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`	Full-sequence CSA2 with local KV, shared compressed KV, and an attention sink.
`DeepseekV41AttentionOutput`	Attention result and the shared state for the next layer.
`DeepseekV41AttentionState`	CSA2 state owned by one full-sequence model forward.
`_Compressor`	Non-overlapping channelwise softmax pooling; ratio one is a projection.
`_CompressorLinear`	Keep pooling weights in FP32 while selecting compute from the input dtype.
`_Indexer`	Frozen released CSA2 indexer with shared keys and hierarchical selection.
`_RotaryEmbedding`	Adjacent-pair RoPE with the official frequency-only YaRN adjustment.

Functions

Name	Description
`_apply_rope`	Rotate the final channels without changing the input storage.
`_select_candidate_blocks`	Keep high-scoring blocks and always retain the latest visible block.

API

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

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.

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 | NoneDefaults to 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 | NoneDefaults to None

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

cp_group

dist.ProcessGroup | NoneDefaults to 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,

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

floatDefaults to 0.02

Standard deviation of projection weights.

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.

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

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

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

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

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

Bases: Linear

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

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

Project […, input] into […, output] with the input’s compute dtype.

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

wq_b

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]

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 | NoneDefaults to None

Optional physical global positions [batch, local_sequence].

packed_seq_ids

torch.Tensor | NoneDefaults to None

Optional document IDs [batch, local_sequence], zero for padding.

cp_group

dist.ProcessGroup | NoneDefaults to 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],

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

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

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

boolDefaults to False

Conjugate the rotation for the attention output.

Returns: torch.Tensor

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

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