[amd-nprotaso]

Motivation

On ROCm, Qwen3.8 QSA decode and speculative verify run the sparse attention as two passes:

1. a Triton gather copies each row's selected KV (up to 2051 positions) into packed scratch;
2. a Triton decode runs one program per (row, KV head).

At decode batch sizes that is an extra round trip through the selected KV, and most of the GPU sits idle. At 32k context with MTP on one MI355X, QSA attention takes about 9.4% of decode GPU time, roughly 120 µs per layer call.

AITER now has a fused FlyDSL kernel for gfx950: aiter.ops.flydsl.qsa_sparse_decode (AITER PR). It reads the selected KV straight from the pool through req_to_token and splits each row's selection across CTAs. This PR serves QSA decode and verify with it.

Modifications

• QwenSparseAttnBackend.forward_decode → _forward_paged_attention takes the fused kernel first on HIP when can_use_flydsl_qsa_sparse_decode accepts the operands: gfx950, bf16 q, bf16 or fp8 KV, head_dim 256, GQA group ≤ 16. That covers decode, target verify and draft extend, since each query row is its own length-one query. Everything else falls through to the existing path unchanged.
• Split count: recommended_num_splits sets the KV split count per batch size, so it is fixed for each captured CUDA graph.
• Workspace: sized once at the captured maximum with workspace_partials, because graphs keep raw pointers to it. A larger eager batch gets a per-call workspace.
• Tests:
  • test/registered/kernels/ops/attention/qsa/test_qsa_flydsl_decode.py checks that the backend gives the kernel the same rows, lengths and per-row request map as the Triton fallback, and that the outputs match.
  • Kernel correctness tests live in AITER (op_tests/test_flydsl_qsa_sparse_decode.py).
  • test/registered/kernels/benchmark/attention/bench_qsa_sparse_decode.py compares the two paths.

Requires AITER with aiter.ops.flydsl.qsa_sparse_decode. No recommended-deployment flags change, so the cookbook page is unchanged.

Accuracy Tests

GSM8K through the chat API: greedy, thinking off, final number scored, 1024 max tokens, MTP on with real acceptance.

• Short: 500 questions, 88-token prompts.
• Long: 48 questions, each preceded by 120 solved problems (~23k-token prompts), so decode selects 2051 of ~23k positions.

Accept length is total completion tokens ÷ total verify steps, from each response's spec_tokens_details.

MI355X, TP1, MTP 3/1/4	Triton (flag off)	FlyDSL (this PR)
GSM8K short (500)	95.2%	96.0%
GSM8K long, ~23k ctx (48)	95.8%	97.9%
MTP accept length, short	3.503	3.506
MTP accept length, long	3.394	3.441

Repeated baseline runs vary by about 1 point on short GSM8K, so the accuracy differences are within noise.

At the kernel level, against an fp64 reference at Qwen3.8 TP1 shapes (Hq=24, Hkv=2, top-k 2051), the fused kernel is more accurate than the Triton path it replaces:

• relative L2 error 1.7e-3 vs 4–6e-3;
• 3–5× lower max error with a peaked softmax.

Speed Tests and Profiling

Setup:

• 1× MI355X: Qwen3.8-Flash-Next Quark MXFP4 (PLE FP8), TP1, --attention-backend aiter --page-size 64, MTP NEXTN 3/1/4, --max-running-requests 16.
• The only difference between the two arms is enabled/disabled FlyDSL QSA kernel
• Perf runs use simulated MTP acceptance (SGLANG_SIMULATE_ACC_LEN=2.32), so both arms do identical verify work; accept length is 2.32 in both.

End to end: ISL 32768 / OSL 32768, 24 prompts at concurrency 8, 8 warmup requests (sglang.benchmark.serving, random dataset).

	Triton	FlyDSL	Δ
Output throughput (tok/s)	977.2	1075.0	+10.0%
Mean TPOT (ms)	8.06	7.31	−9.3%
P99 TPOT (ms)	8.20	7.43	−9.4%
Mean TTFT (ms)	4275	4388	+2.6%
Mean E2E latency (s)	268.3	243.8	−9.1%
Benchmark duration (s)	804.8	731.6	−9.1%

TTFT is prefill, which this PR does not touch; the difference is run-to-run noise.

Decode profile: 40 decode steps at full batch, 8 requests at 32k context, MTP.

	Triton	FlyDSL
QSA attention GPU time	72.0 ms (9.4%)	20.0 ms (2.8%)
per QSA call (600 calls)	120.0 µs	33.3 µs
total GPU kernel time	767.7 ms	705.2 ms (−8.1%)

Triton per call: _sparse_gqa_chunk_prefill 71.8 µs + _compact_kv 39.4 µs + _fa2_valid_counts 4.4 µs + _fa2_prefix_sum 4.4 µs. FlyDSL per call: sparse_decode_attn_kernel 28.9 µs + sparse_decode_reduce_kernel 4.4 µs.

Kernel benchmark: bench_qsa_sparse_decode.py, top-k 2051, cold cache, CUDA graph.

Hkv, group	batch	FlyDSL (µs)	Triton gather + decode (µs)	speedup
2, 12 (TP1)	1	16.6	89.6	5.4×
2, 12	4	17.2	97.6	5.7×
2, 12	16	25.3	112.4	4.4×
2, 12	64	59.7	202.6	3.4×
2, 12	128	100.9	344.4	3.4×

---

CI States

Latest PR Test (Base): ❌ Run #40664232734
Latest PR Test (Extra): ❌ Run #40664232716
Latest PR Test (AMD ROCm 10): ❌ Run #40664232700

[1am9trash]

Hi, @amd-nprotaso
Should _resolve_flydsl_sparse_decode() also check et_bool_env_var("SGLANG_USE_AITER")?

[amd-nprotaso]

I removed SGLANG_ENABLE_QSA_FLYDSL_DECODE and use SGLANG_USE_AITER instead.

[1am9trash]

LGTM. The comments resolved (small nits).
Only impact on amd code path.

[1am9trash]

PR-test-base passed.
Link: https://github.com/sgl-project/sglang/actions/runs/40303135933/job/121505901651?pr=45196

[sogalin]

LGTM with AMD flag to protect the common code path.

Please modify the description as well. It looks like you change the code from comments but haven't modified the descriptions.

[amd-nprotaso]

LGTM with AMD flag to protect the common code path.

Please modify the description as well. It looks like you change the code from comments but haven't modified the descriptions.

Updated

[amd-nprotaso]

ROCm/aiter#6670 was merged