[ROCm] Enable FP8 KV-cache and relax constraints for RDNA4 custom paged attention

[laudney]

Summary

Enable FP8 KV-cache support for the gfx12 (RDNA4) custom paged attention kernel via software dequantization, and fix several overly-restrictive constraints that prevented the Navi kernel from being used in valid configurations.

FP8 KV-cache halves the KV-cache memory footprint (16-bit → 8-bit per element), roughly doubling token capacity. On an AMD Radeon AI PRO R9700 (32GB), this increases KV-cache from ~132K tokens to ~264K tokens for a Qwen3-Coder-30B AWQ model — the difference between 1.4x and 2.75x concurrent 96K-context requests.

Kernel changes (csrc/rocm/attention.cu)

• convert_b8x8_to_b16x8(): Portable FP8 → FP16/BF16 vector conversion using HIP runtime fp8::vec_conversion (works across gfx11/gfx12 without arch-specific intrinsics)
• FP8 dequant wired into gfx12 QKV kernel: When kv_cache_dtype == "fp8", V-cache loads go through software dequant using per-block KV cache scale factors before the dot-product accumulation
• GQA ratio 1–2 support on gfx12: The gfx12 kernel correctly handles GQA ratios down to 1 (MHA), but the Python guard restricted to gqa_ratio >= 3 — now gqa_ratio >= 1

Platform guard fixes (vllm/platforms/rocm.py)

Guard	Before (bug)	After (fix)
_ON_GFX12	didn't exist	New flag to distinguish gfx12 (FP8-capable) from gfx11 (no FP8)
block_size	== 16 only	== 16 on gfx11, == 16 or == 32 on gfx12 (VBLOCKS_PER_LANE=1)
head_size	64 or 128 on Navi	128 only — kernel computes vhead_elem = warpid * 16 + lane16id without bounds check, reads past head slice for 64-wide heads
kv_cache_dtype	"auto" only	"auto" or "fp8"/"fp8_e4m3" on gfx12

Test changes (tests/kernels/attention/test_attention.py)

• FP8 KV-cache tests now run on supports_fp8() platforms instead of blanket-skipping all FP8 on Navi

Related PRs (RDNA4/gfx12 series)

• [ROCm] Enable wvSplitK skinny GEMM kernel for RDNA4/gfx1x decode #34709 — [ROCm] Enable wvSplitK/wvSplitKQ skinny GEMM kernels for RDNA4
• [ROCm] Add MXFP4 inline dequant Triton kernel for RDNA4/gfx12 #34632 — [ROCm] Add MXFP4 inline dequant Triton kernel for RDNA4/gfx12
• [ROCm] Make Whisper causal attention backend-agnostic #34631 — [ROCm] Make Whisper causal attention backend-agnostic

Test plan

• Qwen3-Coder-30B-A3B AWQ-4bit with kv-cache-dtype: fp8 on gfx1201 — server reports "Using fp8 data type to store kv cache", 263K tokens KV capacity, correct inference output
• Qwen3-14B-FP8 with kv-cache-dtype: fp8_e4m3 on gfx1201 — 150K tokens, correct output
• GPT-OSS-20B MXFP4 with kv-cache-dtype: fp8 on gfx1201 — correct output
• No regression on gfx11 (FP8 path gated behind _ON_GFX12)
• No regression on gfx9/MI-series (separate kernel under __HIP__GFX9__ guard)
• test_paged_attention FP8 variants pass on gfx12

[gemini-code-assist]

Code Review

This pull request introduces well-structured and robust changes to enable FP8 KV-cache support for RDNA4 (gfx12) and relaxes several constraints for the Navi kernel. The implementation of software dequantization in the attention kernel is clean and leverages if constexpr effectively to handle different data types. The refactoring of Q-loading logic to be independent of the cache data type improves correctness and readability. The Python-level guards have been appropriately updated to reflect the new capabilities and fixes. The test suite has also been correctly modified to include the new FP8 configurations on supported platforms. Overall, this is a high-quality contribution that significantly enhances performance and capabilities on ROCm hardware.

[mergify]

Hi @laudney, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[tjtanaa]

won't this affect the condition of gfx11x?

[tjtanaa]

The custom paged attention kernels implementation for gfx12x and gfx11x are different.

[tjtanaa]

I thought on gfx12 it supports block size of 32 as stated in

block_size_ok = block_size == 16 or (_ON_GFX12 and block_size == 32)

[laudney]

Rebased onto latest main and addressed review feedback:

Conflict resolution:

• rocm.py had a conflict in the navi else branch — resolved cleanly (upstream removed VLLM_ROCM_CUSTOM_PAGED_ATTN env var, so we dropped that line)

Fixes for review comments (new commit):

1. @tjtanaa's concern about gfx11 behavior — gqa_ratio >= 1 was incorrectly applied to all gfx1x. Now scoped: gqa_min = 1 if _ON_GFX12 else 3, so gfx11 retains >= 3.
2. block_size=32 test coverage — test skip now allows block_size=32 on gfx12 instead of unconditionally requiring 16.
3. k_scale/v_scale at 1.0 — kept as-is because the reference computation doesn't apply dequant scales. Added a NOTE comment documenting this limitation.

Pre-commit hooks all pass. All changes reviewed by GPT-5.4 (Codex) — no blockers found.