[Kernel] Integrate FA4 SM90 FP8 MMA for Gemma-4 full attention

.buildkite/test_areas/lm_eval.yaml

@@ -320,6 +320,22 @@ steps:
       commands:
       - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-watermark-dspark-tp2.txt
 
+# SM90-only path (FA4 FP8 MMA for head_size=512), so no AMD mirror.
+- label: ":nvidia: (H100) LM Eval Hopper FP8 KV"
+  key: lm-eval-hopper-fp8-kv-2xh100
+  device: h100
+  num_devices: 2
+  timeout_in_minutes: 30
+  source_file_dependencies:
+  - cmake/external_projects/vllm_flash_attn.cmake
+  - vllm/v1/attention/backends/fa_utils.py
+  - vllm/v1/attention/backends/flash_attn.py
+  - vllm/vllm_flash_attn/
+  - tests/evals/gsm8k/configs/gemma-4-31B-it-FP8-dynamic-fp8kv-TP2.yaml
+  - tests/evals/gsm8k/configs/models-hopper-fp8-kv.txt
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-hopper-fp8-kv.txt
+
 - label: ":nvidia: (B200) [Nightly] LM Eval Spec Decode"
   key: lm-eval-spec-decode-4xb200
   timeout_in_minutes: 120

cmake/external_projects/vllm_flash_attn.cmake

@@ -39,7 +39,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG b273a02621d3c94bc747df67204f591b67f38fae
+          GIT_TAG 9ed6b1291b2a4ef8cbd2fd4fa7ebad6f280d9ee9
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

tests/evals/gsm8k/configs/gemma-4-31B-it-FP8-dynamic-fp8kv-TP2.yaml

@@ -0,0 +1,17 @@
+# Gemma-4 full-attention layers (head_size=512) with an FP8 KV cache take the
+# FA4 SM90 FP8 MMA path on Hopper; the sliding-window layers use FA3 FP8.
+# The hf-override makes the model text-only. Without it Gemma-4 is an
+# mm-prefix LM, which FLASH_ATTN does not support with FP8 KV on SM90.
+model_name: "RedHatAI/gemma-4-31B-it-FP8-dynamic"
+# Measured 0.970 on 2x H200 with this harness and FA PR 199.
+accuracy_threshold: 0.96
+num_questions: 1319
+num_fewshot: 5
+use_chat_completions: true
+max_tokens: 512
+server_args: >-
+  --tensor-parallel-size 2
+  --max-model-len 4096
+  --kv-cache-dtype fp8
+  --attention-backend FLASH_ATTN
+  --hf-overrides '{"text_config":{"use_bidirectional_attention":null}}'

tests/evals/gsm8k/configs/models-hopper-fp8-kv.txt

@@ -0,0 +1 @@
+gemma-4-31B-it-FP8-dynamic-fp8kv-TP2.yaml

tests/v1/attention/test_flash_attn_sm90_fp8_mma_selection.py

