[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).
  • SGLANG_AITER_ASM_CONTEXT_PREFILL=0 turns it off.
• Draft-model catch-up chunks (EAGLE/MTP) hit the same branch and get the same speedup — they are ordinary extend batches with a prefix.

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.

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.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #33130605150
Latest PR Test (Extra): ❌ Run #33130604897
Latest PR Test (AMD ROCm 7.2): ❌ Run #33130605020