[Kernel] Integrate FA4 SM90 FP8 MMA for Gemma-4 full attention by katec846 · Pull Request #60372 · vllm-project/vllm

katec846 · 2026-10-07T08:53:10Z

Overview

Enable FP8 MMA in the SM90 FA4 full-attention kernel (PR199) for Gemma-4 while retaining FP8 KV-cache storage.

Gemma-4 uses head dimension 512 for full attention and 256 for sliding-window attention. This change compares FP8 MMA against the existing path that reads FP8 KV, dequantizes it, and performs FP16 MMA.

On H200, FP8 full attention uses FA4, while FP8 sliding-window attention continues to use FA3. FA4 FP8 sliding-window attention with head dimension 256 is currently unsupported.

Claims

Integrate the SM90 d512 FP8-MMA kernel while preserving native FP16/BF16 queries and supported FA3 sliding-window routing.
Gate unsupported kernel combinations and normalize descale layouts.
Allow FA4 full-attention layers to select splits automatically, removing the previous CUDA-graph cap of 32 for these layers; retain the cap for FA3.
Report measured gains over the KV-dequantization baseline and the remaining regression against BF16 KV/MMA.

Validation

The MMA labels describe the full-attention path. Both FP8-KV configurations use FA3 FP8 sliding-window attention; the BF16-KV reference uses FA4 BF16 sliding-window attention.

Serving setup:

Gemma-4-31B-it-FP8 weights for all configurations.
Four H200 GPUs on one node; TP4.
Approximately 25K shared prefix plus 3K unique input tokens.
128 requests, 200 output tokens per request.
MTP enabled with three speculative tokens.
Prefix-cache reuse enabled.
Concurrency 1–64; median of three rounds.

Serving performance

The 8k-in-1k-out benchmark of gemma_4_31b_it_h200 on one H200 at TP=1, with GPU clocks locked. Each number is the median of 3 repetitions, as perf-eval aggregates them; no request failed.

	Output tok/s	Mean TPOT (ms)	Median ITL (ms)
Before (`5d35ce3934` + FA `b273a02`)	796.5	126.5	75.1
This PR + vllm-project/flash-attention#199	831.1 (1.04x)	123.2	69.3
Delta	+4.3%	-2.6%	-7.8%

Reproduce

Before is this PR's base, 5d35ce3934, with FlashAttention at Add SM90 FP8 MMA forward for head dimension 512 flash-attention#199's base, b273a02. After is this PR's head, 237043eee6, with Add support for H100 #199's head, 2648f1b.
Neither PR changes C++. Both runs use one editable build of 237043eee6 with VLLM_FLASH_ATTN_SRC_DIR set to the FlashAttention checkout, and differ only in Python sources.
The model is a local copy of RedHatAI/gemma-4-31B-it-FP8-dynamic. Server and benchmark match perf-eval@5f4e445 (lib/server.sh, lib/run_vllm_bench.sh, lib/aggregate_perf.py).

VLLM_DEEP_GEMM_WARMUP=skip VLLM_ENGINE_READY_TIMEOUT_S=3600 \
vllm serve $MODEL --port 8000 \
  --tensor-parallel-size 1 --kv-cache-dtype fp8 \
  --tool-call-parser gemma4 --reasoning-parser gemma4 --enable-auto-tool-choice

# 3 repetitions against the same server; median of each metric
vllm bench serve --backend openai --base-url http://127.0.0.1:8000 --model $MODEL \
  --dataset-name random --random-input-len 8192 --random-output-len 1024 --ignore-eos \
  --num-prompts 512 --max-concurrency 128 --num-warmups 128

Model evaluation

Accuracy scores are percentages. 87.19% for AIME 2025, 85.73% for GPQA Diamond (avg@16), and 86.86% for MMLU-Pro are measured in this PR; the two reference columns are reported baselines from PR #48666.

Benchmark	FP8 KV + FP8 MMA (PR199)	FP8 KV + FP16 MMA (reference FA4)	BF16 TRITON
AIME 2025 (avg@64)	87.19	86.88	87.29
GPQA Diamond (avg@16)	85.73*	86.55	85.42
MMLU-Pro	86.86	86.64	86.44

