[AMD] Quantize AITER unified prefill queries to FP8

[vorapolsiloai]

Motivation

#38757 routes AITER extend/prefill through unified_attention. When the KV cache is FP8, it still passes Q in BF16. AITER then sees a BF16-Q / FP8-KV call, which has two effects:

• AITER does not select its tuned FP8-Q / FP8-KV prefill config ([Triton/Gluon] [Config] Enable unified-attention skip-mask for gfx950 hd256 FP8 prefill ROCm/aiter#5598).
• The gfx950 Gluon kernel is not used, because it requires Q and KV to have the same dtype ([Triton/Gluon] [GFX950] Add Unified Attention Gluon Kernel ROCm/aiter#4614).

The unified decode path already quantizes Q to FP8 with scaled_fp8_quant() and passes the matching q_descale. This PR makes prefill do the same.

Modifications

In _forward_extend_unified(), when kv_cache_dtype is FP8:

• Quantize Q to FP8 with the layer's K scale (layer.k_scale, or the backend default), as decode does.
• Pass the same scale as q_descale to unified_attention.

With a BF16 KV cache, nothing changes.

When this code runs: _forward_extend_unified() is used when SGLANG_USE_AITER_UNIFIED_ATTN=1 is set, and also, automatically, for models with a sliding-window KV pool. In both cases decode already uses FP8 Q, so this change gives prefill the same Q precision as decode. Routing does not change.

Kernel selection check

MI355X, Qwen3.8-27B (24 Q heads, 4 KV heads, head_dim=256), FP8 KV:

Stack	Q / KV dtype	AITER kernel config
Before this PR	BF16 / FP8	generic Triton D_GEQ_256 (BLOCK_M=16, tile 32)
This PR + ROCm/aiter#5598	FP8 / FP8	D_GEQ_256.Q_GEQ_256.DT_fp8_fp8 (BLOCK_M=128, tile 64)
This PR + ROCm/aiter#4614	FP8 / FP8	gfx950 Gluon 2D prefill and 3D decode

Accuracy Tests

• New unit tests in test_aiter_fp8_q_unified_attention.py:
  • FP8 KV: Q is quantized with the K scale, q_descale is passed, cu_seqlens_q is correct, and the output shape and dtype are unchanged.
  • BF16 KV: Q passes through unchanged and q_descale is None.
• The full test file passes on MI355X.
• Re-validated post-rebase on AMD Instinct MI355X (ROCm 10, AITER): all 5 unit tests in test/registered/attention/test_aiter_fp8_q_unified_attention.py pass cleanly.
• Deterministic chat-template checks give the same answer with BF16 Q, FP8 Q, [Triton/Gluon] [Config] Enable unified-attention skip-mask for gfx950 hd256 FP8 prefill ROCm/aiter#5598, and [Triton/Gluon] [GFX950] Add Unified Attention Gluon Kernel ROCm/aiter#4614. This covers a short prompt and a 60K-token needle prompt.

Speed Tests and Profiling

These numbers are for this PR together with ROCm/aiter#5598. This PR is what lets AITER pick the FP8 config, but the speedup comes from the AITER kernel.

MI355X, TP1, Qwen3.8-27B MXFP4, FP8 KV, --chunked-prefill-size 16384. Same request count and token count in every run.

Workload (input/output, concurrency)	Baseline (tok/s)	This PR + aiter#5598 (tok/s)	Change
8K/1K, c4	372.32	377.03	+1.3%
60K/600, c4	156.85	177.33	+13.1%
60K/600, c16	218.84	260.60	+19.1%
128K/1K, c4	105.84	132.14	+24.8%
128K/1K, c16	130.78	173.95	+33.0%

The baseline is the same setup without this PR or aiter#5598. With prefill-heavy requests (output length 8), median TTFT was 52% lower at 60K input and 65% lower at 128K input than the same stack with BF16 Q.

Checklist

• Format code with pre-commit.
• Add unit tests.
• Provide accuracy and speed benchmark results.

---

CI States

Latest PR Test (Base): ❌ Run #37836096134
Latest PR Test (Extra): ❌ Run #37836095849
Latest PR Test (AMD ROCm 10): ❌ Run #37836096212