[QSA] Tune sm120 sparse prefill geometry

[Dev-Jahn]

Depends on #36787 and stays a draft until it merges. #36497 merged into main through #37500 on 2026-09-08, so this branch is now #36787 rebased onto main plus one commit; that commit is what this PR adds (its GPU test moved to test/registered/kernel/qsa/ in the rebase).

#36787 gave sm120 its QSA sparse decode, split-K and MQA scoring paths but left the two sparse prefill kernels (_sparse_gqa_prefill and _sparse_gqa_chunk_prefill) launching with the L20 fallback selection: the table's large-prefill entry supplies BLOCK_N=16 with one warp and two stages, and the head tile comes from max(16, next_power_of_2(group_size)). This PR selects an sm120 launch family by measurement.

Motivation

Prefill for this model routes every full-attention layer through the sparse kernels: the indexer selects token_topk + compress_ratio - 1 = 2051 token slots per query row and the attention kernel streams up to that budget (rows whose visible context is shorter stream less). That makes the cost per call nearly independent of context length once the budget saturates; in the measurements below the chunk kernel moves only from 2.98 ms to 3.05 ms between 32k and 128k context on the previous geometry. The cost sits in launch configuration rather than data volume, so tuning the geometry improves it without touching the arithmetic.

The tuned family, held in _SM120_PREFILL_CONFIGS and selected by measurement:

• Group size 6 (the FP8 TP4/EP4 rank shape, 6 local query heads over 1 KV head): BLOCK_M=8, BLOCK_N=16, 4 warps, 3 stages.
• Group size 12 (the NVFP4 TP2 rank shape, 12 local query heads over 1 KV head): BLOCK_M=16, BLOCK_N=16, 2 warps, 3 stages.
• The tuned tuples apply only inside the measured workload envelope; everything else keeps the existing H20/L20 tables. The selector receives the kernel kind, the top-k, the local KV-head count, and the longest request from both launch sites, and the tuned family requires all of: the 188-SM SM120 configuration the numbers come from (an exact device set, not a threshold; unmeasured SM120 parts stay on the table), top-k 2051 (the only measured index width), head dimension 256, group size 6 with one local KV head or group size 12 with one or two (the measured TP rank shapes), 8,192 or more total query rows for ordinary prefill (its lowest measured point), and for chunk prefill 512 or more rows per request on average plus a grid-utilization bound (max_q * num_requests <= 2 * total_q) so a batch of one long request padded by short ones cannot take a geometry measured on uniform batches. Triton older than 3.7 is a policy floor: the family was measured on 3.7.1, and a 3.6.0 spot check compiled and ran the 8-row dot. sm121 (GB10) is excluded the same way as in [Qwen 3.8 Flash Next] Add sm120 (RTX PRO 6000 Blackwell) support #36787. SGLANG_QSA_PREFILL_GEOMETRY offers table (force the device table) and tuned (skip the device set and row gates while keeping the workload, architecture, and version gates, for measuring an unlisted SM120 part with the included benchmark), in the style of the [Qwen 3.8 Flash Next] Add sm120 (RTX PRO 6000 Blackwell) support #36787 backend overrides. The 170-SM 5090 CI runner exercises correctness through the forced tuned path; its performance was not measured, which is why it is not in the device set.

On the mechanism: in both kernels one Triton program handles one query row, and BLOCK_M tiles the query heads of one KV group, so the fp32 accumulator is [BLOCK_M, 256] per program and a BLOCK_M below the group size is invalid rather than untried (Triton lowers the 8-row dot onto the 16-row MMA instruction shape). The ablations and the register table below carry the evidence: for group 6 either change alone measured slower, so the gain needs the smaller head tile and the schedule together, and the register data rules out the single-factor stories. For group 12, which cannot shrink BLOCK_M below its 12 heads, the 2-warp 3-stage schedule accounts for the full gain.

Modifications

One commit over the stacked history:

• _get_prefill_config in sparse_attn.py with the named _SM120_PREFILL_CONFIGS table, the workload-envelope conditions (kernel kind, top-k, local KV-head count, row thresholds, and the grid-utilization bound), the exact device set, the Triton 3.7 version floor, the SGLANG_QSA_PREFILL_GEOMETRY override, and a bounded cache on the device-table choice.
• Two registered test files (GPU kernel tests and a CPU selector test) and a manual benchmark alongside the existing decode and MQA ones, with providers for the shipped selector, the previous geometry derived from the shipped table logic, forced geometry tuples from the command line, a torch reference, and a gather plus FlashAttention synthetic reference, over both prefill kernels and a configurable query count. The benchmark exits nonzero when a provider exceeds its error bound against the requested torch reference.

