[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 the 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;
• the residency hint requires SM89 or newer, where FP8 tensor-core support is native. SM86 measurements below show a repeatable ~1% regression from the register cap;
• 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 the 48-SM SM121 GPU #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.

Measurements (GeForce RTX 5060 Ti, sm_120, 36 SMs, CUDA 13.0)

Kernel medians from Nsight Systems CUDA activity traces, FP16 activations, M = 8, block_size 128. Each trace contains 100 measured launches after warmup; before and after select the plain and residency-pinned entry points from the same binary.

N	K	grid	before	after	speedup
1168	5120	73	9.088 us	7.936 us	1.15x
1440	5120	90	10.928 us	8.480 us	1.29x
1440	6144	90	12.896 us	10.016 us	1.29x
1440	17408	90	34.704 us	26.816 us	1.29x
1728	5120	108	11.872 us	9.264 us	1.28x

The predicate boundaries behave as intended: N = 1152 (exactly 2 blocks/SM), N = 1744 (just above 3 blocks/SM), and M = 16 all stay on the plain kernel and measure within 1% between A/B arms. A ragged in-window N = 1169 produces bitwise-identical output (max_abs_diff = 0).

Measurements (GeForce RTX 3060, sm_86, 28 SMs, CUDA 13.0)

The same methodology shows that applying the hint on SM86 is a small but repeatable pessimization, which motivates the SM89 minimum:

N	K	grid	before	after	speedup
912	5120	57	71.648 us	72.608 us	0.987x
1024	5120	64	72.096 us	72.959 us	0.988x
1120	5120	70	72.576 us	73.183 us	0.992x
1344	5120	84	73.760 us	74.352 us	0.992x
1024	17408	64	236.814 us	239.135 us	0.990x

Three independent N = 1024, K = 5120 traces with alternating arm order reproduce a 0.988-0.991x speedup. cuobjdump -res-usage shows that SM86 moves from REG:56 STACK:0 for the plain FP16 kernel to REG:40 STACK:8 for the pinned kernel. Boundary and M = 16 controls remain neutral, and a ragged in-window correctness case is bitwise identical. The final dispatch therefore leaves pre-SM89 devices on the plain kernel.

Measurements (GeForce RTX 4090, sm_89, 128 SMs, CUDA 13.3)

Kernel time measured with Nsight Systems CUDA traces, FP16 activations, block_size 128. The baseline and pinned arms are separate provider binaries built from the same merged PR head; the baseline changes only the residency dispatch to select the plain entry point.

M	N	K	Grid	Plain	Pinned	Speedup
8	5120	6144	320	18.43 us	14.56 us	1.27x
8	6144	5120	384	16.45 us	13.76 us	1.20x
8	5120	17408	320	101.83 us	98.11 us	1.04x

Controls outside the pinned dispatch remain neutral within run-to-run variance:

Control	Plain	PR	Ratio
M = 8, N = 1024, K = 5120 (grid below 2 blocks/SM)	5.76 us	5.92 us	0.97x
M = 16, N = 5120, K = 6144 (MTiles = 2)	31.36 us	30.78 us	1.02x

Using the occurrence counts in the Qwen3.8 workload microbenchmark, the three affected shapes save about 407 us per forward. A reverse-order repeat reproduced the main gains (1.24x, 1.22x, and 1.04x respectively).

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, SM86 exclusion, SM89 inclusion, and the KSplit / MTiles exclusions, so the expectations do not move with the test machine.
• GemvTensorCorePinnedResidency ??? 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 pre-SM89 devices and 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.

[Baiju Meswani (baijumeswani)]

The residency-pinned KSplit16 kernel is measurably positive on a 48-SM SM121 GB10 when a matrix falls inside the PR's dispatch window. The strongest and most stable batch-1 decode results show approximately 8-10% lower FP8 GEMV operator latency at M=1 for Qwen-like K dimensions.