[Rockyeast]

Overview

#62956 routes unquantized BF16 linear layers and LM heads through the shape-dynamic CuTe DSL skinny GEMM when skinny_gemm_plans.py has a measured plan for the device, the local (N, K) and the token count. SM90 has no plans, so on H100 / H200 every dense decode GEMM still goes to cuBLAS, which reads the weight at 2.3-3.7 TB/s at 1-4 tokens.

This PR adds SM90_SKINNY_GEMM_PLANS for the TP=1 local shapes of Qwen3-4B / 8B / 14B / 32B and Llama-3.1-8B (QKV, o_proj, fused gate/up, down, LM head), registered under capability (9, 0). No code path changes: shapes or token counts without a plan keep using cuBLAS.

Claims

• E2E, H200, one request: median TPOT -14.9% on Qwen3-8B (6.02 -> 5.12 ms), -12.5% on Qwen3-4B, -11.0% on Llama-3.1-8B, -10.5% on Qwen3-14B, -12.7% on Qwen3-32B (20.24 -> 17.68 ms). Two concurrent requests: -8.3% to -12.2%. Four: -5.4% on Qwen3-8B, -3.9% on Qwen3-32B, -2.3% on Qwen3-14B, within run-to-run spread on Qwen3-4B and Llama-3.1-8B. Eight: unchanged (batch > 4 stays on cuBLAS).
• Qwen/Qwen3-30B-A3B (BF16 MoE, torch.compile), added after the first review round: three dense cells (fused QKV (5120, 2048), o_proj (2048, 4096), LM head (151936, 2048)): median TPOT -5.5% at one request (4.755 -> 4.493 ms), -2.6% at 2, about -3% at 4, unchanged at 8 (no plan).
• Accuracy: GSM8K 5-shot on Qwen3-8B (1319 questions, greedy, 2 concurrent requests so decode runs the planned kernels): 0.8886 without, 0.8893 with the plans.
• Kernels: 52 cells, each 1.07x-1.61x faster than cuBLAS with the weight read from HBM. A cell is listed only when it also wins with a hot L2 and, for 4 tokens, inside a running model (see Details).

Validation

Per-cell GEMM time

H200 NVL, vLLM nightly 0.33.0rc1.dev421 with #62956 applied, BF16, CUDA graph replay, the weight rotated over enough copies that every call reads it from HBM, and each GEMM preceded by a small non-PDL elementwise kernel (as in the model; its time is subtracted). Median of 7 paired repeats, cuBLAS -> plan in us:

projection	(N, K)	M=1	M=2	M=4
Qwen3-4B QKV	(6144, 2560)	12.0 → 8.8 (1.37x)	12.9 → 9.3 (1.39x)	cuBLAS
Qwen3-4B gate/up	(19456, 2560)	32.3 → 23.9 (1.35x)	32.4 → 24.3 (1.33x)	32.4 → 29.0 (1.12x)
Qwen3-4B LM head	(151936, 2560)	193.5 → 178.6 (1.08x)	194.1 → 180.4 (1.08x)	cuBLAS
Qwen3-4B o_proj	(2560, 4096)	11.1 → 6.9 (1.61x)	11.2 → 7.7 (1.45x)	11.0 → 8.7 (1.26x)
Qwen3-4B down	(2560, 9728)	19.6 → 13.2 (1.48x)	19.5 → 13.8 (1.41x)	cuBLAS
Qwen3-8B / Llama-3.1-8B o_proj	(4096, 4096)	14.9 → 9.6 (1.55x)	14.8 → 10.4 (1.42x)	14.7 → 11.4 (1.30x)
Qwen3-8B / Llama-3.1-8B QKV	(6144, 4096)	18.7 → 13.0 (1.43x)	17.5 → 13.8 (1.28x)	cuBLAS
Qwen3-8B gate/up	(24576, 4096)	57.8 → 48.4 (1.20x)	58.1 → 49.4 (1.18x)	58.4 → 52.6 (1.11x)
Qwen3-8B LM head	(151936, 4096)	309.8 → 281.5 (1.10x)	310.6 → 284.4 (1.09x)	cuBLAS
Qwen3-8B down	(4096, 12288)	32.8 → 24.4 (1.35x)	33.0 → 25.9 (1.27x)	33.3 → 27.8 (1.20x)
Llama-3.1-8B gate/up	(28672, 4096)	61.0 → 55.9 (1.09x)	59.7 → 57.3 (1.04x)	cuBLAS
Llama-3.1-8B LM head	(128256, 4096)	257.0 → 238.2 (1.08x)	257.0 → 241.0 (1.07x)	cuBLAS
Llama-3.1-8B down	(4096, 14336)	34.5 → 28.0 (1.23x)	34.5 → 29.3 (1.18x)	cuBLAS
Qwen3-14B o_proj	(5120, 5120)	18.0 → 14.0 (1.28x)	17.4 → 14.2 (1.22x)	cuBLAS
Qwen3-14B QKV	(7168, 5120)	25.4 → 17.6 (1.44x)	24.9 → 19.4 (1.29x)	cuBLAS
Qwen3-14B gate/up	(34816, 5120)	92.2 → 82.7 (1.11x)	93.3 → 85.2 (1.09x)	cuBLAS
Qwen3-14B / 32B LM head	(151936, 5120)	384.7 → 351.3 (1.10x)	385.3 → 354.2 (1.09x)	cuBLAS
Qwen3-14B down	(5120, 17408)	55.1 → 42.6 (1.29x)	54.8 → 42.9 (1.28x)	55.2 → 47.5 (1.16x)
Qwen3-32B QKV	(10240, 5120)	30.9 → 24.9 (1.24x)	30.8 → 26.0 (1.19x)	cuBLAS
Qwen3-32B gate/up	(51200, 5120)	140.0 → 121.4 (1.15x)	140.1 → 123.0 (1.14x)	cuBLAS
Qwen3-32B o_proj	(5120, 8192)	28.7 → 20.1 (1.43x)	29.1 → 21.1 (1.38x)	29.3 → 24.5 (1.20x)
Qwen3-32B down	(5120, 25600)	74.4 → 62.8 (1.18x)	74.0 → 63.1 (1.17x)	74.1 → 68.5 (1.08x)

