Repository navigation
[Perf][Model] Qwen4Exp QSA: sm_90 tuning table for _select_config - #57273
Conversation
Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com>
mergify
Bot
added
performance
gau-nernst
left a comment
gau-nernst
left a comment
There was a problem hiding this comment.
Please remove the benchmark script. Also I think the unit test changes are not necessary
| elif _is_sm90() and (sm90 := _select_sm90_config(base_programs, is_fp8)): | ||
| BLOCK_N, target_splits, num_warps = sm90 |
There was a problem hiding this comment.
why not do BLOCK_N, target_splits, num_warps = _select_sm90_config(base_programs, is_fp8) like everyone else?
| Measured over the full attention + split-K merge + output gate at group | ||
| 3/6/12 with 1 and 2 kv heads, bf16 and fp8 K/V. Two effects drive the table. | ||
| Small batches are latency-bound and the kernel hands tiles to split programs | ||
| stride-wise, so the default (32, 64, 4) leaves one program with a two-tile | ||
| tail (65 tiles over 64 splits) that the whole launch waits for; a 64-wide | ||
| tile makes 33 tiles, the 64 requested splits clamp to 33 and every program | ||
| runs exactly one tile. From a few hundred programs up the kernel is | ||
| gather-bound and wants resident CTAs: BLOCK_N 64 needs 68-76 KB of shared | ||
| memory per CTA (num_stages=2, head_dim 256), 3 CTAs per SM, against 6 for | ||
| BLOCK_N 32 with one warp; extra warps only split an already tiny | ||
| [group, 256] x [256, BLOCK_N] dot. With fp8 caches the default entries for | ||
| bp 25..256 are already the fastest measured, and the bp > 2048 region was | ||
| not re-measured for either dtype. |
Signed-off-by: Shuolei Wang <948904026@qq.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
gau-nernst
left a comment
gau-nernst
left a comment
There was a problem hiding this comment.
LGTM. Wonder whether H20 tuned config translates to H200. But again, probably won't matter much
|
/ci run |
|
❌ This PR is 29 commits behind upstream |
|
/ci run |
|
❌ @ShuoleiWang, A reviewer with write access must run |
|
/ci run |
|
✅ Triggered Buildkite CI #89521 for commit |
|
✅ @ShuoleiWang, CI is now available for this PR.
|
|
/ci run |
|
❌ This PR is 1 commit behind upstream |
|
/ci run --allow-stale |
|
✅ Triggered Buildkite CI #89545 for commit
|
…lm-project#57273) Signed-off-by: Shuolei Wang <shuoleiwang123@gmail.com> Signed-off-by: Shuolei Wang <948904026@qq.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Purpose
Retune one kernel's launch parameters for sm_90 (H100/H200/H20):
_qsa_sparse_paged_gqa_splitk_kernel, the Triton split-K kernel that runs the sparse attention of Qwen3.8-Flash-Next's QSA layers over the 2048 selected tokens per query. Its tile width, split count and warps per CTA come from_select_config()invllm/models/qwen4_exp/nvidia/ops/qsa.py, a table tuned on GB300 (#55557 added an sm_120 table beside it). Measured on H20, the GB300 entries are the wrong geometry in most decode regions, so this PR adds_select_sm90_config()in the same shape as the sm_120 helper. Besides that nothing else changes. Regions where the default entry is already fastest returnNoneand keep it.Each row below is one region of
_select_config's key,base_programs = query rows × kv_heads per rank(at TP8 with MTP3 that is 4 × concurrency, so the second column shows which serving points land in the region). "GB300 entry" is what runs today, "sm_90 entry" what this PR selects on H100/H200/H20; a tuple is (BLOCK_N, target_splits, num_warps). "Latency" is the change in this kernel's time per call when switching from the GB300 entry to the sm_90 entry, measured on one H20 (attention + split-K merge + fused output gate, CUDA-graph replay, median of 3 interleaved orders × 9 replays, hot and L2-evicted), at the three head geometries group 3 / 6 / 12 = query heads per kv head at TP8 / TP4 / TP1-2. Negative is faster.bf16 K/V cache:
base_programsfp8 K/V cache (the path added by #55557):
base_programsWhy: at small
base_programsthe kernel hands tiles to split programs stride-wise, so (32, 64, 4) leaves one program with a two-tile tail (65 tiles over 64 splits) that the launch waits for, whereas a 64-wide tile gives 33 tiles over 33 splits, one per program. From a few hundred programs up the kernel is gather-bound and wants resident CTAs: BLOCK_N 64 needs 68–76 KB of shared memory per CTA (3 CTAs/SM) against 35–41 KB for BLOCK_N 32 with one warp (6 CTAs/SM). With fp8 the in-kernel dequant adds per-tile work, which is why the 64-wide entries lose almost half the time there.No region regresses beyond run-to-run noise on any geometry; every candidate passed the FP32 reference at
rtol=atol=2e-2. This kernel is 3–9 % of a decode step on this model, so end-to-end throughput moves by well under 1 % with bf16 caches; with fp8 caches the change is about 6 % of a C128 decode step. All measurements are H20; H100/H200 were not available.benchmarks/kernels/benchmark_qsa_select_config.py(added here) re-measures one shape in about 10 s, so a maintainer can check the direction there:Duplicate check.
gh pr list --repo vllm-project/vllm --state open --search "<q>"forqsa select_config,qsa sm90,qsa hopper,qsa H20,qwen4_exp sparse attention tuning: no open PR touches the sm_90 entries of this table. The open QSA PRs are #55430 (tile-union prefill kernel on SM121) and #55122 (deterministic persistent top-k), both orthogonal; #55557 (merged) added the sm_120 table this PR sits next to.Test Plan
pytest tests/models/qwen4_exp/test_qsa_reference.py -k "select_config_regions or sm90_" pytest tests/models/qwen4_exp/test_qsa_reference.py python benchmarks/kernels/benchmark_qsa_select_config.py --requests 128 --query-len 4 --group 3 --kv-heads 1 --candidates 64,4,2 32,8,1 32,6,1 python benchmarks/kernels/benchmark_qsa_select_config.py --requests 1 --query-len 4 --group 3 --kv-heads 1 --candidates 32,64,4 64,64,2 python benchmarks/kernels/benchmark_qsa_select_config.py --requests 128 --query-len 4 --group 3 --kv-heads 1 --kv-dtype fp8 --candidates 64,4,2 32,6,1 python benchmarks/kernels/benchmark_qsa_select_config.py --requests 4 --query-len 4 --group 3 --kv-heads 1 --kv-dtype fp8 --candidates 32,64,4 64,64,4In addition, on 8× H20 (TP8/EP8, MTP3, bf16 K/V,
--cudagraph-capture-sizes 1 4 16 64 256 512), stock vs this PR: server start and graph capture,vllm bench serverandom 8k-in/1k-out at C1 (8 prompts) and C128 (256 prompts), greedy generations withlogprobson 16 prompts, and GSM8K through lm-evaluation-harness (0-shot chat template, greedy, first 300 problems, identical prompts on both servers).Test Result
Unit tests (H20, main
fc8132a5tree):-k "select_config_regions or sm90_"Sweep script (H20;
currentis the entry this PR selects):Serving (8× H20, bf16 K/V; one launch per arm, so the tok/s differences are within the 1–3 % stock-vs-stock drift of this box and only show the absence of a regression):
Both servers captured every PIECEWISE and FULL graph size. The patched server's Triton cache holds only the (1 warp, N32) and (2 warps, N64) attention specializations; the stock one also has (4 warps, N32), i.e. the sm_90 table is the one running.
Model evaluation. The split count changes the split-K reduction order, so greedy outputs diverge from stock at near-ties: at every divergence the top-1/top-2 margin is 0–0.75 logprob (0–6 bf16 ULPs) and the other arm's token is inside the top-5; 12 of 16 chat prompts have bit-identical first-token top-5 logprobs.
AI assistance
Claude were used for the implementation, benchmark scripts and this description. The submitter reviewed every changed line and ran the commands above.