[zyongye]

Overview

Folds the three DeepGEMM MegaMoE expert implementations into one DeepGemmMegaMoEExperts layer. The three are DeepSeek V4/V4.1's DeepseekV4MegaMoEExperts, Kimi K3's KimiK3MegaMoEExperts, and the DeepGemmMegaMoEBackend pair under both. The new layer is driven by a feature matrix: routed precision × activation × fusable shared-expert precision. The PR also adds NVFP4 routed experts (DeepGEMM nvfp4_mega_moe) as a new row. It keeps the procedure from #60096: weights load straight into the DeepGEMM layout, with no post-load transform.

Draft: the end-to-end eval has not been run yet (see Validation).

Claims

• One layer. Every DeepGEMM MegaMoE caller uses forward(hidden_states, topk_weights, topk_ids, shared_experts=None).

• Fusion decided at construction. Fusion of the shared expert is decided once, from the shared MLP's weights, and never changes at runtime.

• Supported matrix. It covers every routed precision and activation that the pinned DeepGEMM exposes:

  routed	kernel	dispatch	fusable shared	activations
  mxfp4	fp8_fp4_mega_moe (fp8xfp4)	MXFP8	MXFP8	swiglu, situ
  mxfp8	fp8_fp4_mega_moe (fp8xfp8)	MXFP8	MXFP8	swiglu, situ
  nvfp4	nvfp4_mega_moe	NVFP4, one scale per token	MXFP8, BF16	swiglu
  bf16	bf16_mega_moe	BF16	BF16	swiglu

  swiglu takes optional clamp, alpha and beta. situ (Kimi K3) takes alpha and beta.

• Existing models unchanged. DeepSeek V4/V4.1 and Kimi K3 still build the MXFP4 row and pass the same kernel arguments as before. The mxfp8, nvfp4 and bf16 rows have no model caller yet.

• Size. Code under vllm/ is +360 lines net. The refactor commit is −120, the NVFP4 commit is +472, and the two follow-up commits are −15 and +23.

Validation

GB200 (SM100), vendored DeepGEMM at the current pin:

.venv/bin/python -m pytest tests/kernels/moe/test_deep_gemm_mega_moe.py \
    tests/models/test_deepseek_v4_mega_moe.py \
    tests/model_executor/test_deep_gemm_warmup.py
# 138 passed in 67s

Kernel numerics. The real kernel is compared against an fp32 reference that runs on the same quantized weights and activations. The reference covers dispatch quantization, the activation, the L2 activation quantization and the shared expert. The metric is DeepGEMM's calc_diff; the test requires < 1e-3 and cosine > 0.99. The table shows the worst case over the activations each format supports:

routed	unfused shared	fused MXFP8 shared	fused BF16 shared
mxfp4	2.9e-5	1.8e-4	–
mxfp8	3.1e-5	1.8e-4	–
nvfp4	1.8e-6	1.6e-4	1.4e-6
bf16	2.1e-6	–	1.7e-6

Negative controls. Each of these deliberate breaks made the expected cases fail:

break	cases that failed
wrong weight layout, activation or shared input (MX)	the affected cases
drop NVFP4 per-token x_scales, l1_alpha or l2_alpha	all 9 NVFP4 cases
drop the shared-expert input	all 6 fused cases
wrong shared BLOCK_M	the 3 fused-MXFP8 cases

Input staging. The Triton staging kernels match DeepGEMM's reference quantizers (per_token_cast_to_fp8, per_token_cast_to_nvfp4) bit for bit. This covers values, scales, shared-expert inputs, unrouted padding and the zero-token case.

Not yet done:

• End-to-end DeepSeek-V4.1 gsm8k A/B against main: TP4, --enable-expert-parallel, --kernel-config '{"moe_backend":"deep_gemm_mega_moe"}'. The main baseline is 92.65% flexible / 92.57% strict.
• End-to-end Kimi K3 with moe_backend=deep_gemm_mega_moe.
• Only single-rank kernel numerics have been run; there is no multi-rank EP numerics test.

Details

Feature matrix.

