[CUDA] Speed up the NVFP4 QMoE decode GEMV and enable it for MTP verify - #31159
[CUDA] Speed up the NVFP4 QMoE decode GEMV and enable it for MTP verify#31159tianleiwu wants to merge 8 commits into
Conversation
256c315 to
d73cf35
Compare
There was a problem hiding this comment.
🟡 Not ready to approve
The updated GEMV device kernels introduce alignment-unsafe vector/packed accesses that can cause undefined behavior and should be fixed before approval.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.
Pull request overview
This PR optimizes the CUDA fused FP4 (NVFP4/MXFP4) QMoE decode GEMV path by reducing instruction count and improving scale/weight handling, and expands the GEMV “decode-shaped” support window so multi-token (MTP verify) steps stay on the fast GEMV path instead of falling back to dequantize + grouped GEMM.
Changes:
- Add a packed E2M1 (FP4) dequantization path (
Fp4I2FConverter::decode_quad) and support SM80 pair-interleaved decoding for single-copy MXFP4 weights. - Add an analytic default tiling heuristic (
Fp4MoeGemvDefaultConfig) with an env override (ORT_FP4_GEMV_DEFAULT_TILING) and plumb SM80-pair-interleaved layout through support checks and launchers. - Expand the supported “expanded rows” window (
kMaxProfiledExpandedRows8→64) and add/extend Python tests to cover multi-token GEMV and fallback parity.
File summaries
| File | Description |
|---|---|
| onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py | Updates GEMV support-window docs and adds multi-token (MTP-style) NVFP4 GEMV test coverage plus fallback companion. |
| onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py | Adds SM80 grouped-GEMM single-weight-copy parity and memory-release tests plus env override helper. |
| onnxruntime/contrib_ops/cuda/moe/moe_quantization.h | Adds state for releasing FP4 raw weights and tracking whether GEMV reads SM80 pair-interleaved layout. |
| onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc | Enables raw-weight release in SM80 regime, adds safety guard, improves GEMV config seeding, and supports SM80-pair-interleaved GEMV dispatch. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.h | Raises max profiled expanded rows to keep MTP verify on GEMV. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h | Adds default-tiling helper and extends supported-query/launch signatures for SM80 pair-interleaved layout. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu | Implements default tiling heuristic, SM80-pair-interleaved GEMV details, and updated support/launch logic. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh | Refactors inner-loop to reduce ALU/memory traffic (vectorized scale loads, load/decode scheduling, K-paired accumulation). |
| onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h | Implements packed FP4 decode and adds pair-interleaved FP4 decode mode, forwarding layout selection via ConverterWrapper. |
| docs/contrib_ops/cuda/qmoe_gemv_experiments.md | Documents benchmarks/analysis for packed decode and analytic tiling. |
| docs/contrib_ops/cuda/moe_qmoe.md | Documents SM80 single-copy weights behavior and updated environment variables/shape gates. |
Review details
Comments suppressed due to low confidence (2)
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh:329
- Same issue as the non-SwiGLU kernel: when
ScalesAccessvectorizes,load_scales()writes through afloat4*intovec_scale, butvec_scaleis not guaranteed 16-byte aligned. Add explicit alignment to avoid misaligned vector stores/loads.
TypeA vec_scale[CtaN];
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_device.cuh:348
- Same alignment concern as in
moe_gemv_kernel: the packed FP4 decode path writes totile_wviauint32_t*, which requires 4-byte alignment. Explicitly aligntile_wto avoid undefined behavior.
TypeA tile_w[StepK];
- Files reviewed: 11/11 changed files
- Comments generated: 2
- Review effort level: Low
We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.
d73cf35 to
4a22d73
Compare
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Review summary (5-agent pass: readability, code/test quality, adversarial/critical, deep numerical, cross-module integration)Bottom line: The core bit-trick ( ✅ Verified correct
🔴 Major — two places where "numerically neutral" isn't strictly true
Neither is necessarily a blocking bug — the drift is likely within existing GEMV accuracy tolerances — but the "pure tiling knob, bit-exact" framing in the PR should be corrected, and a cancellation-targeted regression test should be added. Confirming the actual error magnitude would require a full CUDA build + targeted kernel test (happy to help with that if useful, but flagging it as expensive/optional rather than doing it unprompted). 🟡 Major — test coverage gaps
🟢 Minor / Nit
✅ Integration — no cross-module issues found
|
87b9392 to
e25d639
Compare
Two changes to the QMoE FP4 GEMV, both profile-driven on Qwen3.6-35B-A3B MTP decode (H200, hidden 2048, inter 512, 256 experts, top-k 8, block size 16). Vectorize the per-column scale load. A block loads the scales of CtaN adjacent output columns, which for the non-interleaved layout are contiguous, but it did so one scalar at a time. With GroupSize 16 and StepK 8 a warp's lanes cover 16 distinct scale rows, so each of the CtaN loads touched 16 sectors while using two bytes of each: 128 of the 176 L1 sectors a warp requested per K iteration went to scales. One float4 load replaces all CtaN of them. ncu reported 63% excessive sectors and 10.2 of 32 bytes used per sector before, 27.8 after, and L1/TEX utilization fell from 90% to 34%. Pair the products along K instead of across columns. The resulting ALU-bound kernel spent about a ninth of its issued instructions on the prmt sequence that shuffles decoded weights into column pairs for the shared `pack_to_vec2`/`mma` path. Pairing along K needs no shuffle: the activation tile is already in k order and the converters emit k pairs contiguously. The group scale also moves out of the per-weight dequantize and onto the tile sum, since a thread's tile sits inside one scale group -- four fma2 per column and K tile become one. The inner loop drops from 292 to 250 SASS instructions. Batching the CtaN weight loads ahead of the decode is what makes the second change pay off; interleaving each load with its own consumer left one load in flight and made the kernel long-scoreboard bound instead. Measured on the 40-layer MTP decode with N_spec 3: the two GEMV kernels go from 1.657 to 1.174 ms/step and end-to-end decode from 8.998 to 8.539 ms/step (4 interleaved reps per arm, no overlap between arms).
…TP decode Raises kMaxProfiledExpandedRows from 8 to 64 so the fused NVFP4 decode GEMV also covers Qwen-style top_k=8 MTP verification steps: an (N+1)-token verify for num_speculative_tokens=N expands to (N+1)*8 rows, i.e. up to 64 for N=7. Previously those steps fell out of the GEMV support window and back onto the dequantize + grouped-GEMM path. Adds multi-token GEMV coverage (and a gemv_mode="0" fallback companion on the identical shape) to test_qmoe_nvfp4_cuda.py.
load_scales() writes vec_scale through a float4*, the activation iterator writes tile_a through AccessTypeA*, and the packed FP4 decode writes tile_w through uint32_t*. All three were only TypeA-aligned (2 bytes), which is undefined behavior for the wider accesses. Also make the scalar FP4 decode an explicit else branch: with the packed path returning early, nvcc 13 flags the trailing loop as unreachable and -Werror all-warnings turns that into a build failure.
Scale FP4 weights before low-precision multiplication, keep the widened MTP row window FP4-specific, and make tiling use the owning device properties.
…erage Review follow-up. The CtaN/Threads knob was documented as bit-exact, but a block walks K in strides of StepK * Threads and the epilogue reduces across Threads/32 warps, so Threads changes the floating-point summation order. Same for the K-paired inner loop, which is a reassociation of the previous serial per-column chain. Comments in moe_gemv_fp4.h/.cu, moe_quantization.cc, moe_gemv_device.cuh and the experiments doc now say so. Tests add the expanded-rows window boundary (64) and the reject case above it, a k=1024 case that reaches the second clause of Fp4MoeGemvDefaultConfig and doubles as a long-K GEMV-vs-fallback parity check, and a subprocess case for the ORT_FP4_GEMV_DEFAULT_TILING=0 opt-out (the env var is latched in a function-local static, so it needs a fresh process). Also documents the decode_quad prmt selectors byte by byte, the packed path's 4-byte alignment contract, and why the dense fpA_intB GEMV keeps scalar scale loads. No functional change.
e25d639 to
43e7ccb
Compare
Description
Four changes to the fused NVFP4 QMoE decode GEMV. Stacked on #31154 — review only the top four
commits; the base branch is that PR.
Fp4I2FConverter::decode_quad) — decode a whole 32-bit weightword (eight codes) per step instead of one code at a time.
ORT_FP4_GEMV_DEFAULT_TILING— env switch to bypass the autotuner and take the defaulttiling, for A/B and for avoiding autotune cost in short runs.
kMaxProfiledExpandedRows8 -> 64 so MTP verify steps stay on the GEMV path.Motivation and Context
Prior profiling established that this kernel is ALU-pipeline bound, not memory- or
tiling-bound. On the actual Qwen3.6 decode shapes (
hidden=2048,inter=512,E=256,top_k=8, bf16, SwiGLU), ncu reported for the FC1 SwiGLU-fused GEMV:That is why the levers here are instruction-count levers. Two things were measured and
explicitly dropped because of it: smaller
CtaNtiling (the autotuner still picksthreads64/CtaN=8;CtaN=4never wins because the kernel is compute-bound, notoccupancy-bound), and halving scale bandwidth by storing combined scales as 1-byte e4m3 (DRAM is
only ~7%, so it cannot move the needle).
All numbers below: 1x H200 SXM (SM90, 132 SM, ~4.8 TB/s HBM), CUDA 13.0, Qwen3.6-35B-A3B-NVFP4
N=3(verify batchM=4).1. Packed E2M1 dequantize
prmtselects four bytes per instruction, so a 4-element magnitude lookup costs one instructioninstead of four. Bit-identical to the per-element path (same magnitude tables, same sign
handling). The FP4 GEMV kernel SASS shrinks ~30%, and the two QMoE GEMVs drop:
2. Cut memory and ALU traffic — −0.46 ms/step (−5.1%)
The scales of the
CtaNcolumns a block owns sitInterleaveelements apart, so for thenon-interleaved ColumnMajor layout (
Interleave == 1) the wholeCtaN-wide scale vector iscontiguous and can be fetched with one wide access instead of
CtaNscalar ones. This mattersfar more than the byte count suggests: with a groupwise scale (NVFP4
GroupSize = 16) andStepK = 8, a warp's 32 lanes cover 16 distinct scale rows that arenelements apart, soevery scale load touches 16 different sectors —
CtaN * 16sectors, using 2 bytes out of each32-byte sector.
Per-kernel (graph OFF, 40 launches/step each):
moe_gemv_interleaved_swiglu_kernelmoe_gemv_kernelEnd-to-end (4 interleaved
.so-swap reps per arm):8.998 -> 8.539 ms/step. No overlap between the two sets.
3.
kMaxProfiledExpandedRows8 -> 64The fused GEMV rejects
expanded_num_rows > kMaxProfiledExpandedRows. Qwen3.6 is top-8, sosingle-token decode expands to 8 rows (accepted), but an MTP verify does not: an
(N+1)-tokenverify for
num_speculative_tokens = Nexpands to(N+1) * 8rows, i.e. up to 64 for N=7.Those steps fell out of the window and back onto the dequantize + CUTLASS grouped-GEMM path,
which re-dequantizes all 256 experts per token.
The impact of that fallback is large: with the limit at 8, the 2-token verify (expanded 16)
dropped MTP to ~2.4 tok/s; raising the limit put it at ~30–55 tok/s (12–23x). 64 covers
the
N=3shape used today with headroom toN=7.Tests
onnxruntime_provider_testFP4/FP8/QMoE: 18/18 pass.onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py: 22/22 pass, including newmulti-token GEMV cases and a
gemv_mode="0"dequant-fallback companion on the identical shape,so both must match the same exact dequantized reference.
Methodology note
End-to-end deltas are quoted as ms/step from a fixed-step measurement, never tok/s: any
numerics change alters the generated sequence and therefore the MTP acceptance rate, which swamps
the speed delta. Per-kernel durations are taken with CUDA graphs off —
nsys --cuda-graph-trace=nodeinflates durations ~35% globally and up to 3.8x for large-gridkernels.
Important
Fp4I2FConverter::convert()gained aPairInterleavedtemplate parameter in #31154. Thepacked path added here assumes the plain nibble order (nibble
jof the word is logicalelement
j), which is what itsprmtselectors encode, so it is nested insideif constexpr (!PairInterleaved). Please check that guard carefully during review — appliedwithout it, the pair-interleaved SM80 layout would silently decode to the wrong values.