[Kernel] Resubmit PR 48666 - Gemma4 FP8 KV FA4 head dim 512 backend selection (#53175)

Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>

tests/kernels/attention/test_attention_selector.py

@@ -848,7 +848,7 @@ def hopper_selection():
 def test_hopper_mm_prefix_selects_triton_flash_attn(
     use_mm_prefix, flash_attn_version, hopper_selection
 ):
-    """Hopper selects the composite only when its causal route resolves FA4."""
+    """Hopper selects the composite when its causal route resolves FA3 or FA4."""
     from vllm.engine.arg_utils import EngineArgs
 
     config = EngineArgs(
@@ -860,11 +860,7 @@ def test_hopper_mm_prefix_selects_triton_flash_attn(
         backend = get_attn_backend(
             256, torch.bfloat16, None, use_mm_prefix=use_mm_prefix
         )
-    fa4_resolved = config.attention_config.flash_attn_version == 4
-    if use_mm_prefix:
-        expected = "TRITON_FLASH_ATTN" if fa4_resolved else "TRITON_ATTN"
-    else:
-        expected = "FLASH_ATTN"
+    expected = "TRITON_FLASH_ATTN" if use_mm_prefix else "FLASH_ATTN"
     assert backend.get_name() == expected
 
 

tests/v1/attention/test_backend_per_kind.py

@@ -29,7 +29,7 @@ from vllm.v1.kv_cache_interface import KVCacheSpecKind
 def test_largest_kernel_block_within(supported_sizes, expected):
     class Backend:
         @staticmethod
-        def get_supported_kernel_block_sizes():
+        def get_supported_kernel_block_sizes(kv_cache_spec=None):
             return supported_sizes
 
     assert _largest_kernel_block_within(Backend, 1024, 1024 * 1536, 2048) == expected

vllm/model_executor/layers/attention/attention.py

@@ -8,7 +8,10 @@ import torch.nn as nn
 
 import vllm.envs as envs
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
-from vllm.config import CacheConfig, get_current_vllm_config
+from vllm.config import (
+    CacheConfig,
+    get_current_vllm_config,
+)
 from vllm.config.vllm import VllmConfig
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.logger import init_logger
@@ -100,6 +103,7 @@ def _largest_kernel_block_within(
     per_token_bytes: int,
     page_budget: int,
     fallback: int,
+    kv_cache_spec: KVCacheSpec | None = None,
 ) -> int:
     """Largest supported kernel block size whose page fits in ``page_budget``.
 
@@ -111,7 +115,7 @@ def _largest_kernel_block_within(
     """
     from vllm.v1.attention.backend import MultipleOf
 
-    sizes = attn_backend.get_supported_kernel_block_sizes()
+    sizes = attn_backend.get_supported_kernel_block_sizes(kv_cache_spec)
     max_block_size = page_budget // per_token_bytes
     candidates = [s for s in sizes if isinstance(s, int)]
     candidates.extend(
@@ -632,7 +636,7 @@ class Attention(nn.Module, AttentionLayerBase):
             # ``unify`` scales it up by an integer ratio.
             shared_page = vllm_config.cache_config.skip_page_size_padded
             # The backend owns its packing
-            sw_per_token = self.attn_backend.customize_spec(
+            kv_cache_spec = self.attn_backend.customize_spec(
                 SlidingWindowSpec(
                     block_size=1,
                     num_kv_heads=self.num_kv_heads,
@@ -642,10 +646,15 @@ class Attention(nn.Module, AttentionLayerBase):
                     kv_quant_mode=quant_mode,
                     sliding_window=self.sliding_window,
                 )
-            ).real_page_size_bytes
+            )
+            sw_per_token = kv_cache_spec.real_page_size_bytes
             page_budget = shared_page or sw_per_token * block_size
             sw_block_size = _largest_kernel_block_within(
-                self.attn_backend, sw_per_token, page_budget, block_size
+                self.attn_backend,
+                sw_per_token,
+                page_budget,
+                block_size,
+                kv_cache_spec,
             )
             return SlidingWindowSpec(
                 block_size=sw_block_size,

vllm/model_executor/models/config.py

@@ -212,12 +212,15 @@ class Gemma4Config(VerifyAndUpdateConfig):
         """Configure attention for heterogeneous head dimensions.
 
         Gemma4 uses different head dimensions for sliding window vs full attention
-        layers. The default FA3 on Hopper cannot handle head_dim > 256, which causes
-        mixed backend selection and numerical divergence.
+        layers. The default FA3 on Hopper cannot handle head_dim > 256.
 
-        When FA4 is available we force it for ALL layers, giving a uniform kernel path
-        and avoiding the mixed FA3+FA4 penalty. When FA4 is not available we fall back
-        to Triton.
+        On SM90 with FP8 KV cache, use FA3 for supported layers and let the generic
+        FlashAttention selector upgrade larger head dimensions to FA4.
+        The multimodal-prefix composite routes image masks to Triton and causal
+        requests to this per-layer FA3/FA4 selection. For other configurations,
+        force FA4 for all layers to avoid the mixed
+        FA3+FA4 penalty.
+        When FA4 is not available, fall back to Triton.
         """
         model_config = vllm_config.model_config
         arch_config = model_config.model_arch_config
@@ -230,6 +233,7 @@ class Gemma4Config(VerifyAndUpdateConfig):
         if len(set(head_dims.values())) <= 1:
             return
 
+        from vllm.platforms import current_platform
         from vllm.v1.attention.backends.fa_utils import is_fa_version_supported
         from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
@@ -241,12 +245,22 @@ class Gemma4Config(VerifyAndUpdateConfig):
                 and vllm_config.attention_config.backend
                 in (None, AttentionBackendEnum.FLASH_ATTN)
             ):
-                vllm_config.attention_config.flash_attn_version = 4
-                logger.info(
-                    "Gemma4 model has heterogeneous head dimensions %s. Using FA4 for "
-                    "all layers to avoid mixed FA3/FA4 penalty.",
-                    head_dims,
-                )
+                use_per_layer_fa = current_platform.is_device_capability_family(
+                    90
+                ) and vllm_config.cache_config.cache_dtype.startswith("fp8")
+                if use_per_layer_fa:
+                    logger.info(
+                        "Gemma4 model has heterogeneous head dimensions %s. Using "
+                        "per-layer FA3/FA4 selection for FP8 KV cache on SM90.",
+                        head_dims,
+                    )
+                else:
+                    vllm_config.attention_config.flash_attn_version = 4
+                    logger.info(
+                        "Gemma4 model has heterogeneous head dimensions %s. Using FA4 "
+                        "for all layers to avoid mixed FA3/FA4 penalty.",
+                        head_dims,
+                    )
         elif vllm_config.attention_config.backend is None:
             vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
             logger.info(

vllm/models/deepseek_v4/compressor.py

@@ -66,7 +66,7 @@ class CompressorBackend(AttentionBackend):
         return "CompressorBackend"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(1)]
 
     @classmethod

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

@@ -111,7 +111,7 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [256]
 
     @staticmethod

vllm/models/deepseek_v4/sparse_mla.py

@@ -57,7 +57,7 @@ class DeepseekV4SparseMLABackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [256]
 
     @staticmethod

vllm/models/deepseek_v41/compressor.py

@@ -37,7 +37,7 @@ class CompressorBackend(AttentionBackend):
         return "CompressorBackend"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(1)]
 
     @classmethod

vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py

@@ -113,7 +113,7 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [128]
 
     @staticmethod

vllm/models/deepseek_v41/sparse_mla.py

@@ -86,7 +86,7 @@ class DeepseekV4SparseMLABackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64 if current_platform.is_device_capability_family(90) else 128]
 
     @staticmethod
@@ -255,7 +255,7 @@ class FlashMLAMegaAttnBackend(DeepseekV4FlashMLABackend):
         return "FLASHMLA_MEGA_ATTN_DSV41"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [128]
 
     @classmethod

vllm/models/minimax_m3/common/indexer.py

@@ -91,7 +91,7 @@ class MiniMaxM3IndexerBackend(AttentionBackend):
         return [128]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [128]
 
     @classmethod

vllm/models/minimax_m3/common/sparse_attention.py

@@ -159,7 +159,7 @@ class MiniMaxM3SparseBackend(AttentionBackend):
         return [128]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # Page size == sparse block size (one sparse block per KV page).
         return [128]
 

vllm/models/qwen4_exp/amd/qsa.py

@@ -75,7 +75,7 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
         return "QWEN4_EXP_QSA_TRITON"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # QSA consumes manager pages directly and does not use FA4 paged attention.
         return [MultipleOf(16)]
 

vllm/models/qwen4_exp/nvidia/qsa.py

@@ -103,7 +103,7 @@ class Qwen4ExpQSAFlashAttentionBackend(FlashAttentionBackend):
         return "QWEN4_EXP_QSA_TRITON"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # QSA consumes manager pages directly and does not use FA4 paged attention.
         return [MultipleOf(16)]
 

vllm/v1/attention/backend.py

@@ -69,7 +69,9 @@ class AttentionBackend(ABC):
     forward_includes_kv_cache_update: bool = True
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(
+        kv_cache_spec: "KVCacheSpec | None" = None,
+    ) -> list[int | MultipleOf]:
         return [MultipleOf(1)]
 
     @staticmethod

vllm/v1/attention/backends/b12x.py

@@ -183,7 +183,7 @@ class B12xPagedAttentionBackend(AttentionBackend):
         return B12xPagedMetadataBuilder
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return list(_B12X_SUPPORTED_PAGE_SIZES)
 
     @classmethod

vllm/v1/attention/backends/composite.py

@@ -116,14 +116,16 @@ class MMPrefixAttentionRouting(CompositeAttentionRouting):
 
 
 def _intersect_kernel_block_sizes(
-    first: type[AttentionBackend], second: type[AttentionBackend]
+    first: type[AttentionBackend],
+    second: type[AttentionBackend],
+    kv_cache_spec: KVCacheSpec | None = None,
 ):
     from vllm.v1.attention.backend import MultipleOf
 
     result: list[int | MultipleOf] = []
     seen: set[tuple[type, int]] = set()
-    for lhs in first.get_supported_kernel_block_sizes():
-        for rhs in second.get_supported_kernel_block_sizes():
+    for lhs in first.get_supported_kernel_block_sizes(kv_cache_spec):
+        for rhs in second.get_supported_kernel_block_sizes(kv_cache_spec):
             candidate: int | MultipleOf | None = None
             if isinstance(lhs, MultipleOf) and isinstance(rhs, MultipleOf):
                 candidate = MultipleOf(math.lcm(lhs.base, rhs.base))
@@ -378,7 +380,9 @@ def create_composite_attention_backend(
             return CompositeAttentionMetadataBuilder
 
         @staticmethod
-        def get_supported_kernel_block_sizes():
+        def get_supported_kernel_block_sizes(
+            kv_cache_spec: KVCacheSpec | None = None,
+        ):
             if kernel_block_sizes:
                 return [
                     size
@@ -388,7 +392,9 @@ def create_composite_attention_backend(
                         for backend in (general_backend, causal_backend)
                     )
                 ]
-            return _intersect_kernel_block_sizes(general_backend, causal_backend)
+            return _intersect_kernel_block_sizes(
+                general_backend, causal_backend, kv_cache_spec
+            )
 
         @classmethod
         def supports_block_size(cls, block_size):

vllm/v1/attention/backends/cpu_attn.py

@@ -57,7 +57,7 @@ class CPUAttentionBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(32)]
 
     @classmethod

vllm/v1/attention/backends/fa_utils.py

@@ -303,7 +303,11 @@ def flash_attn_supports_kv_cache_dtype(
         kv_cache_block_size=kv_cache_block_size,
         supports_fa4_hd256=supports_fa4_hd256,
     )
-    return (fa_version == 3 and current_platform.is_device_capability_family(90)) or (
+    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)
+    )
+    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

@@ -291,18 +291,56 @@ class FlashAttentionBackend(AttentionBackend):
     ]
     head_size_v: int | None = None
 
