[kimseunghyun-kr]

Overview

Adds A100 (SM80) TP=8 skinny GEMM plans for the Qwen3.8-Flash-Next BF16 decode linears. SM80 had no plans, so every decode linear used cuBLASLt, which on these shapes runs split-K plus a splitKreduce kernel. Follows #64482 and #59214.

The skinny GEMM did not compile on SM80: its epilogue used cute.arch.elect_one(), which CuTe DSL accepts only on SM90+ (CONFIG_UNSUPPORTED_ARCH 'sm_80'). Below SM90 this PR writes the per-warp partial from lane 0 instead (same meaning: one lane per warp writes its partial sum); SM90+ keeps elect_one().

Claims

• 1.06–3.1x over cuBLAS per projection on A100 at M = 1–16 (38 plan entries, 10 shapes).
• End to end on 8 x A100 at concurrency 1, two pairs in opposite arm orders: TPOT P50 −8.3 % and −8.6 %, P95 −6.6 % and −4.2 %, output tok/s +4.9 % and +5.6 %. At concurrency 8: within noise (table in Details).

Validation

1x A100-SXM4-40GB, vLLM v0.32.0 CUDA 13 image + this branch's _skinny_gemm.py and low_latency_gemm.py
pre-commit run --files <3 changed files>                                   → passed
per-plan check (= test_qwen4_exp_sm80_selected_shapes, image has no pytest deps for main's test file):
  plan selected by qwen4_exp_low_latency_gemm, cosine vs F.linear > 0.999,
  max abs error vs FP32 reference <= F.linear's                            → 38/38 passed
sweep: 270 configs per (N, K, M), CUDA graph replay (50 calls x 20 replays), hot L2 and L2-flush (weights rotated over > 160 MB)

Cold-L2 results (the sweep). Plans are kept only where they are at least 3 % faster in both hot and cold L2:

projection	(N, K)	M	cuBLAS → plan (µs, cold L2)	speedup cold / hot
GDN fused QKVZ	(2048, 2560)	1	16.39 → 11.55	1.42x / 2.25x
GDN fused QKVZ	(2048, 2560)	2	14.96 → 12.54	1.19x / 1.45x
GDN fused QKVZ	(2048, 2560)	4	15.00 → 13.23	1.13x / 1.17x
GDN/QSA out	(2560, 768)	1	7.30 → 4.89	1.49x / 1.90x
GDN/QSA out	(2560, 768)	2	11.46 → 5.11	2.24x / 2.71x
GDN/QSA out	(2560, 768)	4	11.48 → 5.85	1.96x / 2.29x
GDN/QSA out	(2560, 768)	8	11.55 → 6.79	1.70x / 1.82x
GDN/QSA out	(2560, 768)	16	11.68 → 10.44	1.12x / 1.06x
GDN fused B/A	(12, 2560)	1	4.29 → 1.75	2.45x / 2.55x
GDN fused B/A	(12, 2560)	2	5.51 → 1.90	2.89x / 3.13x
GDN fused B/A	(12, 2560)	4	5.52 → 2.23	2.47x / 2.69x
GDN fused B/A	(12, 2560)	8	5.86 → 2.89	2.02x / 2.04x
GDN fused B/A	(12, 2560)	16	6.22 → 3.94	1.58x / 1.55x
QSA QKV/gate + indexer Q/K	(2688, 2560)	1	16.18 → 13.43	1.20x / 1.80x
QSA QKV/gate + indexer Q/K	(2688, 2560)	2	15.62 → 14.05	1.11x / 1.43x
QSA QKV/gate + indexer Q/K	(2688, 2560)	4	17.41 → 14.44	1.21x / 1.09x
QSA QKV/gate + indexer Q/K	(2688, 2560)	8	22.98 → 18.90	1.22x / 1.23x
shared gate/up	(160, 2560)	1	5.07 → 2.66	1.90x / 1.70x
shared gate/up	(160, 2560)	2	7.56 → 3.17	2.39x / 2.66x
shared gate/up	(160, 2560)	4	7.58 → 3.86	1.96x / 2.23x
shared gate/up	(160, 2560)	8	7.65 → 4.54	1.68x / 1.72x
router	(512, 2560)	1	6.36 → 4.19	1.52x / 1.61x
router	(512, 2560)	2	8.14 → 4.45	1.83x / 2.10x
router	(512, 2560)	4	8.17 → 5.27	1.55x / 1.69x
router	(512, 2560)	8	8.24 → 6.74	1.22x / 1.17x
shared-expert gate	(1, 2560)	1	3.58 → 1.66	2.16x / 2.26x
shared-expert gate	(1, 2560)	2	4.10 → 1.80	2.28x / 2.37x
shared-expert gate	(1, 2560)	4	4.14 → 2.12	1.95x / 2.01x
shared-expert gate	(1, 2560)	8	4.03 → 2.44	1.65x / 1.61x
shared-expert gate	(1, 2560)	16	4.19 → 3.39	1.23x / 1.22x
LM head	(31040, 2560)	1	125.01 → 110.12	1.14x / 1.13x
LM head	(31040, 2560)	2	121.75 → 111.15	1.10x / 1.10x
LM head	(31040, 2560)	4	122.25 → 115.15	1.06x / 1.09x
final HC down	(320, 10240)	1	10.85 → 7.76	1.40x / 1.68x
final HC down	(320, 10240)	2	12.29 → 9.62	1.28x / 1.54x
HC up	(10240, 320)	1	9.50 → 7.64	1.24x / 1.44x
HC up	(10240, 320)	2	9.95 → 8.02	1.24x / 1.48x
HC up	(10240, 320)	4	9.99 → 8.86	1.13x / 1.13x

Details

• Shapes and points that lose or win by less than 3 % stay on cuBLAS. Examples: QKVZ and QSA at M = 16, LM head at M ≥ 8, and the shared-expert down projection (2560, 80), which has no valid config. I also dropped final HC down at M = 4: it was 1.07x / 1.18x in the sweep but 0.91x in the validation run.

• The local TP=8 shapes come from the model's projection sizes (for example, QKVZ 16384 / 8 = 2048, QSA 768 q + 768 gate + 512 replicated KV head + 640 replicated indexer Q/K = 2688).

• elect_one() stays on SM90+: CuteSkinnyGemm takes use_elect_one, set from current_platform.has_device_capability(90), so SM90+ compiles the same epilogue as main and only SM80/SM86 use the lane-0 store. On H100, all 24 QWEN4_EXP_SM90_GEMM_PLANS entries give bitwise-equal output with main and time within −0.6 % to +0.4 %.

• End to end: 8 x A100-40GB, TP8 + EP, MTP 3, Qwen3.8-Flash-Next-FP8, v0.32.0 CUDA 13 image with these two files mounted, MLPerf agentic dataset, 30 min per arm, A then B in one job. The start-up log confirms that the plans are used (Warming up CuTeDSL compile_units=36 names=('shape-dynamic skinny GEMM',)). 0 failed samples in every arm.

  load (pair)	arm	TTFT P50 / P95 (ms)	TPOT P50 / P95 / P99 (ms)	output tok/s	samples
  C1 (1: A then B)	cuBLAS (main)	305 / 720	5.16 / 6.52 / 7.26	71.5	278
  C1 (1: A then B)	this PR	313 / 690	4.73 / 6.09 / 6.59	75.0	302
  C1 (2: B then A)	cuBLAS (main)	304 / 685	5.21 / 6.39 / 6.74	72.0	272
  C1 (2: B then A)	this PR	301 / 682	4.76 / 6.12 / 6.56	76.0	272
  C8	cuBLAS (main)	593 / 1,102	8.50 / 14.25 / 18.74	512	2,226
  C8	this PR	583 / 1,322	8.46 / 14.78 / 19.26	532	2,065

  At concurrency 1 (decode M = 4 with MTP 3), both pairs agree, in both arm orders: TPOT P50 −8.3 % / −8.6 %, P95 −6.6 % / −4.2 %, output tok/s +4.9 % / +5.6 %. At concurrency 8 the change is within run-to-run noise: decode steps there are dominated by MoE, all-reduce and many small kernels. Accuracy, paired on the turns both arms scored (the benchmark's per-turn score): C1 pair 1, 239 turns, diff −0.011 (95 % bootstrap CI [−0.060, +0.038]); C1 pair 2, 233 turns, diff −0.003 (CI [−0.051, +0.045]); C8, 1,863 turns, diff +0.005 (CI [−0.017, +0.026]). All intervals include zero; generation is sampled, so 26–36 % of turns differ between any two runs. The kernel outputs match F.linear (cosine >= 0.999997, max error not above cuBLAS's).

• CI: the new test runs only on SM80 with CuTe DSL; it skips elsewhere, like the SM90/SM100/SM121 tests in the same file.

---

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.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @kimseunghyun-kr.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[kimseunghyun-kr]

@gau-nernst, a gentle ping on your question about other architectures. Since b26ff89, SM90+ compiles the same elect_one() epilogue as main, and only SM80/SM86 take the new path. On H100, all 24 SM90 plan entries are bitwise equal to main, and their time difference is −0.6 % to +0.4 % (noise); details are in my reply in your thread. Could you take another look when you have time?

(Drafted with AI assistance; I reviewed it.)