[Qwen3.8-Flash-Next] Separate prefill and decode paths for QSA indexer

[gau-nernst]

Purpose

#53896 ships with a simple QSA indexer kernel that is used for both decode and prefill requests. This is not efficient as we can't design kernels specialized for decode and prefill shapes separately. This PR splits a mixed batch into decode and prefill requests, so that efficient decode/prefill kernels can be invoked separately, similar to how other attention backends work in vLLM.

Currently this PR ships 2 decode/prefill-specialized Triton kernels to replace the original generic QSA indexer, but these are not meant to be SOL. Further optimizations are left for future PRs.

Main changes

• Replace _qsa_mqa_paged_kernel with _qsa_mqa_paged_uniform_kernel (decode and spec-decode) and _qsa_mqa_paged_prefill_kernel (prefill). Create a new file qsa_indexer.py under NVIDIA ops.
• Add visible_blocks = (logical_position + 1) // compress_ratio to QSA metadata. Previously this is computed by the indexer score kernel, which is not quite natural.
• Add warmup logic for _qsa_mqa_paged_uniform_kernel. The prefill kernel only has 1 specialization, it doesn't need warmup
• Switch _LOGITS_WORKSPACE_BYTES to env var VLLM_SPARSE_INDEXER_MAX_LOGITS_MB, which raises the default workspace size from 128 MiB to 512 MiB. Decode kernel does not observe this limit, following DeepSeek sparse attention.

The kernels' heuristics are currently tuned on GB300. They might not be optimal for all GPUs.

Microbenchmarks

All measurements are done on GB300. Using page_size=196 (actual page size for non-MTP production)

Decode and Spec-decode

Note: Decode query len (DQL) = 4 corresponds to MTP3 (3 speculative tokens)

Single request

DQL	Context	Baseline µs	PR µs	Baseline GB/s	PR GB/s	Speedup
1	8K	3.984	3.792	134	141	1.05x
1	50K	4.704	4.544	708	733	1.04x
1	100K	5.696	5.600	1,169	1,189	1.02x
4	8K	3.968	3.936	141	143	1.01x
4	50K	5.536	4.800	630	726	1.15x
4	100K	7.632	5.792	913	1,203	1.32x

Multiple requests Context lengths are sampled from log-normal distribution to simulate skewed contexts

DQL	Requests	Context tokens (min / mean / max)	Baseline µs	PR µs	Baseline GB/s	PR GB/s	Speedup
1	4	1,584 / 10,240 / 24,708	4.704	4.608	567	579	1.02x
1	16	1,536 / 10,240 / 43,876	10.912	8.192	977	1,302	1.33x
1	4	6,548 / 51,199 / 102,400	7.664	7.536	1,737	1,767	1.02x
1	16	3,508 / 51,200 / 102,400	18.368	15.552	2,900	3,425	1.18x
4	4	1,584 / 10,240 / 24,708	10.816	4.864	259	576	2.22x
4	16	1,536 / 10,240 / 43,876	14.048	9.120	798	1,229	1.54x
4	4	6,548 / 51,199 / 102,400	16.000	8.160	871	1,709	1.96x
4	16	3,508 / 51,200 / 102,400	31.456	16.528	1,773	3,374	1.90x

Prefill

Query layout	Prefix	Baseline µs	PR µs	Baseline TFLOP/s	PR TFLOP/s	Speedup
1x8K	8K	1,025.100	162.431	25.1	158.6	6.31x
1x8K	50K	2,728.835	360.766	42.5	321.4	7.56x
1x8K	100K	4,752.939	524.830	47.0	425.5	9.06x
1x16K	8K	2,372.762	328.510	29.0	209.2	7.22x
1x16K	50K	5,775.507	691.965	43.1	360.0	8.35x
1x16K	100K	9,815.913	1,047.393	47.3	442.9	9.37x
2x8K	8K	2,039.702	364.038	25.3	141.6	5.60x
2x8K	50K	5,458.671	701.340	42.5	330.7	7.78x
2x8K	100K	9,517.954	1,058.042	46.9	422.2	9.00x
4x4K	8K	1,880.741	453.158	22.8	94.8	4.15x
4x4K	50K	5,295.237	782.044	42.2	285.6	6.77x
4x4K	100K	9,359.874	1,182.170	46.8	370.6	7.92x

Prefill is still very far from SOL (2 PFLOPS). But good enough for a pure Triton baseline implementation. A more optimized kernel can be added in the future.

E2E perf benchmarks

