From 6deb05e0e4f6c298a60975f188cd044773ae6dec Mon Sep 17 00:00:00 2001 From: Luciano Martins <22145370+lucianommartins@users.noreply.github.com> Date: Wed, 10 Jun 2026 02:45:39 +0200 Subject: [PATCH] [Core][Model] Gemma4: Unified FA4 for all layers + FlashAttention mm_prefix support (#42175) Signed-off-by: Luciano Martins Signed-off-by: Lucas Wilkinson Signed-off-by: Matthew Bonanni Co-authored-by: Luciano Martins Co-authored-by: Lucas Wilkinson Co-authored-by: Matthew Bonanni --- vllm/model_executor/models/config.py | 65 +++++++------ vllm/v1/attention/backend.py | 8 ++ vllm/v1/attention/backends/flash_attn.py | 92 +++++++++++++++++- vllm/v1/attention/backends/flex_attention.py | 1 + .../attention/backends/mla/flashattn_mla.py | 1 + .../attention/backends/mla/flashinfer_mla.py | 1 + .../backends/mla/flashinfer_mla_sparse.py | 1 + vllm/v1/attention/backends/mla/flashmla.py | 1 + .../attention/backends/mla/tokenspeed_mla.py | 1 + vllm/v1/attention/backends/triton_attn.py | 50 +++------- vllm/v1/attention/backends/utils.py | 30 ++++++ vllm/v1/spec_decode/dflash.py | 5 + vllm/v1/worker/gpu_model_runner.py | 95 +++++-------------- vllm/vllm_flash_attn/flash_attn_interface.py | 13 +++ 14 files changed, 222 insertions(+), 142 deletions(-) diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index ebd1c53e81..64d606c289 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -57,50 +57,49 @@ class Gemma3TextModelConfig(VerifyAndUpdateConfig): class Gemma4Config(VerifyAndUpdateConfig): @staticmethod def verify_and_update_config(vllm_config: "VllmConfig") -> None: - """Force unified attention backend for models with heterogeneous - head dimensions. + """Configure attention for heterogeneous head dimensions. - Some Gemma4 variants use different head dimensions for - sliding window (head_dim) vs full attention (global_head_dim) layers. - When global_head_dim > 256, FlashAttention rejects those layers - (head_size <= 256 kernel limit), causing vLLM to select a different - backend for each layer type. This mixed-backend execution produces - numerical divergence and output corruption. + Gemma4 uses different head dimensions for sliding window + (head_dim) vs full attention (global_head_dim) layers. The + default FA3 on Hopper cannot handle head_dim > 256, which + causes mixed backend selection and numerical divergence. - The fix detects heterogeneous head dimensions from the model config - and forces TRITON_ATTN (which has no head_size ceiling) for all - layers when the user hasn't explicitly chosen a backend. - - TODO: Heterogeneous head_sizes (head_dim != global_head_dim) - require NixlConnector changes to support per-layer KV transfer - with different head dimensions for prefill-decode disaggregation. + 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. """ hf_text_config = vllm_config.model_config.hf_text_config head_dim = getattr(hf_text_config, "head_dim", None) global_head_dim = getattr(hf_text_config, "global_head_dim", None) - # Only force Triton when head dimensions actually differ AND the - # larger one exceeds FlashAttention's kernel limit (head_size <= 256). - # This avoids unnecessary backend forcing on smaller models where - # the config carries global_head_dim but all layers can still use - # the same FA backend. - max_head_dim = max(head_dim or 0, global_head_dim or 0) - if ( - head_dim is not None - and global_head_dim is not None - and head_dim != global_head_dim - and max_head_dim > 256 - and vllm_config.attention_config.backend is None - ): - from vllm.v1.attention.backends.registry import ( - AttentionBackendEnum, - ) + if head_dim is None or global_head_dim is None or head_dim == global_head_dim: + return + from vllm.v1.attention.backends.fa_utils import is_fa_version_supported + from vllm.v1.attention.backends.registry import AttentionBackendEnum + + max_head_dim = max(head_dim, global_head_dim) + + if is_fa_version_supported(4) and max_head_dim <= 512: + if ( + vllm_config.attention_config.flash_attn_version is None + 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 " + "(head_dim=%d, global_head_dim=%d). Using FA4 for " + "all layers to avoid mixed FA3/FA4 penalty.", + head_dim, + global_head_dim, + ) + elif vllm_config.attention_config.backend is None: vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN logger.info( "Gemma4 model has heterogeneous head dimensions " - "(head_dim=%d, global_head_dim=%d). Forcing TRITON_ATTN " - "backend to prevent mixed-backend numerical divergence.", + "(head_dim=%d, global_head_dim=%d). FA4 not available, " + "forcing TRITON_ATTN backend.", head_dim, global_head_dim, ) diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 4b4a4435b3..32b4b8ab9a 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -267,6 +267,7 @@ class AttentionBackend(ABC): use_mla: bool, has_sink: bool, use_sparse: bool, + use_mm_prefix: bool, device_capability: "DeviceCapability", ) -> str | None: return None @@ -334,6 +335,7 @@ class AttentionBackend(ABC): use_mla, has_sink, use_sparse, + use_mm_prefix, device_capability, ) if combination_reason is not None: @@ -415,6 +417,12 @@ class CommonAttentionMetadata: decode rows (assumes every draft was accepted). Not safe for kernels that need exact per-row context lengths on decode rows.""" + mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None + """PrefixLM bidirectional ranges for multimodal tokens. Maps + request index to list of (start, end) token position ranges + where bidirectional attention should apply. None for text-only + batches or non-PrefixLM models.""" + # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0) _seq_lens_cpu: torch.Tensor | None = None _num_computed_tokens_cpu: torch.Tensor | None = None diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index c56c4ee6e1..d6774a6eb9 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -192,6 +192,10 @@ class FlashAttentionBackend(AttentionBackend): ) return kv_cache_dtype in ["auto", "float16", "bfloat16"] + @classmethod + def supports_mm_prefix(cls) -> bool: + return is_fa_version_supported(4) + @classmethod def supports_sink(cls) -> bool: if not is_flash_attn_varlen_func_available(): @@ -212,10 +216,20 @@ class FlashAttentionBackend(AttentionBackend): use_mla: bool, has_sink: bool, use_sparse: bool, + use_mm_prefix: bool, device_capability: DeviceCapability, ) -> 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 get_flash_attn_version(head_size=head_size, has_sinks=has_sink) != 4 + ): + return ( + "mm_prefix (PrefixLM bidirectional attention) requires " + "FlashAttention v4, which does not resolve for this " + "head_size" + ) return None @@ -255,6 +269,10 @@ class FlashAttentionMetadata: causal: bool = True + # PrefixLM bidirectional ranges for multimodal tokens. + # Shape: (num_seqs, max_ranges, 2) int32, [start, end] per range. + mm_prefix_range_tensor: torch.Tensor | None = None + def _get_sliding_window_configs( vllm_config: VllmConfig, @@ -572,6 +590,19 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad max_num_splits=max_num_splits, causal=causal, ) + + # Compute mm_prefix range tensor if the batch contains + # multimodal tokens with bidirectional ranges. + mm_ranges = common_attn_metadata.mm_req_doc_ranges + if mm_ranges is not None: + from vllm.v1.attention.backends.utils import ( + compute_mm_prefix_range_tensor, + ) + + attn_metadata.mm_prefix_range_tensor = compute_mm_prefix_range_tensor( + mm_ranges, num_reqs, seq_lens.device + ) + return attn_metadata def update_block_table( @@ -793,6 +824,18 @@ class FlashAttentionImpl(AttentionImpl): if self.sliding_window is not None else None ) + mm_prefix_ranges = attn_metadata.mm_prefix_range_tensor + mm_mask_mod = None + mm_aux = None + if ( + mm_prefix_ranges is not None + and attn_metadata.causal + and self.vllm_flash_attn_version == 4 + ): + max_ranges = mm_prefix_ranges.shape[1] + mm_mask_mod = _make_mm_prefix_mask_mod(max_ranges) + mm_aux = [mm_prefix_ranges] + flash_attn_varlen_func( q=query[:num_actual_tokens], k=key_cache, @@ -815,6 +858,8 @@ class FlashAttentionImpl(AttentionImpl): v_descale=v_descale, num_splits=attn_metadata.max_num_splits, s_aux=self.sinks, + mask_mod=mm_mask_mod, + aux_tensors=mm_aux, ) return output @@ -1040,17 +1085,58 @@ class FlashAttentionImpl(AttentionImpl): window_size=sliding_window_size, softcap=self.logits_soft_cap, fa_version=self.vllm_flash_attn_version, - q_descale=layer._q_scale.expand(descale_shape) + q_descale=layer._q_scale.expand(descale_shape) # type: ignore[operator] if self.supports_quant_query_input else None, - k_descale=layer._k_scale.expand(descale_shape), - v_descale=layer._v_scale.expand(descale_shape), + k_descale=layer._k_scale.expand(descale_shape), # type: ignore[operator] + v_descale=layer._v_scale.expand(descale_shape), # type: ignore[operator] num_splits=1 if self.batch_invariant_enabled else 0, ) return output +def _make_mm_prefix_mask_mod(max_ranges: int): + """Build a CuTE-DSL mask_mod implementing (causal OR mm_prefix). + + Returns a @cute.jit callable that evaluates: + keep = (kv_idx <= q_idx) OR + (q_idx in [r_start,r_end] AND kv_idx in [r_start,r_end]) + for each mm_prefix range stored in aux_tensors[0]. + """ + import cutlass + import cutlass.cute as cute + from cutlass import Int32 # type: ignore[attr-defined] + + from vllm.vllm_flash_attn.cute.utils import ( # type: ignore[import-untyped] + scalar_to_ssa, + ) + + @cute.jit + def mm_prefix_mask_mod( + batch_idx: cute.TensorSSA, + head_idx: cute.TensorSSA, + q_idx: cute.TensorSSA, + kv_idx: cute.TensorSSA, + seqlen_info, + aux_tensors, + ): + keep = kv_idx <= q_idx + ranges = aux_tensors[0] + b = batch_idx[0] + for i in cutlass.range_constexpr(max_ranges): # type: ignore[attr-defined] + r_start = scalar_to_ssa(ranges[b, i, 0], Int32) + r_end = scalar_to_ssa(ranges[b, i, 1], Int32) + valid = r_start < r_end + q_in = (q_idx >= r_start) & (q_idx <= r_end) & valid + k_in = (kv_idx >= r_start) & (kv_idx <= r_end) & valid + keep = keep | (q_in & k_in) + return keep + + mm_prefix_mask_mod.use_fast_sampling = True + return mm_prefix_mask_mod + + def use_cascade_attention( common_prefix_len: int, query_lens: np.ndarray, diff --git a/vllm/v1/attention/backends/flex_attention.py b/vllm/v1/attention/backends/flex_attention.py index b870142520..f1299f53ad 100644 --- a/vllm/v1/attention/backends/flex_attention.py +++ b/vllm/v1/attention/backends/flex_attention.py @@ -957,6 +957,7 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat persistent_kv_indices=self.persistent_kv_indices, persistent_kv_num_blocks=self.persistent_kv_num_blocks, persistent_doc_ids=self.persistent_doc_ids, + mm_prefix_range=common_attn_metadata.mm_req_doc_ranges, ) # Pre-build block_mask so it is ready before CUDA graph capture. diff --git a/vllm/v1/attention/backends/mla/flashattn_mla.py b/vllm/v1/attention/backends/mla/flashattn_mla.py index bd947296e8..63daa860fd 100644 --- a/vllm/v1/attention/backends/mla/flashattn_mla.py +++ b/vllm/v1/attention/backends/mla/flashattn_mla.py @@ -82,6 +82,7 @@ class FlashAttnMLABackend(MLACommonBackend): use_mla: bool, has_sink: bool, use_sparse: bool, + use_mm_prefix: bool, device_capability: DeviceCapability, ) -> str | None: if not flash_attn_supports_mla(): diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index e98bee9d79..e3d8637deb 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py @@ -75,6 +75,7 @@ class FlashInferMLABackend(MLACommonBackend): use_mla: bool, has_sink: bool, use_sparse: bool, + use_mm_prefix: bool, device_capability: DeviceCapability, ) -> str | None: # FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128, 192] diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 842153f403..aa6301c13b 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -111,6 +111,7 @@ class FlashInferMLASparseBackend(AttentionBackend): use_mla: bool, has_sink: bool, use_sparse: bool, + use_mm_prefix: bool, device_capability: DeviceCapability, ) -> str | None: # FlashInfer MLA sparse kernel requires qk_nope_head_dim in [128, 192] diff --git a/vllm/v1/attention/backends/mla/flashmla.py b/vllm/v1/attention/backends/mla/flashmla.py index 2f6058d69a..43aa186b51 100644 --- a/vllm/v1/attention/backends/mla/flashmla.py +++ b/vllm/v1/attention/backends/mla/flashmla.py @@ -84,6 +84,7 @@ class FlashMLABackend(MLACommonBackend): use_mla: bool, has_sink: bool, use_sparse: bool, + use_mm_prefix: bool, device_capability: DeviceCapability, ) -> str | None: if use_sparse: diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index 6c8dedd77f..0f819fe8ce 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py @@ -93,6 +93,7 @@ class TokenspeedMLABackend(MLACommonBackend): use_mla: bool, has_sink: bool, use_sparse: bool, + use_mm_prefix: bool, device_capability: DeviceCapability, ) -> str | None: # Surface a clear install hint up front rather than letting a raw diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 716d56e817..92ff08cc0f 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -19,7 +19,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( from vllm.platforms import current_platform from vllm.platforms.interface import DeviceCapability from vllm.utils.math_utils import next_power_of_2 -from vllm.utils.torch_utils import async_tensor_h2d, is_quantized_kv_cache +from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -30,7 +30,10 @@ from vllm.v1.attention.backend import ( CommonAttentionMetadata, MultipleOf, ) -from vllm.v1.attention.backends.utils import get_kv_cache_layout +from vllm.v1.attention.backends.utils import ( + compute_mm_prefix_range_tensor, + get_kv_cache_layout, +) from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd from vllm.v1.attention.ops.triton_reshape_and_cache_flash import ( triton_reshape_and_cache_flash, @@ -89,40 +92,6 @@ class TritonAttentionMetadata: mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None mm_prefix_range_tensor: torch.Tensor | None = None - @staticmethod - def compute_mm_prefix_range_tensor( - mm_prefix_range: dict[int, list[tuple[int, int]]] | None, - num_seqs: int, - device: torch.device, - ) -> torch.Tensor | None: - """Convert mm_prefix_range dict to padded tensor for Triton kernel. - - Returns shape: (num_seqs, max_ranges, 2) with 0-padding for empty ranges. - Empty ranges have start==end==0, which kernel skips via is_valid check. - """ - if mm_prefix_range is None: - return None - - # Collect ranges, using [(0,0)] for empty sequences to ensure uniform dims - range_lists = [ - mm_prefix_range.get(i, [(0, 0)]) or [(0, 0)] for i in range(num_seqs) - ] - - # Return None if all ranges are trivial (only (0,0) placeholders) - if all(r == [(0, 0)] for r in range_lists): - return None - - # Build on CPU first then move to GPU in a single H2D transfer - max_ranges = max(len(r) for r in range_lists) - # Pad all sequences to the same number of ranges - padded = [] - for r in range_lists: - padded_r = list(r) + [(0, 0)] * (max_ranges - len(r)) - padded.append(padded_r) - # Build on pinned CPU memory so the H2D transfer is non-blocking. - padded = async_tensor_h2d(padded, dtype=torch.int32, device=device) - return padded.view(num_seqs, max_ranges, 2) - class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]): _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS @@ -215,6 +184,7 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> TritonAttentionMetadata: + num_reqs = common_attn_metadata.num_reqs num_actual_tokens = common_attn_metadata.num_actual_tokens max_query_len = common_attn_metadata.max_query_len @@ -261,6 +231,14 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet softmax_segm_max=self.softmax_segm_max, softmax_segm_expsum=self.softmax_segm_expsum, ) + + mm_ranges = common_attn_metadata.mm_req_doc_ranges + if mm_ranges is not None: + attn_metadata.mm_prefix_range = mm_ranges + attn_metadata.mm_prefix_range_tensor = compute_mm_prefix_range_tensor( + mm_ranges, num_reqs, seq_lens.device + ) + return attn_metadata diff --git a/vllm/v1/attention/backends/utils.py b/vllm/v1/attention/backends/utils.py index 7bc20fc715..30db5d5f5a 100644 --- a/vllm/v1/attention/backends/utils.py +++ b/vllm/v1/attention/backends/utils.py @@ -17,6 +17,7 @@ from typing_extensions import runtime_checkable from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.utils.math_utils import cdiv +from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec if TYPE_CHECKING: @@ -45,6 +46,35 @@ PAD_SLOT_ID = -1 NULL_BLOCK_ID = 0 +def compute_mm_prefix_range_tensor( + mm_prefix_range: dict[int, list[tuple[int, int]]] | None, + num_seqs: int, + device: torch.device, +) -> torch.Tensor | None: + """Convert mm_prefix_range dict to padded tensor for Triton kernel. + + Returns shape: (num_seqs, max_ranges, 2) with 0-padding for empty ranges. + Empty ranges have start==end==0, which kernel skips via is_valid check. + """ + if mm_prefix_range is None: + return None + + range_lists = [ + mm_prefix_range.get(i, [(0, 0)]) or [(0, 0)] for i in range(num_seqs) + ] + + if all(r == [(0, 0)] for r in range_lists): + return None + + max_ranges = max(len(r) for r in range_lists) + padded = [] + for r in range_lists: + padded_r = list(r) + [(0, 0)] * (max_ranges - len(r)) + padded.append(padded_r) + padded = async_tensor_h2d(padded, dtype=torch.int32, device=device) + return padded.view(num_seqs, max_ranges, 2) + + def is_valid_kv_cache_layout(value: str) -> bool: return value in get_args(KVCacheLayoutType) diff --git a/vllm/v1/spec_decode/dflash.py b/vllm/v1/spec_decode/dflash.py index 72d0f99d07..f76305d085 100644 --- a/vllm/v1/spec_decode/dflash.py +++ b/vllm/v1/spec_decode/dflash.py @@ -73,6 +73,11 @@ class DFlashProposer(SpecDecodeBaseProposer): @override def _create_draft_vllm_config(self) -> VllmConfig: base = super()._create_draft_vllm_config() + # The draft model is text-only — clear the target's multimodal + # flag so flash_attn is not rejected for mm_prefix support. + arch = base.model_config.model_arch_config + if arch.is_mm_prefix_lm: + base.model_config.model_arch_config = replace(arch, is_mm_prefix_lm=False) return replace( base, attention_config=replace( diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 0cbad37a5e..f3f52c75d8 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -2286,6 +2286,30 @@ class GPUModelRunner( seq_lens_cpu = None num_computed_tokens_cpu = None + # Compute mm_prefix bidirectional ranges before building + # attention metadata so builders handle them during build(). + # Ranges exceeding sliding_window are skipped to prevent + # early tokens from attending across the entire image span. + req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None + if self.is_mm_prefix_lm: + req_doc_ranges = {} + hf_text_config = self.model_config.hf_text_config + _bidi_sw = getattr(hf_text_config, "sliding_window", None) + for req_id in self.input_batch.req_ids: + image_doc_ranges = [] + req_state = self.requests[req_id] + for mm_feature in req_state.mm_features: + if mm_feature.modality == "audio": + continue + pos_info = mm_feature.mm_position + img_doc_range = pos_info.extract_embeds_range() + for r in img_doc_range: + if _bidi_sw is not None and (r[1] - r[0] + 1) > _bidi_sw: + continue + image_doc_ranges.append(r) + req_idx = self.input_batch.req_id_to_index[req_id] + req_doc_ranges[req_idx] = image_doc_ranges + cm_base = CommonAttentionMetadata( query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1], query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1], @@ -2302,6 +2326,7 @@ class GPUModelRunner( causal=True, is_prefilling=is_prefilling, positions=self.positions[:num_tokens_padded], + mm_req_doc_ranges=req_doc_ranges, ) if self.dcp_world_size > 1: @@ -2454,36 +2479,6 @@ class GPUModelRunner( else: _build_attn_group_metadata(kv_cache_gid, attn_gid, cm) - if self.is_mm_prefix_lm: - req_doc_ranges = {} - - # Gemma4 bidi: skip ranges that exceed the sliding - # window. When image tokens > sliding_window, bidi causes - # early image tokens to attend to the entire image - # (e.g. 6 → 1092 targets), degrading spatial precision. - # Per-range filtering keeps bidi for small images/video - # frames while skipping oversized images. - hf_text_config = self.model_config.hf_text_config - _bidi_sw = getattr(hf_text_config, "sliding_window", None) - - for req_id in self.input_batch.req_ids: - image_doc_ranges = [] - req_state = self.requests[req_id] - for mm_feature in req_state.mm_features: - if mm_feature.modality == "audio": - continue - pos_info = mm_feature.mm_position - img_doc_range = pos_info.extract_embeds_range() - for r in img_doc_range: - if _bidi_sw is not None and (r[1] - r[0] + 1) > _bidi_sw: - continue - image_doc_ranges.append(r) - req_idx = self.input_batch.req_id_to_index[req_id] - req_doc_ranges[req_idx] = image_doc_ranges - - # Set mm_prefix_range for all attention metadata - self._set_mm_prefix_range_for_metadata(attn_metadata, req_doc_ranges) - if spec_decode_common_attn_metadata is not None and ( num_reqs != num_reqs_padded or num_tokens != num_tokens_padded ): @@ -6891,46 +6886,6 @@ class GPUModelRunner( return self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds) # type: ignore[assignment] - def _set_mm_prefix_range_for_metadata( - self, - attn_metadata: Any, - req_doc_ranges: dict[int, list[tuple[int, int]]], - ) -> None: - """Set mm_prefix_range for all attention metadata objects. - - This method handles both list and non-list attention metadata, - computing mm_prefix_range_tensor once and sharing it across all - metadata objects to avoid redundant host-to-device transfers. - """ - from vllm.v1.attention.backends.triton_attn import ( - TritonAttentionMetadata, - ) - - # Get all metadata objects from either list or dict structure - metadata_list = [] - if isinstance(attn_metadata, list): - for ub_metadata in attn_metadata: - metadata_list.extend(ub_metadata.values()) - else: - metadata_list.extend(attn_metadata.values()) - - # Set mm_prefix_range for all metadata and compute tensor once - shared_tensor = None - for metadata in metadata_list: - metadata.mm_prefix_range = req_doc_ranges # type: ignore[attr-defined] - - # Only compute tensor for TritonAttentionMetadata - if isinstance(metadata, TritonAttentionMetadata): - if shared_tensor is None: - shared_tensor = ( - TritonAttentionMetadata.compute_mm_prefix_range_tensor( - req_doc_ranges, - metadata.seq_lens.shape[0], # type: ignore[attr-defined] - metadata.seq_lens.device, # type: ignore[attr-defined] - ) - ) - metadata.mm_prefix_range_tensor = shared_tensor - def may_reinitialize_input_batch( self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int] ) -> None: diff --git a/vllm/vllm_flash_attn/flash_attn_interface.py b/vllm/vllm_flash_attn/flash_attn_interface.py index 33955bb239..5004ba9c8f 100644 --- a/vllm/vllm_flash_attn/flash_attn_interface.py +++ b/vllm/vllm_flash_attn/flash_attn_interface.py @@ -206,6 +206,9 @@ def flash_attn_varlen_func( cp_world_size=1, cp_rank=0, cp_tot_seqused_k=None, + # FA4 only + mask_mod=None, + aux_tensors=None, ): """dropout_p should be set to 0.0 during evaluation Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads @@ -297,6 +300,10 @@ def flash_attn_varlen_func( raise NotImplementedError("FA2 does not support s_aux") if num_splits > 1: raise NotImplementedError("FA2 does not support num_splits > 1") + if mask_mod is not None: + raise NotImplementedError("FA2 does not support mask_mod") + if aux_tensors is not None: + raise NotImplementedError("FA2 does not support aux_tensors") out, softmax_lse = torch.ops._vllm_fa2_C.varlen_fwd( q, k, @@ -325,6 +332,10 @@ def flash_attn_varlen_func( ) elif fa_version == 3: assert alibi_slopes is None, "Alibi is not supported in FA3" + if mask_mod is not None: + raise NotImplementedError("FA3 does not support mask_mod") + if aux_tensors is not None: + raise NotImplementedError("FA3 does not support aux_tensors") out, softmax_lse, _, _ = torch.ops._vllm_fa3_C.fwd( q, k, @@ -388,6 +399,8 @@ def flash_attn_varlen_func( return_lse=return_softmax_lse, out=out, learnable_sink=s_aux, + mask_mod=mask_mod, + aux_tensors=aux_tensors, ) else: raise ValueError(f"Unsupported FA version: {fa_version}")