Repository navigation
Conversation
With `moe_backend=deep_gemm_mega_moe` and expert parallel, GLM-5.3-Flash's MoE layers (MTP included) run DeepGEMM's FP8xFP8 MegaMoE: one kernel dispatches, runs the routed and the shared expert and combines across the EP group. The official checkpoint's experts are block FP8 with a float32 scale per 128x128 block; the SM100 kernel takes power-of-two scales per row and 32 columns. The weights are requantized once at load. Routing stays outside the kernel (`fused_topk_bias`, the routed scaling factor folded into the top-k weights). Requires tensor_parallel_size=1. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
Move the DeepGEMM-specific part into the shared backend abstraction: `DeepGemmSm100Fp8MegaMoEBackend` (`mma_type="fp8xfp8"`) requantizes block-FP8 routed and shared expert weights to the kernel's UE8M0 1x32 scales, and `get_deep_gemm_mega_moe_backend` selects it by `mma_type`. `Glm5NextMegaMoEExperts` becomes a thin adapter over `DeepseekV4MegaMoEExperts`, as Kimi K3's is: it only declares the checkpoint's loader parameters and picks the FP8 backend. The base weight loader now splits gate/up by half of the parameter's rows, so block scales shard like weights. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
|
I'm not very sure about the regression for accuracy. But do we have another better choice for MegaMoE? |
|
@gcanlin thanks for the bring up, can u run a AA-LCR or AIME2026 to see the accuracy regression level |
🥲I can't access to my machine recently. So if you have free machine, could you help run the accuracy benchmark? And I just noticed that we disabled MegaMoE for Qwen as the gsm8k had degraded. But currently, GLM's gsm8k doesn't happen this regression. So I think this approach looks promising. |
|
This pull request has merge conflicts that must be resolved before it can be |
Port the GLM-5.3-Flash FP8 MegaMoE onto the unified DeepGEMM MegaMoE experts from vllm-project#61538 / vllm-project#60096, which replaced DeepseekV4MegaMoEExperts and the post-load MegaMoE backend transform: - Glm5NextMegaMoEExperts now subclasses DeepGemmMegaMoEExperts with the MEGA_MOE_MXFP8 (fp8xfp8) format. E4M3 shards load straight into the kernel layout; the float32 128x128 block scales are staged and every freshly loaded shard is requantized in place to UE8M0 1x32 scales after load_weights (requant_block_fp8_to_ue8m0 moves to the GLM adapter). - The block-FP8 shared expert is fused the same way: the experts own its weights and the separate shared MLP is dropped. - The SwiGLU clamp moves to MegaMoEActivation; the deep_gemm_mega_moe.py and DeepSeek V4 conflicts take main's version. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ault mypy inferred tuple[int, ...] for weight_block_size; keep it as tuple[int, int]. ruff B008 flagged slice() as a default argument in the test; use a module constant for the local expert range. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
- A weight update can send an expert weight and its block scale in separate load_weights calls. Requantization ran at the end of every call and only weights marked a shard, so a weight arriving before its scale was rounded with the stale scale and the new scale was never applied. Track weights and scales separately and requantize a shard once both have arrived; a weight whose scale was never loaded still raises. - Fuse only a block-FP8 shared expert. The fallback to the base fusion could accept an MXFP8 shared expert that the GLM loader cannot place (gate_proj/up_proj names). - Dequantize with a broadcast block scale and quantize with the shared MXFP8 reference (_mxfp8_e4m3_quantize_torch) instead of a local copy; requant_block_fp8_to_ue8m0 now returns UE8M0 bytes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
Reloading weights in place (layerwise reload, weight transfer) produced garbage output: - Requantization ran at the end of load_weights. During a layerwise reload the layer is still on meta at that point, so it was skipped, and the replayed raw block-FP8 values were copied back with stale kernel scales. - The block-FP8 shared expert was written straight into .data instead of through param.weight_loader, so reload never saw those writes. - The kernel scales derived by requantization are never loaded, so the layer looked incomplete and was deferred to finalize with all of its weights buffered. Requantize from the model's process_weights_after_loading (text model, multimodal wrapper and MTP drafter), load the shared expert through its weight loader, and give the derived scales weight_loader_numel = 0. The finish step rejects any shard whose weight or block scale is still missing, instead of guessing from scale values. DeepGemmMegaMoEExperts gains small hooks so the subclass no longer copies base logic: _shared_linear_fits (shared-expert fusion check), _shared_param_name / _shared_shard_id (also accept gate_proj/up_proj), and _weight_shard_view (kernel-layout view of one shard). DeepSeek V4 behavior is unchanged. Cold load works as before. Reload and weight transfer additionally need vllm-project#57733, which calls the model-level hook at the end of finalize_layerwise_reload. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
…7733 Layerwise reload and weight transfer replay the block-FP8 shards into the MegaMoE experts, but nothing requantizes them until the reload calls the model's process_weights_after_loading (vllm-project#57733), so the model would serve garbage. Like expert substitution, the text model, the multimodal wrapper and the MTP drafter now detect a reload at the start of load_weights, put the previous weights back and raise NotImplementedError. Models without MegaMoE are unaffected. Remove the check once vllm-project#57733 lands. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
Route the routed experts of DeepseekV2MoE through the block-FP8 MegaMoE experts from vllm-project#60895 when --moe-backend deep_gemm_mega_moe is set, for the DSA models (GLM-5.x, DeepSeek-V3.2; GlmMoeDsaForCausalLM / DeepseekV32ForCausalLM and their MTP). - Routing (scoring_func + e_score_correction_bias, renormalized, times routed_scaling_factor) runs before the kernel; the shared expert stays a separate MLP. - Requires EP, TP=1, one expert group and a block-FP8 checkpoint. Works with data parallel and prefill context parallel: every EP rank calls the kernel once per layer per step. - Weights load straight into the kernel layout (routed_experts_prefix ''), are requantized in process_weights_after_loading, and layerwise reload is refused until vllm-project#57733, as in vllm-project#60895. - Other moe backends and other DeepseekV2MoE users are unchanged; DeepSeek-V2/V3 models get a clear NotImplementedError for this backend. Signed-off-by: Muhammad Ibrahim <ibrahimnd2000@gmail.com>
|
This pull request has merge conflicts that must be resolved before it can be |
Route the routed experts of DeepseekV2MoE through the block-FP8 MegaMoE experts from vllm-project#60895 when --moe-backend deep_gemm_mega_moe is set, for the DSA models (GLM-5.x, DeepSeek-V3.2; GlmMoeDsaForCausalLM / DeepseekV32ForCausalLM and their MTP). - Routing (scoring_func + e_score_correction_bias, renormalized, times routed_scaling_factor) runs before the kernel; the shared expert stays a separate MLP. - Requires EP, TP=1, one expert group and a block-FP8 checkpoint. Works with data parallel and prefill context parallel: every EP rank calls the kernel once per layer per step. - Weights load straight into the kernel layout (routed_experts_prefix ''), are requantized in process_weights_after_loading, and layerwise reload is refused until vllm-project#57733, as in vllm-project#60895. - Other moe backends and other DeepseekV2MoE users are unchanged; DeepSeek-V2/V3 models get a clear NotImplementedError for this backend. Signed-off-by: Muhammad Ibrahim <ibrahimnd2000@gmail.com>
Resolve the loader conflicts with vllm-project#61457 (per-layer fused shared expert selection) in the GLM-5.3 model and MTP drafter. Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
vllm-project#64357 runs the model's process_weights_after_loading at the end of every reload, and vllm-project#64336's direct mode reruns it too, so the block-FP8 shards a reload writes are requantized in place. Remove the refusal added until that landed. Test both modes through start_reload / finish_reload. Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Keep the requant math, split weight updates, and reload in both modes, which now also checks the cold load layout. Drop tests that repeat those or restate the shared-expert fit check. Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
DeepGemmMegaMoEExperts needs no hooks for GLM: the adapter overrides the existing shared-expert fusion decision and loaders, and reads a shard's kernel-layout view through _kernel_layout_shard. The backend file is back to main. Under MegaMoE the model no longer builds a shared MLP only to drop it, which removes the name probing around it. Drop checks for states the block-FP8 entry check already rules out and helpers with one caller. Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…kend Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
It has nothing specific to GLM, and it no longer imports a private function from another module. Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Requantizing block-FP8 MoE weights to MXFP8 is exact when every block scale is a power of two, and rounds the weights otherwise, as with GLM-5.3-Flash's checkpoint. Warn once in that case. Signed-off-by: Canlin <canlinguosdu@gmail.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
gau-nernst
left a comment
There was a problem hiding this comment.
LGTM. just a minor nit, non-blocking. I think as long as we make it clear that requant takes place, it should be ok.
cc @JaredforReal for a quick check
|
✅ @gcanlin, CI is now available for this PR.
|
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Canlin <canlinguosdu@gmail.com>
|
/ci run |
|
✅ Triggered Buildkite CI #99906 for commit ℹ️ This PR is 20 commits behind upstream |
Overview
Run GLM-5.3-Flash's MoE layers with DeepGEMM's FP8xFP8 MegaMoE (routed and shared experts in one kernel that also dispatches and combines across the EP group), using the official FP8 checkpoint. Enabled by
--kernel-config '{"moe_backend":"deep_gemm_mega_moe"}'with--tensor-parallel-size 1 --data-parallel-size N --enable-expert-parallel.Related: #55275 (BF16 / MXFP4 MegaMoE for GLM, open), #54049 (FlashInfer NVFP4 mega kernel).
Claims
moe_backend.Validation
4x B200, nightly
8cbd5d03with this change applied, MTP 5 with acceptance fixed bysynthetic_acceptance_length3.61. The three arms ran on the same node and in the same container, one run each; they differ only in parallel layout andmoe_backend.AgentX replay (InferenceX-pinned aiperf,
semianalysis_cc_traces_weka_062126), concurrency 96, 1 hour:vllm bench serve, random prompts, concurrency 128, on the same servers:Accuracy:
GSM8K ran on nightly
43b4aaea(TP4) and the servers above (MegaMoE). AIME ran on nightlydf3bd094with this branch, both arms in the same container (MTP 3, fp8 KV, 131K output cap) with DeepInfer'seval/aimeprompt and rule-based judge, temperature 1.0, top_p 0.95; the generic arm runsFLASHINFER_TRTLLMon the checkpoint's own scales. With the full 1M context MegaMoE reaches 0.977 ± 0.005; the remaining loss is mostly problem 15 (0.72), where 12 of 32 samples never stop and loop until the context is full.Weight reload, 4x B200, nightly
e0d00e4b8ewith this branch, DEP4, MTP 3. The experts are zeroed before each reload, which breaks every greedy output; all results are compared with the cold start:The same holds for a second reload, a reload of the MTP drafter, and an update that sends expert weights and block scales in separate
load_weightscalls. Long prompts (about 14K tokens) already vary between identical runs without a reload, on both MoE backends, so they are not compared. No real NCCL or IPC trainer was run.Unit tests:
pytest tests/models/glm5next/test_mega_moe.py, 5 passed.Commands
AIME 2026:
The client sends DeepInfer's
eval/aimeprompts (build_records) to/v1/chat/completionswithtemperature=1.0, top_p=0.95, 32 samples per problem, and scoresreasoning + "</think>" + contentwith itsscore_rollout. Samples were drawn with a 131K output cap first; the 82 that hit it were continued from their own text to the 1M context, which samples the same distribution as the larger cap.The baselines drop
moe_backend(DEP4, generic) or use--tensor-parallel-size 4alone with--intra-node-data-parallel-size 1on the router (TP4). The generic path pickedFLASHINFER_TRTLLMfor the FP8 experts. The random-prompt rows were measured on the same servers before the replay withvllm bench serve --dataset-name random --max-concurrency 128.Details
requant_block_fp8_to_ue8m0).fused_topk_bias(sigmoid, score correction bias), withrouted_scaling_factorfolded into the top-k weights. DeepGEMM's Mega Gate only supportssqrtsoftplus.Glm5NextMegaMoEExpertssubclassesDeepGemmMegaMoEExperts, which is unchanged, with theMEGA_MOE_MXFP8format. E4M3 shards load straight into the kernel layout; their float32 block scales are staged in*_weight_scale_inv, and the model'sprocess_weights_after_loadingcallsrequantize_block_fp8(), which rounds every newly loaded shard to UE8M0 1x32 scales in place (main model, multimodal wrapper and MTP drafter).load_weightscalls; by the end of the load both must be in, orrequantize_block_fp8()raises.weight_loader_numel = 0, so a layerwise reload replays a layer as soon as its checkpoint tensors are in.--prefix-match-unitmatters for this model independently of this PR: the KDA state page makes the KV block 8,960 tokens under DEP4, and without the flag prefix hits land on whole blocks.tensor_parallel_size=1only; SM100 only; no sharded RDT or sparse NCCL updates. Under DEP4 each rank holds 8.8M KV tokens.Pull Request Checklist
I used vLLM's
/pr-checklistskill. (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.
🤖 Generated with Claude Code