Baseline e16b5e5. Results obtained on GB300 TP4. Qwen/Qwen3.8-Flash-Next

vllm serve {model}
--port {port}
-tp 4
--no-enable-prefix-caching
--no-enable-flashinfer-autotune
--enable-auto-tool-choice
--tool-call-parser qwen3_coder
--reasoning-parser qwen3

VLLM_USE_RUST_BENCH=1

vllm bench serve
--backend openai-chat
--base-url http://{remote_host}:8000
--model {model}
--dataset-name speed-bench
--speed-bench-config throughput_8k
--speed-bench-max-input-len 8192
--speed-bench-output-len 1024
--ignore-eos
--num-warmups 5
--sweep-max-concurrency 1,4,16,64
--sweep-num-prompts-factor 10
--percentile-metrics "ttft,tpot,itl,e2el"
--metric-percentiles 50,90,99
--result-dir "{log_dir}/results"
--save-result

Non-MTP (P50 statistics)

Concurrency	TTFT before	TTFT after	1/TPOT before	1/TPOT after
1	195.0 ms	163.6 ms (-16.1%)	211.7 tok/s	212.5 tok/s (+0.4%)
4	393.8 ms	396.8 ms (+0.8%)	151.8 tok/s	156.1 tok/s (+2.9%)
16	681.2 ms	616.1 ms (-9.6%)	91.4 tok/s	95.0 tok/s (+3.9%)
64	759.5 ms	689.9 ms (-9.2%)	46.6 tok/s	47.8 tok/s (+2.6%)

MTP3 (--speculative-config '{"method":"mtp","num_speculative_tokens":3}', P50 statistics)

Concurrency	TTFT before	TTFT after	1/TPOT before	1/TPOT after
1	198.0 ms	162.1 ms (-18.1%)	243.8 tok/s	227.2 tok/s (−6.8%)
4	226.1 ms	191.6 ms (-15.2%)	172.2 tok/s	176.8 tok/s (+2.7%)
16	230.8 ms	202.6 ms (-12.2%)	96.0 tok/s	101.6 tok/s (+5.9%)
64	262.1 ms	228.3 ms (-12.9%)	43.9 tok/s	47.4 tok/s (+8.1%)

Test Plan

Unit test

.venv/bin/python -m pytest -v -s \
  tests/models/qwen4_exp/test_qsa_pre_indexer.py \
  tests/models/qwen4_exp/test_qsa_reference.py

E2E accuracy test

TP4 on GB300 with MTP3. Qwen/Qwen3.8-Flash-Next

vllm serve {model}
--port {port}
-tp 4
--enable-prefix-caching
--no-enable-flashinfer-autotune
--speculative-config '{"method": "mtp", "num_speculative_tokens": 3}'
--enable-auto-tool-choice
--tool-call-parser qwen3_coder
--reasoning-parser qwen3

Test	Sampling config	main (e16b5e5)	this PR
GSM8K	temp=0, max_token=1024, think=false	0.9674	0.9682
MMMU_PRO	max_tokens=65536, think=true	0.7671	0.7746

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @gau-nernst.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[gau-nernst]

/ci run

[github-actions]

✅ Triggered Buildkite CI #86528 for commit db191d62fb13.

[ZJY0516]

Do we need to optimize this for loop? or you want to leave it for future work?

[gau-nernst]

The loop follows existing DSv4 sparse indexer pattern to limit logits workspace size

vllm/vllm/model_executor/layers/sparse_attn_indexer.py

Line 449 in 55aa766

for chunk in prefill_metadata.chunks:

The current main also has this chunking logic

[ZJY0516]

vLLM has many place using top_k, could we unify them?

The condition looks strange

current_platform.has_device_capability(90)
 and not current_platform.is_device_capability_family(120)

[gau-nernst]

I'm following the logic for DSv4 sparse indexer here as well

vllm/vllm/model_executor/layers/sparse_attn_indexer.py

Lines 617 to 629 in 55aa766

	use_cooperative_topk = (
	current_platform.is_cuda()
	and topk_tokens in (512, 1024, 2048)
	and num_rows <= 64
	and logits.stride(0) % 4 == 0 # TMA 16-byte alignment
	and current_platform.has_device_capability(90)
	and not current_platform.is_device_capability_family(120)
	)
	use_persistent_topk = current_platform.is_cuda() and topk_tokens in (
	512,
	1024,
	2048,
	)

[gau-nernst]

/ci run

[github-actions]

✅ Triggered Buildkite CI #86625 for commit 34ecc53f22b3.