-    @classmethod
-    def _get_fa4_hd256_block_size(cls) -> int | None:
-        vllm_config = get_current_vllm_config_or_none()
-        if vllm_config is None or vllm_config.model_config is None:
-            return None
+    @staticmethod
+    def _get_sm90_fa4_fp8_kv_block_size(
+        kv_cache_spec: KVCacheSpec | None = None,
+    ) -> int | None:
+        if kv_cache_spec is None:
+            vllm_config = get_current_vllm_config_or_none()
+            if vllm_config is None or vllm_config.model_config is None:
+                return None
+            head_size = vllm_config.model_config.get_head_size()
+            uses_fp8_kv = vllm_config.cache_config.cache_dtype in (
+                "fp8",
+                "fp8_e4m3",
+            )
+        else:
+            if not isinstance(kv_cache_spec, AttentionSpec):
+                return None
+            head_size = kv_cache_spec.head_size
+            uses_fp8_kv = kv_cache_spec.kv_quant_mode == KVQuantMode.FP8_PER_TENSOR
 
-        head_size = vllm_config.model_config.get_head_size()
         if (
-            uses_fa4_hd256_kernel(head_size, cls.head_size_v)
+            current_platform.is_device_capability_family(90)
+            and uses_fp8_kv
+            and head_size == 512
+            and get_flash_attn_version(head_size=head_size) == 4
+        ):
+            # The SM90 FP8-KV-dequant kernel uses a 64-token TMA tile/page.
+            return 64
+        return None
+
+    @classmethod
+    def _get_fa4_hd256_block_size(
+        cls, kv_cache_spec: KVCacheSpec | None = None
+    ) -> int | None:
+        if kv_cache_spec is None:
+            vllm_config = get_current_vllm_config_or_none()
+            if vllm_config is None or vllm_config.model_config is None:
+                return None
+            head_size = vllm_config.model_config.get_head_size()
+            head_size_v = cls.head_size_v
+        else:
+            if not isinstance(kv_cache_spec, AttentionSpec):
+                return None
+            head_size = kv_cache_spec.head_size
+            head_size_v = kv_cache_spec.head_size_v
+
+        if (
+            uses_fa4_hd256_kernel(head_size, head_size_v)
             and get_flash_attn_version(
                 head_size=head_size,
-                head_size_v=cls.head_size_v,
+                head_size_v=head_size_v,
                 supports_fa4_hd256=True,
             )
             == 4
@@ -311,8 +349,12 @@ class FlashAttentionBackend(AttentionBackend):
         return None
 
     @classmethod
-    def get_supported_kernel_block_sizes(cls) -> list[int | MultipleOf]:
-        if block_size := cls._get_fa4_hd256_block_size():
+    def get_supported_kernel_block_sizes(
+        cls, kv_cache_spec: KVCacheSpec | None = None
+    ) -> list[int | MultipleOf]:
+        if block_size := cls._get_sm90_fa4_fp8_kv_block_size(kv_cache_spec):
+            return [block_size]
+        if block_size := cls._get_fa4_hd256_block_size(kv_cache_spec):
             # Sliding-window specs select the smallest advertised size.
             return [block_size]
         return [MultipleOf(16)]
@@ -321,6 +363,8 @@ class FlashAttentionBackend(AttentionBackend):
 
     @classmethod
     def get_preferred_block_size(cls, default_block_size: int) -> int:
+        if block_size := cls._get_sm90_fa4_fp8_kv_block_size():
+            return max(default_block_size, block_size)
         if block_size := cls._get_fa4_hd256_block_size():
             return max(default_block_size, block_size)
         if current_platform.is_xpu():
@@ -419,6 +463,13 @@ class FlashAttentionBackend(AttentionBackend):
     ) -> str | None:
         if has_sink and device_capability < DeviceCapability(9, 0):
             return "sink not supported on compute capability < 9.0"