Kernel measurements

Produced by the included benchmark on one GPU of the host below: batch 1, chunk 8192, 2,051 selected token slots per query row, head dimension 256, one KV head, BF16, CUDA events, 10 warmups, median of 100 iterations. Timed providers alternate inside one shared window because sequential per-provider runs measured up to 6.8% higher on the first provider from clock state. A provider's absolute time still depends on which providers share its window, so milliseconds are comparable within one table and the ratios are the result; repeats of one cell within a fixed window agree to about 0.1%. The torch reference is computed on every timing run, every provider is checked against it, and the benchmark exits nonzero past a 0.02 max-abs bound. The torch and synthetic columns are timed in their own windows at three iterations, so they are context rather than part of the interleaved comparison. The selected indices come from a deterministic affine walk over the visible range; adjacent rows therefore share fewer selected tokens than a real top-k output would, and the tables below measure the streaming cost, at most a pessimistic view of inter-row cache reuse. Iterations reuse the same K/V tensors, so at 8k and 32k the working set stays L2-resident across iterations; the 128k rows exceed it. A cold-cache or streaming variant (rotating K/V and index buffers between iterations) was not measured, so the cutoffs are established in the regime described here. The previous provider derives the pre-tuning geometry from the shipped table logic, so the delta reproduces on a single checkout, and every cell of every table can be regenerated with a forced-geometry provider from the benchmark's command line.

The chunk kernel, which serves prefill with a cached prefix and every chunk after the first (a prefix-free single-chunk prompt takes the non-chunked kernel):

Rank shape	KV length	Previous	Tuned	Ratio	Torch reference	Gather+FlashAttention (synthetic)
6 Q / 1 KV	8,192	2.556 ms	2.146 ms	0.840x	81.8 ms	27.4 ms
6 Q / 1 KV	32,768	2.978 ms	2.334 ms	0.784x	81.7 ms	31.2 ms
6 Q / 1 KV	131,072	3.058 ms	2.343 ms	0.766x	81.1 ms	31.2 ms
12 Q / 1 KV	8,192	2.621 ms	2.436 ms	0.929x	81.2 ms	27.4 ms
12 Q / 1 KV	32,768	2.984 ms	2.757 ms	0.924x	81.3 ms	29.7 ms
12 Q / 1 KV	131,072	3.079 ms	2.837 ms	0.921x	82.1 ms	29.7 ms

The synthetic column materializes packed K/V and calls FlashAttention on the result; the backend has no such prefill route, so it shows what a packed implementation would cost, and past 32k its gather materializes about 17 GB per iteration and dominates it.

The non-chunked kernel, retuned by the same selector, measured with query rows equal to the KV length (torch-checked run):

Rank shape	KV length	Previous	Tuned	Ratio
6 Q / 1 KV	8,192	2.386 ms	2.200 ms	0.922x
6 Q / 1 KV	32,768	11.142 ms	10.035 ms	0.901x
6 Q / 1 KV	131,072	45.407 ms	40.873 ms	0.900x
12 Q / 1 KV	8,192	2.574 ms	2.486 ms	0.966x
12 Q / 1 KV	32,768	11.067 ms	10.687 ms	0.966x
12 Q / 1 KV	131,072	45.156 ms	43.430 ms	0.962x

The TP1 rank shape (24 query heads over 2 KV heads, group size 12) was measured separately and the tuned tuple wins there as well: 4.2 to 5.0% on the chunk kernel across the three KV lengths and 1.1% on the non-chunked kernel at 8k, at 8,192 query rows. The per-program tile is KV-head independent (the KV-head count only scales the grid), which the measurement confirms.

Query-count sweep and the 512 cutoff

The benchmark takes a query count and a forced geometry, so the same harness measured the fixed sm120 tuples against the geometry the previous selector picks at each total_q (KV 8,192):

