core.transformer.moe.megakernel.mok.runtime

MOK functional runtime and its MCore autograd bridge.

Module Contents

Classes

_MoKWeightArguments	Non-owning references in the external MOK forward/backward ABI order.
_MoKAutograd	Autograd bridge from MCore parameters to MoK’s functional API.

Functions

_gate_up_weight_arguments	Adapt MCore-owned combined FC1 weights to MOK’s gate/up API.
_gate_up_main_grad_arguments	Adapt canonical FC1/FC2 main-grad buffers to MOK’s six-buffer order.

API

core.transformer.moe.megakernel.mok.runtime._gate_up_weight_arguments(

shared_fc1: torch.Tensor,

routed_fc1,

intermediate_size: int,

)

Adapt MCore-owned combined FC1 weights to MOK’s gate/up API.

Shared FC1 is a dense two-dimensional parameter, so gate and up are ordinary zero-copy row views. Routed FC1 may span multiple experts or descriptor-backed allocations; passing the same combined payload for gate and up preserves its storage metadata, and MOK applies the existing logical gate/up offsets when the two arguments alias.

core.transformer.moe.megakernel.mok.runtime._gate_up_main_grad_arguments(

main_grads,

main_grad_storage_tables,

intermediate_size: int,

)

Adapt canonical FC1/FC2 main-grad buffers to MOK’s six-buffer order.

class core.transformer.moe.megakernel.mok.runtime._MoKWeightArguments

Bases: typing.NamedTuple

Non-owning references in the external MOK forward/backward ABI order.

Routed arguments may be tensors, MXFP8 tuples, or split-weight descriptors. This transient tuple neither copies payloads nor registers/caches parameters.

shared_gate: torch.Tensor

None

shared_up: torch.Tensor

None

shared_down: torch.Tensor

None

routed_gate: object

None

routed_up: object

None

routed_down: object

None

classmethod from_combined(

shared_fc1,

shared_fc2,

routed_fc1,

routed_fc2,

intermediate_size,

)

Split only shared FC1 while retaining the caller’s routed weight views.

class core.transformer.moe.megakernel.mok.runtime._MoKAutograd

Bases: torch.autograd.Function

Autograd bridge from MCore parameters to MoK’s functional API.

static forward(

ctx,

module: megatron.core.transformer.moe.megakernel.mok.backend.MoKMegakernel,

x: torch.Tensor,

router_weights: torch.Tensor,

top_experts: torch.Tensor,

shared_output_gate_weight: torch.Tensor | None,

*parameters: torch.Tensor,

) → torch.Tensor

Run MOK forward and retain the state needed by backward.

static backward(ctx, grad_output: torch.Tensor)

Run MOK backward and return gradients for its autograd inputs.