* Note that in a paird GPQA Diamond comparison using the same 198 questions and 16 seeds, FP8 MMA scored 85.73% versus 85.32% for the FP8 KV + FP16 MMA. The difference was +0.41 percentage points, with a paired question-bootstrap 95% interval of [−0.51, +1.36] points. Both runs had zero invalid or truncated answers. The table’s 86.55% reference is imported from PR #48666 and is separate from this matched control.

MTP speculative decoding

Mean acceptance length (AL), median of three rounds, from the same sweep as the serving-performance table above. All configurations use FP8 model weights, TP4 on four H200 GPUs, three speculative tokens, and prefix-cache reuse enabled. Each measurement completed 128/128 requests without errors. AL includes the bonus token emitted per speculative step.

Configuration	cc1	cc2	cc4	cc8	cc16	cc32	cc64
FP8 KV + FP8 MMA	2.9830	2.9932	3.0002	2.9640	2.9849	2.9860	2.9793
FP8 KV + FP16 MMA	2.9924	2.9691	3.0130	2.9995	2.9993	2.9857	2.9979
BF16 KV + BF16 MMA	3.0053	2.9823	2.9940	2.9935	2.9995	3.0002	2.9982

Reproduction commands

The Hopper FP8 MMA path supports text-only workloads. The commands use --hf-overrides '{"text_config":{"use_bidirectional_attention":null}}' to disable Gemma-4’s bidirectional/mm-prefix attention. Multimodal Gemma-4 with FP8 KV uses Triton instead of this kernel. FP8 sliding-window layers use FA3, so do not force FA4 globally.

Serve and benchmark: 25K shared prefix, TP4, MTP3, cache reuse on

FP8 KV + FP8 MMA (this PR):

vllm serve "$MODEL" \
  --max-num-seqs 64 --max-num-batched-tokens 8192 --max-model-len 32768 \
  --host 127.0.0.1 --port 8000 --tensor-parallel-size 4 \
  --gpu-memory-utilization 0.9 --trust-remote-code --async-scheduling \
  --enable-prefix-caching --attention-backend FLASH_ATTN \
  --max-cudagraph-capture-size 256 --kv-cache-dtype fp8 \
  --compilation-config '{"use_inductor_graph_partition":1,"pass_config":{"eliminate_noops":true,"fuse_allreduce_rms":true}}' \
  --hf-overrides '{"text_config":{"use_bidirectional_attention":null}}' \
  --speculative-config "{\"model\":\"$ASSISTANT\",\"num_speculative_tokens\":3}"

