> 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.inkling.layers

Native feed-forward and short-convolution layers for Inkling.

The equations follow the published Transformers 5.14 Inkling implementation,
while the routed experts use NeMo AutoModel's grouped-GEMM and expert-parallel
implementations.

## Module Contents

### Classes

| Name                                                                                                            | Description                                                                        |
| --------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------------- |
| [`InklingCorrectionBias`](#nemo_automodel-components-models-inkling-layers-InklingCorrectionBias)               | Own and apply Inkling's trained fp32 router correction bias.                       |
| [`InklingDenseMLP`](#nemo_automodel-components-models-inkling-layers-InklingDenseMLP)                           | Dense Inkling MLP retaining the checkpoint's fused interleaved layout.             |
| [`InklingGate`](#nemo_automodel-components-models-inkling-layers-InklingGate)                                   | Top-k router matching `transformers` `InklingTopkRouter`.                          |
| [`InklingMoE`](#nemo_automodel-components-models-inkling-layers-InklingMoE)                                     | Inkling MoE feed-forward: custom router + shared experts + grouped routed experts. |
| [`InklingSharedExperts`](#nemo_automodel-components-models-inkling-layers-InklingSharedExperts)                 | Always-on shared experts, matching `transformers` `InklingSharedExperts`.          |
| [`InklingShortConvolution`](#nemo_automodel-components-models-inkling-layers-InklingShortConvolution)           | Keep short-convolution weights in a model-owned fp32 holder.                       |
| [`_InklingShortConvolutionFP32`](#nemo_automodel-components-models-inkling-layers-_InklingShortConvolutionFP32) | Run one short convolution inside a dtype-uniform fp32 FSDP unit.                   |

### Functions

| Name                                                                                                    | Description                                                             |
| ------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`_causal_conv1d`](#nemo_automodel-components-models-inkling-layers-_causal_conv1d)                     | Apply a causal depthwise convolution to a full sequence.                |
| [`_causal_conv1d_update`](#nemo_automodel-components-models-inkling-layers-_causal_conv1d_update)       | Update cached depthwise-convolution state for one decoding step.        |
| [`_mask_padding_states`](#nemo_automodel-components-models-inkling-layers-_mask_padding_states)         | Zero padded tokens before a short convolution.                          |
| [`build_inkling_moe_config`](#nemo_automodel-components-models-inkling-layers-build_inkling_moe_config) | Build the routed-expert :class:`MoEConfig` from an `InklingTextConfig`. |
| [`inkling_swiglu`](#nemo_automodel-components-models-inkling-layers-inkling_swiglu)                     | Apply SwiGLU to Inkling's interleaved `[gate, up]` projection layout.   |

### API

```python
class nemo_automodel.components.models.inkling.layers.InklingCorrectionBias(
    num_experts: int
)
```

**Bases:** `Module`

Own and apply Inkling's trained fp32 router correction bias.

**`e_score_correction_bias`**

---

```python
nemo_automodel.components.models.inkling.layers.InklingCorrectionBias.forward(
    routed_scores: torch.Tensor
) -> torch.Tensor
```

Add the correction bias in the routing score dtype.

```python
class nemo_automodel.components.models.inkling.layers.InklingDenseMLP(
    config
)
```

**Bases:** `Module`

Dense Inkling MLP retaining the checkpoint's fused interleaved layout.

**`down_proj`**

---

**`gate_up_proj`**

---

**`global_scale`** `= nn.Parameter(torch.ones(1, dtype=model_dtype))`

---

```python
nemo_automodel.components.models.inkling.layers.InklingDenseMLP.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor
```

Run the dense MLP over the fused interleaved gate/up projection.

**Parameters:**

**`hidden_states`** `torch.Tensor`

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

---

**Returns:** `torch.Tensor`

torch.Tensor: Tensor of shape `[..., hidden]`. `gate_up_proj` maps

```python
nemo_automodel.components.models.inkling.layers.InklingDenseMLP.init_weights(
    init_std: float = 0.02
) -> None
```

```python
class nemo_automodel.components.models.inkling.layers.InklingGate(
    config,
    gate_precision: torch.dtype | None = torch.float32
)
```

**Bases:** `Module`

Top-k router matching `transformers` `InklingTopkRouter`.

The gate weight bank additionally holds `n_shared_experts` rows whose logits
participate in the softmax normalization; the resulting shared weights
(`gammas`) scale the always-on shared experts. Selection uses sigmoid scores
plus a per-expert correction bias, while the returned weights come from a
log-sigmoid softmax over the selected routed logits concatenated with the
shared logits, scaled by `route_scale * global_scale`.

**`_fp32_params`** `= InklingCorrectionBias(self.num_experts)`

---

**`e_score_correction_bias`** `Parameter`

Expose the correction bias under the Transformers-compatible name.

---

**`global_scale`** `= nn.Parameter(torch.ones(1, dtype=model_dtype))`

---

**`hidden_dim`** `= config.hidden_size`

---

**`n_shared_experts`** `= config.n_shared_experts`

---

**`n_total_experts`** `= self.num_experts + self.n_shared_experts`

---

**`num_experts`** `= config.n_routed_experts`

---

**`route_scale`** `= config.route_scale`

---

**`top_k`** `= config.num_experts_per_tok`

---

**`weight`**

---

```python
nemo_automodel.components.models.inkling.layers.InklingGate.forward(
    hidden_states: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]
```

Route tokens to experts.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Input tokens of shape `[num_tokens, hidden_dim]`.

---

**Returns:** `torch.Tensor`

`(topk_weights, topk_indices, shared_gammas)` where

```python
nemo_automodel.components.models.inkling.layers.InklingGate.update_bias() -> None
```

No-op: Inkling's correction bias is a trained parameter, not an aux-loss update.

```python
class nemo_automodel.components.models.inkling.layers.InklingMoE(
    config,
    backend: nemo_automodel.components.models.common.utils.BackendConfig,
    moe_config: nemo_automodel.components.moe.config.MoEConfig | None = None
)
```

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

Inkling MoE feed-forward: custom router + shared experts + grouped routed experts.

Subclasses :class:`~nemo_automodel.components.moe.layers.MoE` so the distributed
parallelizer (which discovers MoE blocks via `isinstance` and shards
`self.experts` across the expert-parallel mesh) treats it uniformly, but the
generic `__init__` / `forward` are replaced to match Inkling's routing.

**`gate`** `= InklingGate(config, gate_precision=gate_precision)`

---

**`moe_config`**

---

**`shared_experts`** `= InklingSharedExperts(config)`

---

```python
nemo_automodel.components.models.inkling.layers.InklingMoE.forward(
    hidden_states: torch.Tensor,
    padding_mask: torch.Tensor | None = None
) -> torch.Tensor
```

Route tokens through the shared and routed experts and recombine.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Input of shape `[..., hidden_dim]` (the trailing
dimension is the model dim; leading dims are flattened to tokens).

---

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

Optional boolean mask of padding positions; padded tokens
are excluded from routed-expert dispatch.

---

**Returns:** `torch.Tensor`

torch.Tensor: Output with the same shape as `hidden_states`.

```python
nemo_automodel.components.models.inkling.layers.InklingMoE.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

Initialize MoE parameters in-place (used for from-scratch training).

```python
class nemo_automodel.components.models.inkling.layers.InklingSharedExperts(
    config
)
```

**Bases:** `Module`

Always-on shared experts, matching `transformers` `InklingSharedExperts`.

The fused gate/up projection stays in the checkpoint's interleaved layout so
distributed loading only needs transpose views. The per-token `gammas`
(from :class:`InklingGate`) weight each expert's output before summation.

**`down_proj`**

---

**`gate_up_proj`**

---

**`n_shared_experts`** `= config.n_shared_experts`

---

```python
nemo_automodel.components.models.inkling.layers.InklingSharedExperts.forward(
    hidden_states: torch.Tensor,
    gammas: torch.Tensor
) -> torch.Tensor
```

Apply the shared experts.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Input tokens of shape `[num_tokens, hidden_dim]`.

---

**`gammas`** `torch.Tensor`

Per-token shared weights of shape `[num_tokens, n_shared_experts]`.

---

**Returns:** `torch.Tensor`

torch.Tensor: Output of shape `[num_tokens, hidden_dim]`.

```python
class nemo_automodel.components.models.inkling.layers.InklingShortConvolution(
    hidden_size: int,
    conv_kernel_size: int,
    layer_idx: int,
    conv_idx: int
)
```

**Bases:** `Module`

Keep short-convolution weights in a model-owned fp32 holder.

**`_fp32_params`**

---

```python
nemo_automodel.components.models.inkling.layers.InklingShortConvolution.forward(
    hidden_states: torch.Tensor,
    past_key_values: typing.Any | None = None,
    conv_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Run the short convolution in fp32 and cast the result back.

**Parameters:**

**`hidden_states`** `torch.Tensor`

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

---

**`past_key_values`** `Any | None` — default: None

Optional cache holding the per-layer conv state.

---

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

Optional boolean padding mask of shape `[batch, sequence]`.

---

**Returns:** `torch.Tensor`

torch.Tensor: Tensor of shape `[batch, sequence, hidden]` in the input

```python
nemo_automodel.components.models.inkling.layers.InklingShortConvolution.init_weights(
    init_std: float
) -> None
```

Initialize the fp32 depthwise-convolution weights.

```python
class nemo_automodel.components.models.inkling.layers._InklingShortConvolutionFP32(
    hidden_size: int,
    conv_kernel_size: int,
    layer_idx: int,
    conv_idx: int
)
```

**Bases:** `Module`

Run one short convolution inside a dtype-uniform fp32 FSDP unit.

**`weight`**

---

```python
nemo_automodel.components.models.inkling.layers._InklingShortConvolutionFP32.forward(
    hidden_states: torch.Tensor,
    past_key_values: typing.Any | None = None,
    conv_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor
```

Apply the causal short convolution with a residual connection.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape `[batch, sequence, hidden]`. Internally
transposed to `[batch, hidden, sequence]` for the depthwise
`conv1d` and transposed back before the residual add.

---

**`past_key_values`** `Any | None` — default: None

Optional cache holding the per-layer conv state; updated
in place during incremental decoding.

---

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

Optional boolean padding mask of shape `[batch, sequence]`
that zeroes padded positions before the convolution.

---

**Returns:** `torch.Tensor`

torch.Tensor: Tensor of shape `[batch, sequence, hidden]` (input plus

```python
nemo_automodel.components.models.inkling.layers._causal_conv1d(
    hidden_states: torch.Tensor,
    weight: torch.Tensor
) -> torch.Tensor
```

Apply a causal depthwise convolution to a full sequence.

**Parameters:**

**`hidden_states`** `torch.Tensor`

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

---

**`weight`** `torch.Tensor`

Depthwise weights of shape `[hidden, 1, kernel]`.

---

**Returns:** `torch.Tensor`

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

```python
nemo_automodel.components.models.inkling.layers._causal_conv1d_update(
    hidden_states: torch.Tensor,
    conv_state: torch.Tensor,
    weight: torch.Tensor
) -> torch.Tensor
```

Update cached depthwise-convolution state for one decoding step.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape `[batch, hidden, sequence]` containing new tokens.

---

**`conv_state`** `torch.Tensor`

Mutable cache tensor of shape `[batch, hidden, kernel]`.

---

**`weight`** `torch.Tensor`

Depthwise weights of shape `[hidden, 1, kernel]`.

---

**Returns:** `torch.Tensor`

Tensor of shape `[batch, hidden, sequence]`. `conv_state` is updated in place.

```python
nemo_automodel.components.models.inkling.layers._mask_padding_states(
    hidden_states: torch.Tensor,
    attention_mask: torch.Tensor | None
) -> torch.Tensor
```

Zero padded tokens before a short convolution.

**Parameters:**

**`hidden_states`** `torch.Tensor`

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

---

**`attention_mask`** `torch.Tensor | None`

Optional boolean tensor of shape `[batch, sequence]` where
`True` marks a valid token.

---

**Returns:** `torch.Tensor`

Tensor of shape `[batch, sequence, hidden]`. The output aliases the input

```python
nemo_automodel.components.models.inkling.layers.build_inkling_moe_config(
    text_config,
    backend: nemo_automodel.components.models.common.utils.BackendConfig
) -> nemo_automodel.components.moe.config.MoEConfig
```

Build the routed-expert :class:`MoEConfig` from an `InklingTextConfig`.

Only the fields consumed by :class:`GroupedExperts` matter here (routing and
shared experts are handled by :class:`InklingGate` / :class:`InklingSharedExperts`),
so `n_shared_experts` is set to 0 and the gate-specific fields are left at
neutral defaults.

**Parameters:**

**`text_config`**

Local `InklingTextConfig`.

---

**`backend`** `BackendConfig`

Backend configuration selecting the expert compute/dispatch kernels.

---

**Returns:** `MoEConfig`

Configuration for the routed grouped experts.

```python
nemo_automodel.components.models.inkling.layers.inkling_swiglu(
    hidden_states: torch.Tensor,
    routing_weights: torch.Tensor
) -> torch.Tensor
```

Apply SwiGLU to Inkling's interleaved `[gate, up]` projection layout.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape `[..., 2 * intermediate]` with arbitrary
leading dimensions, whose last axis interleaves gate and up channels
as `[g0, u0, g1, u1, ...]` (`::2` selects gate, `1::2` selects up).

---

**`routing_weights`** `torch.Tensor`

Tensor broadcastable to `[..., intermediate]` that scales
the activated output.

---

**Returns:** `torch.Tensor`

torch.Tensor: Activated tensor of shape `[..., intermediate]` in