total Q	Rank shape	Table geometry	Table ms	sm120 family ms	Ratio
32	6 Q / 1 KV	(16,32,8,2)	0.1081	0.1198	1.108x
64	6 Q / 1 KV	(16,64,8,2)	0.0830	0.1199	1.444x
128	6 Q / 1 KV	(16,64,4,2)	0.0814	0.1199	1.474x
256	6 Q / 1 KV	(16,32,4,2)	0.1485	0.1512	1.018x
512	6 Q / 1 KV	(16,32,4,2)	0.2924	0.1798	0.615x
1,024	6 Q / 1 KV	(16,16,1,2)	0.3797	0.3764	0.991x
8,192	6 Q / 1 KV	(16,16,1,2)	2.4172	2.0153	0.834x
32	12 Q / 1 KV	(16,32,8,2)	0.1080	0.1349	1.249x
64	12 Q / 1 KV	(16,64,8,2)	0.0830	0.1360	1.638x
128	12 Q / 1 KV	(16,64,4,2)	0.0806	0.1448	1.796x
256	12 Q / 1 KV	(16,32,4,2)	0.1467	0.1670	1.138x
512	12 Q / 1 KV	(16,32,4,2)	0.2944	0.1860	0.632x
1,024	12 Q / 1 KV	(16,16,1,2)	0.3886	0.3619	0.931x
8,192	12 Q / 1 KV	(16,16,1,2)	2.6221	2.4357	0.929x

The table wins short prefills and the tuned family wins at 512 and above. The 512 cutoff is the lowest measured row count at which the family wins for both shapes; the observed crossover sits in (256, 512]. The sweep points below 8,192 hold KV at 8,192 with a prefix, so every row streams the full selected budget; the 8,192 point is prefix-free with ramping visibility, and it is the only measured ordinary-prefill point, which is why the ordinary path requires 8,192 rows while the sweep-backed cutoff applies to the chunk path. Production total_q is a batch sum. A measured batch of 16 requests of 64 rows (1,024 rows total) ran 4.2% slower on the group-6 tuple and 6.3% faster on the group-12 tuple than the table; in that benchmark shape the per-request context also drops to about 512 tokens, so the batching effect and the shorter-context effect are not separated by this one point. The gate requires 512 rows per request on average, which keeps every measured win and sends that mixed regime to the table, and the added grid-utilization bound keeps batches whose shape the sweep did not cover off the tuned path as well; the benchmark now takes per-request KV lengths so the two effects can be measured separately later. The gap is not monotonic above the cutoff: at 1,024 rows the win shrinks to 0.9% (group 6) and 6.9% (group 12) where the previous selector switches to its large-total_q entry, then widens again at 8,192. The family never regresses at or above 512 in these measurements, and the CPU selector test pins both tuples at every measured query count plus the 511/512 boundary.

Geometry ablations

Group 6 (chunk kernel, 8,192 query rows, same tensors and indices across tuples):

KV length	(16,16,1,2) previous	(8,16,1,2)	(16,16,4,3)	(8,16,4,3) tuned
8,192	2.4177 ms	2.6044 ms	2.5406 ms	2.0138 ms
32,768	2.9819 ms	3.4373 ms	3.1568 ms	2.3370 ms
131,072	3.0229 ms	3.4725 ms	3.1616 ms	2.3496 ms

Group 12 (BLOCK_M cannot go below the 12 heads, so the schedule is the only axis):

KV length	1 warp / 2 stages previous	2 warps / 3 stages tuned	4 warps / 3 stages
8,192	2.6210 ms	2.4352 ms	2.6461 ms
32,768	2.9812 ms	2.7546 ms	3.0025 ms
131,072	3.0645 ms	2.8119 ms	3.0349 ms

Compiled-kernel attributes for the ablation tuples:

Group	Tuple	Registers/thread	Spilled values
6	(16,16,1,2) previous	254	0
6	(8,16,4,3) tuned	128	2
6	(8,16,1,2)	255	12
6	(16,16,4,3)	96	0
12	(16,16,1,2) previous	254	0
12	(16,16,2,3) tuned	164	0
12	(16,16,4,3)	96	0

End-to-end serving

NVFP4 TP2 on two GPUs of the host below, PLE host offload, CUDA graphs, chunked prefill 8,192, MTP off; random inputs at ISL 8,192 / OSL 128, 64 prompts, this branch against its exact base commit. Concurrency 1 was run three times per side with the radix cache flushed before every run:

ISL	Revision	Median TTFT, three runs	Spread
8,192	branch	505.6 / 505.1 / 505.2 ms	0.5 ms
8,192	base	511.3 / 501.6 / 504.2 ms	9.7 ms
131,072	branch	8,293.4 / 8,321.9 / 8,321.2 ms	28.5 ms
131,072	base	8,340.8 / 8,356.7 / 8,338.1 ms	18.6 ms

