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.
`_SelectionRecord`	The value-only state fields one layer produced, replayed to its activation-checkpoint recompute.

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.
`csa2_selection_region`	Run the frozen indexer and the sparse-index planning outside any active TorchDispatch mode.
`use_triton_rope`	Route `_apply_rope` through the fused Triton kernel.
`validate_attention_inputs`	Run the host-synchronising checks on the per-forward inputs once.

Data

_ROPE_IMPL

csa2_selection_replay

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

reuse_selection_on_recompute

= bool(config.reuse_sparse_selection_on_recompute)

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._plan_sparse_indices(
    next_state: nemo_automodel.components.models.deepseek_v41.attention.DeepseekV41AttentionState,
    kv_sequence: int,
    positions: torch.Tensor,
    position_ids: torch.Tensor,
    valid_tokens: torch.Tensor,
    global_valid: torch.Tensor,
    global_seq_ids: torch.Tensor | None,
    packed_seq_ids: torch.Tensor | None,
    global_sequence: int,
    all_rows_nonempty: bool
) -> nemo_automodel.components.models.deepseek_v4.optimized_kernels.SparseIndexPlan

Build the sparse slot indices of one selection lineage and prepare them for the backend.

Slots keep the released order: the sliding window first, then the selected compressed positions offset past the window keys. The result depends only on per-forward constants and on the current selection, so every layer that consumes the same selection shares it.

Parameters:

next_state

DeepseekV41AttentionState

State after this layer’s compressor and indexer; topk_indices is read on compressed layers.

kv_sequence

int

Number of key positions per batch row of the K/V this plan is for.

positions

torch.Tensor

Integer tensor of shape [sequence] with this rank’s global query positions.

position_ids

torch.Tensor

Integer tensor of shape [batch, sequence] or [1, sequence], document-local when packed.

valid_tokens

torch.Tensor

Boolean tensor of shape [batch, sequence]; padded queries get no slots.

global_valid

torch.Tensor

Boolean tensor of shape [batch, global_sequence] over the window keys.

global_seq_ids

torch.Tensor | None

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

packed_seq_ids

torch.Tensor | None

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

global_sequence

int

Number of window key positions per batch row.

all_rows_nonempty

bool

Whether every query has at least one valid slot.

Returns: SparseIndexPlan

The backend plan built from slot indices [batch, sequence, padded_slots] (window slots, then

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,
    inputs_validated: bool = False
) -> 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.

inputs_validated

boolDefaults to False

True when the caller already ran validate_attention_inputs on these position_ids / attention_mask / packed_seq_ids (the model does, once per forward); the layer then skips its own host-synchronising checks.

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

Tensor of shape [batch, sequence, hidden].

state

DeepseekV41AttentionState

Shared tensor layouts documented by 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,
    sparse_plan: nemo_automodel.components.models.deepseek_v4.optimized_kernels.SparseIndexPlan | 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

Optional candidate mask of shape [batch, sequence, compressed].

compressed_kv

Tensor | None = None

Rotated KV tensor of shape [batch, compressed, head_dim].

compressed_seq_ids

Tensor | None = None

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

compressed_valid

Tensor | None = None

Optional valid-group mask of shape [batch, compressed].

compression_ratio

int = 0

Number of tokens represented by a compressed position.

index_keys

Tensor | None = None

Rotated index keys of shape [batch, compressed, index_head_dim].

sparse_plan

SparseIndexPlan | None = None

Optional SparseIndexPlan of the current selection for the TileLang / cuDNN backends, built by the first layer that consumes the selection and reused by every later layer of the same lineage; value-only, dropped whenever the selection changes.

topk_indices

Tensor | None = None

Compressed-relative indices of shape [batch, sequence, topk], with -1 denoting absent positions.

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 using the input’s explicit compute dtype.

Parameters:

values

torch.Tensor

Tensor of shape […, input] with arbitrary leading dimensions.

Returns: torch.Tensor

Tensor of shape […, output] in the input 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

class nemo_automodel.components.models.deepseek_v41.attention._SelectionRecord(
    index_keys: torch.Tensor | None,
    topk_indices: torch.Tensor | None,
    candidates: torch.Tensor | None,
    sparse_plan: nemo_automodel.components.models.deepseek_v4.optimized_kernels.SparseIndexPlan | None
)

Dataclass

The value-only state fields one layer produced, replayed to its activation-checkpoint recompute.

candidates

Tensor | None

Optional candidate mask of shape [batch, sequence, compressed].

index_keys

Tensor | None

Rotated index keys of shape [batch, compressed, index_head_dim], or None.

sparse_plan

SparseIndexPlan | None

The lineage’s SparseIndexPlan, or None for the dense backends.

topk_indices

Tensor | None

Compressed-relative selection of shape [batch, sequence, topk], -1 absent, or None.

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

nemo_automodel.components.models.deepseek_v41.attention.csa2_selection_region() -> collections.abc.Iterator[None]

Run the frozen indexer and the sparse-index planning outside any active TorchDispatch mode.

Selective activation checkpointing records every op of the checkpoint forward through a TorchDispatch mode and indexes the recompute against that record. Replaying the recorded selection skips these ops in the recompute, so they must be invisible to the mode in the forward as well (as Qwen3.8-Flash-Next’s qsa_route_selection_region does). Everything inside runs under torch.no_grad; hiding the region only affects checkpoint bookkeeping.

nemo_automodel.components.models.deepseek_v41.attention.use_triton_rope(
    enabled: bool
) -> None

Route _apply_rope through the fused Triton kernel.

The Triton path is opt-in and off by default; the eager path stays the reference.

Parameters:

enabled

bool

Use apply_rope_triton when True, the eager implementation when False.

nemo_automodel.components.models.deepseek_v41.attention.validate_attention_inputs(
    position_ids: torch.Tensor,
    attention_mask: torch.Tensor | None,
    packed_seq_ids: torch.Tensor | None,
    sequence: int,
    cp_group: torch.distributed.ProcessGroup | None = None
) -> None

Run the host-synchronising checks on the per-forward inputs once.

Each check reads a device tensor on the host (torch.equal / torch.all / torch.any), which drains the CUDA launch queue. The model calls this once per forward and passes inputs_validated=True to every layer; a layer called directly (inputs_validated=False) still validates its own inputs. Shape checks stay in the layer because they cost nothing.

nemo_automodel.components.models.deepseek_v41.attention._ROPE_IMPL = 'torch'

nemo_automodel.components.models.deepseek_v41.attention.csa2_selection_replay: RecomputeReplay[_SelectionRecord] = RecomputeReplay('CSA2 selection')

Record/replay channel for the frozen indexer’s selection and the sparse-index plan across AC recompute.

DeepseekV41Block.nemo_checkpoint_context_fn binds it around each checkpointed block call when DeepseekV41TextConfig.reuse_sparse_selection_on_recompute is set: the checkpoint forward records what a layer adds to the state, the recompute takes it back instead of re-running the indexer and the index prologue. Both are integer / frozen-score work autograd never reads, so gradients are unchanged.