[Kernel] Gemma-4 FA4 FP8 Kernel (#48666)

Signed-off-by: Qixiang Lin <qixiangl@nvidia.com>
Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
Co-authored-by: Qixiang Lin <qixiangl@nvidia.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>

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 f3e1a4f74c99145c0717709860bf765de1703779
+          GIT_TAG 617264c1c7955c9e84817654ebeedff069f3c5f1
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

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
@@ -94,6 +97,7 @@ def should_load_quant_weights(quant_method: QuantizeMethodBase | None) -> bool:
 
 def _largest_kernel_block_within(
     attn_backend: "type[AttentionBackend]",
+    vllm_config: VllmConfig,
     per_token_bytes: int,
     page_budget: int | None,
     fallback: int,
@@ -108,7 +112,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_for_config(vllm_config)
     candidates = [s for s in sizes if isinstance(s, int)]
     if not candidates:
         candidates = [s.base for s in sizes if isinstance(s, MultipleOf)]
@@ -629,7 +633,11 @@ class Attention(nn.Module, AttentionLayerBase):
                 sliding_window=self.sliding_window,
             ).real_page_size_bytes
             sw_block_size = _largest_kernel_block_within(
-                self.attn_backend, sw_per_token, shared_page, block_size
+                self.attn_backend,
+                vllm_config,
+                sw_per_token,
+                shared_page,
+                block_size,
             )
             return SlidingWindowSpec(
                 block_size=sw_block_size,

vllm/platforms/interface.py

@@ -612,7 +612,6 @@ class Platform:
         For hybrid models, also aligns block_size with mamba page sizes.
         """
         from vllm.config.cache import CacheConfig
-        from vllm.config.vllm import set_current_vllm_config
 
         cache_config = vllm_config.cache_config
         model_config = vllm_config.model_config
@@ -627,10 +626,9 @@ class Platform:
 
         # Phase 1: Pick block size from backend (skip if user set --block-size)
         if not cache_config.user_specified_block_size:
-            with set_current_vllm_config(vllm_config):
-                preferred = backend_cls.get_preferred_block_size(
-                    CacheConfig.DEFAULT_BLOCK_SIZE
-                )
+            preferred = backend_cls.get_preferred_block_size_for_config(
+                CacheConfig.DEFAULT_BLOCK_SIZE, vllm_config
+            )
             if preferred != CacheConfig.DEFAULT_BLOCK_SIZE:
                 logger.info(
                     "Setting kv cache block size to %d for %s backend.",

vllm/v1/attention/backend.py

@@ -69,6 +69,13 @@ class AttentionBackend(ABC):
     def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
         return [MultipleOf(1)]
 
+    @classmethod
+    def get_supported_kernel_block_sizes_for_config(
+        cls, vllm_config: "VllmConfig"
+    ) -> list[int | MultipleOf]:
+        """Return kernel block sizes for a concrete engine configuration."""
+        return cls.get_supported_kernel_block_sizes()
+
     @staticmethod
     @abstractmethod
     def get_name() -> str:
@@ -201,6 +208,13 @@ class AttentionBackend(ABC):
 
         return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes)
 
+    @classmethod
+    def get_preferred_block_size_for_config(
+        cls, default_block_size: int, vllm_config: "VllmConfig"
+    ) -> int:
+        """Return the preferred block size for a concrete engine config."""
+        return cls.get_preferred_block_size(default_block_size)
+
     @classmethod
     def indexes_kv_by_block_stride(cls) -> bool:
         """Whether the backend reads KV pages by the runtime block stride.

vllm/v1/attention/backends/fa_utils.py

@@ -241,9 +241,9 @@ def flash_attn_supports_kv_cache_dtype(
         head_size_v=head_size_v,
         has_sinks=has_sinks,
     )
-    return (fa_version == 3 and current_platform.is_device_capability_family(90)) or (
-        fa_version == 4 and current_platform.is_device_capability_family(100)
-    )
+    return (
+        fa_version in (3, 4) and current_platform.is_device_capability_family(90)
+    ) or (fa_version == 4 and current_platform.is_device_capability_family(100))
 
 
 def flash_attn_supports_quant_query_input() -> bool:

vllm/v1/attention/backends/flash_attn.py

@@ -85,13 +85,58 @@ class FlashAttentionBackend(AttentionBackend):
     ]
 
     @staticmethod
-    def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
+    def _get_sm90_fa4_fp8_kv_block_size(
+        vllm_config: VllmConfig | None = None,
+    ) -> int | None:
+        if vllm_config 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()
+        if (
+            current_platform.is_device_capability_family(90)
+            and vllm_config.cache_config.cache_dtype in ("fp8", "fp8_e4m3")
+            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_supported_kernel_block_sizes(cls) -> list[int | MultipleOf]:
+        if block_size := cls._get_sm90_fa4_fp8_kv_block_size():
+            # Sliding-window cache specs select the smallest advertised size.
+            # Report the kernel's exact page-size contract instead of the
+            # generic FlashAttention multiple-of-16 capability.
+            return [block_size]
+        return [MultipleOf(16)]
+
+    @classmethod
+    def get_supported_kernel_block_sizes_for_config(
+        cls, vllm_config: VllmConfig
+    ) -> list[int | MultipleOf]:
+        if block_size := cls._get_sm90_fa4_fp8_kv_block_size(vllm_config):
+            return [block_size]
         return [MultipleOf(16)]
 
     forward_includes_kv_cache_update: bool = False
 
     @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 current_platform.is_xpu():
+            return max(default_block_size, 64)
+        return super().get_preferred_block_size(default_block_size)
+
+    @classmethod
+    def get_preferred_block_size_for_config(
+        cls, default_block_size: int, vllm_config: VllmConfig
+    ) -> int:
+        if block_size := cls._get_sm90_fa4_fp8_kv_block_size(vllm_config):
+            return max(default_block_size, block_size)
         if current_platform.is_xpu():
             return max(default_block_size, 64)
         return super().get_preferred_block_size(default_block_size)
@@ -895,7 +940,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
+        )
 
         vllm_config = get_current_vllm_config_or_none()
         dcp_a2a = (

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

@@ -23,6 +23,23 @@ 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. Re-register
+    cloned buffers on the draft layer 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"):
+        target_scale = getattr(target_attn, scale_name)
+        attn.register_buffer(scale_name, target_scale.detach().clone())
+    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:
@@ -76,7 +93,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,15 +109,19 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
 
         target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0)
         num_non_shared = len(target_layer_types) - 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)
