Revert "[Kernel] Gemma-4 FA4 FP8 Kernel" (#52987)
Roger Wang committed
Aug 19, 2026 at 10:48 UTC
755492e37d7d7201d93f0effa200acac56601f38
8 files changed
+29
-152
cmake/external_projects/vllm_flash_attn.cmake
+1
-1
@@ -39,7 +39,7 @@ else()
39
FetchContent_Declare(
40
vllm-flash-attn
41
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
42
- GIT_TAG 617264c1c7955c9e84817654ebeedff069f3c5f1
42
+ GIT_TAG f3e1a4f74c99145c0717709860bf765de1703779
43
GIT_PROGRESS TRUE
44
# Don't share the vllm-flash-attn build between build types
45
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
vllm/model_executor/layers/attention/attention.py
+3
-11
@@ -8,10 +8,7 @@ import torch.nn as nn
8
9
import vllm.envs as envs
10
from vllm.compilation.breakable_cudagraph import eager_break_during_capture
11
-from vllm.config import (
12
- CacheConfig,
13
- get_current_vllm_config,
14
-)
11
+from vllm.config import CacheConfig, get_current_vllm_config
12
from vllm.config.vllm import VllmConfig
13
from vllm.forward_context import ForwardContext, get_forward_context
14
from vllm.logger import init_logger
@@ -97,7 +94,6 @@ def should_load_quant_weights(quant_method: QuantizeMethodBase | None) -> bool:
94
95
def _largest_kernel_block_within(
96
attn_backend: "type[AttentionBackend]",
100
- vllm_config: VllmConfig,
97
per_token_bytes: int,
98
page_budget: int | None,
99
fallback: int,
@@ -112,7 +108,7 @@ def _largest_kernel_block_within(
108
"""
109
from vllm.v1.attention.backend import MultipleOf
110
115
- sizes = attn_backend.get_supported_kernel_block_sizes_for_config(vllm_config)
111
+ sizes = attn_backend.get_supported_kernel_block_sizes()
112
candidates = [s for s in sizes if isinstance(s, int)]
113
if not candidates:
114
candidates = [s.base for s in sizes if isinstance(s, MultipleOf)]
@@ -636,11 +632,7 @@ class Attention(nn.Module, AttentionLayerBase):
632
)
633
).real_page_size_bytes
634
sw_block_size = _largest_kernel_block_within(
639
- self.attn_backend,
640
- vllm_config,
641
- sw_per_token,
642
- shared_page,
643
- block_size,
635
+ self.attn_backend, sw_per_token, shared_page, block_size
636
)
637
return SlidingWindowSpec(
638
block_size=sw_block_size,
vllm/platforms/interface.py
+5
-3
@@ -612,6 +612,7 @@ class Platform:
612
For hybrid models, also aligns block_size with mamba page sizes.
613
"""
614
from vllm.config.cache import CacheConfig
615
+ from vllm.config.vllm import set_current_vllm_config
616
617
cache_config = vllm_config.cache_config
618
model_config = vllm_config.model_config
@@ -626,9 +627,10 @@ class Platform:
627
628
# Phase 1: Pick block size from backend (skip if user set --block-size)
629
if not cache_config.user_specified_block_size:
629
- preferred = backend_cls.get_preferred_block_size_for_config(
630
- CacheConfig.DEFAULT_BLOCK_SIZE, vllm_config
631
- )
630
+ with set_current_vllm_config(vllm_config):
631
+ preferred = backend_cls.get_preferred_block_size(
632
+ CacheConfig.DEFAULT_BLOCK_SIZE
633
+ )
634
if preferred != CacheConfig.DEFAULT_BLOCK_SIZE:
635
logger.info(
636
"Setting kv cache block size to %d for %s backend.",
vllm/v1/attention/backend.py
-14
@@ -69,13 +69,6 @@ class AttentionBackend(ABC):
69
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
70
return [MultipleOf(1)]
71
72
- @classmethod
73
- def get_supported_kernel_block_sizes_for_config(
74
- cls, vllm_config: "VllmConfig"
75
- ) -> list[int | MultipleOf]:
76
- """Return kernel block sizes for a concrete engine configuration."""
77
- return cls.get_supported_kernel_block_sizes()
78
-
72
@staticmethod
73
@abstractmethod
74
def get_name() -> str:
@@ -222,13 +215,6 @@ class AttentionBackend(ABC):
215
216
return min(s.base if isinstance(s, MultipleOf) else s for s in supported_sizes)
217
225
- @classmethod
226
- def get_preferred_block_size_for_config(
227
- cls, default_block_size: int, vllm_config: "VllmConfig"
228
- ) -> int:
229
- """Return the preferred block size for a concrete engine config."""
230
- return cls.get_preferred_block_size(default_block_size)
231
-
218
@classmethod
219
def indexes_kv_by_block_stride(cls) -> bool:
220
"""Whether the backend reads KV pages by the runtime block stride.
vllm/v1/attention/backends/fa_utils.py
+3
-3
@@ -238,9 +238,9 @@ def flash_attn_supports_kv_cache_dtype(
238
head_size_v=head_size_v,
239
has_sinks=has_sinks,
240
)
241
- return (
242
- fa_version in (3, 4) and current_platform.is_device_capability_family(90)
243
- ) or (fa_version == 4 and current_platform.is_device_capability_family(100))
241
+ return (fa_version == 3 and current_platform.is_device_capability_family(90)) or (
242
+ fa_version == 4 and current_platform.is_device_capability_family(100)
243
+ )
244
245
246
def flash_attn_supports_quant_query_input() -> bool:
vllm/v1/attention/backends/flash_attn.py
+2
-57
@@ -85,58 +85,13 @@ class FlashAttentionBackend(AttentionBackend):
85
]
86
87
@staticmethod
88
- def _get_sm90_fa4_fp8_kv_block_size(
89
- vllm_config: VllmConfig | None = None,
90
- ) -> int | None:
91
- if vllm_config is None:
92
- vllm_config = get_current_vllm_config_or_none()
93
- if vllm_config is None or vllm_config.model_config is None:
94
- return None
95
-
96
- head_size = vllm_config.model_config.get_head_size()
97
- if (
98
- current_platform.is_device_capability_family(90)
99
- and vllm_config.cache_config.cache_dtype in ("fp8", "fp8_e4m3")
100
- and head_size == 512
101
- and get_flash_attn_version(head_size=head_size) == 4
102
- ):
103
- # The SM90 FP8-KV-dequant kernel uses a 64-token TMA tile/page.
104
- return 64
105
- return None
106
-
107
- @classmethod
108
- def get_supported_kernel_block_sizes(cls) -> list[int | MultipleOf]:
109
- if block_size := cls._get_sm90_fa4_fp8_kv_block_size():
110
- # Sliding-window cache specs select the smallest advertised size.
111
- # Report the kernel's exact page-size contract instead of the
112
- # generic FlashAttention multiple-of-16 capability.
113
- return [block_size]
114
- return [MultipleOf(16)]
115
-
116
- @classmethod
117
- def get_supported_kernel_block_sizes_for_config(
118
- cls, vllm_config: VllmConfig
119
- ) -> list[int | MultipleOf]:
120
- if block_size := cls._get_sm90_fa4_fp8_kv_block_size(vllm_config):
121
- return [block_size]
88
+ def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
89
return [MultipleOf(16)]
90
91
forward_includes_kv_cache_update: bool = False
92
93
@classmethod
94
def get_preferred_block_size(cls, default_block_size: int) -> int:
128
- if block_size := cls._get_sm90_fa4_fp8_kv_block_size():
129
- return max(default_block_size, block_size)
130
- if current_platform.is_xpu():
131
- return max(default_block_size, 64)
132
- return super().get_preferred_block_size(default_block_size)
133
-
134
- @classmethod
135
- def get_preferred_block_size_for_config(
136
- cls, default_block_size: int, vllm_config: VllmConfig
137
- ) -> int:
138
- if block_size := cls._get_sm90_fa4_fp8_kv_block_size(vllm_config):
139
- return max(default_block_size, block_size)
95
if current_platform.is_xpu():
96
return max(default_block_size, 64)
97
return super().get_preferred_block_size(default_block_size)
@@ -939,17 +894,7 @@ class FlashAttentionImpl(AttentionImpl):
894
"heads in the layer"
895
)
896
942
- # FA4's SM90 FP8-KV path consumes native FP16/BF16 Q and dequantizes
943
- # FP8 K/V in-kernel. Other FA4 paths (notably SM100) still require Q,
944
- # K, and V to have the same FP8 dtype.
945
- uses_sm90_fa4_fp8_kv_dequant = (
946
- self.vllm_flash_attn_version == 4
947
- and current_platform.is_device_capability_family(90)
948
- and self.kv_cache_dtype in ("fp8", "fp8_e4m3")
949
- )
950
- self.supports_quant_query_input = flash_attn_supports_quant_query_input() and (
951
- not uses_sm90_fa4_fp8_kv_dequant
952
- )
897
+ self.supports_quant_query_input = flash_attn_supports_quant_query_input()
898
899
vllm_config = get_current_vllm_config_or_none()
900
dcp_a2a = (
vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py
+12
-43
@@ -23,23 +23,6 @@ from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
23
logger = init_logger(__name__)
24
25
26
-def _copy_target_kv_scales(attn: nn.Module, target_attn: nn.Module) -> None:
27
- """Copy target KV scales while preserving their tensor representation.
28
-
29
- Default attention scales are scalar buffers, while some quantization
30
- methods replace them with length-one or per-head parameters. Re-register
31
- cloned buffers on the draft layer so the shared KV cache is interpreted
32
- with the target's values and shapes without aliasing target parameters.
33
- """
34
- for scale_name in ("_k_scale", "_v_scale"):
35
- target_scale = getattr(target_attn, scale_name)
36
- attn.register_buffer(scale_name, target_scale.detach().clone())
37
- for scale_name in ("_k_scale_float", "_v_scale_float"):
38
- setattr(attn, scale_name, getattr(target_attn, scale_name))
39
- for scale_name in ("_k_scale_cpu", "_v_scale_cpu"):
40
- getattr(attn, scale_name).copy_(getattr(target_attn, scale_name))
41
-
42
-
26
class Gemma4Speculator(AutoRegressiveSpeculator):
27
@property
28
def advance_draft_positions(self) -> bool:
@@ -93,7 +76,7 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
76
model: nn.Module,
77
target_attn_layer_names: set[str],
78
) -> None:
96
- """Wire draft layers to share KV and KV scales with the target model.
79
+ """Wire draft layers to share KV with the target model.
80
81
Each draft decoder layer is mapped to the last non-KV-shared
82
target layer of the same attention type (sliding or full).
@@ -109,19 +92,15 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
92
93
target_num_kv_shared = getattr(target_text_config, "num_kv_shared_layers", 0)
94
num_non_shared = len(target_layer_types) - target_num_kv_shared
112
- target_names_by_index: dict[int, str] = {}
113
- for name in target_attn_layer_names:
114
- _, separator, layer_suffix = name.partition(".layers.")
115
- if not separator:
116
- continue
117
- layer_index, _, _ = layer_suffix.partition(".")
118
- if layer_index.isdigit():
119
- target_names_by_index[int(layer_index)] = name
120
-
121
- type_to_target_names: dict[str, list[str]] = defaultdict(list)
95
+ type_to_target_indices: dict[str, list[int]] = defaultdict(list)
96
for idx, lt in enumerate(target_layer_types[:num_non_shared]):
123
- if target_name := target_names_by_index.get(idx):
124
- type_to_target_names[lt].append(target_name)
97
+ type_to_target_indices[lt].append(idx)
98
+
99
+ target_prefix = "model.layers"
100
+ for name in target_attn_layer_names:
101
+ if ".layers." in name:
102
+ target_prefix = name.split(".layers.")[0] + ".layers"
103
+ break
104
105
draft_layer_types = getattr(draft_text_config, "layer_types", [])
106
for draft_idx, layer in enumerate(model.model.layers):
@@ -136,7 +115,7 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
115
if draft_idx < len(draft_layer_types)
116
else "full_attention"
117
)
139
- candidates = type_to_target_names.get(draft_layer_type, [])
118
+ candidates = type_to_target_indices.get(draft_layer_type, [])
119
if not candidates:
120
logger.warning(
121
"No target layer of type '%s' for draft layer %d",
@@ -145,19 +124,9 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
124
)
125
continue
126
148
- target_layer_name = candidates[-1]
127
+ target_idx = candidates[-1]
128
+ target_layer_name = f"{target_prefix}.{target_idx}.self_attn.attn"
129
attn.kv_sharing_target_layer_name = target_layer_name
150
-
151
- # KV-cache sharing aliases the cache tensor during allocation, but
152
- # the quantization scales live on the Attention modules themselves.
153
- # The BF16 draft model has no quantization config, so its K/V scales
154
- # otherwise remain at the default 1.0 while it reads the target's
155
- # calibrated FP8 cache.
156
- target_attn = self.vllm_config.compilation_config.static_forward_context[
157
- target_layer_name
158
- ]
159
- _copy_target_kv_scales(attn, target_attn)
160
-
130
logger.info(
131
"Gemma4 MTP: draft layer %d (%s) -> %s",
132
draft_idx,
vllm/vllm_flash_attn/flash_attn_interface.py
+3
-20
@@ -404,22 +404,6 @@ def flash_attn_varlen_func(
404
405
from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd
406
407
- # SM90 FA4 fp8-KV path: fp8 e4m3 paged K/V dequantized and the K/V descale folded
408
- # in-kernel; accepts bf16/fp16 Q and writes O in its native dtype (no Q cast, no
409
- # output copy). Only the (batch, num_kv_heads) f32 K/V descales are forwarded.
410
- fa4_fp8_kv_dequant = (
411
- k.dtype == torch.float8_e4m3fn
412
- and torch.cuda.get_device_capability()[0] == 9
413
- )
414
- if fa4_fp8_kv_dequant:
415
- fa4_q_descale = None
416
- fa4_k_descale = k_descale
417
- fa4_v_descale = v_descale
418
- else:
419
- fa4_q_descale = None
420
- fa4_k_descale = None
421
- fa4_v_descale = None
422
-
407
out, softmax_lse, _, _ = _flash_attn_fwd(
408
q,
409
k,
@@ -444,11 +428,10 @@ def flash_attn_varlen_func(
428
block_sparse_tensors=block_sparse_tensors,
429
aux_tensors=aux_tensors,
430
aux_tensor_leading_dims=aux_tensor_leading_dims,
447
- q_descale=fa4_q_descale,
448
- k_descale=fa4_k_descale,
449
- v_descale=fa4_v_descale,
431
+ q_descale=q_descale,
432
+ k_descale=k_descale,
433
+ v_descale=v_descale,
434
output_scale=output_scale,
451
- fp8_kv_dequant=fa4_fp8_kv_dequant,
435
)
436
else:
437
raise ValueError(f"Unsupported FA version: {fa_version}")