At ISL 8,192 the within-side spread reaches 9.7 ms (about 2%), larger than the expected effect, which is about 1.0 ms by construction: a prefix-free single-chunk prompt takes the non-chunked kernel, where the group-12 saving is about 0.09 ms per layer across the twelve full-attention layers. So at that shape the serving-level effect is below what the measurement resolves, in either direction. At ISL 131,072 the prompt spans sixteen chunks and every chunk after the first takes the chunk kernel, where the win is largest; every branch run is faster than every base run and the medians differ by 19.6 ms (0.24%). That is smaller than the sum of the kernel-table savings, which is expected where prefill attention overlaps other per-chunk work, and it is the regime where the change shows at serving level. The kernel tables above are the measurement of the change itself. A single concurrency-16 pair at ISL 8,192 measured TTFT 3,529.8 ms (branch) against 3,564.8 ms (base).

Accuracy Tests

The tuned tuples change the online-softmax tile boundaries (BLOCK_N 32 to 16 at total_q 512, and the schedule elsewhere), so prefill outputs are not bitwise identical to the previous geometry. The registered tests bound the kernel against an FP32 torch reference at atol=rtol=2e-2, and end to end the sglang.test.run_eval GSM8K protocol (5-shot; the five prefix examples are held out of the 1,319-item test split, scoring 1,314; chat API, greedy decoding at temperature 0, thinking enabled) ran on an NVFP4 TP2 server built from this branch and from the exact base commit on the same two GPUs:

Server	GSM8K
This branch	98.10%
Base commit	97.72%

Greedy decoding removes sampling variance; residual differences can still come from request arrival order under 16 client threads, dynamic batching, and kernel reduction order. The 0.38-point difference is in the branch's favor and sits below an approximate standard error for a difference of two runs at this sample count (about 0.55 points, treating the runs as independent; the shared prompt set makes the true uncertainty of the difference smaller than that bound, and one pair of runs cannot settle it), so the right reading is no measured regression. An earlier temperature-1.0 pair of the same protocol moved by the same margin in the other direction, which is consistent with that reading.

Tests

test/registered/kernels/test_qsa_sparse_prefill_triton.py, registered for both the SM120 (RTX 5090) and H100 CI suites so the tuned branch itself runs in CI while the H100 run keeps the fallback path covered:

• Ordinary and chunk prefill against an FP32 torch sparse-attention reference at atol=2e-2, rtol=2e-2: both rank geometries, one and two KV heads, ragged batch 1/2/8, non-multiple-of-four prefixes, negative and out-of-range selected indices, deterministic replay, and side streams.
• A case per tuned group size through the real interface with 2,560 query rows at prefix 0: past the 512-row gate with no selector mock, so on the SM120 runner the real selector picks the tuned tuple, the non-chunked kernel compiles and runs BLOCK_M=8, and rows past index 2,050 stream the full 2,051-slot budget. The case also asserts the selector's returned tuple under the production predicate.
• Chunk-path steady-state cases at prefix 2,050 covering the full 129-tile inner loop with every slot visible and valid.
• A hardware case asserting the geometry the selector returns on the running device, deriving the fallback expectation from _get_table_prefill_config.
• A launch-recorder case intercepting both kernels' launch objects at 513 rows through the real selector and asserting the constexpr geometry they received, so reverting either call site fails the suite on the SM120 runner (the H100 run derives its expectation from the table).

test/registered/unit/layers/attention/test_qsa_sparse_prefill_config.py (CPU): imports _SM120_PREFILL_CONFIGS and pins both tuples at every measured query count plus the 511/512 boundary, covers the SGLANG_QSA_PREFILL_GEOMETRY values, the Triton version gate, and the fallback path with the architecture predicates mocked, and drives the real is_sm120_supported and is_sm121 predicates through mocked CUDA facts.

The existing QSA suites (test_qsa.py, test_qsa_mqa_triton.py, test_qsa_sparse_index_bounds.py, test_qsa_mtp_shared_indexer.py, compressed addressing) pass unchanged.

Compute Sanitizer memcheck, initcheck, racecheck and synccheck are clean on the benchmark's sanitizer driver: batch 1, tuned geometry, both kernels at both 8k and 128k for both rank shapes (eight cases per tool), guard-banded allocations.

Environment

Item	Value
CPU	2x AMD EPYC 9355
GPU	8x NVIDIA RTX PRO 6000 Blackwell Server Edition 96 GB, sm_120a
Driver / CUDA	595.71.05 / 13.2
torch / Triton / FlashInfer	2.13.0+cu130 / 3.7.1 / 0.6.17
OS	Ubuntu 24.04.4, kernel 6.8

Not covered

The indexer scoring and selection chain is unchanged. The chunk wrapper still derives its launch bound with a device read; removing the host reads on that path is a separate change that needs its own measurement.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. The new environment variable is documented in the environment variables reference.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #35072149282
Latest PR Test (Extra): ❌ Run #35072148938
Latest PR Test (AMD ROCm 10): ❌ Run #35072149349

