diff --git a/tests/kernels/attention/test_attention_selector.py b/tests/kernels/attention/test_attention_selector.py index 53867a5edf..c5cd4af08f 100644 --- a/tests/kernels/attention/test_attention_selector.py +++ b/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 diff --git a/tests/v1/attention/test_backend_per_kind.py b/tests/v1/attention/test_backend_per_kind.py index ab77d58825..53dedf3d6b 100644 --- a/tests/v1/attention/test_backend_per_kind.py +++ b/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 diff --git a/vllm/model_executor/layers/attention/attention.py b/vllm/model_executor/layers/attention/attention.py index 79012a4a5c..bec1856d93 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 @@ -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, diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 3a0392c523..4073512d96 100644 --- a/vllm/model_executor/models/config.py +++ b/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( diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index 41af0c7660..e6e0fd959e 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/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 diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 883ef5ecaf..c3b7092fbe 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/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 diff --git a/vllm/models/deepseek_v4/sparse_mla.py b/vllm/models/deepseek_v4/sparse_mla.py index a4000183a0..cdab7d23ae 100644 --- a/vllm/models/deepseek_v4/sparse_mla.py +++ b/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 diff --git a/vllm/models/deepseek_v41/compressor.py b/vllm/models/deepseek_v41/compressor.py index 44440fcb61..42335ae84b 100644 --- a/vllm/models/deepseek_v41/compressor.py +++ b/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 diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 485cdd9ebd..949f383324 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/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 diff --git a/vllm/models/deepseek_v41/sparse_mla.py b/vllm/models/deepseek_v41/sparse_mla.py index 8db6d4ed2b..1130eec090 100644 --- a/vllm/models/deepseek_v41/sparse_mla.py +++ b/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 diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index ac02bca1fc..2459d4ef85 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/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 diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index f385f82067..75aa90100c 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/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] diff --git a/vllm/models/qwen4_exp/amd/qsa.py b/vllm/models/qwen4_exp/amd/qsa.py index bed48a4cba..3c8a61bd54 100644 --- a/vllm/models/qwen4_exp/amd/qsa.py +++ b/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)] diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index e480142a22..b6da89a863 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/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)] diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index a403c9ae39..8b51e39169 100644 --- a/vllm/v1/attention/backend.py +++ b/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 diff --git a/vllm/v1/attention/backends/b12x.py b/vllm/v1/attention/backends/b12x.py index d2b28eed0a..833372ab44 100644 --- a/vllm/v1/attention/backends/b12x.py +++ b/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 diff --git a/vllm/v1/attention/backends/composite.py b/vllm/v1/attention/backends/composite.py index d091a29d87..0ba40e99f2 100644 --- a/vllm/v1/attention/backends/composite.py +++ b/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): diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index 36df340967..b4f25f4194 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/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 diff --git a/vllm/v1/attention/backends/fa_utils.py b/vllm/v1/attention/backends/fa_utils.py index 650fc1dea0..83348bb94a 100644 --- a/vllm/v1/attention/backends/fa_utils.py +++ b/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) ) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index b57af3fd29..129b3e39c3 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/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 diff --git a/vllm/v1/attention/backends/flashinfer.py b/vllm/v1/attention/backends/flashinfer.py index 3113c261dd..2cfaef8166 100755 --- a/vllm/v1/attention/backends/flashinfer.py +++ b/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. diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index 9d0ac08172..1b9a21b655 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/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)] diff --git a/vllm/v1/attention/backends/hpc_attn.py b/vllm/v1/attention/backends/hpc_attn.py index 14545307fe..cbe97da45c 100644 --- a/vllm/v1/attention/backends/hpc_attn.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/amx_mla.py b/vllm/v1/attention/backends/mla/amx_mla.py index 38c8b68af9..16689887af 100644 --- a/vllm/v1/attention/backends/mla/amx_mla.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/cpu_mla.py b/vllm/v1/attention/backends/mla/cpu_mla.py index 91d3cb3339..431f6c3185 100644 --- a/vllm/v1/attention/backends/mla/cpu_mla.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/cutlass_mla.py b/vllm/v1/attention/backends/mla/cutlass_mla.py index 575b9d32a9..b88a4a17a8 100644 --- a/vllm/v1/attention/backends/mla/cutlass_mla.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/flashattn_mla.py b/vllm/v1/attention/backends/mla/flashattn_mla.py index c3729fd06d..b685be102e 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py index 2d1e218640..6048f6b5fd 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index fea49ca154..ca78961eb5 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 142a63ec72..32f67e7e93 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index 830bc52ab9..d39b7343a4 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/flashmla.py b/vllm/v1/attention/backends/mla/flashmla.py index d17ae05c1a..c8cd31d564 100644 --- a/vllm/v1/attention/backends/mla/flashmla.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py index 2e71f0a35c..014358f530 100644 --- a/vllm/v1/attention/backends/mla/flashmla_sparse.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 1ede0442bc..e7eb67b61f 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/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] diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index 3fa5f67bf5..1ab4df0664 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py index f1c3b558be..f6258759d0 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index 7b07dae71c..78727c2eb5 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index 20cddfa4ca..b3d982eac3 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/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 diff --git a/vllm/v1/attention/backends/mla/triton_mla.py b/vllm/v1/attention/backends/mla/triton_mla.py index 3deb996a5b..950d798f22 100644 --- a/vllm/v1/attention/backends/mla/triton_mla.py +++ b/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 diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 4ce8b23d31..3df3b26bba 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/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 diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 711a8f8755..6aa456e22c 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/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 diff --git a/vllm/v1/attention/backends/rocm_attn.py b/vllm/v1/attention/backends/rocm_attn.py index 2671249fb2..76382aa75e 100644 --- a/vllm/v1/attention/backends/rocm_attn.py +++ b/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. diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 1b290c0ae9..8a2fab4cc7 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/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 diff --git a/vllm/v1/attention/backends/triton_flash_attn.py b/vllm/v1/attention/backends/triton_flash_attn.py index 0853b291cd..152a92048e 100644 --- a/vllm/v1/attention/backends/triton_flash_attn.py +++ b/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__, diff --git a/vllm/v1/attention/backends/turboquant_attn.py b/vllm/v1/attention/backends/turboquant_attn.py index 8542e6164f..056e42cf62 100644 --- a/vllm/v1/attention/backends/turboquant_attn.py +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py b/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py index 9a489bcc48..6f76e5055a 100644 --- a/vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py +++ b/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, diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index c10b381e40..3c4c0d8948 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/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}")