[CUDA] QMoE: release raw MXFP4 weight initializers after PrePack - #31154
Merged
Tianlei Wu (tianleiwu) merged 1 commit intoJul 31, 2026
Conversation
Tianlei Wu (tianleiwu)
force-pushed
the
tlwu/20260730/fp4_qmoe_no_weight_copy
branch
from
July 30, 2026 18:29
f59da47 to
c60c50b
Compare
Tianlei Wu (tianleiwu)
force-pushed
the
tlwu/20260730/fp4_qmoe_no_weight_copy
branch
from
July 30, 2026 20:54
c60c50b to
22ff9d0
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
This PR reduces persistent device-memory usage for MXFP4 QMoE on SM80–SM119 by enabling both prefill (SM80 grouped GEMM) and decode (fused GEMV) to share the same SM80 pair-interleaved prepacked weight buffers, and then allowing ORT to release the original raw FP4 weight initializers after PrePack.
Changes:
- Update FP4 GEMV decode to optionally read the SM80 pair-interleaved (prefill) weight layout by in-register un-permutation, avoiding a second decode-layout weight copy when shapes satisfy the interleaved rules.
- Enable releasing raw MXFP4 weight initializers after
PrePackin the SM80 grouped-GEMM regime, with a defensive runtime guard to fail fast if buffers are unexpectedly incomplete. - Add Python CUDA tests validating numerical parity vs. the dequant fallback and validating that raw initializer release returns device memory (when initializers bypass the arena), plus documentation for the single-copy regime.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py | Adds SM80-regime tests for parity vs fallback and for device-memory reduction when raw initializers are released. |
| onnxruntime/contrib_ops/cuda/moe/moe_quantization.h | Adds state/flags tracking raw-initializer release and whether GEMV reads the SM80 layout. |
| onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc | Implements raw-initializer release via PrePack, routes decode GEMV to shared weights when possible, and adds a safety guard if fallback would need released tensors. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h | Extends GEMV APIs to support an SM80 pair-interleaved weight mode and adds a layout-only shape gate helper for PrePack. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu | Adds the SM80-pair GEMV dispatch variant and the layout-only support predicate; threads the new flag through launch paths. |
| onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h | Extends FP4 decode to optionally invert the SM80 pair-interleave at compile time via a templated converter. |
| docs/contrib_ops/cuda/moe_qmoe.md | Documents the new SM80 single-copy weights behavior, shape gate, and the initializer-release interaction with the device allocator / arena. |
Tianlei Wu (tianleiwu)
requested review from
Akshay Sonawane (apsonawane),
Hariharan Seshadri (hariharans29) and
kunal-vaishnavi
July 31, 2026 03:34
Hariharan Seshadri (hariharans29)
approved these changes
Jul 31, 2026
Tianlei Wu (tianleiwu)
added a commit
that referenced
this pull request
Aug 6, 2026
…fy (#31159) ### 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. 1. **Packed E2M1 dequantize** (`Fp4I2FConverter::decode_quad`) — decode a whole 32-bit weight word (eight codes) per step instead of one code at a time. 2. **`ORT_FP4_GEMV_DEFAULT_TILING`** — env switch to bypass the autotuner and take the default tiling, for A/B and for avoiding autotune cost in short runs. 3. **Cut memory and ALU traffic in the decode GEMV.** 4. **`kMaxProfiledExpandedRows` 8 -> 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: > ALU 78.9%, DRAM 7.3%, occupancy 21% (register-limited) That is why the levers here are instruction-count levers. Two things were measured and explicitly **dropped** because of it: smaller `CtaN` tiling (the autotuner still picks `threads64`/`CtaN=8`; `CtaN=4` never wins because the kernel is compute-bound, not occupancy-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 + MTP `N=3` (verify batch `M=4`). ### 1. Packed E2M1 dequantize `prmt` selects four bytes per instruction, so a 4-element magnitude lookup costs one instruction instead 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: | kernel | before | after | |---|---:|---:| | fc1 (SwiGLU-fused) | 33.2 µs | **26.2 µs** | | fc2 | 30.2 µs | **22.2 µs** | ### 2. Cut memory and ALU traffic — −0.46 ms/step (−5.1%) The scales of the `CtaN` columns a block owns sit `Interleave` elements apart, so for the non-interleaved ColumnMajor layout (`Interleave == 1`) the whole `CtaN`-wide scale vector is contiguous and can be fetched with one wide access instead of `CtaN` scalar ones. This matters far more than the byte count suggests: with a groupwise scale (NVFP4 `GroupSize = 16`) and `StepK = 8`, a warp's 32 lanes cover 16 distinct scale rows that are `n` elements apart, so *every* scale load touches 16 different sectors — `CtaN * 16` sectors, using 2 bytes out of each 32-byte sector. Per-kernel (graph OFF, 40 launches/step each): | kernel | before | after | |---|---:|---:| | `moe_gemv_interleaved_swiglu_kernel` | 0.956 ms/step | **0.678 ms/step** | | `moe_gemv_kernel` | 0.700 ms/step | **0.494 ms/step** | | **family total** | **1.657 ms/step** | **1.174 ms/step** | End-to-end (4 interleaved `.so`-swap reps per arm): * before: 8.973 / 9.009 / 8.992 / 9.017 * after: 8.567 / 8.508 / 8.498 / 8.583 **8.998 -> 8.539 ms/step.** No overlap between the two sets. ### 3. `kMaxProfiledExpandedRows` 8 -> 64 The fused GEMV rejects `expanded_num_rows > kMaxProfiledExpandedRows`. Qwen3.6 is top-8, so single-token decode expands to 8 rows (accepted), but an MTP verify does not: an `(N+1)`-token verify for `num_speculative_tokens = N` expands to `(N+1) * 8` rows, 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=3` shape used today with headroom to `N=7`. ### Tests * `onnxruntime_provider_test` FP4/FP8/QMoE: 18/18 pass. * `onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py`: 22/22 pass, including new multi-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=node` inflates durations ~35% globally and up to 3.8x for large-grid kernels. > [!IMPORTANT] > `Fp4I2FConverter::convert()` gained a `PairInterleaved` template parameter in #31154. The > packed path added here assumes the **plain** nibble order (nibble `j` of the word is logical > element `j`), which is what its `prmt` selectors encode, so it is nested inside > `if constexpr (!PairInterleaved)`. Please check that guard carefully during review — applied > without it, the pair-interleaved SM80 layout would silently decode to the wrong values.
Tianlei Wu (tianleiwu)
added a commit
that referenced
this pull request
Aug 6, 2026
…fy (#31159) Four changes to the fused NVFP4 QMoE decode GEMV. Stacked on #31154 — **review only the top four commits**; the base branch is that PR. 1. **Packed E2M1 dequantize** (`Fp4I2FConverter::decode_quad`) — decode a whole 32-bit weight word (eight codes) per step instead of one code at a time. 2. **`ORT_FP4_GEMV_DEFAULT_TILING`** — env switch to bypass the autotuner and take the default tiling, for A/B and for avoiding autotune cost in short runs. 3. **Cut memory and ALU traffic in the decode GEMV.** 4. **`kMaxProfiledExpandedRows` 8 -> 64** so MTP verify steps stay on the GEMV path. 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: > ALU 78.9%, DRAM 7.3%, occupancy 21% (register-limited) That is why the levers here are instruction-count levers. Two things were measured and explicitly **dropped** because of it: smaller `CtaN` tiling (the autotuner still picks `threads64`/`CtaN=8`; `CtaN=4` never wins because the kernel is compute-bound, not occupancy-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 + MTP `N=3` (verify batch `M=4`). `prmt` selects four bytes per instruction, so a 4-element magnitude lookup costs one instruction instead 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: | kernel | before | after | |---|---:|---:| | fc1 (SwiGLU-fused) | 33.2 µs | **26.2 µs** | | fc2 | 30.2 µs | **22.2 µs** | The scales of the `CtaN` columns a block owns sit `Interleave` elements apart, so for the non-interleaved ColumnMajor layout (`Interleave == 1`) the whole `CtaN`-wide scale vector is contiguous and can be fetched with one wide access instead of `CtaN` scalar ones. This matters far more than the byte count suggests: with a groupwise scale (NVFP4 `GroupSize = 16`) and `StepK = 8`, a warp's 32 lanes cover 16 distinct scale rows that are `n` elements apart, so *every* scale load touches 16 different sectors — `CtaN * 16` sectors, using 2 bytes out of each 32-byte sector. Per-kernel (graph OFF, 40 launches/step each): | kernel | before | after | |---|---:|---:| | `moe_gemv_interleaved_swiglu_kernel` | 0.956 ms/step | **0.678 ms/step** | | `moe_gemv_kernel` | 0.700 ms/step | **0.494 ms/step** | | **family total** | **1.657 ms/step** | **1.174 ms/step** | End-to-end (4 interleaved `.so`-swap reps per arm): * before: 8.973 / 9.009 / 8.992 / 9.017 * after: 8.567 / 8.508 / 8.498 / 8.583 **8.998 -> 8.539 ms/step.** No overlap between the two sets. The fused GEMV rejects `expanded_num_rows > kMaxProfiledExpandedRows`. Qwen3.6 is top-8, so single-token decode expands to 8 rows (accepted), but an MTP verify does not: an `(N+1)`-token verify for `num_speculative_tokens = N` expands to `(N+1) * 8` rows, 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=3` shape used today with headroom to `N=7`. * `onnxruntime_provider_test` FP4/FP8/QMoE: 18/18 pass. * `onnxruntime/test/python/transformers/test_qmoe_nvfp4_cuda.py`: 22/22 pass, including new multi-token GEMV cases and a `gemv_mode="0"` dequant-fallback companion on the identical shape, so both must match the same exact dequantized reference. 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=node` inflates durations ~35% globally and up to 3.8x for large-grid kernels. > [!IMPORTANT] > `Fp4I2FConverter::convert()` gained a `PairInterleaved` template parameter in #31154. The > packed path added here assumes the **plain** nibble order (nibble `j` of the word is logical > element `j`), which is what its `prmt` selectors encode, so it is nested inside > `if constexpr (!PairInterleaved)`. Please check that guard carefully during review — applied > without it, the pair-interleaved SM80 layout would silently decode to the wrong values.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
On SM80–SM119 the MXFP4 QMoE node used to keep up to three persistent copies of the expert
weights: the raw
[E, K, N/2]initializers (only ever read by the dequant fallback), the SM80pair-interleaved buffer consumed by the grouped-GEMM prefill, and a GEMV-native buffer consumed
by the fused decode GEMV. For a 20B-class MXFP4 MoE each copy is ~9 GiB, which put post-load
VRAM out of reach of 24 GB consumer cards.
This PR collapses that to a single copy in the default
ORT_FP4_SM80_GEMM=1regime:preprocessor step — the
[e0,e2,e4,e6,e1,e3,e5,e7]nibble pair-interleave applied byinterleave_int4s_inplace_kernel. Inverting it in the decoder is a compile-time index remapof the same eight
decodecalls (Fp4I2FConverter<AType, PairInterleaved>), so there are noextra branches, registers, or memory traffic.
gemv_fp4_fc{1,2}_reads_sm80_layout_recordsper-FC whether the dedicated
gemv_fp4_fc*_weights_decode_copy can be skipped.PrePackreleases the raw initializers. With both prefill and decode served frompre-packed buffers, the dequant fallback — the only consumer of the raw layout — is
unreachable, so
PrePackreportsis_packed = truefor inputs 2/5 and cachesfc*_weights_shape_somoe_helper::CheckInputscan still validate shapes. The staged e8m0block-scale copy is dropped in
TryBuildGemvFp4Scalesfor the same reason.A defensive guard in
ComputeInternalreturns a descriptive error (pointing atORT_FP4_SM80_GEMM=0) if the raw weights were released but the SM80 buffers are somehowincomplete, rather than dereferencing a null initializer.
Measured —
gpt-oss-20bMXFP4, post-load device memoryShape gate and fallback
Reusing the interleaved GEMV (
kInterleave=4,kStepK=32) inherits its rules — MXFP4group_size == 32,n % 16 == 0,k % 64 == 0— checked at pack time byis_moe_gemv_fp4_sm80_layout_supported. Shapes that miss the gate transparently keep thededicated decode-layout copy (logged at INFO). NVFP4 (block 16) always uses the plain
ColToRowlayout and is unaffected.
gpt-oss-20b(k=2880,n=5760/2880) clears the gate.Setting
ORT_FP4_SM80_GEMM=0restores the previous behavior: raw initializers retained, prefillon the dequant fallback, decode on the GEMV-native copy.
Motivation and Context
Makes 20B-class MXFP4 MoE models loadable on 24 GB consumer GPUs, and removes ~9 GiB of dead
device memory on every SM80–SM119 deployment.
Testing
New
TestQMoEFP4Sm80SingleWeightCopyinonnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py(gated on a build with--use_fp4_qmoeand on80 <= SM < 120):test_sm80_parity_vs_dequant_fallback— fp16/bf16 × 1/4/64/256 tokens,hidden = inter = 512, 8 experts, top-2, SwiGLU. Every case is run twice, withORT_FP4_SM80_GEMM=1and=0,and both are compared against a PyTorch reference built from the dequantized weights. The
1/4-token cases exercise the fused decode GEMV un-permuting the pair-interleaved buffer; the
64/256-token cases exercise the SM80 grouped GEMM against the dequant fallback.
test_sm80_release_frees_raw_weight_initializers— 16 experts,hidden = 2048,inter = 1024,session.use_device_allocator_for_initializers=1. A throwaway sessionpre-warms the shared CUDA arena, then session-creation device-memory deltas are compared
between the retaining and releasing configurations.
Results on H200 (SM90), CUDA 13.0 — 8 passed:
Memory test: raw e2m1 initializers 48.0 MiB; retaining session 172.0 MiB vs releasing session
120.0 MiB → 52.0 MiB returned to the device.
Notes for reviewers
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.{cu,h}also change in [CUDA] Address feedback for block-scaled MatMul and MoE GEMV #29910; expecta conflict depending on merge order.