[ROCm][DSV4][Perf] Use FP8 WO_A output projection

[LiuYinfeng01]

Summary

Replace the ROCm DeepSeek V4 BF16 wo_a output-projection path on gfx950 with:

1. AITER inverse_rope_group_quant, which fuses inverse RoPE with per-token
   E8M0 group quantization.
2. AITER batched_gemm_a8w8_mxscale, which consumes the quantized activation
   and the checkpoint's native FP8 wo_a weight.

This removes the BF16 grouped einsum and keeps wo_a in its checkpoint FP8
form. Unsupported weights, scales, devices, and older AITER installations
retain the existing BF16 fallback.

Dependency

This fast path requires AITER >= 0.1.20, which added the MX-scale batched GEMM.
It therefore depends on #52826. Current main remains safe: the capability
probe logs once and falls back to BF16 when the op is unavailable.

Implementation

• Reuse the existing VLLM_ROCM_USE_AITER_FP8BMM gate and restrict the path to
  gfx950.
• Accept native float8_e8m0fnu/uint8 scales, or losslessly convert positive
  power-of-two FP32 scales to biased E8M0 exponent bytes.
• Validate the group-128 weight and scale layouts before enabling the path.
• Cache contiguous BF16 cosine/sine cache views once after weight loading.
• Preserve the existing rocm_inv_rope_einsum path as the fallback.

Performance

Measured on 8 x MI355X, TP1/PP8, one 100,000-token prefill request, 8K prefill
chunks, prefix caching disabled, three runs:

Path	TTFT runs (ms)	Median TTFT	Input throughput
OPUS sparse MLA + tuned GEMMs	2470.2 / 2457.9 / 2459.0	2459.0 ms	40,667.0 tok/s
+ FP8 wo_a	2288.8 / 2276.1 / 2278.4	2278.4 ms	43,891.0 tok/s

Delta: TTFT -7.35% and input throughput +7.93%.

TP8/PP1 prefill concurrency

Measured on 8 x MI355X with TP8/PP1, fixed 8,192-token prompts and one output
token, prefix caching disabled, and 10 requests per concurrency slot. Both arms
use the same image and differ only in the FP8 wo_a gate.

Max concurrency	BF16 wo_a total tok/s	FP8 wo_a total tok/s	Gain
1	16,218.62	16,345.65	+0.78%
2	23,337.44	23,782.97	+1.91%
4	23,618.92	24,004.20	+1.63%
8	23,621.81	24,051.77	+1.82%
16	23,623.62	24,069.49	+1.89%
32	23,624.64	24,069.36	+1.88%
48	23,621.17	24,069.38	+1.90%

All 1,110 requests in each arm completed without errors.

Revisions and runtime:

• vLLM PR commit: a9cf3bc3498eb71bb2d6f5d44c4b0049b62feb7a
• ROCm nightly base digest:
  sha256:cef549da00e0efaeadd9338ac8f351f2b96ff71a5ab8651a99bf989458bf1684
• AITER commit: a6d2b564fd671724a3720b8edf70e8d674e4d694
• FlyDSL: 0.3.1
• Image ID: sha256:0f851f268bc6bdb41c5ed22663a106339b3b269a88afbef17dbda156f2fb7893
• Common settings: TP8/PP1/DCP1/PCP1/EP1, FP8 KV cache, AITER MoE,
  Triton sparse prefill, 8K max batched tokens, no prefix cache
• A/B gates: VLLM_ROCM_USE_AITER_FP8BMM=0/1

The enabled-arm logs load module_inverse_rope_group_quant and
opus_bmm_a8w8_mxscale; neither appears in the disabled-arm logs.

TP1/PP8 prefill concurrency

Measured with the same fixed 8,192-token prompts, one output token, disabled
prefix caching, and 10 requests per concurrency slot. The image, revisions,
runtime settings, and A/B gate are identical to the TP8/PP1 measurements above;
only the topology changes to TP1/PP8.

Max concurrency	BF16 wo_a total tok/s	FP8 wo_a total tok/s	Gain
1	4,119.72	4,271.13	+3.68%
2	6,436.17	6,865.25	+6.67%
4	6,461.36	6,896.31	+6.73%
8	6,460.00	6,891.96	+6.69%
16	6,457.00	6,891.06	+6.72%
32	6,459.55	6,889.21	+6.65%
48	6,460.06	6,889.74	+6.65%

All 1,110 requests in each arm completed without errors. The enabled-arm logs
load module_inverse_rope_group_quant and opus_bmm_a8w8_mxscale on all eight
pipeline stages; neither appears in the disabled-arm logs.

Correctness

Greedy GSM8K, zero invalid responses:

• 100-question gate: 93/100
• Full 1319 questions: 1256/1319 = 95.22%
• Clean fixed-stack reference: 1249/1319 = 94.69%

TP8/PP1 GSM8K

Topology-matched greedy 5-shot evaluation on the TP8 stack above:

Path	Accuracy	Correct	Invalid
BF16 wo_a	93.9348%	1239/1319	0
FP8 wo_a	94.9204%	1252/1319	1

The FP8 candidate first passed the 100-question gate at 93/100 with zero
invalid responses. Full-set accuracy was +0.99 percentage points versus the
topology-matched BF16 reference; one response did not contain a parseable
answer.

Validation

