> This page is for version 0.7.0 · 26.10.
> 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.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/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.

# nemo_automodel.components.models.inkling2.layers

Differentiable Inkling2 primitives and the common MoE integration.

## Module Contents

### Classes

| Name                                                                                     | Description                                                               |
| ---------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`Inkling2LMHead`](#nemo_automodel-components-models-inkling2-layers-Inkling2LMHead)     | Bias-free output projection with FP32 arithmetic and ordinary FSDP hooks. |
| [`Inkling2MoE`](#nemo_automodel-components-models-inkling2-layers-Inkling2MoE)           | Common Gate/GroupedExperts with Inkling2's activation callable.           |
| [`L2Clip`](#nemo_automodel-components-models-inkling2-layers-L2Clip)                     | Per-head L2 clipping, optionally preceded by learned per-channel scales.  |
| [`ResidualQuery`](#nemo_automodel-components-models-inkling2-layers-ResidualQuery)       | Callable FP32 owner for one set of depth-mixture queries.                 |
| [`ResidualState`](#nemo_automodel-components-models-inkling2-layers-ResidualState)       | Immutable depth state; every tensor has shape \[batch, sequence, hidden]. |
| [`ShortConvolution`](#nemo_automodel-components-models-inkling2-layers-ShortConvolution) | Causal width-W convolution with static or token-generated per-head taps.  |

### Functions

| Name                                                                               | Description                                                                 |
| ---------------------------------------------------------------------------------- | --------------------------------------------------------------------------- |
| [`capped_swiglu`](#nemo_automodel-components-models-inkling2-layers-capped_swiglu) | Apply the interleaved capped activation, with routing before the down GEMM. |
| [`residual_mix`](#nemo_automodel-components-models-inkling2-layers-residual_mix)   | Mix depth sources independently at each token in FP32.                      |

### API

```python
class nemo_automodel.components.models.inkling2.layers.Inkling2LMHead()
```

**Bases:** `Linear`

Bias-free output projection with FP32 arithmetic and ordinary FSDP hooks.

```python
nemo_automodel.components.models.inkling2.layers.Inkling2LMHead.forward(
    input: torch.Tensor
) -> torch.Tensor
```

Project through the module boundary so FSDP gathers its weight.

**Parameters:**

**`input`** `torch.Tensor`

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

---

**Returns:** `torch.Tensor`

FP32 tensor of shape \[..., vocab]. Weight storage may be FP32 or BF16.

```python
class nemo_automodel.components.models.inkling2.layers.Inkling2MoE(
    config: nemo_automodel.components.models.inkling2.config.Inkling2TextConfig,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype
)
```

**Bases:** [MoE](/nemo/automodel/v0.7/nemo-automodel/nemo_automodel/components/moe/layers#nemo_automodel-components-moe-layers-MoE)

Common Gate/GroupedExperts with Inkling2's activation callable.

```python
class nemo_automodel.components.models.inkling2.layers.L2Clip(
    heads: int,
    width: int,
    threshold: float,
    learned_scale: bool
)
```

**Bases:** `Module`

Per-head L2 clipping, optionally preceded by learned per-channel scales.

**`scale`**

---

```python
nemo_automodel.components.models.inkling2.layers.L2Clip.forward(
    x: torch.Tensor,
    tau: torch.Tensor | None = None
) -> torch.Tensor
```

Clip head vectors, then apply optional length scaling.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[batch, sequence, heads, head_dim].

---

**`tau`** `torch.Tensor | None` — default: None

Optional FP32 tensor of shape \[batch, sequence], applied after clipping.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, heads, head_dim], in x's dtype.

```python
class nemo_automodel.components.models.inkling2.layers.ResidualQuery(
    hidden_size: int,
    count: int,
    normalize_values: bool
)
```

**Bases:** `Module`

Callable FP32 owner for one set of depth-mixture queries.

**`weight`**

---

```python
nemo_automodel.components.models.inkling2.layers.ResidualQuery.forward(
    state: nemo_automodel.components.models.inkling2.layers.ResidualState
) -> torch.Tensor
```

Mix immutable depth state.

**Parameters:**

**`state`** `ResidualState`

Completed and partial blocks, each \[batch, sequence, hidden].

---

**Returns:** `torch.Tensor`

FP32 tensor of shape \[queries, batch, sequence, hidden].

```python
class nemo_automodel.components.models.inkling2.layers.ResidualState(
    blocks: tuple[torch.Tensor, ...],
    partial: torch.Tensor
)
```

Dataclass

Immutable depth state; every tensor has shape \[batch, sequence, hidden].

**`blocks`** `tuple[Tensor, ...]`

Completed depth blocks, including the initial embedding block.

---

**`partial`** `Tensor`

Current incomplete block. Neither field is detached or mutated.

---

```python
class nemo_automodel.components.models.inkling2.layers.ShortConvolution(
    channels: int,
    width: int,
    dtype: torch.dtype,
    dynamic_input_size: int = 0,
    heads: int = 1
)
```

**Bases:** `Module`

Causal width-W convolution with static or token-generated per-head taps.

**`cp_mesh`** `DeviceMesh | None = None`

---

**`static_w`**

---

**`weight_proj`**

---

```python
nemo_automodel.components.models.inkling2.layers.ShortConvolution.forward(
    x: torch.Tensor,
    valid: torch.Tensor,
    dynamic_input: torch.Tensor | None = None,
    positions: torch.Tensor | None = None
) -> torch.Tensor
```

Apply convolution with exchanged CP history, zero document history and masked padding.

**Parameters:**

**`x`** `torch.Tensor`

Tensor of shape \[batch, sequence, channels].

---

**`valid`** `torch.Tensor`

Boolean tensor of shape \[batch, sequence], true for real tokens.

---

**`dynamic_input`** `torch.Tensor | None` — default: None

Tensor of shape \[batch, sequence, hidden] for dynamic taps.

---

**`positions`** `torch.Tensor | None` — default: None

Logical document positions \[batch, sequence], including CP offsets.
Taps cannot read before the current document start.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, channels], in x's dtype; padding is zero.

```python
nemo_automodel.components.models.inkling2.layers.ShortConvolution.reset_parameters(
    init_std: float
) -> None
```

Initialize every static and dynamic coefficient without checkpoint data.

```python
nemo_automodel.components.models.inkling2.layers.capped_swiglu(
    gate_up: torch.Tensor,
    weights: torch.Tensor,
    gate_softcap: float,
    up_softcap: float
) -> torch.Tensor
```

Apply the interleaved capped activation, with routing before the down GEMM.

**Parameters:**

**`gate_up`** `torch.Tensor`

Tensor of shape \[tokens, 2 \* intermediate], interleaved gate/up.

---

**`weights`** `torch.Tensor`

Tensor of shape \[tokens, 1], containing selected expert weights.

---

**`gate_softcap`** `float`

Positive gate tanh cap; sigmoid still receives the raw gate.

---

**`up_softcap`** `float`

Positive up tanh cap.

---

**Returns:** `torch.Tensor`

Tensor of shape \[tokens, intermediate], in gate_up's dtype.

```python
nemo_automodel.components.models.inkling2.layers.residual_mix(
    state: nemo_automodel.components.models.inkling2.layers.ResidualState,
    queries: torch.Tensor,
    normalize_values: bool
) -> torch.Tensor
```

Mix depth sources independently at each token in FP32.

**Parameters:**

**`state`** `ResidualState`

Completed blocks and partial, each \[batch, sequence, hidden].

---

**`queries`** `torch.Tensor`

Tensor of shape \[queries, hidden], with one or three pseudo-queries.

---

**`normalize_values`** `bool`

Normalize values as well as scoring keys.

---

**Returns:** `torch.Tensor`

Tensor of shape \[queries, batch, sequence, hidden], without mutating state.