[CUDA] Block-scaled MatMul decode GEMV: tensor cores, packed FP4 decode and M-tiling - #31155
Open
tianleiwu wants to merge 6 commits into
Open
[CUDA] Block-scaled MatMul decode GEMV: tensor cores, packed FP4 decode and M-tiling#31155tianleiwu wants to merge 6 commits into
tianleiwu wants to merge 6 commits into
Conversation
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
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.
tianleiwu
force-pushed
the
tlwu/20260730/block_scaled_gemv_decode
branch
from
July 30, 2026 18:59
eb5b749 to
f545074
Compare
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
1 participant
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
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).
Note
Overlaps #29910, which also edits
matmul_block_scaled_fp{4,8}*— whichever lands secondneeds a rebase.