+        if (
+            use_mm_prefix
+            and kv_cache_dtype is not None
+            and is_quantized_kv_cache(kv_cache_dtype)
+            and device_capability == DeviceCapability(9, 0)
+        ):
+            return "SM90 FP8 KV with mm_prefix requires Triton"
         if (
             kv_cache_dtype is not None
             and is_quantized_kv_cache(kv_cache_dtype)
@@ -431,7 +482,10 @@ class FlashAttentionBackend(AttentionBackend):
                 supports_fa4_hd256=True,
             )
         ):
-            return "FP8 KV cache requires FA3 on SM90 or FA4 on SM100"
+            return (
+                "FP8 KV cache requires FA3 on SM90, FA4 with head_size=512 "
+                "on SM90, or FA4 on SM100"
+            )
         if (
             use_mm_prefix
             and get_flash_attn_version(
@@ -514,15 +568,16 @@ def _get_sliding_window_configs(
 ) -> set[tuple[int, int] | None]:
     """Get the set of all sliding window configs used in the model.
 
-    Only inspects FlashAttentionImpl layers. Other backends (e.g.
-    TurboQuant, MLA) use their own metadata builders and are skipped.
+    Only inspects FlashAttentionImpl layers, including implementations wrapped
+    by a composite backend. Other backends (e.g. TurboQuant, MLA) use their own
+    metadata builders and are skipped.
     """
     sliding_window_configs: set[tuple[int, int] | None] = set()
     layers = get_layers_from_vllm_config(vllm_config, Attention)
     for layer in layers.values():
-        if not isinstance(layer.impl, FlashAttentionImpl):
-            continue
-        sliding_window_configs.add(layer.impl.sliding_window)
+        for impl in layer.impl.get_impl_variants():
+            if isinstance(impl, FlashAttentionImpl):
+                sliding_window_configs.add(impl.sliding_window)
     return sliding_window_configs
 
 
@@ -1145,7 +1200,17 @@ class FlashAttentionImpl(AttentionImpl):
                 "heads in the layer"
             )
 
-        self.supports_quant_query_input = flash_attn_supports_quant_query_input()
+        # 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.
+        uses_sm90_fa4_fp8_kv_dequant = (
+            self.vllm_flash_attn_version == 4
+            and current_platform.is_device_capability_family(90)
+            and self.kv_cache_dtype in ("fp8", "fp8_e4m3")
+        )
+        self.supports_quant_query_input = flash_attn_supports_quant_query_input() and (
+            not uses_sm90_fa4_fp8_kv_dequant
+        )
 
         dcp_a2a = (
             vllm_config is not None

vllm/v1/attention/backends/flashinfer.py

@@ -423,7 +423,7 @@ class FlashInferBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # Page sizes >= 128 only run on the trtllm-gen dynamic kernel (GQA/MQA
         # on Blackwell); advertise them only when usable so selection never
         # picks a large kernel block we cannot serve.

vllm/v1/attention/backends/flex_attention.py

@@ -154,7 +154,7 @@ class FlexAttentionBackend(AttentionBackend):
         return []
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(16)]
 
 

vllm/v1/attention/backends/hpc_attn.py

@@ -288,7 +288,7 @@ class HpcAttentionBackend(AttentionBackend):
     forward_includes_kv_cache_update: bool = True
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64]
 
     @staticmethod

