Repository navigation
[Perf] [NVIDIA] DSA: Pick any divisor of the tile count for the TileLang sparse decode split - #45274
JackChuang wants to merge 1 commit into
Conversation
JackChuang
changed the title
a35d350 to
3804c14
Compare
H20 resultsTested on 8×H20 (78 SMs/GPU) with GLM-5.3-Flash, FP8 KV, TileLang DSA, attention DP8 + MTP1. Control and candidate used the same #44262 dependency; the only source difference was this PR. Partial + combine component
Each case used 40 CUDA-event samples after graph warmup and L2 eviction. All outputs were finite, and CUDA Graph replay matched eager bit-exactly. Service A/BThe workload used deterministic LongBench-v2 prompts averaging 8,208.6 rendered input tokens, 1,024 requested output tokens, seed 8102, and no server queue. Throughput below is whole-burst client output throughput, including the 8K prefill.
All six runs completed the requested outputs with CUDA Graph replay and queue=0. There was no systematic change in MTP accept length. QualityMatched 5-shot GSM8K runs with the same legacy completion scorer showed no distinguishable regression: candidate 1203/1314 (91.5525%), retained control 1196/1314 (91.0198%). An earlier control run scored 91.6286%, so control run-to-run variation was larger than the retained A/B difference; no quality improvement is claimed. Overall, the affected H20 partial+combine shapes are about 16–20% faster. The service gain is smaller than on H200, but all three pairs at both concurrency levels move in the same direction. |
3804c14 to
c6675db
Compare
…ang sparse decode split _pick_inner_iter tried only ni, ni // 2, ni // 4, ... tiles per partial group. A top-k width that is not a power-of-two number of tiles (2051 with the KPool tail, padded to 66 tiles at block_I 32) therefore fell back to 2 tiles per group, i.e. 33 partial groups, for 4 to 65 query rows on 132 SMs. On CUDA (FP8 KV), also try the other divisors of the tile count within the same per-SM work limit, and take a larger one when its blocks (query rows x 16-head blocks x groups) keep at least 80% of two block slots per SM busy across the waves they need; with fewer blocks per SM the larger split ran slower than the halving pick. Power-of-two tile counts pick the same value as before, and ROCm callers keep the halving search. Add a CPU unit test for the split selection. Co-authored-by: Han Han <hanhan7630@outlook.com>
c6675db to
99da46e
Compare
Co-authored-by: @HanHan009527
Summary
TL;DR: This PR lets the CUDA TileLang FP8 sparse decode split each query's top-k tiles into groups of any size that divides the tile count, not only halvings of it, as long as the GPU stays busy. For GLM-5.3-Flash, whose 2051-wide top-k pads to 66 tiles, attention DP8 decode with MTP 1/1/2 at 128 concurrent requests now splits each query into 6 groups instead of 33. Fewer groups mean fewer partial results to write and merge. Splits are unchanged for power-of-two tile counts (top-k 2048, as in DeepSeek-V3.2 and GLM-5 to GLM-5.2) and on ROCm.
Motivation
Earlier DSA models select 2048 KV tokens per query; GLM-5.3-Flash selects 2051, so its tile count is no longer a power of two, and for decode-sized calls the split search, which only tries halvings, falls back to many small groups. Its indexer appends up to 3 not-yet-pooled tail tokens (KPool) to the 2048 selected ones, and
_forward_tilelangpads the 2051 columns to the next multiple of 64, 2112: 66 tiles of 32 on CUDA instead of 64.#42957 added the FP8 KV cache to the TileLang DSA backend on CUDA through the partial + combine path of
tilelang_sparse_fwd, which serves both prefill and decode. That path splits thenitop-k KV tiles of each query intoni / inner_itergroups, computes one partial result per group, and merges them with a combine kernel._pick_inner_iterpicksinner_iter, the number of tiles per group, and tries onlyni, ni // 2, ni // 4, .... For 64 tiles these are all the divisors, but halving 66 gives 33, 16, 8, 4 and 2, of which only 33 and 2 divide 66, so the search never tries 3, 6, 11 or 22.The search also keeps at least one block per SM:
seq * ni / inner_iter >= cu * block_per_cu, whereseqis the number of query rows in the call andblock_per_cuis 1 on CUDA. On a 132-SM GPU this allows at mostseq / 2tiles per group. So for 4 to 65 query rows the search falls back to 2 tiles per group: 33 partial results per query, each written out and read back by the combine kernel, although 11 tiles per group (6 groups) fits the same limit at 32 rows. With MTP 1/1/2 (multi-token prediction with--speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2, so one draft token), each decoding request contributes 2 query rows, so 128 concurrent requests on attention DP8 (16 per GPU) give 32 rows per GPU. This PR lets the search also take those other divisors (Modifications).Modifications
python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py:_pick_inner_itergetsany_divisorandhead_blocksoptions. Withany_divisor, it also tries the divisors ofnilarger than the existing pick, within the same work limit, and takes the largest one that passes the fill check below. Only the CUDA FP8 call intilelang_sparse_fwdturns it on. The ROCm calls (FP8 and BF16 KV intilelang_sparse_fwd, and DeepSeek-V4'sdpsk_v4_fp8_attention_fwd) keep the default search.test/registered/unit/layers/attention/test_tilelang_dsa_decode_split.py(new,base-a-test-cpu): checks the picks for 66 tiles on 132 SMs, that power-of-two tile counts pick as before, and that the default search is unchanged. Removing any condition of the new search fails a case.Design notes:
seq * head_blocks * ni / inner_iter: one per 16 heads, query row and group), and each SM runs two of these blocks at once. The existing work limit only guarantees one block per SM, so with few heads per GPU the largest divisor within it can leave slots idle (8 heads at 44 rows: 22 tiles per group gives 132 blocks, one per SM). The check takes a divisor only if its blocks fill at least 80% of the two slots per SM in each wave they need. With it, this PR is at least as fast as main in every case of the kernel test below, and so would be any threshold from 75% to 100%.ni % inner_iter == 0, and main already runs splits that are not powers of two (33 and 66 tiles per group).seq < cu * block_per_cu.Accuracy Tests
Unchanged: GSM8K and the MTP accept length match main, and the kernel error at this PR's splits stays below that of main's splits for 66 or more rows. Each group rounds its softmax probabilities to FP8 against its own running maximum, so a larger group rounds more coarsely; main already uses groups of 33 and 66 tiles for 66 or more rows.
All outputs are finite, and CUDA graph replay is bit-identical to eager.
Accuracy command
Speed Tests and Profiling
Attention DP8, end to end. This measures the layout of the TL;DR. Result: steady decode throughput is 2.5% higher at 128 concurrent requests and 1.2% higher at 64, and the MTP accept length is unchanged (1.90–1.91 tokens per step, including the bonus token). Setup: GLM-5.3-Flash on 8×H200; this PR on main
c5a5bf0(the PR has since been rebased onto71ee2dcwith the same diff), with #44262 (needed for attention DP with MTP on this model) applied to both versions; torch 2.14.1, FlashInfer 0.7.1, CUDA 13.0. Configuration: the command below (TP8/EP8, attention DP8 with 64 heads per GPU, FP8 KV with TileLang DSA prefill and decode, MegaMoE with DeepGEMM, MTP 1/1/2, radix cache off). Workload: 64 or 128 prompts of about 8.2k tokens (LongBench-v2 documents) with 1,024 output tokens each, sent at once, so every GPU decodes 8 or 16 requests with no new arrivals. Steady decode throughput is read from each GPU's decode log lines while all of its requests decode, and summed over the 8 GPUs; ITL comes from the benchmark client. Server restarts alternate between the two versions; the tables show medians.Steady decode, whole model (+% = this PR better)
Over each whole burst, which also includes the 8.2k-token prefills that this PR does not change, the client's output throughput is 1.7% higher at 128 and 0.5% higher at 64.
Where the time goes. This profile of 8 decode steps on every GPU at 128 concurrent requests (32 query rows per call) checks that the gain comes from the two sparse decode attention kernels. Result: they take 27% less GPU time per call (1.36× faster), which accounts for the shorter decode step.
Each decode step makes 12 such calls. On main they take 2.42 ms, 9.6% of the 25.17 ms step; the 0.65 ms they save (2.6% of the step) matches the 0.62 ms shorter step above.
Attention DP8 with MTP 5/1/6, and with MTP 1/1/2 at 256 concurrent requests, end to end. These runs put more query rows into each call while staying in the range where the split changes (up to 65 rows per GPU with 64 heads): MTP 5/1/6 gives 6 query rows per request in each verify call, and 256 concurrent requests with MTP 1/1/2 give 64 rows per GPU. Result: with MTP 5/1/6 the decode step is 3.3% shorter at 64 concurrent requests and 2.0% shorter at 32; with MTP 1/1/2 at 256 it is 4.0% shorter, and steady decode throughput is 4.3% higher (14,733 → 15,370 tok/s). Setup: as above, with
--enable-dp-attention-local-control-broadcastadded for both versions (it removes an all-rank sync on every scheduler iteration), and for MTP 5/1/6 also--enable-linear-replayssm-spec(commands below); 32 to 256 prompts sent at once. Each decode step makes 12 sparse decode attention calls with 6 query rows per request (MTP 5/1/6) or 2 (MTP 1/1/2), and with MTP 5/1/6 also 4 draft calls with 1 row per request (96 and 32 such calls in 8 profiled steps). The MTP accept length is the same within the spread between restarts (3.79–3.99 tokens per step with MTP 5/1/6, 1.90 with MTP 1/1/2), so the table compares decode step times: each step runs the same calls whatever the accept length.Each change holds in every server-restart pair. At 128 concurrent requests with MTP 5/1/6, only the 4 draft calls (16 rows) change split (2 → 6 tiles per group), so the gain is small.
Decode step times (ms)
At 64 concurrent requests with MTP 5/1/6, a profile of 8 decode steps shows where the 0.96 ms shorter step comes from: the 12 calls with 48 query rows per GPU (grid [192, 33] on main and [192, 6] with this PR, i.e. 33 and 6 groups per query) take 28% less GPU time per call in this profile (294.4 → 211.8 µs for partial + combine, 1.39× faster), and the 4 draft calls with 8 rows (33 → 22 groups) take 9% less (54.6 → 49.5 µs); together they save about 1.0 ms per step. Over each whole burst, which also includes the prefills, the client's output throughput is 0.5–0.9% higher with MTP 5/1/6 and 1.6% higher with MTP 1/1/2 at 256.
TP8 low-latency command, end to end. This checks the layout with the fewest heads per GPU (8), where the fill check matters most (Design notes). Result: the decode step is 0.7–1.2% shorter at 2 to 8 concurrent requests, where the split changes, and the same at 1 and 16, where it does not. Setup: as above, with the TP8 low-latency command below (TP8/EP8 without attention DP, MTP 5/1/6, so 6 query rows per request in each verify call); 1 to 16 prompts of about 8.2k tokens with 1,024 output tokens each, sent at once. The decode step is read from the decode log lines while all requests decode.
At 4 concurrent requests the profile shows the split change (grid [24, 33] on main, [24, 11] with this PR) and 9.9 µs less per partial + combine call (41.6 → 31.7 µs). Throughput is not shown for this setup: it also moves with the accept length, which varies by up to 1.8% between restarts of the same version, more than the step-time gain.
Per kernel call (+% = less time with this PR). This test times the CUDA FP8 partial + combine kernels for every head count and row count below, to check that every case is at least as fast as main. Result: wherever the split changes, this PR takes 10.5–26.5% less time; elsewhere it is the same split. Setup: one H200; GLM-5.3-Flash's shape (2051 valid columns padded to 2112, value width 512); each case is a CUDA graph of
tilelang_sparse_fwdtimed after flushing the L2 cache, taking the median over repeated replays.In 89 further cases (8 to 64 heads, 4 to 200 rows) it takes 7.8–26.9% less time where the split changes, and never more.
H20, per kernel call and end to end. This checks a GPU with fewer SMs (78 per GPU instead of 132), since the split depends on the SM count. Result: per call, the partial + combine kernels take 16–20% less GPU time where the split changes (64 heads at 16 and 32 rows, the per-GPU shapes at 64 and 128 concurrent requests below; 8 heads at 12 and 24 rows) and the same where it does not (8 heads at 48 rows); end to end, the mean time per output token (TPOT) is 1.1% lower at 64 concurrent requests and 1.3% lower at 128, lower in every server-restart pair at both, and the MTP accept length is unchanged. Setup: GLM-5.3-Flash on 8×H20, FP8 KV with TileLang DSA, attention DP8 with MTP 1/1/2, #44262 applied to both versions; the same LongBench-v2 prompts as above (about 8.2k tokens, 1,024 output tokens, seed 8102). The tables and commands are in @HanHan009527's comment.
Startup. This table measures the one-time cost of compiling the new kernel pairs. Result: about 35 s once per empty TileLang kernel cache, nothing afterwards.
The compile happens while the prefill CUDA graphs are captured; without them it happens at the first batch that needs each new split.
Commands
Attention DP8 server (the same command for main and this PR):
Attention DP8 server for the MTP 5/1/6 and 256-request runs: the command above with
--enable-dp-attention-local-control-broadcastadded; for MTP 5/1/6,--speculative-num-steps 5 --speculative-eagle-topk 1 --speculative-num-draft-tokens 6 --max-running-requests 128replace the MTP and request-limit flags, and--enable-linear-replayssm-specis added.TP8 low-latency server: the H200 Low Latency command from the GLM-5.3-Flash cookbook, with FP8 KV instead of BF16, and radix cache off:
Load, after two short warm-up bursts (128 prompts shown; each row uses its request count for both
--num-promptsand--max-concurrency):longbench_v2_8k.jsonholds one ShareGPT conversation per LongBench-v2 document of at least 8,000 tokens (GLM-5.3-Flash tokenizer): the document's first 8,000 tokens, its question and choices, and "Explain your reasoning step by step, quoting the relevant parts of the document, then give the answer."Profile: one more burst (128 prompts with MTP 1/1/2, 64 with MTP 5/1/6), then
POST /start_profilewith{"num_steps": 8, "activities": ["GPU"]}once every GPU decodes all of its requests.CI
No test file fails only with this PR in the CPU and single-GPU suites run here against its parent, main
c5a5bf0. The new unit testtest_tilelang_dsa_decode_split.py(base-a-test-cpu) andtest_tilelang_dsa_zero_rope.py, which runs the changed split on GPU, pass on the rebased PR. Multi-GPU suites, Blackwell and ROCm were not run here.Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #39700313502
Latest PR Test (Extra): ❌ Run #39700313399
Latest PR Test (AMD ROCm 10): ❌ Run #39700313519