[AMD] Qwen3.5 ASM FMHA chunked-prefill context attention

[zijiecode]

Motivation

On MI355X, long-input Qwen3.5 serving (70k prompt, chunked prefill) spends most of its prefill attention time in aiter's ck_tile paged mha_batch_prefill. Measured at production shapes (head_dim 256, fp8 KV, 16k-token chunks), that kernel runs at roughly 15% of the GPU's FP8 peak — and paging is not the reason: a fully contiguous KV layout is just as slow, and fp8 inputs are only ~1.2x faster than bf16, so the pipeline never reaches the fp8 rate.

The gfx950 ASM kernel fmha_fwd_hd256_fp8_causal_group (already used by SGLang for the no-prefix first chunk) runs the same shapes 3.5x faster. The only reason context chunks cannot use it today is that it wants contiguous varlen K/V, while the KV cache is paged.

It turns out the gather is nearly free: copying 70k tokens of fp8 KV (~72 MB) into a contiguous buffer takes ~18 µs, against several milliseconds for the attention itself. So we gather and call the ASM kernel.

Modifications

python/sglang/srt/layers/attention/aiter_backend.py, in forward_extend:

• New branch for extend batches with a prefix (the no-prefix case already had an ASM path). It expands the page-granular kv_indptr / kv_indices metadata into per-token pool indices, gathers K/V into a contiguous varlen buffer (via a uint8 view, since index_select is not implemented for fp8), and calls flash_attn_varlen_fp8_pertensor_func with bottom-right causal masking.
• Guard rails:
  • The branch checks that the page metadata actually covers seq_lens before gathering (mixed / speculative batches can disagree); on any mismatch it logs once and falls back to the original paged path. Nothing crashes, worst case is the old speed.
  • Gated on gfx950, bf16 Q, fp8 KV cache, head_dim 256, no sliding window, no sink, no softcap.
  • Disabled when the KV pool uses the vectorized_5d layout (different physical layout).
• Draft-model catch-up chunks (EAGLE/MTP) hit the same branch and get the same speedup — they are ordinary extend batches with a prefix.

Accuracy Tests

GSM8K, 5-shot, all 1319 questions, temperature 0, max new tokens 2048; Qwen3.5-397B-A17B-MXFP4, TP2, MI355X.

Backend	Accuracy	Correct	Invalid
AITER	92.8%	1224/1319	1/1319

Benchmarking and Profiling

Benchmarks were run on a tree with #36330 applied (the current MI355X perf baseline for this model); the change itself is independent of it and this PR is based on main.

Kernel microbenchmark

Single MI355X GPU, causal, head_dim 256, fp8 KV, 16384 new tokens attending to a 32k-token paged context (the production context-chunk shape):

kernel	TP2 shape (16 q-heads / 2 kv-heads)	TFLOP/s	TP4 shape (8 / 1)
ck_tile paged batch_prefill, fp8 KV (current path)	14.24 ms	772 (~15% fp8 MFU)	7.21 ms / 762
ck_tile contiguous (page_size 1), fp8 KV	14.45 ms	761	—
ck_tile paged, bf16 KV	17.48 ms	629	—
ASM fmha_fwd_hd256_fp8_causal_group + gather	4.04 ms	2,725 (~55% fp8 MFU)	2.00 ms / 2,753

Paged and contiguous ck_tile run at the same speed, and fp8 is only ~1.2x faster than bf16 there — so paging is not the bottleneck, kernel efficiency is. Tail chunks show the same ratio (4464 new tokens over a 65k context: 2.44 ms vs 8.99 ms). The KV gather this PR adds costs 17.6 µs at 65k context (~7.6 TB/s effective), negligible against the attention itself.

End-to-end serving

End-to-end serving on MI355X, Qwen3.5-397B-A17B-MXFP4, random-ids 70k input / 300 output, EAGLE (3 steps, top-k 1, 4 draft tokens), warmup + measured requests, otherwise the pinned recipe (--attention-backend aiter --page-size 16 --kv-cache-dtype fp8_e4m3 --chunked-prefill-size 16384):

config	metric	before	after	delta
TP2 c4	mean E2E	13,393 ms	10,057 ms	-24.9%
TP2 c4	mean TTFT	8,274 ms	6,382 ms	-22.9%
TP2 c4	mean TPOT	17.12 ms	12.29 ms	-28.2%
TP4 c4	mean E2E	8,512 ms	6,415 ms	-24.6%
TP4 c4	mean TTFT	4,942 ms	3,638 ms	-26.4%
TP4 c8	mean E2E	16,098 ms	11,724 ms	-27.2%
TP4 c12	mean E2E	23,690 ms	16,933 ms	-28.5%
TP4 c16	mean E2E	31,154 ms	22,585 ms	-27.5%

TPOT improves as a side effect: other requests' prefill chunks block decode for less time. No crashes and zero guard fallbacks across all runs above.

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.
• 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 #33142874228
Latest PR Test (Extra): ❌ Run #33142885073
Latest PR Test (AMD ROCm 7.2): ❌ Run #33142874151

[yichiche]

LGTM.