[lishuoshuo-amd]

Overview

ROCM_ATTN, the default ROCm backend for non-MLA models, runs every decode step through its HIP paged attention, also with AITER enabled. With an fp8 KV cache and 16-token blocks its K/V blocks already match the layout of the AITER asm paged attention, so this PR runs large pure-decode batches on pa_fwd_asm with no KV relayout and keeps the HIP kernel for everything else. ROCM_AITER_FA already uses this kernel, but only when it is selected explicitly with the shuffled KV cache layout.

Claims

• MI355X, VLLM_ROCM_USE_AITER=1, --kv-cache-dtype fp8: output throughput +4.2–6.3% at concurrency 64 and +6.2–8.8% at concurrency 128 on Qwen3-14B-FP8, Llama-3.1-8B-Instruct and Qwen3-8B at TP1, and +2.7–4.4% at TP2. Batches below the cut are unchanged (concurrency 1 and 8 everywhere, and 16 except Qwen3-14B at TP1, which already crosses it there: +2.4%).
• The asm path is taken only for pure-decode batches above 2 * num_compute_units() query heads (512 on MI355X), the cut AITER itself uses between its asm and HIP kernels, and only for GQA ratios up to 8, with the ultra-precision kernel (high_precision=2).
• GSM8K and 3K–32K passkey retrieval match the HIP decode within run-to-run spread.

Validation

MI355X, ROCm 10.0, AITER 0.1.27, vllm/vllm-openai-rocm:nightly-rocm100-a2bcb5657a881a002e79d0adab3c6b5e39436121. In that image rocm_attn.py, _aiter_ops.py and platforms/rocm.py match this PR's base 8af2080e7 except get_preferred_block_size from vllm-project#63231, which only sets block size 16 for ROCM_ATTN, the size used here.

pytest tests/kernels/attention/test_rocm_attn_aiter_asm_decode.py -v      # 15 passed
pytest tests/v1/attention/test_rocm_attention_backends_selection.py -v   # 89 passed, unchanged

VLLM_ROCM_USE_AITER=1 vllm serve <model> --kv-cache-dtype fp8 --max-model-len 40960 --gpu-memory-utilization 0.85 [--tensor-parallel-size 2]
vllm bench serve --backend vllm --model <model> --dataset-name random --random-input-len 1024 --random-output-len 1024 --ignore-eos --max-concurrency <c> --num-prompts <n>

Output tok/s, main -> this PR, fresh server per arm (n = 8 / 32 / 64 / 256 / 256 prompts at concurrency 1 / 8 / 16 / 64 / 128):

model	TP	conc 1	conc 8	conc 16	conc 64	conc 128
Qwen3-14B-FP8	1	128 -> 129	943 -> 947	1756 -> 1798 (+2.4%)	4898 -> 5105 (+4.2%)	6224 -> 6612 (+6.2%)
Llama-3.1-8B-Instruct	1	282 -> 282	1727 -> 1726	3232 -> 3229	9390 -> 9983 (+6.3%)	12661 -> 13670 (+8.0%)
Qwen3-8B	1	255 -> 253	1739 -> 1739	3321 -> 3329	8963 -> 9390 (+4.8%)	12026 -> 13089 (+8.8%)
Qwen3-14B-FP8	2	175 -> 175	1261 -> 1262	2331 -> 2332	6066 -> 6231 (+2.7%)	7750 -> 8066 (+4.1%)
Llama-3.1-8B-Instruct	2	388 -> 388	2314 -> 2322	4241 -> 4239	10540 -> 10926 (+3.7%)	13279 -> 13864 (+4.4%)

Median TPOT drops 4–11% at concurrency 64–128. Repeats agree: at TP1 concurrency 64 / 128, main ran 4852–4899 / 6024–6229 tok/s in four runs of Qwen3-14B-FP8 and the asm decode 5105–5115 / 6612–6635 in five (two at AITER's default precision, which runs at the same speed), and a second Llama pair gave +5.1% / +7.8%. One further PR run had host stalls (p99 ITL 79.7 ms against 15.5–17.1 ms in every other run) and is left out.

GSM8K, 1319 questions, 5-shot chat, greedy, 64 in flight so decode runs above the cut, Qwen3 thinking off; flexible match, mean over runs (each arm: two or three CUDA graph runs and one --enforce-eager run):

model	main	this PR
Qwen3-14B-FP8	0.9239 (0.9227–0.9257)	0.9213 (0.9166–0.9257)
Llama-3.1-8B-Instruct	0.8309 (0.8294–0.8332)	0.8327 (0.8287–0.8378)
Qwen3-8B	0.8948 (0.8901–0.8984)	0.8976 (0.8923–0.9022)

Passkey retrieval at 3K / 8K / 16K / 32K tokens (24 prompts, 30 in flight): 24/24 in every run, both arms, TP1 and TP2. At TP2, GSM8K is 0.8362 -> 0.8347 (Llama) and 0.9189 -> 0.9189 (Qwen3-14B-FP8).

VLLM_GPU_SYNC_CHECK=error with the asm path active (Llama-3.1-8B-Instruct at TP1: GSM8K plus the five concurrencies above; Qwen3-14B-FP8 at TP2: GSM8K): torch's sync debug mode is armed in every worker and no synchronization is reported, no tracebacks, GSM8K 0.8317 and 0.9204. main's Llama run under the same check is clean as well.

Kernel latency, accuracy runs and kernel error

impl.forward in a CUDA graph, µs per call, fp8 cache, all sequences at the given context length (asm at AITER's default precision; the serving runs show high_precision=2 runs at the same speed):

heads (q/kv)	ctx	seqs	HIP (today)	asm	asm / HIP
40/8	1024	16 / 64 / 256	10.69 / 28.58 / 139.59	8.33 / 22.75 / 97.25	0.78 / 0.80 / 0.70
40/8	4096	16 / 64 / 256	26.71 / 130.10 / 448.27	22.47 / 95.05 / 363.04	0.84 / 0.73 / 0.81
32/8	1024	24 / 64 / 256	14.26 / 26.54 / 132.85	10.62 / 22.73 / 96.25	0.74 / 0.86 / 0.72
32/8	4096	24 / 64 / 256	35.05 / 123.93 / 439.31	31.05 / 94.69 / 362.53	0.89 / 0.76 / 0.83

Below the cut, forcing asm is slower on long contexts (40/8, ctx 4096, 1–4 sequences: 20.0–20.6 µs against 8.0–10.7 µs), and AITER's HIP fallback is as fast as ROCM_ATTN's kernel but about 16x less accurate on an fp8 cache (mean absolute error 0.0056 against 0.00035 vs an fp32 reference).

Error vs an fp32 reference above the cut, 4 head layouts × 3 seeds, 1–4000 tokens, max / mean: HIP 0.0035–0.0072 / 0.00013–0.00018; asm at AITER's default precision 0.015–0.026 / 0.00076–0.00108; asm with high_precision=2 0.015–0.026 / 0.00044–0.00078.

GSM8K per run, flexible / strict:

model	main	this PR (high_precision=2)	asm at AITER's default precision
Qwen3-14B-FP8	0.9227 / 0.6528, 0.9257 / 0.6550, eager 0.9234 / 0.6619	0.9249 / 0.6482, 0.9181 / 0.6437, 0.9257 / 0.6550, eager 0.9166 / 0.6581	0.9151, eager 0.9219, ROCM_AITER_FA 0.9189
Llama-3.1-8B-Instruct	0.8332 / 0.8249, 0.8302 / 0.8211, eager 0.8294 / 0.8196	0.8317 / 0.8249, 0.8378 / 0.8256, 0.8287 / 0.8188, eager 0.8324 / 0.8249	0.8234, eager 0.8218, ROCM_AITER_FA 0.8249
Qwen3-8B	0.8984 / 0.8491, 0.8901 / 0.8287, 0.8923 / 0.8400, eager 0.8984 / 0.8408	0.8961 / 0.8438, 0.8923 / 0.8332, 0.8999 / 0.8408, eager 0.9022 / 0.8408	0.8999

"main" includes runs of this branch with VLLM_ROCM_USE_AITER_MHA=0, which turns the path off. Any attention numerics change moves many greedy outputs: main in eager mode against main with CUDA graphs leaves 27% (Qwen3-14B) and 50% (Llama) of the GSM8K outputs identical, and this PR against main 38% / 44% / 49%. At AITER's default precision the asm decode, including the shipped ROCM_AITER_FA path, is 0.2–0.9 points below main's mean in all six Qwen3-14B and Llama runs, which is why this PR asks for the ultra-precision kernel.

Details

• Layout. With an fp8 cache, split_kv_cache gives K as [blocks, kv_heads, head_size / 16, 16, 16] and V as [blocks, kv_heads, head_size, 16]. With 16-token blocks that V is AITER's shuffled [blocks, kv_heads, block_size / 16, head_size, 16] with one group, so unsqueeze(2) is all the asm kernel needs; the interleaved V of the AITER fused RoPE + KV-cache writer has the same bytes. has_native_kv_cache_layout is required because the kernel addresses pages by a dense page stride.
• Scale tables. The asm kernel reads per-slot [blocks, kv_heads, block_size] scales. The metadata builder maps each layer to a table filled with its per-tensor scale once the KV cache is bound, shaped by the KV cache group's spec rather than the model-wide head count, reusing ROCM_AITER_FA's AsmKvScaleTables, so layers with equal scales share one allocation (K and V together about 0.08% of the KV cache for Qwen3-14B) and forward reads nothing on the host.
• Cut and precision. The choice is made in Python on num_actual_tokens * num_heads, so each captured CUDA graph uses one kernel; both kernels are deterministic and batch-invariant on their own, but a request's kernel depends on how many others share its batch, as with AITER's paged_attention_common. ROCM_ATTN does not claim batch invariance, so VLLM_BATCH_INVARIANT=1 never selects it. With high_precision=2 AITER requires the ultra-precision kernel instead of falling back to the lowest-precision one; pa_asm.csv for gfx942 and gfx950 has it (pa_bf16_pertokenFp8_gqa8_2tg_4w_uhp) for GQA up to 8 only, so larger ratios keep the HIP decode. rocm_aiter_ops.pa_fwd_asm gains a high_precision argument with the old default, so ROCM_AITER_FA is unchanged.
• CI. The new test runs in the MI300 and MI355 "Attention Kernels" steps, which collect tests/kernels/attention and depend on vllm/v1/attention and vllm/_aiter_ops.py.
• Related. [ROCm][Attention] Honor VLLM_ROCM_USE_AITER=1 by ranking ROCM_AITER_FA above ROCM_ATTN vllm-project/vllm#61332 (open) would rank ROCM_AITER_FA above ROCM_ATTN when AITER is enabled; the two do not conflict.

Unverified: gfx942 / MI300X outside CI (same kernels, cut 608), TP > 2, GQA 9–16 models (unchanged by design), models other than the three above, and AITER versions other than 0.1.27.

Not a duplicate: no open or recent PR adds an asm decode to ROCM_ATTN (searched rocm_attn pa_fwd_asm, ROCM_ATTN asm, ROCM_ATTN decode aiter, rocm_attn fp8 kv decode). AI assistance was used; the submitter reviewed every changed line and ran the tests above.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

Powered by Hyperloom, with human review, writing and testing.