Add Mixture-of-Kittens MoE backend on main by FDecaYed · Pull Request #8210 · NVIDIA/Megatron-LM

FDecaYed · 2026-10-14T09:33:46Z

What does this PR do?

Enable moe_megakernel_backend=mok on main. MCore supplies routing and owns the expert parameters, optimizer state, and checkpoints; MoK executes dispatch, routed/shared experts, and combine.

Ports #6572 and #7289, building on the earlier main adaptation in #7695. Related to #5795 and #7756. Uses the MoK revision already installed by #7696.

Support BF16 and native MXFP8 routed weights, with BF16 shared experts and existing grouped-weight layouts.
Preserve parameter aliases and their MCore/TE attributes; refresh derived views after device/dtype conversion and checkpoint loading.
Forward packed-sequence metadata and hash-routing inputs, and use main's precision-aware whole-MoE recomputation.
Keep optional imports lazy and reject unsupported configurations. Initial support is TP=1/ETP=1 with fused shared experts. A native shared-expert fallback for TP>1 and the shared-output gate from Support gated shared experts in MOK for Qwen3.5 #7391 are follow-ups; the pinned MoK API has neither a routed-only mode nor the gate argument.

External Dependency

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

Validation

Standard CI passed on 1cdac4aa8d, including the MoE and determinism suites. The review follow-up at 951fd137d5 is running CI.
47 CPU checks passed for configuration, weight/gradient adaptation, parameter metadata, checkpoint ownership, and packed-routing/recompute dispatch.
Required formatter/linter checks, syntax checks, and kernel determinism coverage check passed. Advisory mypy is not clean in the local lint environment.
Blackwell BF16/MXFP8 kernel execution, distributed EP training, CUDA graph replay, and end-to-end checkpoint resume remain to be validated on this main port. The existing GPU tests are included; CPU checks do not qualify those paths.

Draft pending GPU validation. No performance result is claimed.

Port the native-parameter MOK integration and alias metadata preservation from dev onto current main. Forward packed routing metadata and use main's precision-aware recompute entry point. Invalidate derived weight views on parameter conversion and retain the existing optional dependency pin. Based on dev PRs NVIDIA#6572 and NVIDIA#7289 and the earlier main adaptation NVIDIA#7695. Co-authored-by: qizhang <qizhang@nvidia.com> Co-authored-by: Xin Yao <xiny@nvidia.com> Co-authored-by: HaochenYuan <haocheny@nvidia.com> Co-authored-by: Ajay Balasa <abalasa@nvidia.com> Signed-off-by: Deyu Fu <deyuf@nvidia.com>

copy-pr-bot · 2026-10-14T09:33:51Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

FDecaYed · 2026-10-14T09:36:20Z

/ok to test 1cdac4a

zhongbozhu · 2026-10-15T05:25:05Z

Does this PR includes #7391 content?

zhongbozhu · 2026-10-15T05:40:42Z

+                    "MOK does not support overlap_dispatch_backward_with_experts_wgrad; "
+                    "the megakernel never runs the native dispatch path"
+                )
+            if self.tensor_model_parallel_size != 1 or self.expert_tensor_parallel_size != 1:


Does it mean that we cannot do TP2 to dense modules and shared expert while still doing ETP=1 for experts?

If I remember correctly, MOK fuses shared expert in it, but it should be possible to disable shared expert compute and use Mcore's shared expert implementation so that we can do the above operation.

Yes, TP2/ETP1 is currently rejected. The pinned MoK API requires shared-expert weights and has no routed-only switch. A native shared-expert fallback makes sense as a follow-up, with TP/SP routing and backward validation before relaxing this guard.

zhongbozhu · 2026-10-15T05:42:06Z

+                raise ValueError("MOK requires a shared expert")
+            if self.moe_shared_expert_gate or self.moe_shared_expert_overlap:
+                raise ValueError("MOK does not support the MCore shared-expert gate/overlap")
+            if self.moe_latent_size is not None:


This should be possible to support since latent moe just has two projections, one before dispatch and another one after combine?

In MCore those projections wrap only the routed branch; the shared expert consumes the original hidden states. Wrapping the fused MOK call would also project its shared branch and change that behavior. Native latent MoE therefore needs the routed-only/native-shared path first. I kept this guard and added the reason.

zhongbozhu · 2026-10-15T05:43:25Z

+                raise ValueError("MOK does not support the MCore shared-expert gate/overlap")
+            if self.moe_latent_size is not None:
+                raise ValueError("MOK does not support latent MoE")
+            if self.recompute_modules and "shared_experts" in self.recompute_modules:


Like we discussed above, not sure if we have an application where it's reasonable to unfuse shared expert so that we can apply TP or maybe apply mxfp8 to it, since MOK's shared expert is in bf16?

Yes, a separate MCore shared expert could provide TP and MXFP8 shared compute, as well as shared-expert-only recompute. That is a useful follow-up to benchmark, but it needs a routed-only MOK interface. The pinned API requires fused shared weights, so this port keeps the current path.

zhongbozhu · 2026-10-15T05:46:16Z


+        if self.moe_megakernel_backend == "mok":
+            if self.cuda_graph_impl in ("local", "transformer_engine"):
+                if not self.cuda_graph_modules:


this if not self.cuda_graph_modules means full cuda graph?

If that's the case we should be able to actually do full cg if MOK can be integrated with paged stashing?

On the other hand, the amount of tokens before dispatch is static, if we are able to redo the dispatch in backward pass and overlap it with compute, it's also possible to capture cuda graph of MOK even without paged stashing?

Here an empty list means whole-layer capture for local/transformer_engine; full_iteration is already allowed by the config. Per-layer replay needs the saved MOK schedule/activations to survive pending microbatches. Paged stashing or redispatch/recompute could provide that, but the current bridge implements neither. I clarified this distinction and kept the per-layer guard; GPU replay qualification remains pending.

