> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.deepseek_v41.attention

Differentiable CSA2 attention from the official DeepSeek V4.1 inference model.

Full layers publish compressed KV, index keys, and selected positions. Reindex
layers replace only the selection; Reuse layers consume it unchanged. Immutable
per-forward state keeps source gradients intact through activation recomputation.
Weights are dequantized for training while the released FP8 window KV and FP4
compressed KV/indexer representations retain their quantize/dequantize boundaries. The released indexer
weights are frozen because the hard top-k operation provides no language-model
gradient and the inference release does not implement indexer distillation.

## Module Contents

### Classes

| Name                                                                                                                | Description                                                                    |
| ------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------ |
| [`DeepseekV41Attention`](#nemo_automodel-components-models-deepseek_v41-attention-DeepseekV41Attention)             | Full-sequence CSA2 with local KV, shared compressed KV, and an attention sink. |
| [`DeepseekV41AttentionOutput`](#nemo_automodel-components-models-deepseek_v41-attention-DeepseekV41AttentionOutput) | Attention result and the shared state for the next layer.                      |
| [`DeepseekV41AttentionState`](#nemo_automodel-components-models-deepseek_v41-attention-DeepseekV41AttentionState)   | CSA2 state owned by one full-sequence model forward.                           |
| [`_Compressor`](#nemo_automodel-components-models-deepseek_v41-attention-_Compressor)                               | Non-overlapping channelwise softmax pooling; ratio one is a projection.        |
| [`_CompressorLinear`](#nemo_automodel-components-models-deepseek_v41-attention-_CompressorLinear)                   | Keep pooling weights in FP32 while selecting compute from the input dtype.     |
| [`_Indexer`](#nemo_automodel-components-models-deepseek_v41-attention-_Indexer)                                     | Frozen released CSA2 indexer with shared keys and hierarchical selection.      |
| [`_RotaryEmbedding`](#nemo_automodel-components-models-deepseek_v41-attention-_RotaryEmbedding)                     | Adjacent-pair RoPE with the official frequency-only YaRN adjustment.           |

### Functions

| Name                                                                                                            | Description                                                          |
| --------------------------------------------------------------------------------------------------------------- | -------------------------------------------------------------------- |
| [`_apply_rope`](#nemo_automodel-components-models-deepseek_v41-attention-_apply_rope)                           | Rotate the final channels without changing the input storage.        |
| [`_select_candidate_blocks`](#nemo_automodel-components-models-deepseek_v41-attention-_select_candidate_blocks) | Keep high-scoring blocks and always retain the latest visible block. |

### API

```python
class nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41Attention(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

Full-sequence CSA2 with local KV, shared compressed KV, and an attention sink.

The training implementation supports eager, SDPA and TileLang attention with
torch linear layers and eager FP32 or TE RMSNorm. Left padding, KV-cache
decoding and tensor sharding remain unsupported. Context parallelism keeps
contiguous local queries and exchanges window KV and shared compressed KV.
Every local sequence must contain complete compression groups.

**`attention_dropout`** `= config.attention_dropout`

---

**`attn_sink`** `Parameter`

Return the checkpoint's per-head FP32 attention sink parameter.

---

**`compress_ratio`** `= config.compress_ratios[layer_idx]`

---

**`compressor`**

---

**`cp_group`** `ProcessGroup | None = None`

---

**`head_dim`** `= config.head_dim`

---

**`indexer`**

---

**`is_index_source`** `= layer_idx in config.index_source_layer_ids`

---

**`is_kv_source`** `= layer_idx in config.kv_source_layer_ids`

---

**`kv_norm`**

---

**`num_groups`** `= config.o_groups`

---

**`num_heads`** `= config.num_attention_heads`

---

**`q_norm`**

---

**`rotary_emb`**

---

**`sinks_param`**

---

**`window_size`** `= config.sliding_window`

---

**`wkv`**

---

**`wo_a`**

---

**`wo_b`**

---

**`wq_a`**

---

**`wq_b`**

---

```python
nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41Attention._project_output(
    attended: torch.Tensor,
    angles: torch.Tensor,
    valid_tokens: torch.Tensor
) -> torch.Tensor
```

Undo RoPE on \[batch, sequence, heads, head\_dim] and project to hidden width.

`angles` has shape \[batch, sequence, rotary\_pairs]; `valid_tokens`
is boolean \[batch, sequence]. Return \[batch, sequence, hidden] with
padded queries zeroed, preserving the attention output dtype.

```python
nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41Attention.forward(
    hidden_states: torch.Tensor,
    position_ids: torch.Tensor,
    state: nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionState,
    attention_mask: torch.Tensor | None = None,
    cp_group: torch.distributed.ProcessGroup | None = None,
    packed_seq_ids: torch.Tensor | None = None
) -> nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionOutput
```

Apply attention and publish immutable state for the next layer.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

---

**`position_ids`** `torch.Tensor`

Integer tensor of shape \[batch, sequence] or \[1, sequence],
containing contiguous global positions for unpacked input and
document-local positions for packed input.

---

**`packed_seq_ids`** `torch.Tensor | None` — default: None

Optional document IDs \[batch, local\_sequence], zero for padding.
Document starts must be aligned to compression-group boundaries.

---

**`state`** `DeepseekV41AttentionState`

Per-forward tensors documented in DeepseekV41AttentionState.
Consumers must receive the state from their preceding layer.

---

**`attention_mask`** `torch.Tensor | None` — default: None

Optional binary right-padding mask of shape
\[batch, local\_sequence], with one for tokens and zero for padding.

---

**`cp_group`** `dist.ProcessGroup | None` — default: None

Optional CP group, overriding setup\_cp\_attention. Hidden
states and positions are local shards; shared KV state is global.

---

**Returns:** `DeepseekV41AttentionOutput`

Output with hidden\_states \[batch, sequence, hidden] and the new state,

```python
nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41Attention.reset_parameters(
    init_std: float = 0.02
) -> None
```

Initialize every attention parameter after construction or meta materialization.

**Parameters:**

**`init_std`** `float` — default: 0.02

Standard deviation of projection weights.

---

```python
nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41Attention.setup_cp_attention(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> None
```

Configure the model-owned KV transport through the shared CP hook.

```python
class nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionOutput(
    hidden_states: torch.Tensor,
    state: nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionState
)
```

Dataclass

Attention result and the shared state for the next layer.

**`hidden_states`** `Tensor`

---

**`state`** `DeepseekV41AttentionState`

---

```python
class nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionState(
    compressed_kv: torch.Tensor | None = None,
    index_keys: torch.Tensor | None = None,
    topk_indices: torch.Tensor | None = None,
    candidates: torch.Tensor | None = None,
    compressed_valid: torch.Tensor | None = None,
    compression_ratio: int = 0,
    compressed_seq_ids: torch.Tensor | None = None
)
```

Dataclass

CSA2 state owned by one full-sequence model forward.

Under CP, compressed and index-key axes span the global sequence while
sequence axes in topk\_indices and candidates contain only local queries.

Tensor fields retain autograd history and are never modified by consumers.
A model creates an empty state for every forward, including every microbatch.

**`candidates`** `Tensor | None = None`

---

**`compressed_kv`** `Tensor | None = None`

---

**`compressed_seq_ids`** `Tensor | None = None`

---

**`compressed_valid`** `Tensor | None = None`

---

**`compression_ratio`** `int = 0`

---

**`index_keys`** `Tensor | None = None`

---

**`topk_indices`** `Tensor | None = None`

---

```python
class nemo_automodel.components.models.deepseek_v41.attention._Compressor(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    ratio: int,
    dtype: torch.dtype,
    rms_norm: str = 'torch_fp32'
)
```

**Bases:** `Module`

Non-overlapping channelwise softmax pooling; ratio one is a projection.

**`norm`**

---

**`wgate`**

---

**`wkv`**

---

```python
nemo_automodel.components.models.deepseek_v41.attention._Compressor.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Return complete compressed groups before rotary embedding.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, floor(sequence / ratio), head\_dim] in the

```python
class nemo_automodel.components.models.deepseek_v41.attention._CompressorLinear()
```

**Bases:** `Linear`

Keep pooling weights in FP32 while selecting compute from the input dtype.

```python
nemo_automodel.components.models.deepseek_v41.attention._CompressorLinear.forward(
    values: torch.Tensor
) -> torch.Tensor
```

Project \[..., input] into \[..., output] with the input's compute dtype.

```python
class nemo_automodel.components.models.deepseek_v41.attention._Indexer(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    layer_idx: int,
    dtype: torch.dtype,
    rms_norm: str = 'torch_fp32',
    attn_backend: str = 'eager'
)
```

**Bases:** `Module`

Frozen released CSA2 indexer with shared keys and hierarchical selection.

**`candidate_block_size`** `= config.candidate_block_size`

---

**`candidate_topk_blocks`** `= config.candidate_topk_blocks`

---

**`head_dim`** `= config.index_head_dim`

---

**`is_candidate_source`** `= layer_idx == config.candidate_source_layer_id`

---

**`k_norm`**

---

**`num_heads`** `= config.index_n_heads`

---

**`owns_keys`** `= layer_idx in config.kv_source_layer_ids`

---

**`topk`** `= config.index_topk`

---

**`uses_candidates`** `= 0 <= config.candidate_source_layer_id < layer_idx`

---

**`weights_proj`**

---

**`wk`**

---

**`wq_b`**

---

```python
nemo_automodel.components.models.deepseek_v41.attention._Indexer._select_packed(
    scores: torch.Tensor,
    sequence_ids: torch.Tensor,
    state: nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionState
) -> nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionState
```

Run selection on each document's own key interval.

topk tie-breaking depends on the physical row width. Selecting over an entire
pack, even with other documents masked out, can therefore change a document's
attention. Local intervals also reset candidate blocks at document boundaries.

**Parameters:**

**`scores`** `torch.Tensor`

Causal, document-masked scores \[batch, local\_sequence, global\_compressed].

---

**`sequence_ids`** `torch.Tensor`

Document IDs \[batch, local\_sequence], zero for padding.

---

**`state`** `DeepseekV41AttentionState`

Global compressed metadata and keys described by DeepseekV41AttentionState.

---

**Returns:** `DeepseekV41AttentionState`

Immutable state with global-key topk indices \[batch, local\_sequence, topk]

```python
nemo_automodel.components.models.deepseek_v41.attention._Indexer.forward(
    hidden_states: torch.Tensor,
    query_latent: torch.Tensor,
    latent: torch.Tensor | None,
    angles: torch.Tensor,
    compressed_angles: torch.Tensor,
    state: nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionState,
    position_ids: torch.Tensor | None = None,
    cp_group: torch.distributed.ProcessGroup | None = None,
    packed_seq_ids: torch.Tensor | None = None
) -> nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionState
```

Produce index keys when owned, then replace the current selection.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden].

---

**`query_latent`** `torch.Tensor`

Tensor of shape \[batch, sequence, q\_lora\_rank].

---

**`latent`** `torch.Tensor | None`

Optional unrotated KV of shape \[batch, compressed, head\_dim].

---

**`angles`** `torch.Tensor`

FP32 rotary angles of shape \[batch, sequence, rotary\_pairs].

---

**`compressed_angles`** `torch.Tensor`

FP32 angles of shape \[batch, compressed, rotary\_pairs].

---

**`state`** `DeepseekV41AttentionState`

Shared tensors with layouts in DeepseekV41AttentionState.

---

**`position_ids`** `torch.Tensor | None` — default: None

Optional physical global positions \[batch, local\_sequence].

---

**`packed_seq_ids`** `torch.Tensor | None` — default: None

Optional document IDs \[batch, local\_sequence], zero for padding.

---

**`cp_group`** `dist.ProcessGroup | None` — default: None

Optional CP group; latent and compressed\_angles contain
local complete groups, while state index keys span all ranks.

---

**Returns:** `DeepseekV41AttentionState`

New state with index\_keys \[batch, compressed, index\_head\_dim],

```python
class nemo_automodel.components.models.deepseek_v41.attention._RotaryEmbedding(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    compressed: bool
)
```

**Bases:** `Module`

Adjacent-pair RoPE with the official frequency-only YaRN adjustment.

**`beta_fast`**

---

**`beta_slow`**

---

**`dim`** `= config.qk_rope_head_dim`

---

**`factor`**

---

**`original_length`**

---

**`theta`**

---

```python
nemo_automodel.components.models.deepseek_v41.attention._RotaryEmbedding.forward(
    positions: torch.Tensor
) -> torch.Tensor
```

Construct FP32 phase angles without a cast-sensitive frequency buffer.

**Parameters:**

**`positions`** `torch.Tensor`

Integer tensor of shape \[batch, sequence].

---

**Returns:** `torch.Tensor`

FP32 angles of shape \[batch, sequence, rotary\_pairs], where

```python
nemo_automodel.components.models.deepseek_v41.attention._apply_rope(
    values: torch.Tensor,
    angles: torch.Tensor,
    inverse: bool = False
) -> torch.Tensor
```

Rotate the final channels without changing the input storage.

**Parameters:**

**`values`** `torch.Tensor`

Tensor of shape \[batch, sequence, channels] or
\[batch, sequence, heads, channels].

---

**`angles`** `torch.Tensor`

FP32 tensor of shape \[batch, sequence, rotary\_pairs]. The last
2 \* rotary\_pairs channels of values use adjacent-pair rotation.

---

**`inverse`** `bool` — default: False

Conjugate the rotation for the attention output.

---

**Returns:** `torch.Tensor`

Tensor with the shape and dtype of values, in independent storage.

```python
nemo_automodel.components.models.deepseek_v41.attention._select_candidate_blocks(
    scores: torch.Tensor,
    visible_lengths: torch.Tensor,
    topk_blocks: int,
    block_size: int
) -> torch.Tensor
```

Keep high-scoring blocks and always retain the latest visible block.

**Parameters:**

**`scores`** `torch.Tensor`

Causally masked scores of shape \[batch, sequence, compressed].

---

**`visible_lengths`** `torch.Tensor`

Integer tensor of shape \[batch, sequence, 1] counting
visible compressed positions for each query.

---

**`topk_blocks`** `int`

Maximum retained blocks per query.

---

**`block_size`** `int`

Compressed positions per block.

---

**Returns:** `torch.Tensor`

Boolean candidate mask of shape \[batch, sequence, compressed].