• Each MegaMoEWeightFormat row declares five things: the MMA type, the kernel entry point, the routed weight spec (dtype, packing, scale group, scale dtype), the dispatch quant key, and the fusable shared-expert specs.
• The layer validates its MegaMoEActivation against the row at construction.
• A new precision is a new row plus whatever staging it needs. It needs no new expert class.

Shared-expert fusion.

• The layer fuses the shared expert when two conditions hold. First, the shared MLP's weight dtype, scale dtype and shape match one of the row's shared specs. Second, the installed DeepGEMM exposes the shared-expert buffer and kernel API.
• Otherwise forward adds shared_experts(hidden_states) once.
• A fused layer raises if it is still handed a shared module.

Input staging.

• prepare_megamoe.py moved from vllm/models/deepseek_v4/nvidia/ops/ to fused_moe/. The old location caused an import cycle through vllm.models.deepseek_v4.
• It gained a BF16 passthrough and an NVFP4 kernel.
• Shared Triton helpers were factored out of the MX kernel, which is still bit-exact.

NVFP4.

• Dispatch quantizes hidden states with one NVFP4 scale per token, computed at runtime. The checkpoint's w1/w3 input_scale are therefore unused. A model that wires NVFP4 must skip them.
• The kernel takes a single scalar L2 activation scale. The layer uses the largest w2 input_scale across all experts. Every rank reads every expert's scale tensors during loading, so all ranks agree and EPLB expert moves stay valid.
• l2_alpha (the per-expert w2 global scale × the L2 activation scale) is computed inside the staging kernel. This avoids an extra launch and stays correct after EPLB moves or a weight reload.
• The kernel takes one l1_alpha per expert, so the loader raises if w1 and w3 weight_scale_2 differ (relative tolerance 1e-5, as in ModelOpt's allclose check). It compares the two shards of the same load, not the live slot, so a reload that changes the scale does not trip it. ModelOpt's and Quark's NVFP4 MoE only warn here; they also use one scale.
• Meta tensors, which reload dry runs pass, skip this host-side scale state.
• The L2 activation scale is a host scalar in DeepGEMM's API, so CUDA graphs capture its value. A weight reload that changes the w2 input scales needs a graph recapture. In eager mode the layer recomputes the max on reload. Passing it as a device tensor needs a DeepGEMM change and a pin bump, so it is left for a follow-up.

Behaviour changes for existing models:

• DeepSeek V4/V4.1 now also run Kimi K3's one-time first-launch barrier: a device sync plus a gloo barrier on the EP CPU group. It runs on the first forward, which is the eager profile run before graph capture.
• get_deep_gemm_mega_moe_backend() and the backend classes are removed. Their docstring named them as the override point for plugins and SM90. The extension point is now a new MegaMoEWeightFormat row, plus check_deep_gemm_mega_moe_supported.

CI. The new test file and the moved sources are added to the B200 "DeepSeek V4 Kernels" step. Before this PR the SM100 numerics ran only from tests/models/test_deepseek_v4_mega_moe.py. The staging-kernel tests are gated to SM100, where those kernels run. The L4 "MoE Kernels" shards also collect tests/kernels/moe/, and on SM89 Triton casts fp32 to fp8 through fp16 (cvt.rz.f16.f32). That rounds the NVFP4 block scales twice, so the bit-exact test failed there. The loader and fusion-selection tests in the file still run on L4.

Unchanged: the XPU DeepseekV4MegaMoEExperts, which uses a different kernel.

Related PRs:

• Not a duplicate: no open PR unifies the MegaMoE expert classes or wires nvfp4_mega_moe.
• [Model][MoE] Load MegaMoE experts directly into the DeepGEMM layout #60096 (merged): the load-into-kernel-layout procedure this PR keeps.
• [Perf][GLM-5.3-Flash] DeepGEMM FP8 MegaMoE with fused shared experts #60895 (GLM-5.3 block-FP8 MegaMoE): edits the backend classes this PR removes, so the two conflict. After this PR lands it fits the mxfp8 row, plus a load-time block-scale requant.
• [Weight Reload][MegaMoE] Rebuild module-owned packed expert buffers #59003 (weight reload for packed MegaMoE buffers): also targets the removed classes. [Model][MoE] Load MegaMoE experts directly into the DeepGEMM layout #60096 already writes checkpoint shards in place.
• [Kimi K3][Kernel] Fuse BF16 shared experts into latent MegaMoE tail #53556 (Kimi K3 BF16 shared-expert fusion into the latent tail) and [Kernels] Integration for DeepGEMM MegaMoE kernel #40833 (modular-kernel MegaMoE integration, stale draft) take different approaches.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

AI assistance (Claude Code) was used for this change. The submitter reviews every changed line and runs the tests before marking it ready.

🤖 Generated with Claude Code

[hao-aaron]

Checked this against weight reload (layerwise today, and the reload plan in #59502 / #60931). The MXFP4 path DeepSeek V4/V4.1 and Kimi K3 use loads the same way as after #60096: shards are written straight into the kernel layout through the param's loader, and the shared MLP is still dropped in __init__. No reload concerns there.

Two small requested fixes on the new NVFP4 row:

1. Skip the w2_input_scale side-state update for meta tensors (deep_gemm_mega_moe.py:483-484).
Some reload machinery runs load_weights with meta stand-ins (right names and shapes, no data) to learn which loaders get called:

• sharded RDT's plan bake;
• the reload verifier's load-coverage dry run;
• the init-time trace in the planned reload design.

self._w2_input_scales[expert_id] = loaded_weight.max() copies a meta value into a real CPU tensor, which raises, and that aborts the whole dry run. The side state itself is fine for real loads. Adding and not loaded_weight.is_meta to that branch is enough.

2. Check that w1 and w3 weight_scale_2 match.
The kernel takes one L1 alpha per expert, but the loader keeps whichever shard's weight_scale_2 arrives last. A checkpoint or trainer update where they diverge would silently produce wrong alphas. Suggestion: when the second shard's value lands on a slot the first already wrote (nonzero), compare them and raise on mismatch.

Known issue, out of scope for this PR: l2_activation_scale is passed to the kernels as a Python float. CUDA graphs bake it in at capture, so in graph mode a reload that changes the w2 input scales keeps using the old value. The fix is to pass it as a device tensor:

• DeepGEMM's nvfp4_mega_moe takes a pointer instead of a float (kernel, launcher and op schema in vllm-project/DeepGEMM);
• the staging kernel loads it from the same buffer;
• the layer refills that buffer in place on reload.

That needs a DeepGEMM change and a pin bump, so it's better as a follow-up.

This review/comment was written by an AI agent (Claude Code) on my behalf.

[zyongye]

Thanks for checking this against reload. Both requested fixes are in 603db01.

1. Meta tensors. weight_loader now skips the host-side scale state when loaded_weight.is_meta. A new test loads meta stand-ins for w2_input_scale and both w13_weight_scale_2 shards. Without the guard it fails with Cannot copy out of meta tensor.

2. w1/w3 weight_scale_2. The loader now raises when they differ, with a relative tolerance of 1e-5 (the same as ModelOpt's allclose check). One change from the suggestion: it pairs the two shards of the same load instead of comparing with the slot. On a reload that changes the scale, the slot still holds the old value when the first new shard lands, so a slot comparison would raise on a valid reload. The existing reload test loads two checkpoints with different scales in shuffled shard order, and it still passes. ModelOpt's and Quark's NVFP4 MoE paths only warn on this mismatch; this one raises, as you suggested.

CUDA-graph L2 scale. Agreed, this is left for a follow-up together with the DeepGEMM change and pin bump. The PR description notes it.

This reply was written by an AI agent (Claude Code) on behalf of @zyongye.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #94889 for commit 5550180721d0.

[zyongye]

/ci retry

[github-actions]

✅ Queued 6 failed job(s) for retry in Buildkite CI #94889.

[zyongye]

/ci run

[github-actions]

✅ Triggered Buildkite CI #94976 for commit d83c1498aa29.

ℹ️ This PR is 26 commits behind upstream main. Merged upstream main into this branch as d83c1498aa29 before starting CI. Pull the branch before pushing more changes.

[zyongye]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #94976.

[zyongye]

/ci retry

[github-actions]

✅ Queued 3 failed job(s) for retry in Buildkite CI #94976.