> 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. Packed sequences, left padding, KV-cache
decoding, and sequence/context/tensor sharding require additional state rules
and are rejected rather than silently using incorrect compression boundaries.

**`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`**

---

**`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
) -> 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 zero-based positions. Packed/reset and
incremental positions are unsupported.

---

**`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, sequence], with one for tokens and zero for padding.

---

**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
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
)
```

Dataclass

CSA2 state owned by one full-sequence model forward.

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_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'
)
```

**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.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
) -> 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.

---

**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].