-
-        target_prefix = "model.layers"
+        target_names_by_index: dict[int, str] = {}
         for name in target_attn_layer_names:
-            if ".layers." in name:
-                target_prefix = name.split(".layers.")[0] + ".layers"
-                break
+            _, separator, layer_suffix = name.partition(".layers.")
+            if not separator:
+                continue
+            layer_index, _, _ = layer_suffix.partition(".")
+            if layer_index.isdigit():
+                target_names_by_index[int(layer_index)] = name
+
+        type_to_target_names: dict[str, list[str]] = defaultdict(list)
+        for idx, lt in enumerate(target_layer_types[:num_non_shared]):
+            if target_name := target_names_by_index.get(idx):
+                type_to_target_names[lt].append(target_name)
 
         draft_layer_types = getattr(draft_text_config, "layer_types", [])
         for draft_idx, layer in enumerate(model.model.layers):
@@ -115,7 +136,7 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
                 if draft_idx < len(draft_layer_types)
                 else "full_attention"
             )
-            candidates = type_to_target_indices.get(draft_layer_type, [])
+            candidates = type_to_target_names.get(draft_layer_type, [])
             if not candidates:
                 logger.warning(
                     "No target layer of type '%s' for draft layer %d",
@@ -124,9 +145,19 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
                 )
                 continue
 
-            target_idx = candidates[-1]
-            target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn"
+            target_layer_name = candidates[-1]
             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

@@ -404,6 +404,22 @@ 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
+        )
+        if fa4_fp8_kv_dequant:
+            fa4_q_descale = None
+            fa4_k_descale = k_descale
+            fa4_v_descale = v_descale
+        else:
+            fa4_q_descale = None
+            fa4_k_descale = None
+            fa4_v_descale = None
+
         out, softmax_lse, _, _ = _flash_attn_fwd(
             q,
             k,
@@ -428,10 +444,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,
-            k_descale=k_descale,
-            v_descale=v_descale,
+            q_descale=fa4_q_descale,
+            k_descale=fa4_k_descale,
+            v_descale=fa4_v_descale,
             output_scale=output_scale,
+            fp8_kv_dequant=fa4_fp8_kv_dequant,
         )
     else:
         raise ValueError(f"Unsupported FA version: {fa_version}")