[CUDA] Add an opt-in FP8 DeepGEMM MoE decode path for QMoE - #32122
Conversation
Gate a persistent-BF16 SM90 masked grouped GEMM path behind ORT_DSV4_FP4_DEEPGEMM=1. K=6 verify improves from 27.040 to 21.043 ms; real-text decode improves from 19.0 to 17.1 ms/step with zero rank disagreements. Direct FC1/FC2 comparison is bit-exact to cuBLAS.
The decode MoE chain lays tokens out as [32 experts, 64 padded rows, K] but only masked_m[expert] rows are real, typically 1-3 of 64. PackInputKernel and InterleavedSwiGLUKernel both walked the full padded extent, so their cost was constant regardless of how much work there actually was: ncu measured them at a flat 11.0 us and 14.4 us per layer with DRAM SOL of 0.18% and 20.9%. Bound both loops by the per-expert row count instead. The padded rows are now left stale rather than zero-filled, which is safe because the GEMMs are row-independent and UnpackOutputKernel copies back only the first count rows. This was verified bitwise against a poisoned workspace (NaN, +/-Inf, 1e30 across 128 patterns, ragged and zero counts): the compact output is identical, even though 40-60% of the padded FC2 rows do become NaN. Also give SwiGLU a 2D grid over experts and load the adjacent gate/linear pair as one __nv_bfloat162; that part is an exact refactor, bit-identical on valid rows. The two kernels now scale with real work and sit within 0.35 us of the empty kernel launch floor: 1.55 us and 1.51 us at the realistic decode point. Bounding by count removes an implicit clamp the old padded loop bound gave for free, so pin the invariant with a static_assert alongside the alignment one the vectorized load now needs. DeepSeek-V4-Flash, 8xH200, world 8, 1024 prompt / 128 gen: decode 92.12 -> 96.23 tps, 10.86 -> 10.39 ms/token. Acceptance rate 0.329 and 2.65 tokens/step are unchanged, and rank argmax disagreements stay 0.
The DSV4 MoE decode GEMMs are DRAM bound, so the prepacked weight format sets their cost. TryBuildDsv4DeepGemmWeights dequantized the stored MXFP4 weights all the way to bf16, moving 2 bytes per weight. Convert to fp8 e4m3 with a per-[128 N, 128 K] block scale instead and run DeepGEMM's sm90 fp8 masked grouped kernel. The conversion is bit exact. An E2M1 code carries at most two significant bits and e4m3 carries four, so a *power-of-two* block scale only shifts exponents and never disturbs a mantissa. Checked over all 256 experts of all 46 layers: every weight reproduces the fp32 dequantization bitwise, with no underflow and no clipping. The headroom is the group exponent spread within a block, which may reach 14 binades; the measured maximum is 6. Note the conventional amax/448 block scale is *not* usable here: it is not a power of two, so every weight would acquire a full mantissa before being rounded back to three bits (4.8% max relative error). The sm90 fp8 kernel takes both operands in fp8, so activations are quantized too. That work folds into the existing pack and SwiGLU kernels, which already read and write exactly this data, so it is nearly free. Activations use an amax/448 scale per token per 128-channel chunk, where the full e4m3 range is worth more than an exact scale. Measured on 8xH200 at prompt 1024 / generation 512, per decode step: 27.59 ms -> 24.43 ms, a 11.5% reduction, and 32 GiB per rank of weights freed. Standalone the FC1 GEMM goes 70.9 -> 37.9 us and FC2 38.6 -> 22.3 us. Smaller tiles also soften the wave quantization cliff: going from 8 to 9 active experts costs bf16 35% but fp8 only 18%. MMLU-Pro (800 samples) is unchanged at 0.660 vs 0.6625, 2 disagreements, McNemar p=0.48. Rank argmax disagreements remain 0.
… 64) A DSV4 rank owns 256 / world experts, so an eight-rank export has 32 and a four-rank one has 64. kNumExperts was hardcoded to 32, which silently routed every four-rank decode to the generic FP4 grouped GEMM. Both counts are now instantiated and the count travels as a runtime argument. The masked grouped GEMM does not care how many experts there are: its scheduler takes num_m_blocks = ceil_div(masked_m[g], BLOCK_M), so an expert with no rows contributes no tile. Only the workspace, the tensor-map extents and the pack / SwiGLU / unpack grids scale with the count. This does not make 64 experts usable on a 141 GiB H200. DeepGEMM reads an e4m3 mirror of the MXFP4 weights, 1 byte per weight against the checkpoint's 0.5, which at 64 experts is 1,536 MiB per MoE layer and 64.5 GiB per rank over 43 layers -- on top of the SM80-interleaved e2m1 that prefill still reads. The session dies in PrePack. ORT_DSV4_FP4_DEEPGEMM must stay 0 at four ranks until a grouped GEMM reads the 4-bit weights directly.
There was a problem hiding this comment.
Pull request overview
Adds an opt-in FP8 DeepGEMM decode path for fixed-shape CUDA QMoE workloads on Hopper GPUs.
Changes:
- Converts MXFP4 expert weights into persistent E4M3 buffers.
- Adds masked grouped FP8 FC1/FC2 kernels and dispatch.
- Integrates DeepGEMM and documents configuration.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 11 comments.
Show a summary per file
| File | Description |
|---|---|
cmake/external/deep_gemm.cmake |
Fetches DeepGEMM. |
cmake/onnxruntime_providers_cuda.cmake |
Configures bundled CUDA build. |
cmake/onnxruntime_providers_cuda_plugin.cmake |
Configures CUDA plugin build. |
cmake/onnxruntime_cuda_source_filters.cmake |
Classifies the SM90 source. |
docs/contrib_ops/cuda/moe_qmoe.md |
Documents the new path. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/deep_gemm_sm90.cu |
Implements FP8 grouped GEMMs. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/deep_gemm_sm90.h |
Declares fixed-shape interfaces. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu |
Adds workspace and execution dispatch. |
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.h |
Extends runner state and interfaces. |
onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc |
Adds gating, prepack, and routing. |
onnxruntime/contrib_ops/cuda/moe/moe_quantization.h |
Stores configuration and packed buffers. |
onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu |
Implements MXFP4-to-E4M3 conversion. |
onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.h |
Declares the conversion launcher. |
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Factor the arbitrary per-expert global scale out before FP8 quantization so power-of-two scaling remains bit-exact. Fixes the CUDA and TensorRT QMoEFp4ToFp8RoundsScaleUp failure.
|
Review — PR #32122: [CUDA] Add an opt-in FP8 DeepGEMM MoE decode path for QMoE Scope Large PR (17 files, +1121/-9) adding an opt-in specialized decode path for QMoE on Hopper. Behind Author is a trusted CUDA EP contributor and the accuracy validation (MMLU-Pro 800 samples, McNemar p=0.48, 0 rank-argmax disagreements, per-weight bit-exact round-trip check across all 256 × 46 experts) is thorough. Testing story is unusually strong for a change of this size. Reviewing across four axes:
Axis 1 — MXFP4 → E4M3 correctness (the load-bearing insight) The PR description spells out why the conventional
The proof discipline is correct. The runtime The kernel implementation itself in qmoe_kernels.cu:
One micro concern on the launcher's ORT_ENFORCE(n % kQMoEFp8BlockN == 0 && k % kQMoEFp8BlockK == 0 && (n % kQMoEFp8TileN) == 0, ...)Throws Axis 2 — Dispatch gating (both static and runtime) Static gate lives in if (use_fp4_deep_gemm_ && num_rows > 0 && num_rows <= deep_gemm_sm90::kMaxTokensPerExpert &&
hidden_size == deep_gemm_sm90::kHiddenSize && inter_size == deep_gemm_sm90::kInterSize &&
deep_gemm_sm90::NumExpertsSupported(num_experts_per_node) && experts_per_token == 6 &&
activation_type == ActivationType::Swiglu && !use_awq) {
fp4_deep_gemm_workspace_size = deep_gemm_sm90::GetWorkspaceSize(num_experts_per_node);
}The runtime gate in Small inconsistency worth flagging: the sizer gate checks Runtime path: Axis 3 — Third-party dependency (DeepGEMM) New addition to cmake/deps.txt: Concerns:
Axis 4 — Build integration and lint
Documentation quality §9.12 of moe_qmoe.md is unusually thorough:
The one gap: the 64-expert (4-rank) variant is instantiated ( Also flagged in the description: "At four ranks (64 experts) the e4m3 mirror costs 64.5 GiB per rank over 43 layers and the session dies in PrePack, so Test coverage Three new unit tests in qmoe_fp4_to_fp8_kernel_test.cc:
The three tests together validate the arithmetic + safety net. Not covered by unit tests: multi-block and multi-expert quantization (though the kernel is per-block per-expert, so 1-block/1-expert is representative), interaction with the One test-quality nit: API surface additions Two new virtual methods on virtual void setUseFp4DeepGemm(bool /*use_fp4_deep_gemm*/) {}
virtual void setFp4DeepGemmWeightScales(const float* /*fc1_scales*/, const float* /*fc2_scales*/) {}Default no-op in the base, real impl in CI status 88/91 checks OK on Environment variable name change
Summary of items to address before merge Blocking:
Non-blocking but strongly recommended:
Nice to have:
Recommendation Approve — pending #1 (the inexact-flag handling documentation/confirmation) and #2 (cgmanifest.json addition). Everything else is polish. The core theory (power-of-two block scale → bit-exact E2M1 → E4M3 conversion) is correct, the safety net is in place, the gates are conservative, the documentation is thorough, and the accuracy validation is rigorous. Well-executed opt-in specialization. The bit-exact PrePack conversion trick is a genuinely nice insight — it's the kind of numerical-representation reasoning that's worth generalizing. Worth writing up as a design note or blog post separately. |
|
Thanks Hariharan Seshadri (@hariharans29) for the detailed review. I addressed the ten follow-ups in
Validation on the feedback commit:
|
Description
Adds an opt-in DeepGEMM-based MoE decode path for QMoE, behind
ORT_QMOE_FP4_DEEPGEMM(default off). The DSV4 MoE decode GEMMs are DRAM bound, so the prepacked weight format sets their cost; this path stores an fp8 e4m3 mirror of the MXFP4 weights with a power-of-two block scale and runs DeepGEMM's sm90 fp8 masked grouped kernel.Measured 11.5% lower decode step time on 8xH200 and 32 GiB per rank of weights freed relative to the bf16 mirror it replaces.
Summary of Changes
Build
cmake/external/deep_gemm.cmakeFetchContentof deepseek-ai/DeepGEMM.cmake/onnxruntime_providers_cuda.cmakedeep_gemm_sm90.cuto the SM90 TMA target.cmake/onnxruntime_providers_cuda_plugin.cmakecmake/onnxruntime_cuda_source_filters.cmakedeep_gemm_sm90.cuas SM90-only.Kernels and dispatch
onnxruntime/contrib_ops/cuda/llm/moe_gemm/deep_gemm_sm90.{cu,h}onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.{cu,h}onnxruntime/contrib_ops/cuda/moe/moe_quantization.{cc,h}onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.{cu,h}docs/contrib_ops/cuda/moe_qmoe.mdTesting
python test_qmoe_fp4_cuda.py— passes with the path both enabled and disabled.deep_gemm_sm90.cucompiles into the SM90 TMA target.Motivation and Context
Per decode step on 8xH200, prompt 1024 / generation 512:
Smaller tiles also soften the wave-quantization cliff: going from 8 to 9 active experts costs bf16 35% but fp8 only 18%.
Two things reviewers should know:
amax/448block scale is not usable here. It is not a power of two, so every weight would acquire a full mantissa before being rounded back to three bits (4.8% max relative error). The power-of-two scale is load-bearing for exactness, not an optimization. Activations do useamax/448per token per 128-channel chunk, where the full e4m3 range is worth more than an exact scale.moe_qmoe.md.The main reviewer-facing consideration is the new third-party
FetchContentdependency; it is header-only and confined to the SM90 build.Checklist