[CUDA] Block-scaled MatMul decode GEMV: tensor cores, packed FP4 decode and M-tiling - #31155
Conversation
The decode GEMV mapped M onto gridDim.y, so the packed NVFP4 weight was streamed from DRAM once per row of A. Speculative decoding / MTP verify runs M = N_spec + 1 rows and paid M x the weight traffic. Template the kernel on RowsPerBlock (1, 2 or 4): a warp now accumulates several rows at once and the uint4 weight load, the prmt E2M1 decode and the two E4M3 block scales are shared across the tile. Per-row fp32 accumulation order is unchanged, so results are bit-identical to RowsPerBlock == 1. Tiling removes M from gridDim.y, so it is gated on the column grid ceil(N / 8) covering at least one full wave of SMs on its own. H200, M = 4, half: N = 248320 (lm_head) 31040 col blocks 615.8 -> 537.7 us (1.15x) N = 2048 (shared down) 256 col blocks 4.30 -> 3.33 us (1.29x) N = 512 (shared gate) 64 col blocks 3.47 -> 4.27 us (0.81x, gated off) Also pin __launch_bounds__ per tile. Left unconstrained nvcc compiles the 4-row tile to 66 registers, which drops the block from 4 to 3 per SM; the pinned budgets are spill-free and worth ~7% on the large shapes by themselves. Qwen3.6 NVFP4 MTP decode on H200: FP4 dense GEMV 1.099 -> 0.950 ms/step (1.16x), end-to-end 10.448 -> 10.306 ms/step (-1.4%). Set ORT_FP4_GEMV_ROW_TILING=0 to force RowsPerBlock == 1.
The fused FP8 GEMV re-widens both operands to FP32 and does scalar FMAs, so at M > 1 it is issue bound rather than bandwidth bound: at M = 4 a lane runs ~240 instructions per 32 weight bytes and effective bandwidth drops from ~2.35 TB/s (M = 1) to ~1.25 TB/s. Add MatMulBlockScaledFp8MmaGemvKernel, which does the dot products with mma.m16n8k16 and FP32 accumulate. The weight feeds the mma A operand and the activation the mma B operand, so both fragments match the tensors' natural row-major layouts with no transpose or shared-memory staging; the mma "M" extent becomes the output column count (16 per warp) and the mma "N" extent becomes M. Fragment loads are coalesced by permuting the K axis - legal because K is a reduction axis and the same permutation is applied to both operands - which lets each lane load one contiguous uint4 of weight bytes per 64-element K window and feed four mma steps from it. KSplit warps per block take a strided share of the K windows and reduce through shared memory, restoring the memory-level parallelism lost by widening each warp to 16 columns. Used on SM80+ when K % 64 == 0, K >= 256 and block_size % 64 == 0; otherwise the existing kernel runs unchanged. Set ORT_FP8_GEMV_MMA=0 to force it. Not bit-identical to the FMA path but not less accurate: E4M3 to FP16/BF16 is lossless, the products are exact in FP32, and both kernels accumulate in FP32. Scored against an FP64 reference the maximum error is identical for the two. H200, standalone (us), FMA kernel -> mma kernel: 8192x2048 M=1 6.3 -> 5.1 M=4 9.8 -> 5.2 M=8 17.5 -> 5.7 4096x2048 M=1 4.8 -> 4.0 M=4 6.9 -> 4.1 M=8 10.1 -> 4.3 2048x4096 M=1 5.1 -> 4.2 M=4 8.0 -> 4.4 M=8 13.0 -> 4.5 512x2048 M=1 3.3 -> 3.1 M=4 4.7 -> 3.3 M=8 6.4 -> 3.3 Faster at every measured M, so it is preferred whenever its preconditions hold rather than only for M > 1. On a 40-layer Qwen3.6 NVFP4 MTP decode (130 FP8 matmul nodes per step, M = 4) the FP8 GEMV kernel family goes 1.052 -> 0.713 ms/step and the model goes 9.80 -> 9.54 ms/step, with no other kernel family affected.
The fused NVFP4 decode GEMV gave each warp one output column and reduced over K with __shfl_down. That makes the warp re-read the whole A tile once per output column: at RowsPerBlock = 4 it pulls 16 KB of activation for 1 KB of weight, a 16:1 amplification that dominated lm_head (N = 248320, 537.6 us on H200). On SM80+ with K % 128 == 0 and M <= 8, replace the shuffle reduction with mma.m16n8k16 so a warp produces 16 output columns at once. The fragments need no staging because the weight goes in the mma A slot and the activation in the mma B slot: B is [N, K] row-major, exactly the A-row-major fragment, and A is [M, K] row-major, exactly the B-col-major fragment. The K axis is then permuted so each lane's loads are contiguous, which is legal because K is a reduction axis and the same permutation applies to both operands. Since the mma sums across the four lanes holding different 16-element scale blocks, the accumulator cannot be flushed per block; the E4M3 scale is folded into the decoded weight before the mma instead. That is exact in FP16 and BF16 (E2M1 carries 2 significand bits and E4M3 carries 4, so the product needs at most 6) and cannot overflow (max 6 * 448 = 2688, min 0.5 * 2^-9 = 2^-10). KSplit warps per block take a strided share of the K windows and reduce through shared memory, without which the 16x drop in warp count costs more on the small MLP shapes than the extra reuse buys. Measured on H200 (132 SMs, M = 4, FP16), scalar -> tensor core: lm_head N = 248320, K = 2048 537.6 -> 108.3 us 4.96x gate_up_proj N = 512, K = 2048 3.48 -> 2.99 us 1.17x down_proj N = 2048, K = 512 3.32 -> 2.52 us 1.32x Over a Qwen3.6 NVFP4 MTP decode step this is 0.949 -> 0.448 ms/step for the FP4 GEMV family and 9.54 -> 8.99 ms/step end to end. The fp32 accumulation order differs from the scalar path, so results are not bit-identical to it, but max relative error against an fp64 reference is unchanged (2.4e-04 .. 3.5e-04, i.e. NVFP4 quantization noise). Set ORT_FP4_GEMV_MMA=0 to fall back. The new tests vary the weight, the block scales and the activation along K, so an inconsistent k mapping cannot pass; the existing GEMV tests use K = 32/64 and never reach this path.
eb5b749 to
f545074
Compare
There was a problem hiding this comment.
Pull request overview
This PR optimizes the CUDA decode-GEMV fast paths for the block-scaled FP4/FP8 MatMul contrib ops by adding tensor-core (mma.m16n8k16) variants, improving FP4 decode throughput, and adding M-tiling for speculative decode (M>1). It also adds targeted CUDA tests and updates the operator docs to describe the new dispatch behavior and tuning knobs.
Changes:
- FP8: hoist widening work for M>1 tiles and add an SM80+ tensor-core GEMV path gated by
ORT_FP8_GEMV_MMA. - FP4: implement
prmt-based quad decode, add grid-gated row tiling (ORT_FP4_GEMV_ROW_TILING), and add an SM80+ tensor-core GEMV path gated byORT_FP4_GEMV_MMA. - Add new CUDA unit tests for speculative-decode tiling and tensor-core dispatch; update FP4/FP8 operator documentation and experiment notes.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc | Adds coverage for M>1 speculative decode dispatch and the FP8 tensor-core GEMV path (FP16/BF16, ragged tails, bias). |
| onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc | Adds coverage for FP4 row-tiling GEMV and the FP4 tensor-core GEMV path (FP16/BF16, ragged tails, bias). |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h | Extends GEMV launcher signature to accept sm_major and documents the SM80+ tensor-core selection. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu | Implements FP8 M>1 hoisted-widening and introduces an SM80+ mma.m16n8k16 GEMV kernel plus dispatch gating. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc | Passes device sm_major into the FP8 GEMV launcher. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h | Extends GEMV launcher signature to accept sm_major and documents the SM80+ tensor-core sub-path. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu | Implements FP4 prmt quad decode, row tiling logic, and an SM80+ tensor-core GEMV kernel plus dispatch gating. |
| onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc | Passes device sm_major into the FP4 GEMV launcher. |
| docs/contrib_ops/cuda/matmul_block_scaled_fp8.md | Documents the new FP8 tensor-core GEMV sub-path and kill-switch. |
| docs/contrib_ops/cuda/matmul_block_scaled_fp8_experiments.md | Adds a new “Decode GEMV - Tensor Cores” section and renumbers subsequent sections. |
| docs/contrib_ops/cuda/matmul_block_scaled_fp4.md | Updates scalar GEMV description and documents row tiling + tensor-core GEMV sub-path + new env vars. |
Address PR review on the FP4/FP8 tensor-core decode GEMV: document the exact per-lane m16n8k16 fragment ownership (A rows -> output columns, B column -> activation row, D -> output rows 2t/2t+1) in both kernels and both operator docs, and add a one-hot lane-ownership probe test per data type so any swap of the row/column mapping fails immediately instead of cancelling out. Also make the FP8 tensor-core launch gate check M <= 8 explicitly, matching the FP4 gate, so the structural bound is enforced where it is documented.
Automated multi-agent review of PR #31155Five independent review passes (readability, correctness/idiom, adversarial/critical, deep spec, cross-module integration) were run against this diff. Summary below, deduplicated and prioritized. Items requiring GPU execution to confirm are flagged Critical / Major
Minor
Open questions (not blocking, for maintainer/CI follow-up)
Praise (call these out, worth keeping)
Posted by an automated multi-agent review (5 reviewer passes). Findings with |
Addendum: hardware-verified findings (from a second deep-spec pass)A follow-up deep-spec pass actually executed verification code on real SM80 hardware (not just static analysis) and turned up a more serious test-coverage gap than my first pass surfaced. Posting since it's concrete and actionable: Major — FP8 tensor-core tests are numerically vacuous; an all-zeros kernel would pass
The chosen value pattern (weights period-3, activations period-4, small scales) produces near-total cancellation, so
Minor — doc self-contradiction on FP8 mma accuracy
Confirmed on hardware (positive results, reducing risk elsewhere)The same pass fetched the PTX ISA
Two smaller coverage gaps also found: the |
- Parse the GEMV kill switches as bool so ORT_FP4_GEMV_MMA=false and friends actually disable the path instead of silently falling back to the default. - Pass the device properties into the GEMV launchers instead of caching the multiprocessor count in a function-static, which fixed the tiling heuristic to whichever GPU ran first in a heterogeneous multi-GPU process. This also unifies the sm_major derivation between the FP4 and FP8 call sites. - Cap __launch_bounds__ minBlocksPerSM at 4 for pre-SM80, where 1024 resident threads per SM makes 5/6 blocks of 256 unsatisfiable. - Raise the FP8 tensor-core test scales so |expected| is far above the output tolerance; the previous 1/256 scales left an all-zeros output passing. - Cover the FP8 RowsPerWarp = 8 FMA tile (M in [5, 8]) and the FP4 mma lo_ok == false predication (N = 36). - Align the docs: the FP8 mma path is not bit-identical to the FMA path, the M > 1 FMA dispatch re-tune changes the summation order, and FP4 row tiling is superseded by the mma path on SM80+ when K % 128 == 0.
|
Thanks for the very thorough review (both the initial pass and the hardware-verified addendum). Pushed 1ea28c0 which addresses all the actionable findings. Summary below, including two items I'd like to push back on. FixedEnv kill switches were parsed as
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
constexpr int kGemvMaxBlocksPerSm = 4; // 1024 threads/SM = 4 blocks of 256
#else
constexpr int kGemvMaxBlocksPerSm = 8;
#endifand Process-wide SM-count caching — FP8 tensor-core tests were numerically vacuous — you were right, and thanks for actually computing it. With block scales of Untested branches — added Docs — three corrections:
Also removed the now-unused Pushing backFP8 mma scale folding under Tests that flip the kill switches for in-process A/B parity. I'd rather not do this in-process: each switch is read into a function-local Verification
|
Description
Five optimizations to the decode GEMV fast path of the
MatMulBlockQuantizedFp4WeightandMatMulBlockQuantizedFp8Weightcontrib ops. All are behind existing kill switches(
ORT_FP4_GEMV_MMA=0,ORT_FP8_GEMV_MMA=0,ORT_FP4_GEMV_ROW_TILING=0) for A/B.fp8 blocked scaled gemmRowsPerWarp > 1prmt quad decodeprmt.b32instead of oneFP4 dense GEMV: grid-gated row tilingRun FP8 weight-only decode GEMV on tensor coresmma.m16n8k16Run the NVFP4 decode GEMV on tensor coresmma.m16n8k16Motivation and Context
The shipped GEMV kernels were tuned for
M = 1(plain autoregressive decode). Speculative /MTP decoding makes the verify forward
M > 1wide, and atM = 4theM = 1tuning inverts:per row and A once per column, so it goes ALU-bound.
Measured on H200 (µs, shipping
<R,1,1>kernel vs cuBLAS) — flipping dense projections to FP8without touching the kernel was a regression at
M = 4: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. FP8: hoist the fp32 widening for
RowsPerWarp > 1— −0.13 ms/stepM == 1keeps the original scalar path (hoisting measured slightly worse there). ForR > 1,B is widened once per
(col, half)and reused by all rows, A once per(row, half)and reusedby all cols; the 16-element chunk is consumed in two 8-element halves to cap live registers.
Instructions per 16 weight bytes:
8 + 48*R->8*C + 16*R + 16*C + 16*R*C(200 -> 120 percolumn at
R=4, C=2). The fma order is unchanged, so the output is bit-identical.After the fix,
M=4FP8 beats cuBLAS by 1.04–1.59x and the old<R,1,1>by 1.13–1.41x.End-to-end 2x2 (graph ON, 400 steps, 64 warmup, 2 reps), ms/step:
10.80 -> 10.674 = −0.13 ms/step (−1.2%). Acceptance rate unaffected (2.5–2.75 tok/step in
every cell).
2. FP4:
prmtquad decode — −0.26 ms/step (−2.4%)prmt.b32is a 4-byte table lookup in one instruction, so a whole 32-bit word (8 E2M1 codes)is decoded at once instead of per element.
SASS instruction count: half 352 -> 288 (−18%), bf16 376 -> 336 (−11%).
End-to-end (graph ON, 400 steps, 3 interleaved reps): 10.766 -> 10.504 ms/step; every prmt
rep beats every baseline rep. Bit-identical — an exhaustive 256-byte sweep against the old
Raw()*Scale()path gives 0 mismatches for both half and bf16.3. FP4: grid-gated row tiling — −0.14 ms/step (−1.4%)
Process
RowsPerBlockrows of Y at once so theuint4weight load, theprmtdecode and thescale load are amortized across rows. Gated on the column-block count, because at small N the
grid is already narrow and tiling loses:
FP4 GEMV family (graph OFF): 1.099 -> 1.066 (launch bounds) -> 0.950 ms/step (1.157x).
End-to-end (graph ON, 200 steps, 4 interleaved reps): 10.448 -> 10.306 ms/step.
Bit-identical (per-row fp32 accumulation order unchanged).
4. FP8 on tensor cores (
mma.m16n8k16) — −0.27 ms/step (−2.7%)FP8 GEMV family (graph OFF, 130 launches/step): 1.052 -> 0.713 ms/step (1.48x); all kernels
7.368 -> 7.021 ms/step. End-to-end 9.80 -> 9.54 ms/step. Now at 1.28 GB/step = 1.8 TB/s
= 37% of HBM peak.
Not bit-identical, but not worse: E4M3->FP16 is lossless and the
f16 x f16products are exactin fp32, so max error against an FP64 reference is identical for both kernels (2–4e-4).
5. FP4 on tensor cores (
mma.m16n8k16) — −0.55 ms/step (−5.8%)Largest single win of the campaign.
All kernels: 7.018 -> 6.526 ms/step. End-to-end (graph ON, 5 interleaved reps):
9.768 / 9.814 / 9.824 -> 9.558 / 9.584 / 9.466, i.e. 9.54 -> 8.99 ms/step.
Bit-identical (E2M1 decode order and fp32 accumulation order unchanged). The kernel now sits
at 56% of HBM peak, and two further optimizations measured 0%, so this path is considered closed.
Tests
onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cconnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc— includingGemvSpeculativeDecodeTilesFp16(min {2,3,4} xnin {1026,2050,4098,8194}, ragged tails,row- and column-varying data).
Methodology note
All 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 (FP4 lm_head: 2840 µs reported vs 741 µs real).