FDecaYed · 2026-10-15T05:51:20Z

@zhongbozhu #7391 is not included in this port. It adds the Qwen3.5 shared-output gate, which requires a newer MoK API than main's pinned a2ad2d0c. This draft uses ungated fused shared experts; the gate support should come with the matching dependency update in a follow-up.

zhongbozhu · 2026-10-15T05:52:21Z

+                "megatron/core/transformer/moe/README.md"
+            ) from exc
+
+        if not config.gradient_accumulation_fusion:


What's the precision of wgrad accumulation? Assume we have D = A * B + C, does the accumulation precision depend on the dtype of C?

That said, if we keep C dtype to bf16, will the addition precision also be bf16?

The GEMM accumulator is FP32 in both cases. With FP32 main_grad, it stays FP32 through the reduction-store. With BF16 main_grad, each contribution is rounded to BF16 before the BF16 reduction-store, so persistent accumulation is BF16. The shared weight dtype does not select this; main_grad.dtype does. The pinned API also requires shared/routed main-grad buffers to have the same dtype. Documented this from the kernel epilogue.

zhongbozhu · 2026-10-15T06:05:57Z

+    def quantized_routed_weights(self):
+        """Build or refresh cached MOK routed-weight views."""
+        needs_weight_view_update = self._routed_weight_view_cache is None or (
+            self.use_mxfp8_weights and self.is_first_microbatch


Why do we need to update the view under use_mxfp8_weights, which is essentially under fp8 param gather?

The cache includes converted MXFP8 scale buffers. Payloads still alias TE storage; after an optimizer update/parameter sync, the scale values change, so their MOK layout must be refreshed on the first microbatch. This happens in place to preserve descriptor addresses, and later microbatches reuse it. It does not requantize or copy the FP8 weight payload. Added this explanation in the method.

zhongbozhu · 2026-10-15T06:07:33Z

+        self.topk = config.moe_router_topk
+        self.swiglu_limit = config.activation_func_clamp_value
+        self.use_mxfp8_weights = bool(
+            config.fp8 is not None and config.fp8_recipe == "mxfp8" and config.fp8_param


what about reuse_grad_buf_for_mxfp8_param_ag?

That path gathers into shared staging storage, then copy_tensors_to_quantized_params requantizes the native TE parameters before the staging buffer is cleared. MOK aliases the TE payloads. I removed this guard in 951fd137d5: post_all_gather_processing is a no-op for MXFP8 in the pinned TE, so its presence is not the right capability check. Native rowwise/columnwise payload and scale checks remain. Full distributed validation of this combination is still pending.

Signed-off-by: Deyu Fu <deyuf@nvidia.com>

FDecaYed · 2026-10-15T06:29:29Z

/ok to test 951fd13

zhongbozhu · 2026-10-15T07:15:01Z

+    return flat.view(shape)
+
+
+def _swizzle_mxfp8_scale(


we can fuse swizzle into quantization kernel by TE quantizer setting optimize_for_gemm?

Yes, for compute-only MXFP8 tensors; TE's swizzled layout matches MOK. Here we alias native FP8 parameters, and the pinned TE explicitly disables preswizzling for primary FP8 weights. Dequantization and parameter-update paths still expect compact scales, so flipping this flag alone is not safe. I'll keep the zero-copy payload and cached scale conversion in this port; removing that conversion needs support across the native parameter lifecycle.

zhongbozhu · 2026-10-15T07:15:59Z

+                    f"expert={expert}, got={tuple(column_scale.shape)}, "
+                    f"expected={(rows // 32, columns)}"
+                )
+            _swizzle_mxfp8_scale(row_scale.unsqueeze(0), rows=rows, columns=columns, out=row_out)


same comment^, maybe we can leverage optimize_for_gemm in TE quantizer to fuse this swizzle

Same primary-parameter limitation here. This refresh only updates the cached MOK scale bytes on the first microbatch after parameter sync; later microbatches reuse them. Keeping it for this port, with quantize/swizzle fusion as a follow-up once the native parameter lifecycle supports that layout.

FDecaYed added the deepseekv4 DeepSeek V4 PRs label Oct 14, 2026

FDecaYed mentioned this pull request Oct 13, 2026

[Tracking] DeepSeek-v4 HybridModel merge plan #5795

Draft

copy-pr-bot Bot deployed to test October 14, 2026 09:37 Active

copy-pr-bot Bot deployed to public October 14, 2026 09:40 Active

zhongbozhu mentioned this pull request Oct 15, 2026

MegaMOE Integration #6810

Draft

6 tasks

zhongbozhu approved these changes Oct 15, 2026

View reviewed changes

zhongbozhu reviewed Oct 15, 2026

View reviewed changes

zhongbozhu added the MegaMoE MegaMoE development with various possible backends label Oct 15, 2026

Clarify MOK precision and remove the unrelated TE API guard

951fd13

Signed-off-by: Deyu Fu <deyuf@nvidia.com>

copy-pr-bot Bot deployed to test October 15, 2026 06:30 Active

copy-pr-bot Bot deployed to public October 15, 2026 06:33 Active

zhongbozhu reviewed Oct 15, 2026

View reviewed changes

Conversation

FDecaYed commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

External Dependency

Validation

Uh oh!

copy-pr-bot Bot commented Oct 14, 2026

Uh oh!

FDecaYed commented Oct 14, 2026

Uh oh!

zhongbozhu commented Oct 15, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

FDecaYed commented Oct 15, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

FDecaYed commented Oct 15, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

FDecaYed commented Oct 14, 2026 •

edited

Loading