> This page is for version 0.5.0 · 26.06.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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.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`](#nemo_automodel-components-models-deepseek_v41-layers-DeepseekV41HyperConnection) | Predict one sublayer's residual mixing coefficients in FP32.                                        |
| [`DeepseekV41Mix`](#nemo_automodel-components-models-deepseek_v41-layers-DeepseekV41Mix)                         | FP32 coefficients: pre/post \[batch, sequence, streams], comb \[batch, sequence, streams, streams]. |
| [`DeepseekV41RMSNorm`](#nemo_automodel-components-models-deepseek_v41-layers-DeepseekV41RMSNorm)                 | Normalize in FP32 and multiply the scale before casting the result.                                 |

### Functions

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

### Data

[`_HC_COMPILED`](#nemo_automodel-components-models-deepseek_v41-layers-_HC_COMPILED)

[`_NORM_COMPILED`](#nemo_automodel-components-models-deepseek_v41-layers-_NORM_COMPILED)

[`_hc_collapse`](#nemo_automodel-components-models-deepseek_v41-layers-_hc_collapse)

[`_hc_expand`](#nemo_automodel-components-models-deepseek_v41-layers-_hc_expand)

[`_hc_project`](#nemo_automodel-components-models-deepseek_v41-layers-_hc_project)

[`_rms_norm`](#nemo_automodel-components-models-deepseek_v41-layers-_rms_norm)

### API

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

---

**`fn`**

---

**`iterations`** `= config.hc_sinkhorn_iters`

---

**`norm_eps`** `= config.rms_norm_eps`

---

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

---

**`streams`** `= config.hc_mult`

---

```python
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.

```python
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.

```python
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]

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

Initialize all coefficients before checkpoint-free execution.

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

---

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

---

```python
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.

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

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

```python
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.

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

```python
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.

```python
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.

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

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

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

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

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

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