[Kernel] Fix SM90 FP8 decode regression with benchmarked M/K/N routing

[RunFMe]

Motivation

The SM90 row/column-scaled FP8 selector introduced by #34318 can route decode GEMMs to torch._scaled_mm based on K/N alone. The reported H100 W8A8 regression is:

• model: Llama-3.1-8B-Instruct-FP8-dynamic
• test: test/registered/quant/test_w8a8_quantization.py::TestW8A8Fp8.test_throughput
• batch size: 1
• throughput: approximately 215 tok/s before [Kernel] Route large SM90 row/column-scaled FP8 GEMMs to Torch #34318, approximately 190 tok/s after it
• affected down projection: M=1, K=14336, N=4096

Direct profiling of that shape on both H200 and H100 showed that the AOT kernel is substantially faster at decode M, even though K and N are large.

Modifications

Constrain Torch/NVJet routing to the benchmarked large-prefill envelope (doing m >= 6144 is not really worth the gain from my measurements):

m >= 8192 and (
    (k >= 4096 and n >= 6144) or
    (k >= 7168 and n >= 5376)
)

This keeps the following on AOT:

• all decode and small-prefill calls with M < 8192
• the reported K=14336, N=4096 Llama down projection
• narrow TP8 projections K=5376, N=3584 and K=3584, N=5376
• shapes outside the measured projection envelope

Torch remains selected for stable large-prefill wins: wide QKV/gate/up projections and broad down projections.

Speed Tests and Profiling

Common software:

• PyTorch 2.11.0+cu130 (70d99e998b4955e0049d13a98d77ae1b14db1f45)
• CUDA 13.0
• sgl_kernel==0.4.4
• CUDA-event timing with randomized/interleaved AOT and Torch samples

The broad grid used 11 projection families and 12 M values (132 cases): 1, 16, 64, 256, 1024, 2048, 4096, 6144, 8192, 8320, 12288, and 16384. Percentages below are Torch speedups over AOT.

H200 (SM90)

Protocol:

• broad grid: 5 warmups and 15 samples per arm
• two randomized boundary confirmations: 44 cases and 33 cases, with 8–10 warmups and 30–40 samples per arm
• exact-workload conditioning before timing

Reported regression shape:

M	K	N	AOT	Torch	Result
1	14336	4096	26.50 µs	51.12 µs	AOT is 1.93× faster

Routing boundary:

K	N	Projection family	M=8192	M=8320
4096	28672	gate/up	1.74%	3.30%
4096	6144	QKV	0.94%	2.14%
5376	21504	QKV	2.59%	4.64%
5376	28672	gate/up	2.39%	5.17%
5376	14336	gate/up TP2	2.91%	2.62%
7168	5376	down TP4	1.26%	3.29%
14336	5376	down TP2	1.82%	3.25%
28672	5376	down TP1	1.66%	3.09%

Across the H200 broad grid, every one of the 32 measurements selected by the new predicate favored Torch. Their median speedup was 2.50%, with a measured range of 0.17–5.97%. Two independent randomized boundary passes added 48 repeated routed measurements, all of which also favored Torch.

H100 80 GB HBM3 (SM90)

Protocol:

• randomized broad grid: 5 warmups and 20 samples per arm after 0.5 s of exact-workload conditioning per case
• randomized boundary confirmation: 44 cases, with 10 warmups and 40 samples per arm after 2 s of conditioning per case

Reported regression shape:

M	K	N	AOT	Torch	Result
1	14336	4096	32.19 µs	53.55 µs	AOT is 1.66× faster

Routing boundary:

K	N	Projection family	M=8192	M=8320
4096	28672	gate/up	21.02%	43.21%
4096	6144	QKV	5.39%	3.07%
5376	21504	QKV	25.44%	57.38%
5376	28672	gate/up	26.58%	52.16%
5376	14336	gate/up TP2	15.87%	48.31%
7168	5376	down TP4	2.72%	4.43%
14336	5376	down TP2	8.72%	7.08%
28672	5376	down TP1	10.02%	8.57%

Across the H100 broad grid, every one of the 32 measurements selected by the new predicate favored Torch. Their median speedup was 19.48%, with a measured range of 1.36–56.53%. The stronger randomized boundary pass added 24 repeated routed measurements, all favoring Torch, with a 21.02% median and a 2.72–57.38% range.

On both GPUs, the excluded Llama down and narrow TP8 families had weak large-M results that were near parity or changed sign across passes, so this change deliberately leaves them on AOT rather than routing on noise.

---

CI States

Latest PR Test (Base): ✅ Run #33259017128
Latest PR Test (Extra): ❌ Run #33259145867
Latest PR Test (AMD ROCm 7.2): ❌ Run #33259017192

[BBuf]

It seems resonable, approved!

[BBuf]

/tag-run-ci-label extra

[BBuf]

/rerun-test test/registered/quant/test_w8a8_quantization.py

[github-actions]

Results for /rerun-test test/registered/quant/test_w8a8_quantization.py:

🚀 1-gpu-h100 (1 test): ✅ View workflow run

cd test/ && python3 registered/quant/test_w8a8_quantization.py