@@ -0,0 +1,179 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Routing of the FA4 SM90 FP8 MMA path (head_size=512 with an FP8 KV cache).
+
+The device capability is mocked, so no GPU is needed.
+"""
+
+from contextlib import contextmanager
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm.platforms.interface import DeviceCapability
+
+fa_interface = pytest.importorskip("vllm.vllm_flash_attn.flash_attn_interface")
+
+from vllm.v1.attention.backends.fa_utils import get_flash_attn_version  # noqa: E402
+from vllm.v1.attention.backends.flash_attn import (  # noqa: E402
+    FlashAttentionBackend,
+    FlashAttentionImpl,
+)
+
+SM80 = DeviceCapability(8, 0)
+SM90 = DeviceCapability(9, 0)
+SM100 = DeviceCapability(10, 0)
+
+
+@contextmanager
+def _device(capability, *, head_size, kv_cache_dtype, fa_version=None, dcp=1):
+    platform = MagicMock()
+    platform.is_xpu.return_value = False
+    platform.is_rocm.return_value = False
+    platform.is_cuda.return_value = True
+    platform.get_device_capability.return_value = capability
+    platform.is_device_capability_family.side_effect = (
+        lambda family, device_id=0: capability.to_int() // 10 == family // 10
+    )
+    config = MagicMock()
+    config.attention_config.flash_attn_version = fa_version
+    config.model_config.is_diffusion = False
+    config.model_config.get_head_size.return_value = head_size
+    config.cache_config.cache_dtype = kv_cache_dtype
+    config.parallel_config.decode_context_parallel_size = dcp
+    config.parallel_config.dcp_comm_backend = "ag_rs"
+    backends = "vllm.v1.attention.backends"
+    with (
+        patch(f"{backends}.fa_utils.current_platform", platform),
+        patch(f"{backends}.flash_attn.current_platform", platform),
+        patch.object(fa_interface, "is_fa_version_supported", return_value=True),
+        patch("vllm.config.get_current_vllm_config_or_none", return_value=config),
+        patch(
+            f"{backends}.flash_attn.get_current_vllm_config_or_none",
+            return_value=config,
+        ),
+        patch(
+            "vllm.distributed.parallel_state.get_dcp_world_size_and_rank",
+            return_value=(dcp, 0),
+        ),
+    ):
+        yield
+
+
+def _supports_combination(capability, head_size, kv_cache_dtype, block_size=64):
+    return FlashAttentionBackend.supports_combination(
+        head_size=head_size,
+        dtype=torch.bfloat16,
+        kv_cache_dtype=kv_cache_dtype,
+        block_size=block_size,
+        use_mla=False,
+        has_sink=False,
+        use_sparse=False,
+        use_mm_prefix=False,
+        device_capability=capability,
+    )
+
+
+def _impl(head_size, kv_cache_dtype):
+    return FlashAttentionImpl(
+        num_heads=8,
+        head_size=head_size,
+        scale=head_size**-0.5,
+        num_kv_heads=1,
+        alibi_slopes=None,
+        sliding_window=None,
+        kv_cache_dtype=kv_cache_dtype,
+    )
+
+
+@pytest.mark.parametrize(
+    "head_size,kv_cache_dtype,fa_version,sm90_fp8_mma",
+    [
+        (512, "fp8", 4, True),
+        (512, "fp8_e4m3", 4, True),
+        # Gemma-4 sliding-window layers stay on FA3 FP8 with an FP8 query.
+        (256, "fp8", 3, False),
+        (512, "auto", 4, False),
+    ],
+)
+def test_sm90_routing(head_size, kv_cache_dtype, fa_version, sm90_fp8_mma):
+    with _device(SM90, head_size=head_size, kv_cache_dtype=kv_cache_dtype):
+        assert get_flash_attn_version(head_size=head_size) == fa_version
+        assert _supports_combination(SM90, head_size, kv_cache_dtype) is None
+        impl = _impl(head_size, kv_cache_dtype)
+    assert impl.vllm_flash_attn_version == fa_version
+    # The FP8 MMA kernel quantizes Q in-kernel, so vLLM must not quantize it.
+    assert impl.supports_quant_query_input is not sm90_fp8_mma
+
+
+def test_sm90_fp8_mma_block_size():
+    with _device(SM90, head_size=512, kv_cache_dtype="fp8"):
+        assert FlashAttentionBackend.get_supported_kernel_block_sizes() == [64]
+        assert "block size" in _supports_combination(SM90, 512, "fp8", 16)
+    with _device(SM90, head_size=512, kv_cache_dtype="auto"):
+        assert _supports_combination(SM90, 512, "auto", 16) is None
+
+
+@pytest.mark.parametrize(
+    "capability,head_size,fa_version",
+    [
+        (SM90, 256, 4),  # forced FA4 with an FP8 KV cache needs head_size=512
+        (SM100, 512, None),  # FA4 on Blackwell does not support head_size=512
+        (SM80, 512, None),  # FA2 has no FP8 KV cache
+    ],
+)
+def test_fp8_kv_rejected(capability, head_size, fa_version):
+    with _device(
+        capability, head_size=head_size, kv_cache_dtype="fp8", fa_version=fa_version
+    ):
+        assert _supports_combination(capability, head_size, "fp8") is not None
+        with pytest.raises(NotImplementedError, match="fp8 kv-cache"):
+            _impl(head_size, "fp8")
+
+
+@pytest.mark.parametrize(
+    "head_size,kv_cache_dtype,rejected",
+    [(512, "fp8", True), (512, "auto", False), (256, "fp8", False)],
+)
+def test_sm90_fp8_mma_rejects_dcp(head_size, kv_cache_dtype, rejected):
+    with _device(SM90, head_size=head_size, kv_cache_dtype=kv_cache_dtype, dcp=2):
+        if rejected:
+            with pytest.raises(NotImplementedError, match="context parallelism"):
+                _impl(head_size, kv_cache_dtype)
+        else:
+            assert _impl(head_size, kv_cache_dtype).dcp_world_size == 2
+
+
+@pytest.mark.parametrize(
+    "head_size,q_dtype,causal,match",
+    [
+        # FA4 + FP8 KV below d512 raises; no silent switch to FA3.
+        (256, torch.float8_e4m3fn, True, "q=torch.float8_e4m3fn/256"),
+        (128, torch.float8_e4m3fn, False, "head_size=512"),
+        (256, torch.bfloat16, True, "q=torch.bfloat16/256"),
+        (512, torch.bfloat16, False, "noncausal"),
+    ],
+)
+def test_wrapper_rejects_unsupported_sm90_fp8(head_size, q_dtype, causal, match):
+    q = torch.zeros(4, 8, head_size, dtype=q_dtype)
+    k = torch.zeros(1, 64, 1, head_size, dtype=torch.float8_e4m3fn)
+    fa3 = MagicMock()
+    with (
+        patch("torch.cuda.get_device_capability", return_value=(9, 0)),
+        patch.object(torch.ops, "_vllm_fa3_C", fa3, create=True),
+        pytest.raises(NotImplementedError, match=match),
+    ):
+        fa_interface.flash_attn_varlen_func(
+            q,
+            k,
+            k,
+            max_seqlen_q=4,
+            cu_seqlens_q=torch.tensor([0, 4], dtype=torch.int32),
+            max_seqlen_k=4,
+            seqused_k=torch.tensor([4], dtype=torch.int32),
+            block_table=torch.zeros(1, 1, dtype=torch.int32),
+            causal=causal,
+            fa_version=4,
+        )
+    fa3.fwd.assert_not_called()

vllm/v1/attention/backends/fa_utils.py

@@ -305,7 +305,7 @@ def flash_attn_supports_kv_cache_dtype(
     )
     is_sm90 = current_platform.is_device_capability_family(90)
     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)
     )
     return (sm90_fp8_kv_supported and is_sm90) or (
         fa_version == 4 and current_platform.is_device_capability_family(100)

vllm/v1/attention/backends/flash_attn.py

@@ -471,6 +471,16 @@ def supports_combination(
         use_mm_prefix: bool,
         device_capability: DeviceCapability,
     ) -> str | None:
+        if (
+            device_capability == DeviceCapability(9, 0)
+            and kv_cache_dtype is not None
+            and is_quantized_kv_cache(kv_cache_dtype)
+            and head_size == 512
+        ):
+            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"
         if has_sink and device_capability < DeviceCapability(9, 0):
             return "sink not supported on compute capability < 9.0"
         if (
@@ -486,16 +496,13 @@ def supports_combination(
             and not flash_attn_supports_kv_cache_dtype(
                 kv_cache_dtype,
                 head_size=head_size,
-                head_size_v=head_size,
+                head_size_v=cls.head_size_v,
                 has_sinks=has_sink,
                 kv_cache_block_size=block_size,
                 supports_fa4_hd256=True,
             )
         ):
-            return (
-                "FP8 KV cache requires FA3 on SM90, FA4 with head_size=512 "
-                "on SM90, or FA4 on SM100"
-            )
+            return "FP8 KV requires FA3 (SM90), equal-d512 FA4 (SM90), or FA4 (SM100)."
         if (
             use_mm_prefix
             and get_flash_attn_version(
@@ -705,14 +712,14 @@ def __init__(
         self.block_size = kv_cache_spec.block_size
 
         self.max_num_splits = 0  # No upper bound on the number of splits.
-        self.aot_schedule = get_flash_attn_version() == 3
         head_size_v = getattr(kv_cache_spec, "head_size_v", None)
         fa_version = get_flash_attn_version(
             head_size=self.headdim,
             head_size_v=head_size_v,
             kv_cache_block_size=self.block_size,
             supports_fa4_hd256=True,
         )
+        self.aot_schedule = fa_version == 3
         self.fa4_hd256 = fa_version == 4 and uses_fa4_hd256_kernel(
             self.headdim, head_size_v
         )
@@ -1152,9 +1159,8 @@ def __init__(
                 "heads in the layer"
             )
 
-        # FA4's SM90 FP8-KV path consumes native FP16/BF16 Q and dequantizes
-        # FP8 K/V in-kernel. Other FA4 paths (notably SM100) still require Q,
-        # K, and V to have the same FP8 dtype.
+        # FA4's SM90 d512 FP8-KV path consumes native FP16/BF16 Q. Other FA4
+        # FP8 paths (notably SM100) require Q, K, and V to share an FP8 dtype.
         uses_sm90_fa4_fp8_kv_dequant = (
             self.vllm_flash_attn_version == 4
             and current_platform.is_device_capability_family(90)
@@ -1163,6 +1169,13 @@ def __init__(
         self.supports_quant_query_input = flash_attn_supports_quant_query_input() and (
             not uses_sm90_fa4_fp8_kv_dequant
         )
+        if uses_sm90_fa4_fp8_kv_dequant and self.dcp_world_size > 1:
+            # DCP context attention is non-causal, which the SM90 FP8 MMA
+            # kernel rejects at dispatch time.
+            raise NotImplementedError(
+                "FA4 SM90 FP8 KV cache does not support decode context "
+                "parallelism; use an unquantized KV cache."
+            )
 
         dcp_a2a = (
             vllm_config is not None

vllm/v1/attention/backends/flash_attn_diffkv.py

@@ -11,6 +11,7 @@
 )
 from vllm.v1.attention.backend import AttentionType
 from vllm.v1.attention.backends.fa_utils import (
+    flash_attn_supports_kv_cache_dtype,
     get_flash_attn_version,
     is_flash_attn_varlen_func_available,
 )
@@ -86,6 +87,19 @@ def __init__(self, *args, **kwargs) -> None:
             head_size_v=FlashAttentionDiffKVBackend.head_size_v,
             has_sinks=self.sinks is not None,
         )
+        # An explicitly requested backend can bypass automatic selection.
+        if is_quantized_kv_cache(
+            self.kv_cache_dtype
+        ) and not flash_attn_supports_kv_cache_dtype(
+            self.kv_cache_dtype,
+            requires_alibi=self.alibi_slopes is not None,
+            head_size=self.head_size,
+            head_size_v=FlashAttentionDiffKVBackend.head_size_v,
+            has_sinks=self.sinks is not None,
+        ):
+            raise NotImplementedError(
+                "Unsupported FP8 Diff-KV configuration; use an unquantized KV cache."
+            )
         self.fa4_hd256 = False
 
     def do_kv_cache_update(