Repository navigation
Add MoE dispatch/combine fake quantization for vLLM's DeepGEMM Mega-MoE - #2767
Conversation
Give vLLM's DeepGEMM Mega-MoE experts (--moe-backend deep_gemm_mega_moe; DeepSeek-V4/V4.1 and Kimi-K3) the same dispatch_quantizer and combine_quantizer as RoutedExperts. The kernel fuses dispatch, expert GEMMs and combine, so: - dispatch_quantizer fake-quantizes the experts' input tokens, which only feed the kernel's FP8 staging; - combine_quantizer fake-quantizes the router-weighted per-(token, top-k slot) rows the kernel leaves at the end of its symmetric buffer, then sums them again in the kernel's order (ascending slot, FP32 from zero, one BF16 rounding; a fused shared expert's slot last and unquantized) and replaces the kernel output. On the first batch with routed rows, each buffer checks that its rows reproduce the kernel output bit for bit, so a changed DeepGEMM layout raises. Serving needs a calibrated or constant combine scale (each rank re-quantizes rows from all expert ranks), and quantizing the dispatch needs the shared expert unfused. No host syncs after that check, so it runs under CUDA graphs. Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to No actionable merge-blocking issue remains after normal checks. Pre-merge checks ·
  |
| Check name | Status | Explanation | Resolution |
|---|---|---|---|
| Docstring Coverage |  |
Docstring coverage is 27.78% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 18 functions across 3 files. | Write docstrings for the functions missing them to satisfy the coverage threshold. |
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title clearly and concisely describes the primary change: adding MoE dispatch and combine fake quantization for vLLM's DeepGEMM Mega-MoE backend. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Security Anti-Patterns No listed security anti-pattern was introduced. The changed modelopt Python files add only vLLM plugin imports and tensor/quantizer logic; they contain no torch.load, numpy.load(..., allow_pickle=True…
- Fix all pre-merge checks with AI
✨ Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Changes requested: CUDA graph capture can permanently bypass the combine-buffer safety check and silently use an unverified layout.
Needs action:
- Fix unchecked-buffer handling in
vllm_mega_moe.pybefore CUDA graph serving; see the inline comment. - Add a real CUDA capture/replay regression covering padding-only warmup, constant scales, and corrupted combine rows in
test_vllm_mega_moe.py.
No action needed:
- New headers match
LICENSE_HEADER; the change reuses the existing quantizer registry and communication recipes. - Tests were inspected, not executed.
| valid = (symm_buffer.topk_idx[: output.shape[0]] >= 0).T.unsqueeze(-1) | ||
| num_topk = symm_buffer.num_topk | ||
| checked = id(symm_buffer) in _checked_buffers | ||
| if not checked and not torch.cuda.is_current_stream_capturing() and valid.any(): |
There was a problem hiding this comment.
Bot comment.
This skips validation during capture but still captures the re-sum and output replacement. If eager warmup has only padding rows (as the new test documents), the buffer remains unchecked. CUDA graph replay does not execute this Python branch, so the first real routed batch—and every subsequent replay—uses unverified rows without ever raising on a changed layout or debug-zeroed buffer. Constant-scale recipes can also skip the calibration forward entirely via max_calibrate(..., skip_forward_without_activation_calib=True), so calibration cannot guarantee a prior check. Ensure an actual routed eager validation happens before capturing this replacement, or fail closed/defer graph serving until the buffer is verified. Add a real capture/replay regression for this startup sequence rather than only the eager padding test.
There was a problem hiding this comment.
Fixed in 32796ff: a graph captured before a buffer is checked now records the check in the graph, so every replay compares the rows with the kernel output on the device and fails with a device-side assert on a mismatch. test_mega_moe_graph_checks_layout_on_replay captures and replays a real graph after a padding-only warmup with a constant scale, with the correct and with a changed buffer layout.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2767 +/- ##
==========================================
+ Coverage 71.67% 77.84% +6.17%
==========================================
Files 644 648 +4
Lines 71470 72760 +1290
==========================================
+ Hits 51225 56639 +5414
+ Misses 20245 16121 -4124
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
A CUDA graph captured before any eager forward with routed rows would never check the combine rows, since replays skip the Python check. Such a capture now records the check in the graph: every replay compares the rows with the kernel output on the device and fails with a device-side assert on a mismatch. Eager forwards still check each buffer once and raise a RuntimeError. Signed-off-by: Shiyang Chen <shiychen@nvidia.com>
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-6.1-sol) — DM the bot to share feedback.
Approved: the prior CUDA-graph validation bypass and missing replay regression are resolved in the current diff.
No action needed:
- ✔️ Resolved since the last review: unchecked buffers now capture a device-side layout assertion, and a real capture/replay test covers padding-only warmup, constant scales, and corrupted rows.
- The replay test records the assertion condition without destroying the CUDA context.
- New headers match
LICENSE_HEADER; the implementation extends the existing quantizer registry and communication recipes. - Tests were inspected, not executed.
2 participants
What does this PR do?
Type of change: new feature
Follow-up to #2660. Adds the
dispatch_quantizerandcombine_quantizerto vLLM's DeepGEMM Mega-MoE experts (--moe-backend deep_gemm_mega_moe, DeepSeek-V4/V4.1 and Kimi-K3). #2660 rejected this backend: the kernel dispatches, runs the expert GEMMs and combines in one call, so the prepare/finalize hooks ofRoutedExpertsnever run.How it works (new
plugins/vllm_mega_moe.py, ~170 lines):dispatch_quantizer: fake-quantizes the experts' input tokens before the kernel's FP8 staging, which only these tokens feed.combine_quantizer: after the kernel, reads the router-weighted per-(token, top-k slot) BF16 rows it leaves at the end of its symmetric buffer, fake-quantizes them, and sums them again in the kernel's order (valid slots ascending, FP32 from zero, one BF16 rounding, a fused shared expert's slot last and unquantized). The result replaces the kernel output.Safeguards:
DG_COMM_KERNEL_DEBUG=1, which zeroes the buffer) raises instead of silently quantizing the wrong data. After that, there are no host syncs, so CUDA graphs work.algorithm: max) or constant global scale, since each rank re-quantizes rows from all expert ranks.VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION=1), because the fused shared expert reads the same tokens. The combine works either way.No per-backend patches to vLLM or DeepGEMM; the buffer offset comes from DeepGEMM's own
get_symm_buffer_size_for_mega_moe.Usage
Same recipes as #2660, e.g. with the
moe_dispatch_nvfp4/moe_combine_nvfp4units:Testing
vLLM 0.30.0 on B200.
tests/gpu_vllm/torch/quantization/test_vllm_mega_moe.py(new, 1 GPU, stand-in kernel and buffer): 7 passed. Covers dispatch + combine with and without a fused shared expert, the three raise paths (fused-shared dispatch, per-call combine scale, changed buffer layout), registration of both expert classes, and that the layout check waits for a batch with routed rows (warmup batches are all padding).deep_gemm_mega_moe, in three setups: FP8 dispatch + combine with the shared expert unfused; FP8 combine only with the shared expert fused; NVFP4 dispatch + combine. On both ranks and both MoE layers: the layout check passed on real routed rows; the calibrated amax equaled the max of what each quantizer saw and matched across ranks; the output equaled the re-sum of the QDQ'd rows, differed from the unquantized output, and was finite. A run with CUDA graphs also completed.tests/gpu_vllm/torch/quantization/test_vllm_dynamic_modules.py: 39 passed; the one error (test_tiny_deepseek_fp8_kv_quant_cfg) also fails on main.Not tested: accuracy on real DeepSeek-V4 weights, Kimi-K3 end to end, dual-batch overlap.
Before your PR is "Ready for review"
Make sure you read and follow Contributor guidelines and your commits are signed (
git commit -s -S).Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded
trust_remote_code=True,torch.load(..., weights_only=False),pickle, etc.).CONTRIBUTING.md: N/AAdditional Information
Builds on #2660 (merged).
🤖 Generated with Claude Code
Summary by CodeRabbit