vllm/v1/attention/backends/mla/amx_mla.py

@@ -130,7 +130,7 @@ class AMXMLABackend(MLACommonBackend):
         return [576]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(32)]
 
     @classmethod

vllm/v1/attention/backends/mla/cpu_mla.py

@@ -64,7 +64,7 @@ class CPUMLABackend(MLACommonBackend):
         return [576]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [16]
 
     @classmethod

vllm/v1/attention/backends/mla/cutlass_mla.py

@@ -46,7 +46,7 @@ class CutlassMLABackend(MLACommonBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [128]
 
     @staticmethod

vllm/v1/attention/backends/mla/flashattn_mla.py

@@ -49,7 +49,7 @@ class FlashAttnMLABackend(MLACommonBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(16)]
 
     @staticmethod

vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

@@ -40,7 +40,7 @@ class FlashAttnMLASparseBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64]
 
     @staticmethod

vllm/v1/attention/backends/mla/flashinfer_mla.py

@@ -201,7 +201,7 @@ class FlashInferMLABackend(MLACommonBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [32, 64]
 
     @staticmethod

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -83,7 +83,7 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [32, 64]
 
     @staticmethod
@@ -172,7 +172,7 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
         return "FLASHINFER_MLA_SPARSE_SM120"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64, 256]
 
     @staticmethod

vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py

@@ -74,7 +74,7 @@ class FlashInferMLASparseSM90Backend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(64)]
 
     @staticmethod

vllm/v1/attention/backends/mla/flashmla.py

@@ -55,7 +55,7 @@ class FlashMLABackend(MLACommonBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64]
 
     @staticmethod

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -139,7 +139,7 @@ class FlashMLASparseBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64]
 
     @staticmethod

vllm/v1/attention/backends/mla/indexer.py

@@ -200,7 +200,7 @@ class DeepseekV32IndexerBackend(AttentionBackend):
         return "DEEPSEEK_V32_INDEXER"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [1, MultipleOf(16)] if current_platform.is_rocm() else [64]
 
     @classmethod
@@ -228,7 +228,7 @@ class KpoolTailBackend(DeepseekV32IndexerBackend):
         return []
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(1)]
 
     @staticmethod
@@ -257,7 +257,7 @@ class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
         return (KVCacheLayout.BLHNC, KVCacheLayout.BLNHC)
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # Block sizes count uncompressed tokens: C4 indexer pages hold 64 rows.
         return [256]
 
@@ -274,7 +274,7 @@ class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend):
         return "DEEPSEEK_V41_INDEXER"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [64 if current_platform.is_device_capability_family(90) else 128]
 
 

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

