nemo_automodel.components.moe.mok_experts

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

Module Contents

Classes

Name	Description
`GroupedExpertsMoK`	AutoModel expert parameters executed by Mixture-of-Kittens.
`_MoKAutogradFunction`	Connect MoK’s explicit backward API to PyTorch autograd.
`_MoKRuntime`	Own one layer’s MoK configuration and expert-parallel process group.

Functions

Name	Description
`_flatten_mok_tensor_dataclass`	Flatten a MoK schedule/context without retaining tensor references in metadata.
`_load_mok_functional`	Import MoK after distributed setup has selected this process’s GPU.
`_local_tensor`	Return the rank-local tensor while preserving its autograd connection.
`_unflatten_mok_tensor_dataclass`	Reconstruct a MoK dataclass from checkpoint-managed saved tensors.

Data

_MOK_IMPORT_MESSAGE

_mok_functional

API

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

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.

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.

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.

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:

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

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.

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

floatDefaults to 0.02

Standard deviation of the normal weight initialization.

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

Bases: Function

Connect MoK’s explicit backward API to PyTorch autograd.

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

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.

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

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

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

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,

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:

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,

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.

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.

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.

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.

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.

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

nemo_automodel.components.moe.mok_experts._mok_functional = None