• git diff --check: passed
• pre-commit hooks for both changed files: passed
• tests/models/test_deepseek_v4_rocm_wo_a.py: 7 passed
• TP1/PP8 100K prefill: 3/3 requests completed
• GSM8K-1319: zero invalid responses

Related work

#45103 fused inverse RoPE and cached a BF16 wo_a; this PR retains that path as
fallback and uses the native FP8 weight on supported gfx950 configurations.

AI assistance

Cursor assisted with implementation, testing, benchmark analysis, and drafting.
The human submitter reviewed the resulting change and is responsible for it.

---

Re-measurement on the older DeepSeek-V4-Pro checkpoint

Scope note. Everything above this heading was measured on the
0813 checkpoint. Everything below was re-measured on the older
DeepSeek-V4-Pro snapshot
(models--deepseek-ai--DeepSeek-V4-Pro/snapshots/b5968e9190ef611bbf34a7229255be88a0e937c1,
staged locally as DeepSeek-V4-Pro-old). The two sets are not
interchangeable; they are reported side by side so the gate can be judged on
both checkpoints. Nothing above has been edited.

Setup

Item	Value
Hardware	8 x MI355X (gfx950), single node
Image	vllm/vllm-openai-rocm:nightly-d9105ea8001e0a6d77a96327d17515bb5791fb36
vLLM	0.28.1rc1.dev472+gd9105ea80 + this PR applied to vllm/models/deepseek_v4/amd/rocm.py
ROCm / torch	7.2.3 / 2.12.0+git6bbd260
Model	DeepSeek-V4-Pro (older snapshot), FP4
A/B gate	VLLM_ROCM_USE_AITER_FP8BMM=0 vs =1, same image and same patched file
Tuned GEMMs	AITER_CONFIG_GEMM_A8W8_BLOCKSCALE pointed at a merged a8w8-blockscale tuning table
Workload	Identical to the tables above: fixed 8,192-token prompts, one output token, prefix caching disabled, 10 requests per concurrency slot

Fast-path confirmation: the enabled arm's server log loads
module_inverse_rope_group_quant and the mxscale BMM; the disabled arm's log
contains neither (verified per run).

TP8/PP1 prefill concurrency

Max concurrency	BF16 wo_a total tok/s	FP8 wo_a total tok/s	Gain
1	15,544.93	15,623.10	+0.50%
2	20,474.39	21,122.52	+3.17%
4	20,883.86	20,965.92	+0.39%
8	20,963.71	21,157.26	+0.92%
16	20,959.80	21,099.28	+0.67%
32	20,845.46	21,224.69	+1.82%
48	20,848.19	21,178.18	+1.58%

Mean gain across the sweep: +1.29%. All 1,110 requests per arm completed successfully.

TP1/PP8 prefill concurrency

Max concurrency	BF16 wo_a total tok/s	FP8 wo_a total tok/s	Gain
1	4,228.93	4,403.25	+4.12%
2	6,179.04	6,528.63	+5.66%
4	6,176.99	6,595.96	+6.78%
8	6,189.02	6,591.95	+6.51%
16	6,188.72	6,591.11	+6.50%
32	6,190.38	6,591.60	+6.48%
48	6,190.03	6,596.08	+6.56%

Mean gain across the sweep: +6.09%. All 1,110 requests per arm completed successfully.

TP1/PP8 single 100,000-token prefill

Three runs per arm, 8K prefill chunks, prefix caching disabled. The first run of
each arm carries a cold-start spike (~13.8 s on both arms) and is excluded from
the median.

Path	TTFT runs (ms)	Median TTFT	Input throughput
BF16 wo_a	13878.3 / 4309.2 / 4328.5	4328.5 ms	23,100.5 tok/s
FP8 wo_a	13757.7 / 4096.0 / 4203.3	4203.3 ms	23,789.1 tok/s

Delta: TTFT -2.89%, input throughput +2.98%.

GSM8K (full 1,319, greedy, 5-shot)

Addresses the request for a full-set run. Both arms use the same image, the same
patched file, and the same topology (TP8/PP1, DP attention, FP8 KV cache); only
the gate differs.

Path	flexible-extract	strict-match
BF16 wo_a	0.9621 ± 0.0053	0.9629 ± 0.0052
FP8 wo_a	0.9651 ± 0.0051	0.9651 ± 0.0051

The FP8 arm scores 0.9651, i.e. above the 0.96 bar, and +0.30 pp
(flexible-extract) over the topology-matched BF16 reference on this checkpoint.
Both arms exited 0 with zero invalid responses.

Summary on this checkpoint

Measurement	Result
TP8/PP1 throughput	+1.29% mean (+0.39% to +3.17%)
TP1/PP8 throughput	+6.09% mean (+4.12% to +6.78%)
100K prefill TTFT	-2.89%
GSM8K full 1319	0.9651 vs 0.9621 BF16

The direction and magnitude match the 0813 numbers above: a small gain at
TP8/PP1 and a larger one at TP1/PP8, with accuracy at or above the BF16
reference.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[Fangzhou-Ai]

please do a full 1319 gsm8k test, current result seems to be lower than the baseline. We expect the full test score near 0.96

[LiuYinfeng01]

please do a full 1319 gsm8k test, current result seems to be lower than the baseline. We expect the full test score near 0.96

Path	flexible-extract	strict-match
BF16 wo_a	0.9621 ± 0.0053	0.9629 ± 0.0052
FP8 wo_a	0.9651 ± 0.0051	0.9651 ± 0.0051