[ROCm][Perf][GLM-5.3-Flash] Enable AITER sparse MLA on gfx942, including FP8 KV

[jin-amd]

Overview

Lets VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA (#53492) take effect on gfx942 (MI300X/MI325X) for rope-free sparse MLA, GLM-5.3-Flash's layout, with a BF16 or FP8 KV cache. It needs ROCm/aiter#5721 (merged, not yet in an AITER release) and, for FP8, ROCm/aiter#6199 (in review), so it stays a draft until those ship.

Claims

• On gfx942, GLM-5.3-Flash's sparse MLA can run on AITER's sparse_mla_fwd instead of vLLM's Triton kernel. On our MI325X image, calling the same kernel with the same arguments: output throughput +7.1% (7/7 concurrencies), TPOT 5-17% lower, decode attention 7.7x faster per call, gsm8k unchanged.
• --kv-cache-dtype fp8 runs on that kernel too, with FP8 dots. Today this combination reaches AITER's asm MLA decode on gfx942, which is built for 576-wide rows and faults. On our image: +6.9% throughput against BF16 KV, KV capacity 12.2M to 23.1M tokens, gsm8k and needle-in-a-haystack unchanged.
• Nothing changes when the flag is off, on gfx950, or for DeepSeek V4/V4.1. With an AITER that lacks the gfx942 kernel, such as today's v0.1.24.post1 pin, the flag is ignored with a warning, as it is now.

Validation

Unit tests, on MI325X (gfx942): pytest tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py, with this change applied to vllm/vllm-openai-rocm:nightly-43b4aaea3e40e20ef53ce2bfecfa1b72060bf2b5. 17 passed. The new cases cover:

• the gfx942 opt-in, plus the warning and fallback when AITER has no gfx942 kernel;
• the gfx942 scope: rope-free only, and FP8 only when AITER can read it;
• forward_mqa with an FP8 cache on gfx942: q is passed unquantized, with dot_precision="fp8".

pre-commit (ruff, mypy 3.12 and the local hooks) is clean on the changed files.

End to end: not yet run on this branch. No AITER release has #5721 yet, and #6199 is in review. The numbers below come from our GLM-5.3-Flash MI325X images. Their vLLM calls the same kernel with the same arguments through its own dispatch, on AITER 0.1.22.post1 with #5721 and #6199 applied. I'll rerun them on this branch once an AITER build has both PRs.

Setup: GLM-5.3-Flash FP8, TP4 on 4x MI325X, vllm bench serve random 131072 in / 1024 out, 20 prompts per concurrency, medians, no MTP. gsm8k is 5-shot over chat completions, all 1319 questions.

AITER sparse MLA vs vLLM's Triton kernel, BF16 KV (gsm8k strict 0.9719 vs 0.9712, TTFT flat):

conc	tok/s Triton	tok/s AITER	TPOT Triton (ms)	TPOT AITER (ms)
2	80.92	91.36 (+12.9%)	16.57	13.82
4	110.96	121.09 (+9.1%)	22.57	19.51
6	116.93	125.69 (+7.5%)	32.26	28.66
8	129.61	137.74 (+6.3%)	39.59	37.05
10	138.99	145.75 (+4.9%)	49.40	46.63
12	141.29	147.94 (+4.7%)	54.16	51.51
16	143.24	149.89 (+4.6%)	63.89	60.74

FP8 KV vs BF16 KV, both on AITER sparse MLA, geomean over concurrency 2-16:

	BF16 KV	FP8 KV, FP8 dots
Output tok/s	155.1	165.7 (+6.9%)
TPOT	28.1 ms	26.4 ms
gsm8k strict (±0.005)	0.974	0.973
Needle-in-a-haystack, 4K-128K	60/60	60/60
KV cache capacity	12.2M tokens	23.1M tokens

To reproduce on this branch, once AITER has #5721 and #6199 (drop --kv-cache-dtype fp8 for BF16 KV):

VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1 \
vllm serve zai-org/GLM-5.3-Flash --tensor-parallel-size 4 --kv-cache-dtype fp8 \
  --gpu-memory-utilization 0.85 --max-model-len 262144 \
  --no-enable-prefix-caching --trust-remote-code

Details

• Why the opt-in is per caller. is_triton_sparse_mla_enabled() also gates DeepSeek V4/V4.1, whose packed caches are gfx950-only in AITER. A gfx942_ok keyword lets only ROCM_AITER_MLA_SPARSE accept gfx942. It still needs the flag, and it checks that the installed AITER's sparse_mla lists gfx942 in SUPPORTED_ARCHS.
• Rope-free only on gfx942. Layouts with RoPE (DeepSeek V3.2, GLM-5.x) keep AITER's asm MLA decode on gfx942; enabling the kernel there needs its own measurement. _aiter_sparse_mla_unsupported_reason reports this, and the backend falls back.
• FP8 KV on gfx942.
  • AITER's gfx942 kernel reads the cache as e4m3fnuz but takes no fp8 q. So forward_mqa skips quantizing q, _forward_mla_aiter passes dot_precision="fp8", and the kernel quantizes q per tile. gfx950 keeps the pre-quantized q.
  • triton_sparse_mla_fwd gains an optional dot_precision. When it is omitted, the precision is chosen from q's dtype as before.
  • The backend requires an AITER that lists gfx942 in FP8_SCALAR_ARCHS, which [Triton/Gluon] [gfx942] Support fp8 KV caches and fp8 dots in sparse_mla_fwd on gfx942 ROCm/aiter#6199 adds; otherwise it warns and falls back.
• Numerics. FP8 dots put the kernel 3.4-4.5% off an FP32 reference (relative L2), almost all of it from quantizing q; BF16 dots are at 0.2-0.3%. This didn't move gsm8k or needle-in-a-haystack. I haven't exposed a BF16-dots option, to avoid a new env var, but can add one.
• Relation to [ROCm][Bugfix] Support FP8 KV cache for NoPE sparse MLA #57134. [ROCm][Bugfix] Support FP8 KV cache for NoPE sparse MLA #57134 serves rope-free FP8 KV on vLLM's Triton kernel when this flag is off. Both PRs change the q-quantization branch in forward_mqa, so whichever lands second rebases.
• Limitations.
  • Attention sinks are passed through but not exercised on gfx942, since GLM-5.3-Flash has none.
  • No CI job sets this flag on gfx942. The unit tests above cover the dispatch, and AITER's gfx942 op tests cover the kernel.

---

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.