Repository navigation
[Perf][Qwen4Exp] Add SM121 TP=1 skinny-GEMM plans - #59753
Merged
benchislett merged 4 commits intoOct 2, 2026
Merged
Conversation
On DGX Spark (SM121) at TP=1 the Qwen4Exp LM head (248320 x 2560, BF16) has no low-latency plan and falls back to cuBLAS, which streams the weight at about 175 GB/s for a single row. With MTP it runs once per verify step and once per draft step, so it is a large share of decode. Add an SM121 plan table for the TP=1 LM head at M=1, 2, 4, 8 and 16, selected by a FlashInfer CUPTI sweep on GB10, and route capability (12, 1) to it. M > 16 keeps cuBLAS (outside the skinny GEMM range). Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
stecasta
requested review from
AndreasKaratzas,
WoosukKwon,
gau-nernst,
mgoin,
tlrmchlsmth,
yewentao256 and
zyongye
as code owners
October 2, 2026 08:05
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
5 tasks done
6 tasks done
Add GB10 skinny-GEMM plans for the TP=1 router, shared-expert gate and projections, HC projections, GDN projections, QSA projections and the GDN/QSA output projection. Configs are the fastest valid ones from a CUDA-graph sweep against cuBLAS; cells under 1.03x are left on cuBLAS. Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>
stecasta
changed the title
|
✅ @stecasta, CI is now available for this PR.
|
benchislett
enabled auto-merge (squash)
October 2, 2026 15:47
Member
|
/ci run |
|
✅ Triggered Buildkite CI #92659 for commit |
mgoin
approved these changes
Oct 2, 2026
Collaborator
CI selector (shadow): 2 test steps (2 jobs) instead of 36 (49 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (2)
Would skip (today's rules run them) (34)
Would add (today's rules do not run them) (0)none AMD mirrors: would skip (32)
AMD mirrors: would add (0)none 2 changed files · base |
6 tasks done
This was referenced Oct 4, 2026
aijanai
pushed a commit
to aijanai/vllm
that referenced
this pull request
Oct 5, 2026
_gemm_plans() returns {} on anything that is not sm_103, sm_100 or sm_90, so the
CuTe skinny GEMM never runs on sm_120 parts. The shipped plans are also keyed by
the TP-local (N, K) and were measured at TP=4, so even with the gate opened the
TP=1 and TP=2 shapes would miss the table.
Add benchmarks/kernels/qwen38fn_skinny_gemm/, a tuner adapted from
benchmark_k3_cutedsl_residual.py: parameterised over (N, K), no SM gate, driving
the installed shape_dynamic_skinny_gemm, scoring every candidate against
F.linear, and sharding the search across nodes with a merge step that selects the
global winners. Add the resulting sm_120 plan table and route _gemm_plans() to it.
Measured on an RTX 5090 over 31046 points: 58 shape/M points across TP=1 and
TP=2, every one with the weights rotated past the 96 MiB L2.
The gains sit where MTP puts the verification GEMM at batch size 1, since it runs
at M = (k + 1) * batch_size: M=2 at k=1 and M=4 at k=3. They come from the
baseline rather than from this kernel -- at M=1 cuBLAS dispatches a dedicated
GEMV, and from M=2 it pads a tiled GEMM whose minimum tile height the work never
fills. The shared-expert projection shows it plainly: F.linear jumps from 0.0068
to 0.0145 ms between M=1 and M=2 and stays flat to M=16, while the skinny kernel
climbs smoothly from 0.0042 to 0.0076.
The LM head is the exception at 1.06x, flat across M: it is bandwidth-bound, so
cuBLAS wastes nothing and there is nothing to win.
Complementary to vllm-project#59632, which adds plans for sm_121 only despite its SM12x
title, and to vllm-project#59753, which adds the sm_121 LM head.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Alberto Massidda <amassidda@nvidia.com>
6 tasks done
michael-gpu
pushed a commit
to practical-inference/vllm-ampere-extended
that referenced
this pull request
Oct 20, 2026
Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com>
This was referenced Oct 21, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Qwen3.8-Flash-Next decode ITL -9.5% at concurrency 1 on a single DGX Spark (SM121, TP=1) by adding SM121 skinny-GEMM plans for the TP=1 BF16 shapes, which today fall back to cuBLAS. The LM head
(248320, 2560)gives most of it (-8.3% alone); the remaining rows add -1.3% onlocal-inference-lab/Qwen3.8-Flash-Next-NVFP4, where the router, HC projections and shared-expert gate are BF16 and the other dense projections are MXFP8. Onnvidia/Qwen3.8-Flash-Next-NVFP4attention, GDN and shared experts are BF16 as well and use the same rows; that checkpoint does not fit a single Spark, so it is not measured end to end.Complements #59632 (SM121 plans for TP=2 shapes, from #59605): this PR only adds TP=1 keys to the same mechanism. Happy to fold this into #59632 or rebase on it, whichever lands first.
Claims
local-inference-lab/Qwen3.8-Flash-Next-NVFP4). Almost all of the end-to-end gain comes from the three M=1 calls (about 2.2 ms saved each).Validation
Kernel, GB10 (p4242-0063, about 2.5 GHz), FlashInfer CUPTI, CUDA graph, cold L2, median of 30; cosine 1.000000 vs
F.linearfor every config:Other TP=1 shapes, same sweep (selected config vs cuBLAS; cells under 1.03x stay on cuBLAS):
These are cold-L2 numbers. Weights of 7 MB or less stay in L2 in the server, where an in-server torch profile showed little gain for the HC, router and shared-gate rows; the end-to-end numbers below are the claim.
A bake-off on the LM-head shape (FlashInfer
mm_bf16cudnn / cublaslt / tinygemm, b12xbf16_vocab_projectiondefault config) found nothing faster than the skinny GEMM at M=2 to 16; at M=1 b12x was 4% faster than the skinny config but would need a new dependency.End to end, same node (p4242-0074),
local-inference-lab/Qwen3.8-Flash-Next-NVFP4, TP=1, MTP k=3 with forced acceptance length 3, FP8 KV,--max-num-batched-tokens 4096, SPEED-Bench 2K in / 1K out, two interleaved runs per arm, mean:Accuracy: GSM8K, first 256 test questions, greedy, thinking off, speculative decoding off, 8 concurrent requests, two runs each: base 249 and 247, LM-head rows only 245 and 247, full table 246 and 246 (of 256). Batch composition varies between runs, so identical configs differ by a few questions.
Disclaimers:
af7f9488c(the image used for Flash-Next serving) with this PR's table overlaid onlow_latency_gemm.py.(13312, 2560)is the unfused QSA QKV/gate shape; after [Perf][Qwen4Exp] Merge QSA QKVG and indexer QK projections #59533 lands, the fused(13952, 2560)key needs a row (sweep data available).Tests
pytest tests/kernels/test_bf16_skinny_gemm.py -k "qwen4_exp" -vOn GB10 (p4242-0070): the Qwen4Exp capability-routing test and the SM121 selected-shape test (cosine > 0.999 vs
F.linearfor every shape and M in the table), run against this branch'slow_latency_gemm.py: 64 passed. The rest of the file needs GLM-5.2 modules that are not in the Spark serving image.Details
(12, 1), keyed by the local weight shape, configs selected by a CUDA graph sweep against cuBLAS.AI assistance (Claude) was used for the sweep scripts, measurements and this description; the change and the numbers were reviewed by the submitter. Duplicate check: #59632 is the only open SM121 plan PR for Qwen4Exp and covers TP=2 shapes only.