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

forward_context

object

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