FP8 KV + FP16 MMA (PR #164 baseline kernel environment):
(The two FP8-KV serving commands are identical because MMA selection depends on the FlashAttention kernel source, not a CLI flag. For the control, we kept the same vLLM integration but replaced its vendored CuTe sources with the PR #164 dequantizing kernel, including the baseline safety fixes, then restarted the server. --kv-cache-dtype fp8 selects FP8 cache storage in both cases)

vllm serve "$MODEL" \
  --max-num-seqs 64 --max-num-batched-tokens 8192 --max-model-len 32768 \
  --host 127.0.0.1 --port 8000 --tensor-parallel-size 4 \
  --gpu-memory-utilization 0.9 --trust-remote-code --async-scheduling \
  --enable-prefix-caching --attention-backend FLASH_ATTN \
  --max-cudagraph-capture-size 256 --kv-cache-dtype fp8 \
  --compilation-config '{"use_inductor_graph_partition":1,"pass_config":{"eliminate_noops":true,"fuse_allreduce_rms":true}}' \
  --hf-overrides '{"text_config":{"use_bidirectional_attention":null}}' \
  --speculative-config "{\"model\":\"$ASSISTANT\",\"num_speculative_tokens\":3}"

BF16 KV + BF16 MMA:

vllm serve "$MODEL" \
  --max-num-seqs 64 --max-num-batched-tokens 8192 --max-model-len 32768 \
  --host 127.0.0.1 --port 8000 --tensor-parallel-size 4 \
  --gpu-memory-utilization 0.9 --trust-remote-code --async-scheduling \
  --enable-prefix-caching --attention-backend FLASH_ATTN \
  --max-cudagraph-capture-size 256 --kv-cache-dtype bfloat16 \
  --compilation-config '{"use_inductor_graph_partition":1,"pass_config":{"eliminate_noops":true,"fuse_allreduce_rms":true}}' \
  --hf-overrides '{"text_config":{"use_bidirectional_attention":null}}' \
  --speculative-config "{\"model\":\"$ASSISTANT\",\"num_speculative_tokens\":3}"

These recorded runs left the FA version unset. The dequantized control uses the PR #164 baseline kernel environment.

Benchmark:

aiperf profile \
  --model "$MODEL" --tokenizer "$MODEL" --tokenizer-trust-remote-code \
  --endpoint-type chat --endpoint /v1/chat/completions --streaming \
  --use-legacy-max-tokens --url http://127.0.0.1:8000 \
  --input-file "$DATASET" \
  --extra-inputs max_tokens:200 --extra-inputs min_tokens:200 \
  --extra-inputs ignore_eos:true --extra-inputs skip_special_tokens:false \
  --random-seed 1234 --use-server-token-count \
  --concurrency 16 --request-count 128 --num-dataset-entries 128

Run at concurrency 1, 2, 4, 8, 16, 32, and 64; report the median of three rounds.

Evaluation commands

Evaluation server on four H200 GPUs. AIME and GPQA use --max-num-seqs 16; MMLU-Pro uses 32. No MTP.

vllm serve "$MODEL" --served-model-name gemma \
  --host 0.0.0.0 --port 8000 \
  --tensor-parallel-size 1 --data-parallel-size 4 \
  --max-model-len 262144 --max-num-batched-tokens 8192 \
  --enable-chunked-prefill --max-num-seqs 16 \
  --kv-cache-memory-bytes 101716968203 --enforce-eager \
  --trust-remote-code --kv-cache-dtype fp8 \
  --attention-backend FLASH_ATTN \
  --attention-config.flash_attn_version=3 --block-size 64 \
  --kernel-config.enable_flashinfer_autotune=false \
  --kernel-config.enable_cutedsl_warmup=false \
  --hf-overrides '{"text_config":{"use_bidirectional_attention":null}}'

AIME 2025, avg@64

Use aime.yaml with 64 samples per problem, temperature 1.0, top-p 0.95, and the GPT-4.1 judge configuration. Export INFERENCE_API_KEY.

docker run --rm --network host --shm-size=100g \
  -e DUMMY_API_KEY=dummy -e INFERENCE_API_KEY \
  -v "$PWD/aime.yaml:/config.yaml:ro" \
  -v "$PWD/aime-results:/results" \
  --entrypoint bash nvcr.io/nvidia/eval-factory/simple-evals:26.03 \
  -lc 'runner=$(command -v nemo-evaluator || command -v eval-factory); exec "$runner" run_eval --run_config /config.yaml'

GPQA Diamond, avg@16

Use gpqa.yaml with all 198 questions, the eval/aai/mcq-4choices prompt, temperature 1.0, and top-p 0.95. Run seeds 0–15.

pids=()
for seed in $(seq 0 15); do
  docker run --rm --network host --shm-size=4g \
    -e DUMMY_API_KEY=dummy -e HF_TOKEN \
    -e NEMO_EVALUATOR_TELEMETRY_LEVEL=0 -e OMP_NUM_THREADS=1 \
    -v "$PWD/gpqa.yaml:/config.yaml:ro" \
    -v "$PWD/gpqa-results/seed-$seed:/results" \
    --entrypoint bash "$EVAL_IMAGE" \
    -lc 'runner=$(command -v nemo-evaluator || command -v eval-factory); exec "$runner" run_eval --run_config /config.yaml "$@"' \
    -- --overrides "config.params.extra.args=++prompt_config=eval/aai/mcq-4choices ++inference.random_seed=$seed" &
  pids+=("$!")
done

status=0
for pid in "${pids[@]}"; do
  wait "$pid" || status=1
done
test "$status" -eq 0

MMLU-Pro

Restart the evaluation server with --max-num-seqs 32. Use mmlu.yaml with all 12,032 questions, one sample per question, the eval/aai/mcq-10choices-boxed prompt, temperature 1.0, and top-p 0.95.

mkdir -p mmlu-results
chmod 777 mmlu-results

docker run --rm --network host --shm-size=4g \
  -e DUMMY_API_KEY=dummy -e HF_TOKEN \
  -e NEMO_EVALUATOR_TELEMETRY_LEVEL=0 -e OMP_NUM_THREADS=1 \
  -v "$PWD/mmlu.yaml:/config.yaml:ro" \
  -v "$PWD/mmlu-results:/results" \
  --entrypoint bash "$EVAL_IMAGE" \
  -lc 'runner=$(command -v nemo-evaluator || command -v eval-factory); exec "$runner" run_eval --run_config /config.yaml'

Details

Why FP8 is slower than BF16

Nsight shows that faster FP8 full attention does not necessarily produce faster end-to-end execution.

Sliding-window attention and KV-cache operations offset full-attention gains

The profiling setup uses 25K inputs, TP4, MTP enabled, prefix reuse enabled, and offered concurrency 64.

Target-model kernel work	FP8 KV + FP8 MMA	BF16 KV + BF16 MMA	FP8 time reduction vs. BF16
Full attention	6.681 ms	7.584 ms	+11.9%
Sliding-window attention	5.482 ms	4.206 ms	−30.3%
KV-cache writes	0.261 ms	0.157 ms	−66.2%
Attention preparation	0.105 ms	0	N/A
Whole-cycle GPU span	29.879 ms	29.391 ms	−1.7%

Time reduction is 100 × (BF16 time − FP8 time) / BF16 time, calculated from the displayed times. Positive values mean less time (faster); negative values mean more time (slower).

Full attention is approximately 12% faster, but sliding-window attention is approximately 30% slower, with additional KV-write, preparation, and combine work.

KV-write overhead refers to the kernels that reshape and write newly generated keys and values into the paged cache, including the FP8 conversion/quantization path. Attention reads and prefix-cache reuse accounting are separate from this kernel category.

Reuse-off profiles show the same pattern.
The following comparison uses cc16, 25K inputs, TP4, MTP enabled, and prefix-cache reuse disabled. These matched pure-decode cycles process 16 active requests per target step. Values are medians across four GPU ranks of each rank's per-cycle medians, with 69 matched cycles per rank for BF16 and 87 for FP8.

Target-model kernel work	FP8 KV + FP8 MMA	BF16 KV + BF16 MMA	FP8 time reduction vs. BF16
Full attention	1.9670 ms	2.4358 ms	+19.2%
Sliding-window attention	1.9593 ms	1.3867 ms	−41.3%
KV-cache writes	0.2551 ms	0.1413 ms	−80.5%
Attention preparation	0.1098 ms	0	N/A
Whole-cycle GPU span	14.6594 ms	14.2903 ms	−2.6%

Full attention is faster, but SWA, KV-cache writes, and attention preparation/combine offset the gain. The whole-cycle GPU span is 2.6% longer for FP8. This is a profiling comparison, not an unprofiled end-to-end throughput measurement.

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.

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.

jhaotingc · 2026-10-14T23:53:16Z

          vllm-flash-attn
          GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG b273a02621d3c94bc747df67204f591b67f38fae
+          GIT_TAG 6d71c3f914eb7943f49647ea4828c6b2fe9176c6


note: this needs to change when FA PR lands

jhaotingc

Some general comments about: why there are more restriction to use this kernel after switching to FP8 MMA "within the kernel".
I'm curious about whether these restrictions will be traversed by code for actual use case (Gemma4)
Test seems too extensive.

jhaotingc · 2026-10-15T00:00:47Z

+    # specialization supports equal d512 heads, including sinks and dynamic causal.
    sm90_fp8_kv_supported = fa_version == 3 or (
-        fa_version == 4 and head_size in (None, 512)
+        fa_version == 4 and head_size in (None, 512) and head_size_v in (None, 512)


Is there a reason why this restriction needs to be stricken? or it's just missing from previous PR.
Will removing this change breaks Gemma4 + FP8 KV + FP8 MMA for hd 512?

This check was missing from the previous vLLM integration. PR164 already requires equal K/V head dimensions, but the vLLM support check only checked the Q/K dimension.

Removing this addition would not break Gemma-4 with FP8 KV + FP8 MMA, since its full-attention Q/K/V dimensions are all 512. It would, however, allow unsupported Diff-KV shapes through backend selection.

jhaotingc · 2026-10-15T00:03:09Z

+        if sm90_fp8_kv and head_size == 512:
+            # Framework blocks can be split into the advertised 64-token pages.
+            if block_size is not None and block_size % 64 != 0:
+                return "SM90 d512 FP8 KV requires a block size divisible by 64"
+            if use_sparse:
+                return "SM90 d512 FP8 KV does not support block sparsity"


Are these restrictions specific to after your change? Flash attention PR 164 does not need these restrictions?

The page-size requirement already existed in PR 164 (page_size == tile_n, which is 64 for d512). The vLLM check ensures framework blocks can be split into those kernel pages.

Block-sparsity rejection is specific to the current FP8 MMA implementation. Its producer traverses consecutive KV blocks, and its buffer/barrier schedule is not integrated with sparse block-list traversal. Neither restriction rejects Gemma-4’s standard text workload.

No need to create a sm90_fp8_kv if not used elsewhere, just fold in if statement.

jhaotingc · 2026-10-15T16:59:28Z

+        if sm90_fp8_kv and head_size == 512:
+            # Framework blocks can be split into the advertised 64-token pages.
+            if block_size is not None and block_size % 64 != 0:
+                return "SM90 d512 FP8 KV requires a block size divisible by 64"
+            if use_sparse:
+                return "SM90 d512 FP8 KV does not support block sparsity"


No need to create a sm90_fp8_kv if not used elsewhere, just fold in if statement.

simon-veitner-redhat · 2026-10-21T09:56:08Z

Thanks @katec846, this is a nice, compact integration. Overall the vLLM wiring looks good to me.

Before merge: GIT_TAG in cmake/external_projects/vllm_flash_attn.cmake points at 6d71c3f9, the head of FA PR 199. That commit is only reachable as refs/pull/199/head in vllm-project/flash-attention, so source builds without VLLM_FLASH_ATTN_SRC_DIR can't fetch it. So this should wait until FA PR 199 is merged and the tag is re-pinned to the merged commit on main (a squash merge gives a new SHA).

Fixes you can cherry-pick

I put a few small fixes on a branch on top of your head (5446e6c): https://github.com/simon-veitner-redhat/vllm/tree/pr60372-review-fixes

Each commit applies on its own, except the test, which expects the first two.

5b06e28 Remove the FA4 -> FA3 FP8 fallback in flash_attn_varlen_func, because no in-tree caller can reach it. Every caller that can pass fa_version=4 with an FP8 K on SM90 is gated earlier to q/k/v head_size 512 with BF16/FP16 Q, so it always takes the FP8 MMA path. The main gates are the FP8 KV check in FlashAttentionImpl.__init__ and the DiffKV init check, plus supports_quant_query_input=False on this path. The wrapper now raises a NotImplementedError that names the q/k/v dtypes and head sizes. The commit message lists every caller with file:line evidence.
43475aa Reject DCP with FA4 SM90 FP8 KV at impl init. DCP context attention is non-causal, so today a DCP server starts and then crashes on the first prefill with cached context.
8cb6afa CPU selector test for the new routing: SM90 + FP8 + d512 goes to FA4 FP8 MMA, while d256, BF16 KV, non-SM90 and DCP get the expected fallback or error. It also checks that the wrapper raises for FA4 + FP8 K below d512 without calling FA3. It runs in the existing V1 Attention step.
f52595b GSM8K config for Gemma-4-31B-it-FP8 with --kv-cache-dtype fp8 (TP2, text-only override) plus a 2x H100 lm-eval step, because nothing in CI covers this path yet. It measured 0.970 on 2x H200, and the threshold is 0.96.

PR description

Could you update the description with the following?

Exact commands for serve, bench and evals, including --hf-overrides '{"text_config":{"use_bidirectional_attention":null}}'. Without it the FP8 MMA kernel doesn't run.
The text-only limitation. Gemma-4 is an mm-prefix LM by default, and on Hopper FLASH_ATTN with FP8 KV rejects mm-prefix, so multimodal Gemma-4 with FP8 KV uses Triton, not this kernel.
Numbers that disagree:
- The text says AIME 2025 is 85.78, but the table shows 87.19.
- The stated ranges don't match the table. The table gives +0.2% to +2.6% over the dequant control (text: 0.5-2.2%) and -2.1% to -6.9% against BF16 (text: 3.7-9.8%).
- GPQA is 1.39 points below the dequant reference (85.16 vs 86.55). A stderr or a rerun would help, and so would saying where the GPQA and MMLU-Pro values for this PR come from.
- The control is labeled "FP16 MMA" in one table and "BF16 MMA" in the others.
The aot_schedule change (flash_attn.py:717) is a behavior change, not a "preserve". On main with FP8 KV, the FA4 d512 layers were capped at 32 splits under CUDA graphs. Now FA4 picks the split count, and flash_attn_max_num_splits_for_cuda_graph no longer applies to these layers. Most of the high-concurrency gain comes from this change (numbers below), so it should get the credit in the performance section.
The checklist boxes. The commits are co-authored by Codex, but the "AI assistance was used" and /pr-checklist boxes are unchecked.

Our independent results

All runs used 4x H200 with RedHatAI/gemma-4-31B-it-FP8-dynamic and the text-only override. We compared main (FA dequant kernel) against this PR (FA PR 199), and the FA4 path was confirmed with a call tracer.

GSM8K (lm-eval, 5-shot, chat template, 1319 questions; stderr is about 0.5). With FP8 KV: 95.68 with this PR vs 96.44 on main. With BF16 KV: 96.13 on both. The repo's own gsm8k harness gives 0.970 on the branch above.
Single-needle NIAH from 4k to 248k tokens: 70/70 for every config.
Throughput (TP4, no MTP, output tok/s, this PR vs main with FP8 KV):

Workload cc 1 high cc

prefill 16k/64 +1.0% +0.9% (cc32)

prefill 32k/64 +1.8% +1.2% (cc8)

mixed 4k/512 -0.4% +5.7% (cc128)

decode 1k/1k -0.7% +11.5% (cc128)

FP8 KV is still 1-10% slower than BF16 KV, in line with what you report.

Most of the high-concurrency gain comes from removing the 32-split cap. We checked this end to end by running this PR's vLLM code with main's FA, which keeps the dequant kernel but drops the cap. On decode at cc32 / cc128 that gives +4.5% / +11.3% over main, which is 94-100% of the full PR's +4.4% / +12.0%. On mixed it gives +2.9% / +4.8%, which is 81-84% of the PR's +3.4% / +6.0%. The FP8 MMA kernel itself adds about 0-1%.
BF16 is unchanged: BF16 KV decode is within ±0.2% of main at cc 1-128. Main already uses FA4 for all Gemma-4 layers with BF16 KV, so the split change doesn't touch that path.

One small observation, not a blocker: in a harder multi-value retrieval test (6 values plus 10 distractor keys) at 196k tokens, mean recall was slightly lower with this PR. It was 0.86 vs 0.94 on main, over 36 trials (p≈0.06). The results were identical up to 64k and similar at 254k. I'm mentioning it for awareness only.

_{This comment was drafted with Claude Code and reviewed by a human.}

mergify · 2026-10-23T02:59:24Z

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

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

MatthewBonanni

The new kernel doesn't support non-causal attention, right? So if we swap over entirely don't we lose DCP support? Do we need to run DCP through the old kernel?

simon-veitner-redhat · 2026-11-04T20:00:33Z

Right that the new kernel can't serve DCP's non-causal context call: the SM90 FP8 MMA kernel from FA#199 only builds its split QK accumulation for causal and local masks. But no working setup is lost, because DCP never worked on this path in vLLM, before or after this PR.

The FP8 MMA path only applies to head-size-512 layers with an FP8 KV cache on SM90, which means Gemma-4. On main (merge-base 5d35ce3, FA pin b273a02), DCP fails for every Gemma-4 before it reaches this kernel:

vLLM rejects DCP for any model with sliding-window layers ("DCP not support sliding window." at vllm/v1/kv_cache_interface.py:950), and every Gemma-4 has them.
31B and 26B-A4B also need TP greater than their 16 and 8 KV heads.
E2B/E4B KV-shared layers pass key=None into _forward_with_dcp and raise TypeError.
FA3 sliding layers with an FP8 KV cache crash under DCP with "query and key must have the same dtype".

The old fp16-dequant FA4 kernel computed the non-causal call correctly when called directly (rel L2 1.60e-3 on a captured Gemma-4 31B layer, H100), but no real config could reach it, and it no longer exists at 9ed6b12.

This PR makes the limitation explicit: FlashAttentionImpl.__init__ raises NotImplementedError for SM90 FA4 FP8 KV with dcp_world_size > 1, and the wrapper raises on non-causal FP8 MMA calls. With the guard bypassed the kernel runs but is less accurate (rel L2 2.62e-3, LSE error 0.062 vs 0.029 causal), so keeping the guard is right.

For reference, if DCP ever gains sliding-window support: running the context call as local attention with a window covering every key gives the same mask and uses the accurate local build. It matched the BF16-cache DCP path in a DCP=2 emulation on the same layer (rel L2 2.38e-3 vs 2.29e-3 decode). It's a workaround that relies on FA building the QK split for local masks; the cleaner fix would be in FA.

-        if uses_sm90_fa4_fp8_kv_dequant and self.dcp_world_size > 1:
-            raise NotImplementedError(...)
+        self.dcp_context_window = list(self.sliding_window)
+        if (
+            uses_sm90_fa4_fp8_kv_dequant
+            and self.dcp_world_size > 1
+            and self.sliding_window == (-1, -1)
+        ):
+            max_model_len = vllm_config.model_config.max_model_len
+            self.dcp_context_window = [max_model_len, max_model_len]
 ...
                 causal=False,
-                window_size=sliding_window_size,
+                window_size=self.dcp_context_window,

This review was prepared with Claude Code and reviewed by a human (@simon-veitner-redhat).

MatthewBonanni

LGTM, thanks!

MatthewBonanni · 2026-11-04T20:16:44Z

/ci run

github-actions · 2026-11-04T20:18:04Z

✅ Triggered Buildkite CI #100019 for commit 0ffdc07ec70c.

github-actions · 2026-11-04T20:18:10Z

✅ @katec846, CI is now available for this PR.

/ci run starts upstream CI; /amd-ci run starts AMD CI only.
Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
/ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
/amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
/ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

pytest 9 only skips ModuleNotFoundError in importorskip by default, so the ImportError raised by vllm.vllm_flash_attn on ROCm broke collection. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

MatthewBonanni · 2026-11-04T21:31:39Z

/ci run

github-actions · 2026-11-04T21:32:02Z

✅ Triggered Buildkite CI #100033 for commit f2d825998d8e.

DarkLight1337 · 2026-11-05T04:30:12Z

/ci retry

github-actions · 2026-11-05T04:32:12Z

✅ Triggered Buildkite CI #100131 for commit 78e295c3b21e, running 3 failed step(s) from Buildkite CI #100033.

mergify · 2026-11-05T04:36:02Z

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

uv pip install pre-commit>=4.5.1
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.

The conversion helper no longer accepts input_dtype. Remove the stale keyword from the Zen CPU width-logging test so the call matches the API and passes mypy. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Kate Cheng <yunhsuanc@nvidia.com>

DarkLight1337 · 2026-11-05T05:58:48Z

/ci run

github-actions · 2026-11-05T05:59:10Z

✅ Triggered Buildkite CI #100154 for commit feba6078050b.

Run Gemma-4-31B-FP8 from models-large-hopper.txt instead of a dedicated 2xH100 step that ran on every attention-backend change. The lm-eval harness gains an hf_overrides key, needed for the text-only override that routes Gemma-4's full-attention layers to the SM90 FP8 MMA path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Measured on 4x H100 at TP4 with this branch and FA 9ed6b12; the SM90 FP8 MMA kernel ran for the head_size=512 layers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

MatthewBonanni · 2026-11-05T16:06:26Z

/ci run

github-actions · 2026-11-05T16:06:49Z

✅ Triggered Buildkite CI #100296 for commit f2e9e8f337b8.

mergify Bot added the ci/build label Oct 7, 2026

katec846 force-pushed the fa4-fp8mma-draft-20261007 branch from 8fd2432 to 72868c7 Compare October 7, 2026 21:35

jhaotingc reviewed Oct 8, 2026

View reviewed changes

katec846 force-pushed the fa4-fp8mma-draft-20261007 branch 2 times, most recently from 7e83853 to 98b7608 Compare October 8, 2026 17:43

katec846 marked this pull request as ready for review October 9, 2026 17:26

katec846 requested review from Harry-Chen, LucasWilkinson, MatthewBonanni and tlrmchlsmth as code owners October 9, 2026 17:26

claude Bot reviewed Oct 9, 2026

View reviewed changes

katec846 force-pushed the fa4-fp8mma-draft-20261007 branch 2 times, most recently from 52d381d to 9650161 Compare October 9, 2026 19:47

katec846 requested review from AndreasKaratzas, WoosukKwon, mgoin, yewentao256 and zyongye as code owners October 9, 2026 19:47

katec846 force-pushed the fa4-fp8mma-draft-20261007 branch from 00a97d6 to c913b8d Compare October 14, 2026 23:45

jhaotingc reviewed Oct 14, 2026

View reviewed changes

jhaotingc reviewed Oct 15, 2026

View reviewed changes

katec846 force-pushed the fa4-fp8mma-draft-20261007 branch from 2adc688 to 5446e6c Compare October 15, 2026 18:38

katec846 requested review from ZJY0516, khluu and vadiklyutiy as code owners October 23, 2026 02:49

mergify Bot added the needs-rebase label Oct 23, 2026

katec846 force-pushed the fa4-fp8mma-draft-20261007 branch from e71fe29 to 2b76f07 Compare October 23, 2026 03:15

zyongye mentioned this pull request Oct 29, 2026

[Perf][FlashAttention] FA4: pass no FP8 descale for a unit per-tensor scale #63232

Merged

2 of 6 tasks

MatthewBonanni reviewed Nov 3, 2026

View reviewed changes

mgoin mentioned this pull request Nov 4, 2026

[Build] Bump vllm-flash-attn to compress fatbins with --compress-mode=size #65518

Open

6 tasks done

MatthewBonanni approved these changes Nov 4, 2026

View reviewed changes

MatthewBonanni added the ready ONLY add when PR is ready to merge/full CI is needed label Nov 4, 2026

MatthewBonanni enabled auto-merge (squash) November 4, 2026 20:16

Merge branch 'main' into fa4-fp8mma-draft-20261007

0ffdc07

MatthewBonanni and others added 2 commits November 4, 2026 16:23

Merge remote-tracking branch 'upstream-main-now' into fix-60372-ci

f2d8259

Merge branch 'main' into fa4-fp8mma-draft-20261007

78e295c

auto-merge was automatically disabled November 5, 2026 05:33
Head branch was pushed to by a user without write access

Merge branch 'main' into fa4-fp8mma-draft-20261007

feba607

DarkLight1337 enabled auto-merge (squash) November 5, 2026 05:58

MatthewBonanni and others added 3 commits November 5, 2026 10:08

Set Gemma-4 FP8 KV GSM8K reference to the measured 0.937

7894b61

Measured on 4x H100 at TP4 with this branch and FA 9ed6b12; the SM90 FP8 MMA kernel ran for the head_size=512 layers. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Merge branch 'main' into fa4-fp8mma-draft-20261007

f2e9e8f

Conversation

katec846 commented Oct 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Claims

Validation

Serving performance

Model evaluation

MTP speculative decoding

Reproduction commands

Details

Why FP8 is slower than BF16

Sliding-window attention and KV-cache operations offset full-attention gains

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

jhaotingc Oct 14, 2026

Choose a reason for hiding this comment

Uh oh!

jhaotingc left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jhaotingc Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

katec846 Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jhaotingc Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

katec846 Oct 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jhaotingc Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

katec846 Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jhaotingc Oct 15, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

simon-veitner-redhat commented Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Fixes you can cherry-pick

PR description

Our independent results

Uh oh!

mergify Bot commented Oct 23, 2026

Uh oh!

MatthewBonanni left a comment

Choose a reason for hiding this comment

Uh oh!

simon-veitner-redhat commented Nov 4, 2026

Uh oh!

MatthewBonanni left a comment

Choose a reason for hiding this comment

Uh oh!

MatthewBonanni commented Nov 4, 2026

Uh oh!

katec846 commented Oct 7, 2026 •

edited

Loading

katec846 Oct 15, 2026 •

edited

Loading

katec846 Oct 15, 2026 •

edited

Loading

simon-veitner-redhat commented Oct 21, 2026 •

edited

Loading