[jhaotingc]

Overview

On SM90 with an FP8 KV cache, this PR runs attention on FA4's native FP8 forward from vllm-project/flash-attention#206, which computes attention on E4M3 Q/K/V with FP8 WGMMA:

• FA3 → FA4 upgrade. get_flash_attn_version upgrades FA3 to FA4 on SM90 when the KV cache is FP8 E4M3 and FA4 supports the head size: 64, 96, 128, 192 or 256, with the same V head size, no softcap and no batch invariance. An explicitly set flash_attn_version is kept.
• Gemma4 on FA4 for all layers. With an FP8 KV cache on SM90, main runs Gemma4's sliding-window layers on FA3 and its full-attention layers on FA4. This PR uses FA4 for all layers, as Gemma4 already does in other configurations:
  • Sliding-window layers (head dim 256) quantize Q to FP8, as FA3 does, and run FA4's native FP8 forward.
  • Full-attention layers (head dim 512) keep main's FA4 FP8-KV path with BF16 Q.

The FlashAttention pin is bumped once vllm-project/flash-attention#206 merges.

Claims

• Faster. On gemma-4-31B-it-FP8 on H200, output throughput is 1.10x main on the 8k-in-1k-out benchmark of gemma_4_31b_it_h200. This PR runs with [Cute,Fwd,Sm90] Native FP8 (E4M3 Q/K/V) forward for Hopper flash-attention#206, which is based on FlashAttention main. There the head-dim-512 FP8-KV kernel that both main and this PR use for the full-attention layers is about 40% slower than at vLLM's current pin, so this PR's number includes that slowdown and main's does not.
• Accuracy-neutral. GPQA Diamond, AIME 2025, MMLU-Pro and GSM8k all match main and BF16 within noise.

Validation

Serving. The 8k-in-1k-out benchmark of gemma_4_31b_it_h200 on one H200 at TP=1. Each row is the median of 3 runs on the same GPU; no request failed.

	Output tok/s	Mean TPOT (ms)	Median ITL (ms)
main	839.8	133.6	73.1
This PR	921.2 (1.10x)	123.5	60.7

Accuracy. gemma-4-31B-it-FP8 on 4x H200 at TP=1/DP=4. Sampling is temperature 1.0 and top_p 0.95, with thinking on.

	GPQA Diamond avg@16	AIME 2025 avg@64	MMLU-Pro	GSM8k
main	85.54 ± 0.70	86.93 ± 0.72	86.54 ± 0.31	94.01 ± 0.65
This PR	86.05 ± 0.63	87.45 ± 0.68	86.19 ± 0.31	94.31 ± 0.64
BF16 (model card)	85.42	87.29	86.44	94.62

After rebasing on main with #60372, this PR's GSM8k is 94.54 ± 0.63.

Reproduce

Accuracy server. The evals call it directly:

• GPQA Diamond and AIME 2025 use simple_evals. The GPQA answer regex also accepts Answer: $X. AIME is graded by an LLM judge.
• MMLU-Pro uses simple_evals with one pass.
• GSM8k uses lm-eval gsm8k_cot_zeroshot over chat completions, scored by flexible-extract. The system prompt asks for a final line of The answer is X.
• All of them sample with temperature 1.0 and top_p 0.95, with up to 131,072 output tokens.

vllm serve $MODEL --served-model-name gemma \
  --tensor-parallel-size 1 --data-parallel-size 4 \
  --max-model-len 262144 --max-num-batched-tokens 8192 --enable-chunked-prefill --max-num-seqs 16 \
  --trust-remote-code --kv-cache-dtype fp8 --attention-backend FLASH_ATTN --block-size 64 \
  --hf-overrides '{"text_config":{"use_bidirectional_attention":null}}' \
  --chat-template $THINKING_CHAT_TEMPLATE

---

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.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jhaotingc.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[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.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jhaotingc.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork