diff --git a/cmake/external_projects/vllm_flash_attn.cmake b/cmake/external_projects/vllm_flash_attn.cmake index 2a8e9cc781..b70db3adeb 100644 --- a/cmake/external_projects/vllm_flash_attn.cmake +++ b/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 diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 8eb8aaa866..e5960e68a4 100644 --- a/vllm/model_executor/layers/attention/attention.py +++ b/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, diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 6781a9a86e..939094d4c7 100644 --- a/vllm/platforms/interface.py +++ b/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.", diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index e3f05df8dc..12d19221fb 100644 --- a/vllm/v1/attention/backend.py +++ b/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. diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 22f9e6a14d..ae02c30805 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/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: diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index e7ef03d99c..92eab051b4 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/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 = ( diff --git a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py index dfa2c68010..34fc5f4bef 100644 --- a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py +++ b/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, diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 81768e1209..aba0d29727 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/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}")