Kept: >= 5% faster with a cold L2 and >= 3% faster with a hot L2 (single weight copy), plus the in-model check below for 4 tokens. cuBLAS cells did not pass. At 8 and 16 tokens the skinny GEMM loses on every shape (0.4-0.9x), so no plans are listed there.

E2E

H200 NVL, TP=1, vllm serve <model> --max-model-len 4096 --no-enable-prefix-caching, vllm bench serve --dataset-name random --random-input-len 1024 --random-output-len 256 --ignore-eos --seed 0; 32 prompts at concurrency 1, 8x the concurrency otherwise. Both arms run the same nightly image with #62956's files applied; the PR arm adds this plan file. VLLM_DISABLE_COMPILE_CACHE=1 in both arms (see Details). Two runs per arm, interleaved on two GPUs with the GPUs swapped. Median TPOT (ms), both runs:

model	conc.	without plans	with plans	change
Qwen3-8B	1	6.010 / 6.022	5.128 / 5.116	-14.9%
	2	6.105 / 6.050	5.350 / 5.325	-12.2%
	4	6.178 / 6.183	5.843 / 5.845	-5.4%
	8	6.402 / 6.395	6.396 / 6.396	0
Qwen3-4B	1	4.049 / 4.052	3.547 / 3.541	-12.5%
	2	4.119 / 4.125	3.674 / 3.666	-11.0%
	4	4.258 / 4.267	4.260 / 4.245	0
	8	4.450 / 4.458	4.462 / 4.449	0
Llama-3.1-8B-Instruct	1	5.421 / 5.345	4.794 / 4.787	-11.0%
	2	5.487 / 5.440	5.012 / 5.008	-8.3%
	4	5.554 / 5.511	5.448 / 5.494	-1.1%
	8	5.737 / 5.738	5.730 / 5.815	within spread
Qwen3-14B	1	9.346 / 9.364	8.382 / 8.370	-10.5%
	2	9.575 / 9.584	8.659 / 8.651	-9.7%
	4	9.603 / 9.629	9.401 / 9.390	-2.3%
	8	10.025 / 9.909	10.038 / 9.860	within spread
Qwen3-32B	1	20.227 / 20.255	17.668 / 17.693	-12.7%
	2	20.598 / 20.621	18.241 / 18.262	-11.4%
	4	20.634 / 20.697	19.874 / 19.841	-3.9%
	8	21.137 / 21.183	21.191 / 21.548	within spread (output tok/s 326.5 / 326.0 -> 325.9 / 325.9)

