> 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.moe.mok_experts

Mixture-of-Kittens expert backend for AutoModel's shared MoE component.

## Module Contents

### Classes

| Name                                                                                      | Description                                                          |
| ----------------------------------------------------------------------------------------- | -------------------------------------------------------------------- |
| [`GroupedExpertsMoK`](#nemo_automodel-components-moe-mok_experts-GroupedExpertsMoK)       | AutoModel expert parameters executed by Mixture-of-Kittens.          |
| [`_MoKAutogradFunction`](#nemo_automodel-components-moe-mok_experts-_MoKAutogradFunction) | Connect MoK's explicit backward API to PyTorch autograd.             |
| [`_MoKRuntime`](#nemo_automodel-components-moe-mok_experts-_MoKRuntime)                   | Own one layer's MoK configuration and expert-parallel process group. |

### Functions

| Name                                                                                                            | Description                                                                     |
| --------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`_flatten_mok_tensor_dataclass`](#nemo_automodel-components-moe-mok_experts-_flatten_mok_tensor_dataclass)     | Flatten a MoK schedule/context without retaining tensor references in metadata. |
| [`_load_mok_functional`](#nemo_automodel-components-moe-mok_experts-_load_mok_functional)                       | Import MoK after distributed setup has selected this process's GPU.             |
| [`_local_tensor`](#nemo_automodel-components-moe-mok_experts-_local_tensor)                                     | Return the rank-local tensor while preserving its autograd connection.          |
| [`_unflatten_mok_tensor_dataclass`](#nemo_automodel-components-moe-mok_experts-_unflatten_mok_tensor_dataclass) | Reconstruct a MoK dataclass from checkpoint-managed saved tensors.              |

### Data

[`_MOK_IMPORT_MESSAGE`](#nemo_automodel-components-moe-mok_experts-_MOK_IMPORT_MESSAGE)

[`_mok_functional`](#nemo_automodel-components-moe-mok_experts-_mok_functional)

### API

```python
class nemo_automodel.components.moe.mok_experts.GroupedExpertsMoK(
    config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)
```

**Bases:** `Module`

AutoModel expert parameters executed by Mixture-of-Kittens.

Routed parameters use MoK-native contiguous layouts during training while
state-dict hooks preserve AutoModel's established combined expert keys.

**`down_proj_bias`** `None`

Return `None` because MoK does not support expert bias.

---

**`down_projs`** `Tensor`

Return AutoModel's virtual down-projection state-dict tensor.

---

**`ep_mesh`** `DeviceMesh | None = None`

---

**`ep_rank`** `= 0`

---

**`gate_and_up_projs`** `Tensor`

Return AutoModel's virtual combined GateUp state-dict tensor.

---

**`gate_up_proj_bias`** `None`

Return `None` because MoK does not support expert bias.

---

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

---

**`routed_down_weights`**

---

**`routed_gate_weights`**

---

**`routed_up_weights`**

---

**`runtime`**

---

```python
nemo_automodel.components.moe.mok_experts.GroupedExpertsMoK._load_from_state_dict(
    state_dict: dict[str, torch.Tensor],
    prefix: str,
    local_metadata: dict[str, object],
    strict: bool,
    missing_keys: list[str],
    unexpected_keys: list[str],
    error_msgs: list[str]
) -> None
```

Load established AutoModel expert tensors into MoK-native parameters.

```python
nemo_automodel.components.moe.mok_experts.GroupedExpertsMoK._save_to_state_dict(
    destination: collections.OrderedDict[str, torch.Tensor],
    prefix: str,
    keep_vars: bool
) -> None
```

Save virtual AutoModel expert tensors instead of MoK-native parameters.

```python
nemo_automodel.components.moe.mok_experts.GroupedExpertsMoK._validate_model_mok_config(
    config: nemo_automodel.components.moe.config.MoEConfig
) -> None
```

staticmethod

Reject MoE variants outside MoK's current fused-kernel contract.

```python
nemo_automodel.components.moe.mok_experts.GroupedExpertsMoK.forward(
    x: torch.Tensor,
    weights: torch.Tensor,
    indices: torch.Tensor,
    shared_gate_weights: torch.Tensor,
    shared_up_weights: torch.Tensor,
    shared_down_weights: torch.Tensor
) -> torch.Tensor
```

Execute fused shared and routed SwiGLU experts.

**Parameters:**

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

BF16 tensor of shape \[tokens, hidden].

---

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

Tensor of shape \[tokens, activated\_experts].

---

**`indices`** `torch.Tensor`

Int64 tensor of shape \[tokens, activated\_experts].

---

**`shared_gate_weights`** `torch.Tensor`

BF16 tensor of shape \[expert\_intermediate, hidden].

---

**`shared_up_weights`** `torch.Tensor`

BF16 tensor of shape \[expert\_intermediate, hidden].

---

**`shared_down_weights`** `torch.Tensor`

BF16 tensor of shape \[hidden, expert\_intermediate].

---

**Returns:** `torch.Tensor`

BF16 tensor of shape \[tokens, hidden] containing the fused shared and

```python
nemo_automodel.components.moe.mok_experts.GroupedExpertsMoK.init_token_dispatcher(
    ep_mesh: torch.distributed.device_mesh.DeviceMesh
) -> None
```

Initialize MoK on an expert-parallel device mesh.

**Parameters:**

**`ep_mesh`** `DeviceMesh`

One-dimensional device mesh named `ep`.

---

```python
nemo_automodel.components.moe.mok_experts.GroupedExpertsMoK.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None
```

Initialize MoK-native weights from the canonical expert layout.

Draw random values in the same tensor shapes and order as
:class:`GroupedExperts` before transposing them into MoK's native
layouts.  Otherwise switching only the dispatcher changes a
random-initialized model even when every process uses the same seed.

**Parameters:**

**`buffer_device`** `torch.device`

Device on which initialization kernels execute.

---

**`init_std`** `float` — default: 0.02

Standard deviation of the normal weight initialization.

---

```python
class nemo_automodel.components.moe.mok_experts._MoKAutogradFunction()
```

**Bases:** `Function`

Connect MoK's explicit backward API to PyTorch autograd.

```python
nemo_automodel.components.moe.mok_experts._MoKAutogradFunction.backward(
    ctx: object,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor | None, ...]
```

staticmethod

Run the fused MoK backward pass.

**Parameters:**

**`ctx`** `object`

Autograd-owned context containing the forward schedule and tensors.

---

**`grad_output`** `torch.Tensor`

BF16 tensor of shape \[tokens, hidden].

---

**Returns:** `torch.Tensor | None`

Gradients matching every :meth:`forward` input. The runtime and integer

```python
nemo_automodel.components.moe.mok_experts._MoKAutogradFunction.forward(
    ctx: object,
    runtime: nemo_automodel.components.moe.mok_experts._MoKRuntime,
    x: torch.Tensor,
    router_weights: torch.Tensor,
    top_experts: torch.Tensor,
    shared_gate_weights: torch.Tensor,
    shared_up_weights: torch.Tensor,
    shared_down_weights: torch.Tensor,
    routed_gate_weights: torch.Tensor,
    routed_up_weights: torch.Tensor,
    routed_down_weights: torch.Tensor
) -> torch.Tensor
```

staticmethod

Run a fused MoK forward pass.

**Parameters:**

**`ctx`** `object`

Autograd-owned context.

---

**`runtime`** `_MoKRuntime`

Initialized MoK runtime for this layer's EP group.

---

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

Contiguous BF16 tensor of shape \[tokens, hidden].

---

**`router_weights`** `torch.Tensor`

Contiguous FP32 tensor of shape \[tokens, activated\_experts].

---

**`top_experts`** `torch.Tensor`

Contiguous int64 tensor of shape \[tokens, activated\_experts].

---

**`shared_gate_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape \[expert\_intermediate, hidden].

---

**`shared_up_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape \[expert\_intermediate, hidden].

---

**`shared_down_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape \[hidden, expert\_intermediate].

---

**`routed_gate_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape
\[local\_experts, expert\_intermediate, hidden].

---

**`routed_up_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape
\[local\_experts, expert\_intermediate, hidden].

---

**`routed_down_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape
\[local\_experts, hidden, expert\_intermediate].

---

**Returns:** `torch.Tensor`

BF16 tensor of shape \[tokens, hidden].

```python
class nemo_automodel.components.moe.mok_experts._MoKRuntime(
    mok_config: nemo_automodel.components.models.common.MoKBackendConfig,
    swiglu_limit: float
)
```

Own one layer's MoK configuration and expert-parallel process group.

**`config`** `object | None = None`

---

**`ep_group`** `ProcessGroup | None = None`

---

**`swiglu_limit`**

---

```python
nemo_automodel.components.moe.mok_experts._MoKRuntime.backward(
    schedule: object,
    forward_context: object,
    grad_output: torch.Tensor,
    x: torch.Tensor,
    router_weights: torch.Tensor,
    shared_gate_weights: torch.Tensor,
    shared_up_weights: torch.Tensor,
    shared_down_weights: torch.Tensor,
    routed_gate_weights: torch.Tensor,
    routed_up_weights: torch.Tensor,
    routed_down_weights: torch.Tensor
) -> tuple[torch.Tensor, ...]
```

Run MoK's manual backward.

**Parameters:**

**`schedule`** `object`

Opaque schedule returned by :meth:`forward`.

---

**`forward_context`** `object`

Opaque activation context returned by :meth:`forward`.

---

**`grad_output`** `torch.Tensor`

Contiguous BF16 tensor of shape \[tokens, hidden].

---

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

Contiguous BF16 tensor of shape \[tokens, hidden].

---

**`router_weights`** `torch.Tensor`

Contiguous FP32 tensor of shape \[tokens, activated\_experts].

---

**`shared_gate_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape \[expert\_intermediate, hidden].

---

**`shared_up_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape \[expert\_intermediate, hidden].

---

**`shared_down_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape \[hidden, expert\_intermediate].

---

**`routed_gate_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape
\[local\_experts, expert\_intermediate, hidden].

---

**`routed_up_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape
\[local\_experts, expert\_intermediate, hidden].

---

**`routed_down_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape
\[local\_experts, hidden, expert\_intermediate].

---

**Returns:** `torch.Tensor`

Tuple of gradients for `x`, router weights, three routed weights,

```python
nemo_automodel.components.moe.mok_experts._MoKRuntime.forward(
    x: torch.Tensor,
    router_weights: torch.Tensor,
    top_experts: torch.Tensor,
    shared_gate_weights: torch.Tensor,
    shared_up_weights: torch.Tensor,
    shared_down_weights: torch.Tensor,
    routed_gate_weights: torch.Tensor,
    routed_up_weights: torch.Tensor,
    routed_down_weights: torch.Tensor
) -> tuple[torch.Tensor, object, object]
```

Run MoK's manual forward and retain its backward context.

**Parameters:**

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

Contiguous BF16 tensor of shape \[tokens, hidden].

---

**`router_weights`** `torch.Tensor`

Contiguous FP32 tensor of shape \[tokens, activated\_experts].

---

**`top_experts`** `torch.Tensor`

Contiguous int64 tensor of shape \[tokens, activated\_experts].

---

**`shared_gate_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape \[expert\_intermediate, hidden].

---

**`shared_up_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape \[expert\_intermediate, hidden].

---

**`shared_down_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape \[hidden, expert\_intermediate].

---

**`routed_gate_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape
\[local\_experts, expert\_intermediate, hidden].

---

**`routed_up_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape
\[local\_experts, expert\_intermediate, hidden].

---

**`routed_down_weights`** `torch.Tensor`

Contiguous BF16 tensor of shape
\[local\_experts, hidden, expert\_intermediate].

---

**Returns:** `torch.Tensor`

Tuple containing output of shape \[tokens, hidden], an opaque MoK schedule,

```python
nemo_automodel.components.moe.mok_experts._MoKRuntime.initialize(
    ep_mesh: torch.distributed.device_mesh.DeviceMesh,
    n_routed_experts: int
) -> None
```

Attach the runtime to the expert-parallel group.

**Parameters:**

**`ep_mesh`** `DeviceMesh`

One-dimensional device mesh named `ep`.

---

**`n_routed_experts`** `int`

Global number of routed experts, divided evenly over
the EP mesh.

---

```python
nemo_automodel.components.moe.mok_experts._flatten_mok_tensor_dataclass(
    value: object
) -> tuple[tuple[torch.Tensor, ...], tuple[type, tuple[tuple[str, int], ...]]]
```

Flatten a MoK schedule/context without retaining tensor references in metadata.

MoK's dataclasses contain either tensors or tuples of tensors (the latter for
MXFP8 values).  Saving the flattened tensors with `ctx.save_for_backward`
makes them visible to activation-checkpoint saved-tensor hooks.  Storing the
original dataclass directly on `ctx` would keep every layer's macro-sized
context alive until backward.

**Parameters:**

**`value`** `object`

MoK schedule or forward-context dataclass.

---

**Returns:** `tuple[tuple[torch.Tensor, ...], tuple[type, tuple[tuple[str, int], ...]]]`

Flat tensors and tensor-free reconstruction metadata.

```python
nemo_automodel.components.moe.mok_experts._load_mok_functional()
```

Import MoK after distributed setup has selected this process's GPU.

Importing MoK loads its CUDA extension.  Doing that at module-import time is
too early for torchrun workers: every local worker still has CUDA device 0
selected, leaving one stray CUDA context per worker on the first GPU.  MoK
is only needed when expert parallelism is initialized, which happens after
`initialize_distributed` calls `torch.cuda.set_device`.

```python
nemo_automodel.components.moe.mok_experts._local_tensor(
    tensor: torch.Tensor
) -> torch.Tensor
```

Return the rank-local tensor while preserving its autograd connection.

**Parameters:**

**`tensor`** `torch.Tensor`

Tensor or DTensor of arbitrary shape. A DTensor must be sharded or
replicated on the current rank with a materialized local value.

---

**Returns:** `torch.Tensor`

Tensor with the DTensor's rank-local shape, or the original plain tensor.

```python
nemo_automodel.components.moe.mok_experts._unflatten_mok_tensor_dataclass(
    spec: tuple[type, tuple[tuple[str, int], ...]],
    tensors: tuple[torch.Tensor, ...],
    offset: int
) -> tuple[object, int]
```

Reconstruct a MoK dataclass from checkpoint-managed saved tensors.

```python
nemo_automodel.components.moe.mok_experts._MOK_IMPORT_MESSAGE = "dispatcher='mok' requires a built Mixture-of-Kittens installation or an importa...
```

```python
nemo_automodel.components.moe.mok_experts._mok_functional = None
```