[CUDA] Address feedback for block-scaled MatMul and MoE GEMV - #29910
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR follows up on prior CUDA reviews by consolidating duplicated CUDA type-conversion/packing helpers into a shared header, updating the FP4/FP8 block-scaled MatMul kernels and SM120 native path to use the shared utilities, and applying small readability/const-correctness cleanups in the MoE GEMV code. It also adds FP4 edge-case unit tests to lock in behavior for trailing scale blocks and K==0.
Changes:
- Centralize CUDA scalar/packed conversions (half/bf16, vec2 traits, bit-cast, and raw E4M3 helpers) in
cuda_type_helper.cuhand migrate FP4/FP8 kernels to use them. - Tighten SM120 native FP4 launcher contracts and adjust PrePack/launch call sites accordingly.
- Add additional FP4 unit tests for partial trailing scale-block handling, K%32 fast-path gating, and K==0 output semantics.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/core/providers/cuda/cu_inc/cuda_type_helper.cuh | Adds shared CUDA conversion/packing helpers and makes them reusable across TUs. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu | Switches FP4 kernels to use shared helpers and updates comments/invariants. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu | Switches FP8 kernels to use shared helpers and updates vector-path commentary. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4_sm120.cu | Aligns SM120 quantization/conversion code with shared helper APIs and refines contracts. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc | Adjusts PrePack stream usage and updates native SM120 launch argument handling. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h | Clarifies launcher contracts and documents SM120 path scratch/prepacked buffers. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.cu | Const-correctness/readability updates and clarifying synchronization comments. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv.h | Const-correctness updates in GEMV launcher declarations. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu | Const-correctness updates and minor style alignment in FP4 GEMV entry points. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h | Const-correctness updates in FP4 GEMV launcher declarations. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu | Adds clarifying comments around skip-expand invariants and related guards. |
| onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh | Small const-ordering/style alignment for touched parameter declarations. |
| onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc | Adds FP4 edge-case unit tests (partial scale tail, K%32 dispatch gating, and K==0 semantics). |
tianleiwu
enabled auto-merge (squash)
July 29, 2026 00:27
kunal-vaishnavi
previously approved these changes
Jul 29, 2026
This was referenced Jul 30, 2026
kunal-vaishnavi
approved these changes
Jul 30, 2026
tianleiwu
added a commit
that referenced
this pull request
Jul 31, 2026
) ### 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 SM80 pair-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=1` regime: 1. **The decode GEMV learns to read the prefill layout.** The two layouts differ by exactly one preprocessor step — the `[e0,e2,e4,e6,e1,e3,e5,e7]` nibble pair-interleave applied by `interleave_int4s_inplace_kernel`. Inverting it in the decoder is a compile-time index remap of the same eight `decode` calls (`Fp4I2FConverter<AType, PairInterleaved>`), so there are no extra branches, registers, or memory traffic. `gemv_fp4_fc{1,2}_reads_sm80_layout_` records per-FC whether the dedicated `gemv_fp4_fc*_weights_decode_` copy can be skipped. 2. **`PrePack` releases the raw initializers.** With both prefill and decode served from pre-packed buffers, the dequant fallback — the only consumer of the raw layout — is unreachable, so `PrePack` reports `is_packed = true` for inputs 2/5 and caches `fc*_weights_shape_` so `moe_helper::CheckInputs` can still validate shapes. The staged e8m0 block-scale copy is dropped in `TryBuildGemvFp4Scales` for the same reason. A defensive guard in `ComputeInternal` returns a descriptive error (pointing at `ORT_FP4_SM80_GEMM=0`) if the raw weights were released but the SM80 buffers are somehow incomplete, rather than dereferencing a null initializer. ### Measured — `gpt-oss-20b` MXFP4, post-load device memory | configuration | post-load | |---|---| | before (3 copies) | 32908 MiB | | release initializers only (2 copies) | 23164 MiB | | **single copy (this PR, default)** | **13996 MiB** | > Released initializers only shrink the process's *device* footprint when initializers bypass the > BFC arena, i.e. with the session option `session.use_device_allocator_for_initializers = 1`. > Otherwise the freed bytes are recycled inside the arena for later activation/KV allocations. > This is documented in `docs/contrib_ops/cuda/moe_qmoe.md` §9.11. ### Shape gate and fallback Reusing the interleaved GEMV (`kInterleave=4`, `kStepK=32`) inherits its rules — MXFP4 `group_size == 32`, `n % 16 == 0`, `k % 64 == 0` — checked at pack time by `is_moe_gemv_fp4_sm80_layout_supported`. Shapes that miss the gate transparently keep the dedicated decode-layout copy (logged at INFO). NVFP4 (block 16) always uses the plain `ColToRow` layout and is unaffected. `gpt-oss-20b` (`k=2880`, `n=5760`/`2880`) clears the gate. Setting `ORT_FP4_SM80_GEMM=0` restores the previous behavior: raw initializers retained, prefill on 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 `TestQMoEFP4Sm80SingleWeightCopy` in `onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py` (gated on a build with `--use_fp4_qmoe` and on `80 <= 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, with `ORT_FP4_SM80_GEMM=1` and `=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 session pre-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**: | case | SM80-on vs ref | SM80-off vs ref | cross | |---|---|---|---| | fp16 decode 1 / 4 tok | 0.031 / 0.033 | 0.018 / 0.018 | 0.031 / 0.037 | | bf16 decode 1 / 4 tok | 0.125 | 0.125 | 0.000 | | fp16 prefill 64 / 256 tok | 0.016 | 0.016 | 0.016 | | bf16 prefill 64 tok | 0.125 | 0.125 | 0.125 | 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 #29910; expect a conflict depending on merge order.
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
This PR addresses follow-up review feedback from recently merged CUDA PRs covering block-scaled FP4/FP8 MatMul and MoE GEMV paths. The update consolidates duplicated CUDA type helpers into a shared header, tightens kernel contracts and comments, and adds edge-case validation tests for FP4 block-scaled MatMul behavior.
The goal is to improve maintainability and correctness without changing intended operator semantics.
Summary of Changes
CUDA helper consolidation and kernel cleanup
FP4/NVFP4 behavior and API cleanups
MoE GEMV and related style/documentation updates
Tests
Testing
Motivation and Context
In #29818 and #29887 and #29896, reviewer feedback identified duplicate helper logic and opportunities to tighten contracts/comments across the CUDA block-scaled MatMul and MoE GEMV code paths. Consolidating type/packing helpers in one shared CUDA header reduces duplication and future drift, while the added tests lock in edge-case expectations for FP4 block-scaled behavior.
Checklist