Output throughput moves with TPOT (Qwen3-8B at concurrency 1: 162.7 -> 190.2 tok/s). TTFT is unchanged within run-to-run spread.

Qwen/Qwen3-30B-A3B (BF16), same setup, four launches on one GPU in the order base, PR, PR, base; the routed experts run the Triton fused MoE in both arms, the plans only cover the dense projections and the LM head:

conc.	without	with this PR	change
1	4.755 / 4.755	4.493 / 4.493	-5.5%
2	5.514 / 5.512	5.370 / 5.370	-2.6%
4	6.941 / 6.669	6.728 / 6.565	-2% to -3%
8	8.670 / 8.353	8.356 / 8.631	0 (no plan)

Per-cell (H200, cold L2, cuBLAS -> plan, us): o_proj 9.42 -> 5.53 (1.70x) / 10.31 -> 5.74 (1.80x) / 10.19 -> 6.64 (1.54x) at 1 / 2 / 4 tokens; fused QKV 8.81 -> 6.69 (1.32x) / 8.49 -> 7.23 (1.17x), 4 tokens left to cuBLAS (1.03x); LM head 155.4 -> 142.2 (1.09x) / 156.4 -> 143.2 (1.09x), 4 tokens left to cuBLAS (1.01x).

Tests

pytest tests/kernels/test_bf16_skinny_gemm.py on H200 (nightly image with #62956 and this change applied): 369 passed, 438 skipped (other architectures).

• test_skinny_gemm_plans_are_valid (existing, parametrized over every capability) covers the new entries.
• New test_skinny_gemm_plans_run_and_match_linear, parametrized over every capability table (so it also covers the SM121 entries on a GB10): calls shape_dynamic_skinny_gemm directly with the plan's config (a silent F.linear fallback cannot make it pass), asserts the kernel was compiled and compares with F.linear; N is capped at 4096 so the LM-head cases do not allocate 1.5 GB weights.
• test_maybe_use_skinny_gemm is now parametrized over the capability tables with one planned shape each.

Details

• 4 tokens are checked inside the model. For 4 tokens the isolated benchmark does not always predict the model. Each 4-token candidate was checked with an nsys trace of 4 concurrent requests (GEMM time per projection per decode step, cuBLAS time including its split-K reduce). Kept: Qwen3-8B o_proj 511 -> 427 us, gate/up 2071 -> 1922 us, down 1158 -> 1041 us; Qwen3-4B o_proj 387 -> 320 us, gate/up 1137 -> 1103 us; Llama-3.1-8B o_proj 452 -> 383 us; Qwen3-14B down 2077 -> 1898 us; Qwen3-32B o_proj 1811 -> 1553 us, down 4867 -> 4393 us. Dropped although they win in isolation: Llama-3.1-8B down (4096, 14336) (1.09x isolated, 1075 -> 1349 us in the model, made concurrency 4 3% slower) and Qwen3-14B QKV (7168, 5120) (1.15x isolated, 970 -> 1025 us in the model).
• A/B testing plan changes needs a cold compile cache. Whether a layer calls the skinny op is decided while torch.compile traces the model, and the plan file is not part of the compile cache key. Two servers of the same vLLM build, one with and one without a plan, can therefore load each other's cached graph. All E2E numbers above use VLLM_DISABLE_COMPILE_CACHE=1; without it, a second run of the PR arm reproduced the baseline.
• Why the gain is largest at 1 token. cuBLAS picks small-M nvjet tiles (often split-K with a separate reduce kernel) that reach 2.3-3.7 TB/s on these weights; the skinny GEMM streams the weight once with vectorized loads, reaches 3.6-4.3 TB/s, and removes the reduce kernel.
• Startup cost. The planned configs are compiled by the CuTe DSL warmup before CUDA graph capture: 12 compile units for Qwen3-8B (9 for Llama-3.1-8B), about 5-10 s in the server logs, once per process (no disk cache for CuTe DSL kernels today, same as the existing SM121 plans).
• H100 / H800 / H20. The plans are keyed by capability (9, 0) and were measured on an H200 only (132 SMs, 4.8 TB/s). The kernel is bandwidth bound at these sizes, so the ranking should hold on H100, but the margins were not measured there and the table header says so. I can re-measure on an H100 if a reviewer has one; otherwise the existing per-capability keying is what the SM121 table also relies on.
• Scope. TP=1 dense models only; TP>1 local shapes and MoE models are not covered.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.