[detain]

nice work on your PRs!!!

https://github.com/detain/sglang/tree/stack-36644-core

combined this PR with #38144 + #38209 + #36644 on 4× RTX PRO 6000 Workstation Edition (sm_120, 96 GB)

Running this PR as the base of a combined stack against latest main (post-#37500, incl. #36806/#37068), stacked in dependency order: #36787 → #37275 → #38144 → #38209 → #36644.

On top of the five PRs we carry three local additions:

• fix(sm120): read CUDA graph config from exec bag in Qwen4 JIT prewarm — the JIT prewarm read the CUDA-graph config from a stale seed instead of the runtime-context exec bag (get_exec().graph.cuda_graph_config), so kernels were prewarmed under the wrong capture settings after the context refactor; fixed to read the live bag.
• feat(sm120): allow pooled bf16 GDN state with FlashInfer bf16-state kernel — relaxes this PR's exact-sm120 "fp32 unpooled state" contract when the installed FlashInfer wheel ships the bf16-state GDN decode kernel, so --linear-attn-decode-backend flashinfer --mamba-ssm-dtype bfloat16 works with pooled state; falls back to the loud error+guidance when the kernel is absent.
• test(sm120): set is_sm120 on the FlashInfer alignment fixture to keep the fixture honest under the new state-pool logic.

Plus rebase adaptations (no functional change): the NUMA-interleaved PLE pinning now attaches to upstream's allocate_ple_host_table() backend abstraction (pinned path only, file backend untouched), and the environ.py additions were merged against the new SGLANG_QWEN4_PLE_FILE_* block.

Notes: all auto-selection env knobs (SGLANG_QSA_DECODE_BACKEND / SGLANG_QSA_MQA_BACKEND / SGLANG_QSA_PREFILL_GEOMETRY) left at auto; the tuned prefill geometry requires the 188-SM Server Edition under auto (used tuned to force otherwise). #38144 was the difference between stable and silently-wrong long-context top-k for us.

it runs pretty awesome.

my cmd line is

SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION=false \
SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 \
SGLANG_QSA_PREFILL_GEOMETRY=tuned \
sglang serve \
  --model-path Qwen/Qwen3.8-Flash-Next \
  --tp 4 \
  --ep 4 \
  --mem-fraction-static 0.87 \
  --chunked-prefill-size 8192 \
  --kv-cache-dtype fp8_e4m3 \
  --page-size 128 \
  --linear-attn-prefill-backend triton \
  --linear-attn-decode-backend flashinfer \
  --mamba-ssm-dtype bfloat16 \
  --reasoning-parser auto \
  --tool-call-parser auto \
  --host 0.0.0.0 \
  --port 30000 \
  --json-model-override-args '{"text_config": {"rope_parameters": {"mrope_interleaved": true, "mrope_section": [11, 11, 10], "rope_type": "yarn", "rope_theta": 10000000, "partial_rotary_factor": 0.25, "factor": 4.0, "original_max_position_embeddings": 262144}}}' \
  --context-length 1000000 \
  --ple-offload-embedding \
  --speculative-algorithm NEXTN \
  --speculative-num-steps 3 \
  --speculative-eagle-topk 1 \
  --speculative-num-draft-tokens 4 \
  --enable-metrics \
  --enable-mfu-metrics

(Note that SGLANG_QSA_PREFILL_GEOMETRY=tuned line needs dropped for anyone using the Server Edition of the RTX 6000)

[Dev-Jahn]

Thanks for the run and the write-up. That is the first Workstation Edition result I have seen for this stack.

#36497 merged into main on 09-08, so I have rebased #36787 and this PR onto main; the rebase notes are in the #36787 description. Two of your findings are folded in: the prewarm now reads the graph config and the draft width from the runtime context bags, and the fixture change is not needed any more because the GDN commit is gone. On that commit: main now converts the prefill state to fp32 on sm12x by itself, and FlashInfer 0.6.18 ships the pooled bf16 decode kernel, so the exact-sm120 fp32 contract this stack used to carry is dropped rather than relaxed. Your launch line should work unchanged.

On SGLANG_QSA_PREFILL_GEOMETRY=tuned: the device set is exact because every number in this PR comes from the 188-SM Server Edition. If you can run test/manual/kernels/benchmark/attention/bench_qsa_sparse_prefill.py on the Workstation part and post the table, I will add it to the set so the override is not needed there.