> 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

Expert-parallel Mixture-of-Experts layer for the Inkling model.

Only the MoE feed-forward is reimplemented here; every other Inkling submodule
(attention, short convolutions, relative-position logits, norms, vision/audio
towers) is reused verbatim from HuggingFace transformers. The routing math and
shared-expert formulation mirror `transformers.models.inkling.modeling_inkling`
exactly, while the routed experts run through NeMo AutoModel's
:class:`~nemo_automodel.components.moe.experts.GroupedExperts`, which provides
grouped-GEMM compute and expert-parallel (EP) sharding.

## 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 an existing 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                                                             |
| ------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------- |
| [`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.

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

```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:**

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

Expose the correction bias under the Transformers-compatible name.

```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:**

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.

```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:**

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

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.

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

Apply the shared experts.

**Parameters:**

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

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(
    source: torch.nn.Module
)
```

**Bases:** `Module`

Keep short-convolution weights in an existing model-owned fp32 holder.

```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:**

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

Optional cache holding the per-layer conv state.

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

**Returns:** `torch.Tensor`

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

```python
class nemo_automodel.components.models.inkling.layers._InklingShortConvolutionFP32(
    source: torch.nn.Module
)
```

**Bases:** `Module`

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

```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:**

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

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

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.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:**

HuggingFace `InklingTextConfig`.

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:**

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).

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

**Returns:** `torch.Tensor`

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