nemo_automodel.components.models.deepseek_v41.layers

Single-pass mHC and modality-aware routing for DeepSeek V4.1.

The coefficient handoff follows section 2.4.1 of the official technical report: each sublayer consumes the pre-mix produced by the preceding sublayer.

Module Contents

Classes

Name	Description
`DeepseekV41HyperConnection`	Predict one sublayer’s residual mixing coefficients in FP32.
`DeepseekV41Mix`	FP32 coefficients: pre/post [batch, sequence, streams], comb [batch, sequence, streams, streams].
`DeepseekV41RMSNorm`	Normalize in FP32 and multiply the scale before casting the result.

Functions

Name	Description
`_hc_collapse_core`	-
`_hc_expand_core`	-
`_hc_project_core`	Coefficient projection up to the Sinkhorn input: pre, post and the [streams, streams] combination logits.
`_rms_norm_core`	-
`compile_hc_cores`	Wrap the mHC projection, collapse and expand cores with `torch.compile` (`BackendConfig.compile_hc`).
`compile_norm_core`	Wrap the fp32 RMSNorm core with `torch.compile` (`BackendConfig.compile_norm`); runs once per process.

Data

API

class nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41HyperConnection(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    sinkhorn_backend: typing.Literal['torch', 'tilelang'] = 'torch'
)

Bases: Module

Predict one sublayer’s residual mixing coefficients in FP32.

base

eps

= config.hc_eps

iterations

= config.hc_sinkhorn_iters

norm_eps

= config.rms_norm_eps

scale

= nn.Parameter(torch.ones(3, dtype=(torch.float32)))

streams

= config.hc_mult

nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41HyperConnection.collapse(
    hidden_states: torch.Tensor,
    pre_mix: torch.Tensor
) -> torch.Tensor

staticmethod

Collapse streams using the preceding sublayer’s coefficients.

Parameters:

hidden_states

torch.Tensor

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

pre_mix

torch.Tensor

FP32 tensor of shape [batch, sequence, streams].

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden], with the input dtype.

nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41HyperConnection.expand(
    output: torch.Tensor,
    residual: torch.Tensor,
    mix: nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41Mix
) -> torch.Tensor

staticmethod

Mix the sublayer output and residual in the source’s coefficient orientation.

Parameters:

output

torch.Tensor

Tensor of shape [batch, sequence, hidden].

residual

torch.Tensor

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

mix

DeepseekV41Mix

FP32 pre/post tensors of shape [batch, sequence, streams] and comb of shape [batch, sequence, input_streams, output_streams].

Returns: torch.Tensor

Tensor of shape [batch, sequence, streams, hidden], with output’s dtype.

nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41HyperConnection.forward(
    hidden_states: torch.Tensor
) -> nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41Mix

Predict coefficients, preserving projection-before-RMS arithmetic.

Parameters:

hidden_states

torch.Tensor

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

Returns: DeepseekV41Mix

Coefficients with pre/post tensors of shape [batch, sequence, streams]

nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41HyperConnection.reset_parameters(
    std: float = 0.02
) -> None

Initialize all coefficients before checkpoint-free execution.

class nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41Mix()

Bases: NamedTuple

FP32 coefficients: pre/post [batch, sequence, streams], comb [batch, sequence, streams, streams].

comb

Tensor

post

Tensor

pre

Tensor

class nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41RMSNorm(
    dim: int,
    eps: float,
    dtype: torch.dtype
)

Bases: Module

Normalize in FP32 and multiply the scale before casting the result.

weight

= nn.Parameter(torch.ones(dim, dtype=dtype))

nemo_automodel.components.models.deepseek_v41.layers.DeepseekV41RMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Apply the released reference’s RMS normalization.

Parameters:

hidden_states

torch.Tensor

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

Returns: torch.Tensor

Tensor of shape […, hidden], with the input dtype.

nemo_automodel.components.models.deepseek_v41.layers._hc_collapse_core(
    hidden_states: torch.Tensor,
    pre_mix: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v41.layers._hc_expand_core(
    output: torch.Tensor,
    residual: torch.Tensor,
    post: torch.Tensor,
    comb: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v41.layers._hc_project_core(
    hidden_states: torch.Tensor,
    fn: torch.Tensor,
    base: torch.Tensor,
    scale: torch.Tensor,
    streams: int,
    norm_eps: float,
    eps: float
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Coefficient projection up to the Sinkhorn input: pre, post and the [streams, streams] combination logits.

nemo_automodel.components.models.deepseek_v41.layers._rms_norm_core(
    hidden_states: torch.Tensor,
    weight: torch.Tensor,
    eps: float
) -> torch.Tensor

nemo_automodel.components.models.deepseek_v41.layers.compile_hc_cores() -> None

Wrap the mHC projection, collapse and expand cores with torch.compile (BackendConfig.compile_hc).

Runs once per process, static shapes (dynamic=False: the fusion of expand needs concrete shapes; a new sequence length compiles the three cores again). No fullgraph: past dynamo’s recompile limit (many distinct sequence lengths) a core runs eager instead of raising FailOnRecompileLimitHit; the unit tests assert the cores compile without graph breaks. The Sinkhorn normalisation keeps its own backend (TileKernels or torch): compiling the 20-iteration torch loop is slower than the fused TileKernels kernel.

nemo_automodel.components.models.deepseek_v41.layers.compile_norm_core() -> None

Wrap the fp32 RMSNorm core with torch.compile (BackendConfig.compile_norm); runs once per process.

Static shapes: one specialization per norm site (block, attention, compressor and indexer norms) and grad mode. The dynamic-shape variant (Kimi K3’s setting) was measured 15 % slower per call and moved the first-step loss of the 12-layer mini benchmark by 9e-4, while the static kernels reproduced the eager loss to four decimals; no fullgraph, so a shape past dynamo’s recompile limit falls back to eager instead of raising.

nemo_automodel.components.models.deepseek_v41.layers._HC_COMPILED = False

nemo_automodel.components.models.deepseek_v41.layers._NORM_COMPILED = False

nemo_automodel.components.models.deepseek_v41.layers._hc_collapse = _hc_collapse_core

nemo_automodel.components.models.deepseek_v41.layers._hc_expand = _hc_expand_core

nemo_automodel.components.models.deepseek_v41.layers._hc_project = _hc_project_core

nemo_automodel.components.models.deepseek_v41.layers._rms_norm = _rms_norm_core