nemo_automodel.components.models.deepseek_v4.layers

Bases: Module

Sliding-window attention + Compressor + Indexer + attention sink.

Single-head KV (num_key_value_heads=1), grouped low-rank output via DeepseekV4GroupedLinear. compress_ratio == 0 layers skip the compressor / indexer and run pure SWA.

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4Attention._project_query(
    q_residual: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4Attention.forward(
    hidden_states: torch.Tensor,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    attention_mask: torch.Tensor | None = None,
    position_embeddings_compress: tuple[torch.Tensor, torch.Tensor] | None = None,
    rotary_compress: torch.nn.Module | None = None,
    start_pos: int = 0,
    position_ids: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, torch.Tensor | None]

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4Attention.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4Attention.setup_cp_attention(
    cp_mesh
) -> None

Bases: Module

HF PR 45616 port. Long-range KV branch. Pools compress_ratio tokens into one compressed KV; when ratio == 4 the Indexer narrows the pool.

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4Compressor._compute_fsdp_group_has_complete_hca_window(
    local_has_complete_hca_window: bool,
    device: torch.device
) -> bool

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4Compressor._set_hca_param_sync_group(
    process_group
) -> None

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4Compressor.forward(
    hidden_states: torch.Tensor,
    q_residual: torch.Tensor | None,
    rotary: torch.nn.Module,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    cache: nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4TrainCache,
    layer_idx: int,
    start_pos: int,
    enable_hca_fsdp_graph_alignment: bool = False,
    position_ids: torch.Tensor | None = None,
    packed_seq_ids: torch.Tensor | None = None,
    cp_group = None
) -> tuple[torch.Tensor, torch.LongTensor | None, torch.Tensor | None, torch.Tensor | None]

Bases: Module

Callable holder for fp32 tensors that need their own FSDP unit.

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4FP32Parameter.forward(
    reference: torch.Tensor | None = None
) -> torch.Tensor

Bases: Linear

Block-diagonal grouped linear (HF PR 45616 port).

weight parameter has the standard nn.Linear shape [out_features, in_features_per_group] so quantizers keyed on nn.Linear.weight still find it; forward does per-group bmm.

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4GroupedLinear.forward(
    x: torch.Tensor
) -> torch.Tensor

Bases: Module

Per-site HyperConnection mixer (attention or FFN). Ported from transformers/src/transformers/models/deepseek_v4/modular_deepseek_v4.py class DeepseekV4HyperConnection.

Owns fn (packed linear), base (bias), and scale (scalar per-head gains). compute_weights produces three mixer tensors:

• pre [B, S, H] : sigmoid-gated collapse weights
• post [B, S, H] : sigmoid-gated expand weights
• comb [B, S, H, H] : doubly-stochastic combination matrix from Sinkhorn-normalising sigmoid gates

All math runs in fp32 regardless of the outer cast policy; parameters cast themselves via .float() on each forward. HF lists these params in _keep_in_fp32_modules_strict — the KAutomodel adapter does the same via submodule-name matching.

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4HyperConnection._expand_fused(
    output: torch.Tensor,
    residual: torch.Tensor,
    post: torch.Tensor,
    comb: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4HyperConnection._project(
    hidden_streams: torch.Tensor,
    fn: torch.Tensor,
    base: torch.Tensor,
    scale: torch.Tensor,
    hc: int,
    norm_eps: float,
    hc_eps: float
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4HyperConnection.collapse(
    hidden_streams: torch.Tensor,
    pre: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4HyperConnection.compute_weights(
    hidden_streams: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4HyperConnection.expand(
    output: torch.Tensor,
    residual: torch.Tensor,
    post: torch.Tensor,
    comb: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4HyperConnection.forward(
    hidden_streams: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4HyperConnection.init_weights(
    init_std: float
) -> None

Bases: Module

Final HC-stream collapse before the shared RMSNorm + lm_head. Ported from modular_deepseek_v4.py class DeepseekV4HyperHead.

Sigmoid-weighted sum over the hc_mult streams (no Sinkhorn). Used once at the end of DeepseekV4Model.forward to go from [B, S, H, D] back to [B, S, D].

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4HyperHead.forward(
    x: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4HyperHead.init_weights(
    init_std: float
) -> None

Bases: Module

HF PR 45616 port. Picks the top-k compressed positions per query when compress_ratio == 4. Owns its own pool at index_head_dim plus a query projection + weights_proj head-mixer.

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4Indexer._select_keys(
    hidden_states: torch.Tensor,
    q_residual: torch.Tensor,
    rotary: torch.nn.Module,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    cache: nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4TrainCache,
    layer_idx: int,
    start_pos: int,
    position_ids: torch.Tensor | None = None,
    packed_seq_ids: torch.Tensor | None = None,
    cp_group: torch.distributed.ProcessGroup | None = None
) -> torch.LongTensor

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4Indexer.forward(
    hidden_states: torch.Tensor,
    q_residual: torch.Tensor,
    rotary: torch.nn.Module,
    position_embeddings: tuple[torch.Tensor, torch.Tensor],
    cache: nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4TrainCache,
    layer_idx: int,
    start_pos: int,
    position_ids: torch.Tensor | None = None,
    packed_seq_ids: torch.Tensor | None = None,
    cp_group: torch.distributed.ProcessGroup | None = None
) -> torch.LongTensor

Bases: Module

V4 rotary embedding. Produces (cos, sin) sized to qk_rope_head_dim (via partial_rotary_factor = qk_rope_head_dim / head_dim), matching HF.

YaRN: when rope_scaling is a YaRN-typed dict ({"type": "yarn", "factor": F, "original_max_position_embeddings": L0, "beta_fast": ..., "beta_slow": ...}), modify inv_freq per dsv4flash/inference/model.py:precompute_freqs_cis — frequency interpolation with a smooth linear ramp between beta_fast/beta_slow correction dims. Used by the compress-rope (theta=160000) on layers with compress_ratio > 0. The main rope (theta=10000, used only on sliding-window layers) gets rope_scaling=None because the reference builds it with original_seq_len=0 for those layers.

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4RotaryEmbedding.forward(
    x: torch.Tensor,
    position_ids: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Training-only cache shim mirroring the three methods DeepseekV4Compressor / DeepseekV4Indexer call on DeepseekV4Cache.

KAutomodel training forward is stateless — we never persist KV or compressor windows across calls. Each DeepseekV4Attention.forward creates a fresh cache instance, which holds per-layer scratch dicts for the duration of the call. When a full window hasn’t accumulated yet we return an empty tensor and let the downstream code handle it.

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4TrainCache._branch_state(
    state_key: str,
    layer_idx: int
) -> dict

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4TrainCache.accumulate_windows(
    kv: torch.Tensor,
    gate: torch.Tensor,
    layer_idx: int,
    state_key: str,
    ratio: int,
    start_pos: int
) -> tuple[torch.Tensor, torch.Tensor, int]

nemo_automodel.components.models.deepseek_v4.layers.DeepseekV4TrainCache.update_pool(
    new_pooled: torch.Tensor,
    layer_idx: int,
    state_key: str
) -> torch.Tensor

Bases: Function

Per-group matmul that reads and writes token-major activations in place.

torch.bmm wants the group axis first. Permuting a token-major [tokens, groups, d_in] activation is a free view on the input side, but the [groups, tokens, d_out] result had to be permuted back and copied (flatten / reshape by the caller) and, in backward, the permuted input gradient was copied again by the reshape that produced the grouped view (ViewBackward): two 0.5 GB copies per attention layer at the DeepSeek-V4.1 benchmark shape. cuBLAS handles a strided output, so the forward and the input gradient are written straight into token-major buffers through out= and the copies disappear.

nemo_automodel.components.models.deepseek_v4.layers._GroupedLinearTokenMajor.backward(
    ctx: nemo_automodel.components.models.deepseek_v4.layers._GroupedProjectionContext,
    grad_y: torch.Tensor
) -> tuple[torch.Tensor | None, torch.Tensor | None]

nemo_automodel.components.models.deepseek_v4.layers._GroupedLinearTokenMajor.forward(
    ctx: nemo_automodel.components.models.deepseek_v4.layers._GroupedProjectionContext,
    x: torch.Tensor,
    w: torch.Tensor
) -> torch.Tensor

Split x along its last dim into nope (first) and rope (last rope_head_dim) slices, rotate only the rope slice with INTERLEAVED pair-RoPE (pairs (2k, 2k+1)), concat back.

The DSV4-Flash released checkpoint uses interleaved RoPE end-to-end (see dsv4flash/inference/model.py:apply_rotary_emb — complex multiplication on view_as_complex of pairs). HF transformers PR 45616 / PR 45643 ship a Llama-style rotate_half here instead, which pairs (d, d+rd/2). Same algebra but a different dim-to-frequency mapping — the released weights expect the interleaved layout, so the Llama-style helper produces wrong activations on the released checkpoint (verified empirically: kv_post_rope cosine drops from 0.9999 to 0.866 after one block under Llama-style; matches at >0.999 under interleaved).

Parameters:

Returns: torch.Tensor

Tensor with x’s shape and dtype. The leading channels - rope_head_dim

Interleaved RoPE on the last rope_head_dim dims of x (pairs are (2k, 2k+1)). Matches the DeepSeek inference reference’s complex-mul formulation in dsv4flash/inference/model.py:apply_rotary_emb: the released DSV4-Flash weights were trained with this layout, NOT the Llama-style rotate_half layout HF transformers PR 45616/45643 still uses (pairs (d, d+rd/2)).

Inverse rotation: pass -sin instead of sin (caller’s responsibility — same as our existing inverse-rope call site).

Parameters:

Build the additive compressed-position mask for Indexer-selected pool IDs.

Use TileLang DSV4 sparse kernels only when the attention backend requests them.

Use optional TileKernels sinkhorn when available, otherwise torch fallback.

Reshape [B, S, ratio, 2*head_dim] -> [B, S, 2*ratio, head_dim] with the cross-window overlap from the DeepSeek inference reference (Compressor.overlap_transform in dsv4flash/inference/model.py:307-314).

Window N consumes:

• positions [ratio:] of the new tensor: the second half of the feature dim of window N (current block).
• positions [:ratio] of the new tensor: the first half of the feature dim of window N-1 (previous block, i.e. the overlap into the past).

Window 0 has no previous block, so its [:ratio] slice is left at fill_value (0 for the kv tensor, -inf for the score tensor so softmax masks it out).

Softmax-gated sum-pool over ratio consecutive tokens.

Non-overlap mode (HF PR 45616 layout, ratio==128 in V4-Flash): Input kv/gate of shape [B, length, head_dim]. Reshape to [B, length/ratio, ratio, head_dim] and pool over the ratio axis.

Overlap mode (DeepSeek inference reference layout, ratio==4 in V4-Flash): Input kv/gate of shape [B, length, 2*head_dim] (wkv/wgate project to 2*head_dim so each window can carry both its own kv and a half-overlap into the next window). Reshape to [B, length/ratio, ratio, 2*head_dim], apply _overlap_transform to remap to [B, length/ratio, 2*ratio, head_dim], then pool over the 2*ratio axis. Each compressed token thus aggregates 2*ratio = 8 raw tokens — the ratio tokens of the current window plus the ratio tokens of the previous window — giving smoother compression boundaries that the released checkpoint was trained under.

HF PR 45616 omits the overlap path entirely; the released DSV4-Flash safetensors have ape/wkv/wgate shapes that only match the overlap layout ([ratio, 2*head_dim] and [2*head_dim, hidden]), so we must support it here to load the released weights.

RMS-normalize the last dim without materializing an x.square() tensor.

Rotate half the hidden dims of the input (Llama / GPT-NeoX style).

Port of transformers.models.llama.modeling_llama.apply_rotary_pos_emb.

Build a 4D additive causal+padding (+optional sliding-window) mask compatible with eager_attention_with_sink.

Mirrors HF’s create_sliding_window_causal_mask (used in DeepseekV4Model.forward): each query at position i attends only to keys at positions [max(0, i - sliding_window + 1), i]. The DSV4-Flash weights were trained with this banding on every layer, so dropping it makes the softmax see a different distribution than training and degrades loss.

Inputs

attention_mask: 2D [B, S] tensor with 1=valid, 0=padding (HF convention), or already-4D additive mask, or None. sliding_window: if not None, mask out keys further back than this many positions from the query (in addition to causal masking).

Returns: [B, 1, S, S] additive mask of dtype (0 where keep, large negative where mask).

Build a 4D additive block-causal mask from packed-sequence lengths.

Eager attention with per-head sink: appends an extra softmax column whose logit is module.sinks[h] and whose value-slot is zero. Ported verbatim from HF PR 45616.

Port of transformers.models.llama.modeling_llama.repeat_kv.