@@ -282,7 +282,7 @@ class AiterMLABackend(MLACommonBackend):
         return []
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # The aiter MLA decode kernel always operates with page_size=1
         # internally (the wrapper flattens kv_buffer via .view(-1, 1, 1, H)).
         # We support any kernel_block_size by expanding block-level indices

vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py

@@ -287,7 +287,7 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [1, MultipleOf(16)]
 
     @staticmethod

vllm/v1/attention/backends/mla/sparse_swa.py

@@ -148,7 +148,7 @@ class DeepseekSparseSWABackend(AttentionBackend):
         return "DEEPSEEK_SPARSE_SWA"
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(32)]
 
     @classmethod

vllm/v1/attention/backends/mla/tokenspeed_mla.py

@@ -87,7 +87,7 @@ class TokenspeedMLABackend(MLACommonBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [32, 64]
 
     @staticmethod

vllm/v1/attention/backends/mla/triton_mla.py

@@ -107,7 +107,7 @@ class TritonMLABackend(MLACommonBackend):
         return []
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(16)]
 
     @classmethod

vllm/v1/attention/backends/rocm_aiter_fa.py

@@ -852,7 +852,7 @@ class AiterFlashAttentionBackend(AttentionBackend):
         return attn_type in (AttentionType.DECODER,)
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         if not rocm_aiter_ops.is_shuffle_kv_cache_enabled():
             return [16, 32]
         # Only gluon serves 128-token pages; the pa_fwd_asm/ll4mi decode

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

@@ -64,7 +64,7 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(16)]
 
     @classmethod

vllm/v1/attention/backends/rocm_attn.py

@@ -176,7 +176,7 @@ class RocmAttentionBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         # ROCM paged attention native C++ kernel only supports block sizes 16 and 32
         # due to shared memory (LDS) constraints on AMD GPUs.
         # See csrc/rocm/attention.cu CALL_CUSTOM_LAUNCHER_BLK macro.

vllm/v1/attention/backends/triton_attn.py

@@ -304,7 +304,7 @@ class TritonAttentionBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [MultipleOf(16)]
 
     @classmethod

vllm/v1/attention/backends/triton_flash_attn.py

@@ -15,7 +15,7 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
 from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
 
 
-class _FA4Backend(FlashAttentionBackend):
+class _FABackend(FlashAttentionBackend):
     @classmethod
     def supports_combination(
         cls,
@@ -42,22 +42,20 @@ class _FA4Backend(FlashAttentionBackend):
         )
         if reason is not None:
             return reason
