feat(moe): add a megakernel backend with MOK integration

[QiZhangNV]

Summary

This PR introduces an MoE megakernel backend and integrates Mixture-of-Kittens (MOK) as the first backend.

MOK replaces the post-router MoE execution path, including dispatch, routed and shared expert computation, and combine.

The MOK implementation used by this PR is maintained on the following integration
branch:

https://github.com/QiZhangNV/mixture-of-kittens/tree/qizhang/mcore-integration

Design

Pluggable megakernel backend

A generic MegakernelBackend boundary is added so that future megakernel
implementations can be selected through:

moe_megakernel_backend: mok

MOK-specific implementation details are isolated under the MOK backend directory
instead of being added directly to the normal MCore MoE execution path.

MCore-owned parameters

MCore remains the canonical owner of routed and shared expert parameters.

The adapter exposes MCore/Transformer Engine storage to MOK through tensor views or descriptor tables without registering duplicate optimizer parameters.

Both routed-weight layouts are supported:

• moe_single_grouped_weight=True
• Per-expert non-single weights

For MXFP8 weights, the adapter builds MOK-compatible data and scale views over the native TE-managed storage. Converted scale layouts are refreshed once per optimizer iteration while retaining CUDA-Graph-stable addresses.

Authoritative routing maps

MOK consumes the routing map produced by the MCore router. A Triton adapter converts the dense routing map into MOK's fixed-width routing representation and gathers the corresponding router probabilities. Missing routes are represented by expert ID -1, allowing the backend to preserve routing decisions when a token has fewer than top-k valid routes.

The backward adapter scatters router-weight gradients back into the original dense probability tensor.

Gradient accumulation

MOK writes routed and shared expert gradients into MCore-owned main_grad buffers. Both single-grouped and per-expert non-single storage are supported.

This initial integration requires:

gradient_accumulation_fusion: true

Non-fused gradient accumulation is intentionally rejected and left as follow-up work.

Checkpointing

The checkpoint remains expressed through the native MCore modules and parameter names. MOK parameter aliases do not introduce additional checkpoint ownership.

Distributed checkpoint save/resume has been validated for the supported single-weight and non-single-weight configurations.

Checkpoint conversion between different physical weight layouts is not included in this PR.

Supported configurations

The current integration has been validated with:

• GB300
• BF16 routed experts
• MXFP8 routed experts
• Single-grouped and per-expert non-single routed weights
• BF16 shared experts
• Fused gradient accumulation
• FP8 parameter gather and gradient-buffer reuse
• Full-iteration CUDA Graph
• MoE layer recompute
• Force-balanced and non-force-balanced routing
• Distributed checkpoint save and resume

Current MOK limitations include:

• Routed and shared experts must use the same intermediate size
• Shared expert parameters must remain BF16
• MCore's optional shared-expert output gate is not supported
• Non-fused gradient accumulation is not supported

Validation

Proxy model

The integration was validated on a 32×GB300 Proxy configuration with:

• TP1 / PP2 / EP16
• Sequence length 4096
• 96 routed experts
• BF16 and MXFP8
• Single-grouped and non-single weights
• Continuous training and distributed-checkpoint save/resume
• Full-iteration CUDA Graph and eager execution

All tested cases completed without skipped or NaN iterations. The refactored single/non-single paths retained the previously validated loss behavior and performance.

For the current MXFP8 non-single Proxy runs:

Implementation	Median TFLOPS/GPU	Iter-50 loss
MCore FusedMLP baseline	939.3	1.074571
MOK	988.2	1.051451

These jobs used separate allocations, so the result should be treated as an E2E performance indication rather than a controlled kernel-only comparison.

DSv4-Pro Full Model

The formal Full Model experiment used:

• 256×GB300
• 61 layers
• 384 routed experts, top-k 6
• TP1 / PP4 / EP64 / VPP4
• Sequence length 4096
• MBS1 / GBS4096
• MXFP8 routed experts with BF16 main gradients
• Non-single routed weights
• Full-iteration CUDA Graph
• Selective recompute and activation offload

Implementation	Median latency	Median TFLOPS/GPU	Iter-50 loss	PP0 total device memory
MCore HybridEP + FusedMLP	22.376 s	944.75	0.2207791	281,381 MB
MOK non-single	20.857 s	1013.50	0.2167743	268,027 MB

This corresponds to a 1.0728× median latency speedup and approximately 13.4 GB less PP0 physical device memory. The jobs used separate allocations.

A later post-refactor regression run achieved 1068.1 median TFLOPS/GPU.

Test plan

• MOK unit and operator tests
• BF16 single-grouped Proxy training
• BF16 non-single Proxy training
• MXFP8 single-grouped Proxy training
• MXFP8 non-single Proxy training
• Force-balanced routing
• Non-force-balanced routing with missing route slots
• Distributed checkpoint save and resume
• Full-iteration CUDA Graph
• MoE layer recompute
• 50-iteration 256-GPU DSv4-Pro Full Model training
• Pre/post-refactor loss, throughput, and memory comparison

Possible feature work

• Batch the per-expert MXFP8 scale swizzle to reduce weight-preparation overhead.
• Let megakernels consume MCore routing results directly to remove the standalone conversion.
• Add support for additional shared-expert configurations and precisions.
• Support non-fused gradient accumulation and zero_out_wgrad if required.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[QiZhangNV]

/ok to test 46003cb