> 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.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                                                                     |
| --------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------- |
| [`_copy_local_init_weight`](#nemo_automodel-components-moe-mok_experts-_copy_local_init_weight)                             | Copy a canonical initialization into the local MOK parameter storage.           |
| [`_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.             |
| [`_load_mok_ops`](#nemo_automodel-components-moe-mok_experts-_load_mok_ops)                                                 | Import MoK's weight-quantization ops after selecting this process's GPU.        |
| [`_local_tensor`](#nemo_automodel-components-moe-mok_experts-_local_tensor)                                                 | Return the rank-local tensor while preserving its autograd connection.          |
| [`_mxfp8_weight_both`](#nemo_automodel-components-moe-mok_experts-_mxfp8_weight_both)                                       | Prequantize both MoK MXFP8 layouts in one kernel launch.                        |
| [`_unflatten_mok_tensor_dataclass`](#nemo_automodel-components-moe-mok_experts-_unflatten_mok_tensor_dataclass)             | Reconstruct a MoK dataclass from checkpoint-managed saved tensors.              |
| [`enable_mok_mxfp8_optimizer_step_cache`](#nemo_automodel-components-moe-mok_experts-enable_mok_mxfp8_optimizer_step_cache) | Retain MXFP8 weights within a step and invalidate them after optimizer updates. |

### Data

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

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

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

### 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 unsharded
`GroupedExperts` for this EP rank's experts, then transpose into
MoK's native layouts and select the local FSDP feature slices. Keeping
full feature widths in the draw preserves the random stream and the
initialized values when FSDP splits the parameters.

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

**`_mxfp8_cache_generation`** `= 0`

---

**`_mxfp8_cached_generation`** `int | None = None`

---

**`_mxfp8_weights`** `tuple[tuple[Tensor, Tensor, Tensor, Tensor], ...] | None = None`

---

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

---

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

---

**`swiglu_limit`**

---

```python
nemo_automodel.components.moe.mok_experts._MoKRuntime._get_mxfp8_weights(
    routed_gate_weights: torch.Tensor,
    routed_up_weights: torch.Tensor,
    routed_down_weights: torch.Tensor
) -> tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], ...]
```

Prequantize all routed-weight layouts once per optimizer generation.

**Parameters:**

**`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:** `tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]`

Gate, up, and down projection layouts. Each entry contains the normal

```python
nemo_automodel.components.moe.mok_experts._MoKRuntime._get_workspace(
    x: torch.Tensor,
    topk: int
) -> object
```

Return MoK's workspace for one local activation shape.

**Parameters:**

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

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

---

**`topk`** `int`

Number of activated experts per token.

---

**Returns:** `object`

Opaque MoK workspace shared by schedule construction and kernels.

```python
nemo_automodel.components.moe.mok_experts._MoKRuntime._invalidate_mxfp8_cache() -> None
```

Advance the BF16-weight generation and drop its quantized tensors.

```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,
    mxfp8_weights: tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], ...] | None
) -> tuple[torch.Tensor, ...]
```

Run MoK's manual backward.

**Parameters:**

**`schedule`** `object`

Opaque schedule returned by `forward`.

---

**`forward_context`** `object`

Opaque activation context returned by `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].

---

**`mxfp8_weights`** `tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], ...] | None`

Prequantized gate, up, and down projections. Each is a
tuple `(weight_fp8, weight_scale, weight_t_fp8, weight_t_scale)`
in MoK's opaque E4M3/E8M0 layouts. Must be non-`None` for MXFP8
and `None` for BF16.

---

**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, tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], ...] | None]
```

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 the BF16 output of shape \[tokens, hidden], an opaque

```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._copy_local_init_weight(
    destination: torch.Tensor,
    source: torch.Tensor
) -> None
```

Copy a canonical initialization into the local MOK parameter storage.

**Parameters:**

**`destination`** `torch.Tensor`

Parameter of shape \[experts, output\_features, input\_features].
A DTensor may shard experts over EP and feature axes over FSDP.
Its local storage is mutated in place without replacing the parameter.

---

**`source`** `torch.Tensor`

Tensor of shape \[local\_experts, output\_features, input\_features],
already restricted to this EP rank's experts but containing the full
feature axes. It may be a noncontiguous transposed view and is not mutated.

---

```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._load_mok_ops()
```

Import MoK's weight-quantization ops after selecting this process's GPU.

```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._mxfp8_weight_both(
    weight: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
```

Prequantize both MoK MXFP8 layouts in one kernel launch.

**Parameters:**

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

Contiguous BF16 tensor of shape
\[local\_experts, output\_features, input\_features].

---

**Returns:** `torch.Tensor`

Normal E4M3 weight, its E8M0 block scales, transposed E4M3 weight,

```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.enable_mok_mxfp8_optimizer_step_cache(
    model_parts: list[torch.nn.Module],
    optimizers: list[torch.optim.Optimizer]
) -> None
```

Retain MXFP8 weights within a step and invalidate them after optimizer updates.

**Parameters:**

**`model_parts`** `list[nn.Module]`

Pipeline model parts paired one-to-one with `optimizers`.

---

**`optimizers`** `list[torch.optim.Optimizer]`

Optimizers whose post-step hooks own cache invalidation.

---

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

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