-        if (
-            get_flash_attn_version(
-                head_size=head_size,
-                has_sinks=has_sink,
-                kv_cache_block_size=block_size,
-                supports_fa4_hd256=True,
-            )
-            != 4
-        ):
-            return "causal route requires FlashAttention v4"
+        fa_version = get_flash_attn_version(
+            head_size=head_size,
+            has_sinks=has_sink,
+            kv_cache_block_size=block_size,
+            supports_fa4_hd256=True,
+        )
+        if fa_version not in (3, 4):
+            return "causal route requires FlashAttention v3 or v4"
         return None
 
 
 TritonFlashAttentionBackend = create_composite_attention_backend(
     TritonAttentionBackend,
-    _FA4Backend,
+    _FABackend,
     name="TritonFlashAttentionBackend",
     backend_name="TRITON_FLASH_ATTN",
     module=__name__,

vllm/v1/attention/backends/turboquant_attn.py

@@ -161,7 +161,7 @@ class TurboQuantAttentionBackend(AttentionBackend):
         return (KVCacheLayout.LBNHC,)
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
         return [16, 32, 64, 128]
 
     @classmethod

vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py

@@ -24,6 +24,27 @@ from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
 logger = init_logger(__name__)
 
 
+def _copy_target_kv_scales(attn: nn.Module, target_attn: nn.Module) -> None:
+    """Copy target KV scales while preserving their tensor representation.
+
+    Default attention scales are scalar buffers, while some quantization
+    methods replace them with length-one or per-head parameters. Preserve the
+    draft layer's registration type so the shared KV cache is interpreted with
+    the target's values and shapes without aliasing target parameters.
+    """
+    for scale_name in ("_k_scale", "_v_scale"):
+        draft_scale = getattr(attn, scale_name)
+        target_scale = getattr(target_attn, scale_name)
+        scale = target_scale.detach().clone()
+        if isinstance(draft_scale, nn.Parameter):
+            scale = nn.Parameter(scale, requires_grad=False)
+        setattr(attn, scale_name, scale)
+    for scale_name in ("_k_scale_float", "_v_scale_float"):
+        setattr(attn, scale_name, getattr(target_attn, scale_name))
+    for scale_name in ("_k_scale_cpu", "_v_scale_cpu"):
+        getattr(attn, scale_name).copy_(getattr(target_attn, scale_name))
+
+
 class Gemma4Speculator(AutoRegressiveSpeculator):
     @property
     def advance_draft_positions(self) -> bool:
@@ -77,7 +98,7 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
         model: nn.Module,
         target_attn_layer_names: set[str],
     ) -> None:
-        """Wire draft layers to share KV with the target model.
+        """Wire draft layers to share KV and KV scales with the target model.
 
         Each draft decoder layer is mapped to the last non-KV-shared
         target layer of the same attention type (sliding or full).
@@ -92,7 +113,10 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
             return
 
         target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0)
-        num_non_shared = len(target_layer_types) - target_num_kv_shared
+        target_num_layers = getattr(
+            target_text_config, "num_hidden_layers", len(target_layer_types)
+        )
+        num_non_shared = target_num_layers - target_num_kv_shared
         type_to_target_indices: dict[str, list[int]] = defaultdict(list)
         for idx, lt in enumerate(target_layer_types[:num_non_shared]):
             type_to_target_indices[lt].append(idx)
@@ -128,6 +152,17 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
             target_idx = candidates[-1]
             target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn"
             attn.kv_sharing_target_layer_name = target_layer_name
+
+            # KV-cache sharing aliases the cache tensor during allocation, but
+            # the quantization scales live on the Attention modules themselves.
+            # The BF16 draft model has no quantization config, so its K/V scales
+            # otherwise remain at the default 1.0 while it reads the target's
+            # calibrated FP8 cache.
+            target_attn = self.vllm_config.compilation_config.static_forward_context[
+                target_layer_name
+            ]
+            _copy_target_kv_scales(attn, target_attn)
+
             logger.info(
                 "Gemma4 MTP: draft layer %d (%s) -> %s",
                 draft_idx,

vllm/vllm_flash_attn/flash_attn_interface.py

@@ -434,6 +434,14 @@ def flash_attn_varlen_func(
 
         from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd
 
+        # SM90 FA4 fp8-KV path: fp8 e4m3 paged K/V dequantized and the K/V descale folded
+        # in-kernel; accepts bf16/fp16 Q and writes O in its native dtype (no Q cast, no
+        # output copy). Only the (batch, num_kv_heads) f32 K/V descales are forwarded.
+        fa4_fp8_kv_dequant = (
+            k.dtype == torch.float8_e4m3fn
+            and torch.cuda.get_device_capability()[0] == 9
+        )
+
         out, softmax_lse, _, _ = _flash_attn_fwd(
             q,
             k,
@@ -458,10 +466,11 @@ def flash_attn_varlen_func(
             block_sparse_tensors=block_sparse_tensors,
             aux_tensors=aux_tensors,
             aux_tensor_leading_dims=aux_tensor_leading_dims,
-            q_descale=q_descale,
+            q_descale=None if fa4_fp8_kv_dequant else q_descale,
             k_descale=k_descale,
             v_descale=v_descale,
             output_scale=output_scale,
+            fp8_kv_dequant=fa4_fp8_kv_dequant,
         )
     else:
         raise ValueError(f"Unsupported FA version: {fa_version}")