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.
`_load_mok_ops`	Import MoK’s weight-quantization ops after selecting this process’s GPU.
`_local_tensor`	Return the rank-local tensor while preserving its autograd connection.
`_mxfp8_weight_both`	Prequantize both MoK MXFP8 layouts in one kernel launch.
`_unflatten_mok_tensor_dataclass`	Reconstruct a MoK dataclass from checkpoint-managed saved tensors.
`enable_mok_mxfp8_optimizer_step_cache`	Retain MXFP8 weights within a step and invalidate them after optimizer updates.

Data

_MOK_IMPORT_MESSAGE

_mok_functional

_mok_ops

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.

_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

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

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:

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.

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

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

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

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,

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:

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

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

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

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._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,

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

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

nemo_automodel.components.moe.mok_experts._mok_ops = None