Pin FP8 GEMV residency for grids just past two blocks per SM

[Tianlei Wu (tianleiwu)]

Description

Stacked on #32409. That PR's commit is the first commit here; only the second commit (Pin FP8 GEMV residency ...) belongs to this PR. Please merge #32409 first — GitHub will then collapse this diff to the residency change alone.

The tensor-core FP8 GEMV launches ceil(N / 16) blocks. A 16-warp block only fits twice per SM, so an N just above 32 * sm_count spills into a second, nearly empty wave: on H200 N = 5120 launches 1.21 waves and ncu measures 66% active cycles.

This adds a second __global__ entry point over the same kernel body carrying __launch_bounds__(32 * KSplit, 3), and routes the shapes in that window to it. The hint caps registers at 40, which makes three blocks resident and collapses those shapes to a single wave. cuobjdump -res-usage confirms the new entry point at REG:40 SHARED:9216, so 40 * 512 * 3 = 61440 registers and 27 KiB of shared memory per SM — three blocks fit.

The window is narrow and the hint is a pessimization outside it, so Fp8MmaGemvPinsResidency in matmul_block_scaled_fp8_tiling.h gates on all four conditions that measurement showed matter:

• a grid at or below 2 blocks per SM is already one wave (N = 1024, 64 blocks, measured 0.96x);
• a grid above 3 blocks per SM stays multi-wave either way;
• 8-warp blocks (KSplit 8, taken from N >= 8192) lose 1.05-1.08x from any explicit __launch_bounds__ — declaring it replaces nvcc's implicit bounds even when the register cap is unchanged. 32-warp blocks (the KSplit 32 arm from Add FP8 GEMV KSplit32 scheduling for client SM12x GPUs #32409) cannot host 3 blocks per SM at all;
• only one row tile meets the 40-register cap. M = 16 (two tiles) measures 0.74x and M = 32 (four tiles) 0.24x, both from spills.

The plain kernel is unchanged and stays the default for everything outside the window, so no currently-selected shape changes code path.

Also adds ORT_FP8_GEMV_KSPLIT to override the K-split heuristic for A/B sweeps, matching the override already present in the FP4 GEMV.

Measurements (H200, sm_90, 132 SMs, CUDA 13.0)

Kernel time, FP16 activations, block_size 128:

N	K	before	after	speedup
5120	6144	13.25 us	11.04 us	1.20x
6144	5120	13.18 us	9.63 us	1.37x
5120	17408	52.00 us	36.51 us	1.42x

On a Qwen3 27B NVFP4 weights + FP8 KV decode workload the FP8 GEMV total drops from 12150 to 11712 us per forward. End-to-end decode throughput (median of repeats, both arms built as real binaries):

config	speedup
batch 1, speculative	1.030 - 1.036x
batch 1, non-speculative	~1.04x
batch 4, speculative (M = 32, gated out)	~1.00x

Acceptance rate is identical to four decimals in every cell.

Tests

onnxruntime_provider_test --gtest_filter='*MatMulBlockQuantizedFp8Weight*' — 14/14 pass.

Two new tests:

• GemvTensorCorePinnedResidencyBoundaries — pure predicate test at a fixed sm_count of 132, asserting the (264, 396] column-block window and the KSplit / MTiles exclusions, so the expectations do not move with the test machine.
• GemvTensorCorePinnedResidencyFp16 — runs the hinted kernel. Which N selects it depends on the device's SM count, so the shape is derived from cudaGetDeviceProperties (N = 16 * (2 * sm_count + 1), plus a ragged + 5) and cross-checked against PickFp8MmaKSplit. Skips on devices whose window lands above N = 8192, where the launcher drops to 8 warps and stops hinting.

Note for #32409

cuobjdump -res-usage shows MatMulBlockScaledFp8MmaGemvKernel<32, 4, ...> at SHARED:66560 (65 KiB). PickFp8MmaKSplit only returns 32 when m <= 8, so MTiles is always 1 there and that instantiation can never be launched — but it is still compiled and emitted. Calling launch_mma.template operator()<32, 1>() directly instead of going through launch_for_ksplit<32>() would drop it. Not changed here since it belongs to #32409.

[Copilot]

🔵 Needs a closer look

Hardware-specific CUDA residency tuning and its unmerged stacked dependency require final human validation.

Pull request overview

Pins FP8 tensor-core GEMV residency for grids between two and three blocks per SM. No actionable issues found.

Changes:

• Adds a launch-bounded kernel entry point and dispatch predicate.
• Adds a K-split override for benchmarking.
• Adds boundary and device-dependent correctness tests.

File summaries

File	Description
matmul_block_scaled_fp8.cu	Adds pinned-kernel dispatch and K-split override.
matmul_block_scaled_fp8_tiling.h	Defines tiling and residency-selection predicates.
matmul_block_scaled_fp8_test.cc	Tests selection boundaries and pinned FP16 execution.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 0
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.