Repository navigation
[Perf][GLM-5.3-Flash] Add SM90 skinny GEMM plans for the BF16 decode projections - #65306
Open
ischencheng wants to merge 5 commits into
Open
ischencheng wants to merge 5 commits into
ischencheng wants to merge 5 commits into
Conversation
Register a capability (9, 0) table for the TP4 and TP8 rank shapes of GLM-5.3-Flash's unquantized BF16 linears (KDA, MLA, indexer, LM head) at M = 1..4. Entries are kept only where the skinny GEMM beats cuBLAS by at least 5% under CUDA graph replay on both H100 and H200. Add a test that runs every SM90 entry against F.linear, eagerly and under graph replay, and cover (9, 0) in the selection test. Signed-off-by: Chen Cheng <85005591+ischencheng@users.noreply.github.com>
Remove the four H200 cells that measured under 1.08x against cuBLAS in the graph-replay A/B: (6416, 4096) M=3, (19360, 4096) M=2 and M=3, and (38720, 4096) M=4. Identical code on different buffers moves these LM head and KDA in_proj cells by up to 4% on H200, so they could fall under the 5% rule; those token counts go back to cuBLAS. Signed-off-by: Chen Cheng <85005591+ischencheng@users.noreply.github.com>
tests/kernels/test_bf16_skinny_gemm.py only runs in the B200 Miscellaneous Kernels step, where the SM90 plan tests skip. Add an H100 step that runs the SM90 plan, plan validity and selection tests when the plans, the skinny GEMM or its dispatch change. Signed-off-by: Chen Cheng <85005591+ischencheng@users.noreply.github.com>
…/glm53-flash-sm90-skinny-plans Signed-off-by: Chen Cheng <85005591+ischencheng@users.noreply.github.com>
The H100 step now also runs when linear.py or vocab_parallel_embedding.py change, since both call maybe_use_skinny_gemm. The table comment records that every cell was measured on both H100 and H200. Signed-off-by: Chen Cheng <85005591+ischencheng@users.noreply.github.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
ischencheng
requested review from
AndreasKaratzas,
Harry-Chen,
WoosukKwon,
khluu,
mgoin,
tlrmchlsmth,
yewentao256 and
zyongye
as code owners
November 3, 2026 05:22
26 of 45 tasks
This branch has not been deployed
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
On H100 / H200, GLM-5.3-Flash's unquantized BF16 decode projections all run cuBLAS, because
skinny_gemm_plans.pyonly has an SM121 table (#62956). This adds an SM90 table for the TP=4 and TP=8 local shapes (KDA in/out projections and f_b/g_b, the MLA projections vLLM dequantizes at load, the indexer wq_b / wk, the LM head), so they take the existing CuTe DSL skinny GEMM at 1-4 tokens. No kernel or dispatch change.Part of the GLM 5.3 perf tracker #57406 (same mechanism as #62956 on SM121) and the Q4 roadmap's GLM 5.3 model performance item #61024.
Claims
zai-org/GLM-5.3-Flashofficial FP8, 4x H200, TP=4: C1 TPOT -5.7% (6.09 -> 5.75 ms). C2 moves within run spread, C4 and C8 are unchanged.Validation
Component, H200 and H100, M=1. Speedup is the A/B described below (a chain with a small kernel between GEMMs, so it is lower than the GEMM-only ratio). The times come from the config sweeps, GEMM only, with that kernel subtracted. "% of HBM peak" is weight bytes / time / peak (4.8 TB/s H200, 3.35 TB/s H100).
The A/B covers every shape at M=1-4 on H200, and on H100 the TP8 shapes (GLM's FP8 checkpoint needs TP8 there). Geomean over those 92 cases is 1.184x. The lowest planned cell is 1.07x. Unplanned cells are 0.98x to 1.02x, which is the spread identical code shows on different buffers. A separate H100 run of every shape covers the TP4-only rows too, and its lowest planned cell is 1.06x.
How it was measured: each case is a chain over copies of the weight adding up to 1 GiB or more (so every GEMM reads HBM, not L2), captured in a CUDA graph, timed interleaved against cuBLAS over 11 trials of 300 replays. Outputs are compared elementwise to a float64 reference with |out - ref| <= 7.5e-5 + 2^-8 |ref| (cuBLAS needs up to 3.8e-5 beyond bf16 rounding on these shapes, the skinny configs under 1e-6), on normal, outlier and scaled/sparse inputs. M=8 and M=64 (no plan) match main.
Full model A/B:
zai-org/GLM-5.3-Flashofficial FP8, 4x H200, TP=4, one container, the two sides differ only in this table, server starts main / PR / main / PR.vllm serve --tensor-parallel-size 4 --language-model-only --disable-flashinfer-autotune --max-model-len 16384 --max-num-seqs 64 --no-enable-prefix-caching.vllm bench servewith 2048 random input tokens, 256 output tokens, greedy, two runs per start, median of the four.C1 is tight (main 6.07-6.10, PR 5.74-5.76). C2 is noisy (main 8.98-9.38, PR 8.65-9.42). 7168-token single-request TTFT is 279.4 -> 283.2 ms, with runs from 277 to 287 ms.
GSM8K: 5-shot completions, greedy, seed 42, in the same setup. Which cells run depends on M, so it ran at three concurrencies.
The concurrency 1 and 2 rows are smaller because of the GPU budget. At M=4 the TP4 table only has the MLA qkv_a and o_proj and the indexer wk, so those two rows are the ones that cover the M=1 and M=2 cells. M=3 cells run only incidentally at concurrency 4; the float64 test covers them at kernel level. Greedy text does change: at concurrency 1, with no batching, it matches on 101 of 266 questions and the final answer on 250. I did not run main vs main.
TP=8 has no full-model run. On a one-rank dummy model with TP=8 shapes (1x H100, no all-reduce, 512-token prompts, 192 greedy tokens, median of base / plans / plans / base with two runs per arm) the step is C1 5.277 -> 4.982 ms (-5.6%), C2 -5.4%, C3 and C4 -3.8%, C8 -0.2%. That dummy model pads its vocab to 19392, so its LM head misses the plan; a real TP=8 rank has 19360 rows. A dummy rank has no all-reduce, so expect less on the real model at M >= 2, as the C2 and C4 rows above suggest.
Tests, 1x H100, on this head:
The new test runs every SM90 cell, eagerly and under CUDA graph replay, against a float64 reference. It skips on other GPUs.
Details
cuBLAS picks 64x8 tiles for these shapes and launches about N/64 CTAs (64 for N=4096, 104 for the TP4 KDA in_proj on 132 SMs), so about 14% of warps are active and DRAM reaches 38-67% of peak (ncu, H200, M=1, LM head 80%). The MLA fused_qkv_a, MLA o_proj, indexer wk and TP8 in_proj also go through splitK plus a reduce kernel. The skinny GEMM launches N / outputs_per_block CTAs with no reduce: for the MLA o_proj, DRAM goes from 58% to 68% of peak and warps active from 14% to 85%. Summed over one TP4 decode step, these projections read 3.55 GB and take 1749 us with cuBLAS and 1460 us with the plans (ncu, serial, cold L2).
The table is keyed by capability (9, 0) and local (N, K). So H100 and H200 share it, H20 / H800 take it too (not measured), and any other model with the same shape gets these configs at M=1-4 and the JIT compile at startup, for example the o_proj of Llama-3.1-8B at TP=1 (4096, 4096) and of Qwen3-30B-A3B (2048, 4096). Each config comes from sweeps on both GPUs: the best worst-case speedup, kept only if it holds in the A/B above. TP4 f_b/g_b (2048, 128) loses on H200, and several M=3/4 cells barely win, so those stay on cuBLAS.
In eager mode a planned layer goes through #62956's custom op, which adds 4-7 us of host time per call (measured on single calls). Decode runs under CUDA graphs on Hopper. Prefill steps run eager, and the TTFT above stays within run spread. #64417's eager bypass would work for this op too, as a follow-up.
MTP: draft steps at C1 run M=1 and use the MLA and LM head plans, and verify at MTP5 runs M>=6, which has no plan. I did not measure MTP; all numbers here are with spec decoding off, the recipe default.
CI: the existing NVIDIA step for this test file runs on B200, where the SM90 cases skip, so this adds a small H100 step.
Not a duplicate: #64979 adds SM90 plans for dense Qwen3 / Llama-3.1 TP=1 shapes, and two of its keys, (2048, 4096) and (4096, 4096), are GLM shapes too. Whichever lands second merges the two tables. No open PR adds GEMM plans for GLM-5.3-Flash on SM90 (#57585 tunes its KDA recurrent kernel there, a different kernel).
AI assistance (Claude Code) was used. I reviewed the change and ran the tests above.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (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.