A high-throughput and memory-efficient inference and serving engine for LLMs
[Core][Model] Gemma4: Unified FA4 for all layers + FlashAttention mm_prefix support (#42175)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Luciano Martins committed
Jun 10, 2026 at 02:45 UTC
6deb05e0e4f6c298a60975f188cd044773ae6dec
14 files changed
+224
-144
vllm/model_executor/models/config.py
+34
-35
@@ -57,50 +57,49 @@ class Gemma3TextModelConfig(VerifyAndUpdateConfig):
57
class Gemma4Config(VerifyAndUpdateConfig):
58
@staticmethod
59
def verify_and_update_config(vllm_config: "VllmConfig") -> None:
60
- """Force unified attention backend for models with heterogeneous
61
- head dimensions.
62
-
63
- Some Gemma4 variants use different head dimensions for
64
- sliding window (head_dim) vs full attention (global_head_dim) layers.
65
- When global_head_dim > 256, FlashAttention rejects those layers
66
- (head_size <= 256 kernel limit), causing vLLM to select a different
67
- backend for each layer type. This mixed-backend execution produces
68
- numerical divergence and output corruption.
69
-
70
- The fix detects heterogeneous head dimensions from the model config
71
- and forces TRITON_ATTN (which has no head_size ceiling) for all
72
- layers when the user hasn't explicitly chosen a backend.
73
-
74
- TODO: Heterogeneous head_sizes (head_dim != global_head_dim)
75
- require NixlConnector changes to support per-layer KV transfer
76
- with different head dimensions for prefill-decode disaggregation.
60
+ """Configure attention for heterogeneous head dimensions.
61
+
62
+ Gemma4 uses different head dimensions for sliding window
63
+ (head_dim) vs full attention (global_head_dim) layers. The
64
+ default FA3 on Hopper cannot handle head_dim > 256, which
65
+ causes mixed backend selection and numerical divergence.
66
+
67
+ When FA4 is available we force it for ALL layers, giving a
68
+ uniform kernel path and avoiding the mixed FA3+FA4 penalty.
69
+ When FA4 is not available we fall back to Triton.
70
"""
71
hf_text_config = vllm_config.model_config.hf_text_config
72
head_dim = getattr(hf_text_config, "head_dim", None)
73
global_head_dim = getattr(hf_text_config, "global_head_dim", None)
74
82
- # Only force Triton when head dimensions actually differ AND the
83
- # larger one exceeds FlashAttention's kernel limit (head_size <= 256).
84
- # This avoids unnecessary backend forcing on smaller models where
85
- # the config carries global_head_dim but all layers can still use
86
- # the same FA backend.
87
- max_head_dim = max(head_dim or 0, global_head_dim or 0)
88
- if (
89
- head_dim is not None
90
- and global_head_dim is not None
91
- and head_dim != global_head_dim
92
- and max_head_dim > 256
93
- and vllm_config.attention_config.backend is None
94
- ):
95
- from vllm.v1.attention.backends.registry import (
96
- AttentionBackendEnum,
97
- )
75
+ if head_dim is None or global_head_dim is None or head_dim == global_head_dim:
76
+ return
77
78
+ from vllm.v1.attention.backends.fa_utils import is_fa_version_supported
79
+ from vllm.v1.attention.backends.registry import AttentionBackendEnum
80
+
81
+ max_head_dim = max(head_dim, global_head_dim)
82
+
83
+ if is_fa_version_supported(4) and max_head_dim <= 512:
84
+ if (
85
+ vllm_config.attention_config.flash_attn_version is None
86
+ and vllm_config.attention_config.backend
87
+ in (None, AttentionBackendEnum.FLASH_ATTN)
88
+ ):
89
+ vllm_config.attention_config.flash_attn_version = 4
90
+ logger.info(
91
+ "Gemma4 model has heterogeneous head dimensions "
92
+ "(head_dim=%d, global_head_dim=%d). Using FA4 for "
93
+ "all layers to avoid mixed FA3/FA4 penalty.",
94
+ head_dim,
95
+ global_head_dim,
96
+ )
97
+ elif vllm_config.attention_config.backend is None:
98
vllm_config.attention_config.backend = AttentionBackendEnum.TRITON_ATTN
99
logger.info(
100
"Gemma4 model has heterogeneous head dimensions "
102
- "(head_dim=%d, global_head_dim=%d). Forcing TRITON_ATTN "
103
- "backend to prevent mixed-backend numerical divergence.",
101
+ "(head_dim=%d, global_head_dim=%d). FA4 not available, "
102
+ "forcing TRITON_ATTN backend.",
103
head_dim,
104
global_head_dim,
105
)
vllm/v1/attention/backend.py
+8
@@ -267,6 +267,7 @@ class AttentionBackend(ABC):
267
use_mla: bool,
268
has_sink: bool,
269
use_sparse: bool,
270
+ use_mm_prefix: bool,
271
device_capability: "DeviceCapability",
272
) -> str | None:
273
return None
@@ -334,6 +335,7 @@ class AttentionBackend(ABC):
335
use_mla,
336
has_sink,
337
use_sparse,
338
+ use_mm_prefix,
339
device_capability,
340
)
341
if combination_reason is not None:
@@ -415,6 +417,12 @@ class CommonAttentionMetadata:
417
decode rows (assumes every draft was accepted). Not safe for kernels
418
that need exact per-row context lengths on decode rows."""
419
420
+ mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None
421
+ """PrefixLM bidirectional ranges for multimodal tokens. Maps
422
+ request index to list of (start, end) token position ranges
423
+ where bidirectional attention should apply. None for text-only
424
+ batches or non-PrefixLM models."""
425
+
426
# WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)
427
_seq_lens_cpu: torch.Tensor | None = None
428
_num_computed_tokens_cpu: torch.Tensor | None = None
vllm/v1/attention/backends/flash_attn.py
+89
-3
@@ -192,6 +192,10 @@ class FlashAttentionBackend(AttentionBackend):
192
)
193
return kv_cache_dtype in ["auto", "float16", "bfloat16"]
194
195
+ @classmethod
196
+ def supports_mm_prefix(cls) -> bool:
197
+ return is_fa_version_supported(4)
198
+
199
@classmethod
200
def supports_sink(cls) -> bool:
201
if not is_flash_attn_varlen_func_available():
@@ -212,10 +216,20 @@ class FlashAttentionBackend(AttentionBackend):
216
use_mla: bool,
217
has_sink: bool,
218
use_sparse: bool,
219
+ use_mm_prefix: bool,
220
device_capability: DeviceCapability,
221
) -> str | None:
222
if has_sink and device_capability < DeviceCapability(9, 0):
223
return "sink not supported on compute capability < 9.0"
224
+ if (
225
+ use_mm_prefix
226
+ and get_flash_attn_version(head_size=head_size, has_sinks=has_sink) != 4
227
+ ):
228
+ return (
229
+ "mm_prefix (PrefixLM bidirectional attention) requires "
230
+ "FlashAttention v4, which does not resolve for this "
231
+ "head_size"
232
+ )
233
return None
234
235
@@ -255,6 +269,10 @@ class FlashAttentionMetadata:
269
270
causal: bool = True
271
272
+ # PrefixLM bidirectional ranges for multimodal tokens.
273
+ # Shape: (num_seqs, max_ranges, 2) int32, [start, end] per range.
274
+ mm_prefix_range_tensor: torch.Tensor | None = None
275
+
276
277
def _get_sliding_window_configs(
278
vllm_config: VllmConfig,
@@ -572,6 +590,19 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
590
max_num_splits=max_num_splits,
591
causal=causal,
592
)
593
+
594
+ # Compute mm_prefix range tensor if the batch contains
595
+ # multimodal tokens with bidirectional ranges.
596
+ mm_ranges = common_attn_metadata.mm_req_doc_ranges
597
+ if mm_ranges is not None:
598
+ from vllm.v1.attention.backends.utils import (
599
+ compute_mm_prefix_range_tensor,
600
+ )
601
+
602
+ attn_metadata.mm_prefix_range_tensor = compute_mm_prefix_range_tensor(
603
+ mm_ranges, num_reqs, seq_lens.device
604
+ )
605
+
606
return attn_metadata
607
608
def update_block_table(
@@ -793,6 +824,18 @@ class FlashAttentionImpl(AttentionImpl):
824
if self.sliding_window is not None
825
else None
826
)
827
+ mm_prefix_ranges = attn_metadata.mm_prefix_range_tensor
828
+ mm_mask_mod = None
829
+ mm_aux = None
830
+ if (
831
+ mm_prefix_ranges is not None
832
+ and attn_metadata.causal
833
+ and self.vllm_flash_attn_version == 4
834
+ ):
835
+ max_ranges = mm_prefix_ranges.shape[1]
836
+ mm_mask_mod = _make_mm_prefix_mask_mod(max_ranges)
837
+ mm_aux = [mm_prefix_ranges]
838
+
839
flash_attn_varlen_func(
840
q=query[:num_actual_tokens],
841
k=key_cache,
@@ -815,6 +858,8 @@ class FlashAttentionImpl(AttentionImpl):
858
v_descale=v_descale,
859
num_splits=attn_metadata.max_num_splits,
860
s_aux=self.sinks,
861
+ mask_mod=mm_mask_mod,
862
+ aux_tensors=mm_aux,
863
)
864
return output
865
@@ -1040,17 +1085,58 @@ class FlashAttentionImpl(AttentionImpl):
1085
window_size=sliding_window_size,
1086
softcap=self.logits_soft_cap,
1087
fa_version=self.vllm_flash_attn_version,
1043
- q_descale=layer._q_scale.expand(descale_shape)
1088
+ q_descale=layer._q_scale.expand(descale_shape) # type: ignore[operator]
1089
if self.supports_quant_query_input
1090
else None,
1046
- k_descale=layer._k_scale.expand(descale_shape),
1047
- v_descale=layer._v_scale.expand(descale_shape),
1091
+ k_descale=layer._k_scale.expand(descale_shape), # type: ignore[operator]
1092
+ v_descale=layer._v_scale.expand(descale_shape), # type: ignore[operator]
1093
num_splits=1 if self.batch_invariant_enabled else 0,
1094
)
1095
1096
return output
1097
1098
1099
+def _make_mm_prefix_mask_mod(max_ranges: int):
1100
+ """Build a CuTE-DSL mask_mod implementing (causal OR mm_prefix).
1101
+
1102
+ Returns a @cute.jit callable that evaluates:
1103
+ keep = (kv_idx <= q_idx) OR
1104
+ (q_idx in [r_start,r_end] AND kv_idx in [r_start,r_end])
1105
+ for each mm_prefix range stored in aux_tensors[0].
1106
+ """
1107
+ import cutlass
1108
+ import cutlass.cute as cute
1109
+ from cutlass import Int32 # type: ignore[attr-defined]
1110
+
1111
+ from vllm.vllm_flash_attn.cute.utils import ( # type: ignore[import-untyped]
1112
+ scalar_to_ssa,
1113
+ )
1114
+
1115
+ @cute.jit
1116
+ def mm_prefix_mask_mod(
1117
+ batch_idx: cute.TensorSSA,
1118
+ head_idx: cute.TensorSSA,
1119
+ q_idx: cute.TensorSSA,
1120
+ kv_idx: cute.TensorSSA,
1121
+ seqlen_info,
1122
+ aux_tensors,
1123
+ ):
1124
+ keep = kv_idx <= q_idx
1125
+ ranges = aux_tensors[0]
1126
+ b = batch_idx[0]
1127
+ for i in cutlass.range_constexpr(max_ranges): # type: ignore[attr-defined]
1128
+ r_start = scalar_to_ssa(ranges[b, i, 0], Int32)
1129
+ r_end = scalar_to_ssa(ranges[b, i, 1], Int32)
1130
+ valid = r_start < r_end
1131
+ q_in = (q_idx >= r_start) & (q_idx <= r_end) & valid
1132
+ k_in = (kv_idx >= r_start) & (kv_idx <= r_end) & valid
1133
+ keep = keep | (q_in & k_in)
1134
+ return keep
1135
+
1136
+ mm_prefix_mask_mod.use_fast_sampling = True
1137
+ return mm_prefix_mask_mod
1138
+
1139
+
1140
def use_cascade_attention(
1141
common_prefix_len: int,
1142
query_lens: np.ndarray,
vllm/v1/attention/backends/flex_attention.py
+1
@@ -957,6 +957,7 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
957
persistent_kv_indices=self.persistent_kv_indices,
958
persistent_kv_num_blocks=self.persistent_kv_num_blocks,
959
persistent_doc_ids=self.persistent_doc_ids,
960
+ mm_prefix_range=common_attn_metadata.mm_req_doc_ranges,
961
)
962
963
# Pre-build block_mask so it is ready before CUDA graph capture.
vllm/v1/attention/backends/mla/flashattn_mla.py
+1
@@ -82,6 +82,7 @@ class FlashAttnMLABackend(MLACommonBackend):
82
use_mla: bool,
83
has_sink: bool,
84
use_sparse: bool,
85
+ use_mm_prefix: bool,
86
device_capability: DeviceCapability,
87
) -> str | None:
88
if not flash_attn_supports_mla():
vllm/v1/attention/backends/mla/flashinfer_mla.py
+1
@@ -75,6 +75,7 @@ class FlashInferMLABackend(MLACommonBackend):
75
use_mla: bool,
76
has_sink: bool,
77
use_sparse: bool,
78
+ use_mm_prefix: bool,
79
device_capability: DeviceCapability,
80
) -> str | None:
81
# FlashInfer MLA kernel requires qk_nope_head_dim in [64, 128, 192]
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+1
@@ -111,6 +111,7 @@ class FlashInferMLASparseBackend(AttentionBackend):
111
use_mla: bool,
112
has_sink: bool,
113
use_sparse: bool,
114
+ use_mm_prefix: bool,
115
device_capability: DeviceCapability,
116
) -> str | None:
117
# FlashInfer MLA sparse kernel requires qk_nope_head_dim in [128, 192]
vllm/v1/attention/backends/mla/flashmla.py
+1
@@ -84,6 +84,7 @@ class FlashMLABackend(MLACommonBackend):
84
use_mla: bool,
85
has_sink: bool,
86
use_sparse: bool,
87
+ use_mm_prefix: bool,
88
device_capability: DeviceCapability,
89
) -> str | None:
90
if use_sparse:
vllm/v1/attention/backends/mla/tokenspeed_mla.py
+1
@@ -93,6 +93,7 @@ class TokenspeedMLABackend(MLACommonBackend):
93
use_mla: bool,
94
has_sink: bool,
95
use_sparse: bool,
96
+ use_mm_prefix: bool,
97
device_capability: DeviceCapability,
98
) -> str | None:
99
# Surface a clear install hint up front rather than letting a raw
vllm/v1/attention/backends/triton_attn.py
+14
-36
@@ -19,7 +19,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
19
from vllm.platforms import current_platform
20
from vllm.platforms.interface import DeviceCapability
21
from vllm.utils.math_utils import next_power_of_2
22
-from vllm.utils.torch_utils import async_tensor_h2d, is_quantized_kv_cache
22
+from vllm.utils.torch_utils import is_quantized_kv_cache
23
from vllm.v1.attention.backend import (
24
AttentionBackend,
25
AttentionCGSupport,
@@ -30,7 +30,10 @@ from vllm.v1.attention.backend import (
30
CommonAttentionMetadata,
31
MultipleOf,
32
)
33
-from vllm.v1.attention.backends.utils import get_kv_cache_layout
33
+from vllm.v1.attention.backends.utils import (
34
+ compute_mm_prefix_range_tensor,
35
+ get_kv_cache_layout,
36
+)
37
from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd
38
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
39
triton_reshape_and_cache_flash,
@@ -89,40 +92,6 @@ class TritonAttentionMetadata:
92
mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None
93
mm_prefix_range_tensor: torch.Tensor | None = None
94
92
- @staticmethod
93
- def compute_mm_prefix_range_tensor(
94
- mm_prefix_range: dict[int, list[tuple[int, int]]] | None,
95
- num_seqs: int,
96
- device: torch.device,
97
- ) -> torch.Tensor | None:
98
- """Convert mm_prefix_range dict to padded tensor for Triton kernel.
99
-
100
- Returns shape: (num_seqs, max_ranges, 2) with 0-padding for empty ranges.
101
- Empty ranges have start==end==0, which kernel skips via is_valid check.
102
- """
103
- if mm_prefix_range is None:
104
- return None
105
-
106
- # Collect ranges, using [(0,0)] for empty sequences to ensure uniform dims
107
- range_lists = [
108
- mm_prefix_range.get(i, [(0, 0)]) or [(0, 0)] for i in range(num_seqs)
109
- ]
110
-
111
- # Return None if all ranges are trivial (only (0,0) placeholders)
112
- if all(r == [(0, 0)] for r in range_lists):
113
- return None
114
-
115
- # Build on CPU first then move to GPU in a single H2D transfer
116
- max_ranges = max(len(r) for r in range_lists)
117
- # Pad all sequences to the same number of ranges
118
- padded = []
119
- for r in range_lists:
120
- padded_r = list(r) + [(0, 0)] * (max_ranges - len(r))
121
- padded.append(padded_r)
122
- # Build on pinned CPU memory so the H2D transfer is non-blocking.
123
- padded = async_tensor_h2d(padded, dtype=torch.int32, device=device)
124
- return padded.view(num_seqs, max_ranges, 2)
125
-
95
96
class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]):
97
_cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
@@ -215,6 +184,7 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
184
common_attn_metadata: CommonAttentionMetadata,
185
fast_build: bool = False,
186
) -> TritonAttentionMetadata:
187
+ num_reqs = common_attn_metadata.num_reqs
188
num_actual_tokens = common_attn_metadata.num_actual_tokens
189
max_query_len = common_attn_metadata.max_query_len
190
@@ -261,6 +231,14 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
231
softmax_segm_max=self.softmax_segm_max,
232
softmax_segm_expsum=self.softmax_segm_expsum,
233
)
234
+
235
+ mm_ranges = common_attn_metadata.mm_req_doc_ranges
236
+ if mm_ranges is not None:
237
+ attn_metadata.mm_prefix_range = mm_ranges
238
+ attn_metadata.mm_prefix_range_tensor = compute_mm_prefix_range_tensor(
239
+ mm_ranges, num_reqs, seq_lens.device
240
+ )
241
+
242
return attn_metadata
243
244
vllm/v1/attention/backends/utils.py
+30
@@ -17,6 +17,7 @@ from typing_extensions import runtime_checkable
17
18
from vllm.config import VllmConfig, get_layers_from_vllm_config
19
from vllm.utils.math_utils import cdiv
20
+from vllm.utils.torch_utils import async_tensor_h2d
21
from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
22
23
if TYPE_CHECKING:
@@ -45,6 +46,35 @@ PAD_SLOT_ID = -1
46
NULL_BLOCK_ID = 0
47
48
49
+def compute_mm_prefix_range_tensor(
50
+ mm_prefix_range: dict[int, list[tuple[int, int]]] | None,
51
+ num_seqs: int,
52
+ device: torch.device,
53
+) -> torch.Tensor | None:
54
+ """Convert mm_prefix_range dict to padded tensor for Triton kernel.
55
+
56
+ Returns shape: (num_seqs, max_ranges, 2) with 0-padding for empty ranges.
57
+ Empty ranges have start==end==0, which kernel skips via is_valid check.
58
+ """
59
+ if mm_prefix_range is None:
60
+ return None
61
+
62
+ range_lists = [
63
+ mm_prefix_range.get(i, [(0, 0)]) or [(0, 0)] for i in range(num_seqs)
64
+ ]
65
+
66
+ if all(r == [(0, 0)] for r in range_lists):
67
+ return None
68
+
69
+ max_ranges = max(len(r) for r in range_lists)
70
+ padded = []
71
+ for r in range_lists:
72
+ padded_r = list(r) + [(0, 0)] * (max_ranges - len(r))
73
+ padded.append(padded_r)
74
+ padded = async_tensor_h2d(padded, dtype=torch.int32, device=device)
75
+ return padded.view(num_seqs, max_ranges, 2)
76
+
77
+
78
def is_valid_kv_cache_layout(value: str) -> bool:
79
return value in get_args(KVCacheLayoutType)
80
vllm/v1/spec_decode/dflash.py
+5
@@ -73,6 +73,11 @@ class DFlashProposer(SpecDecodeBaseProposer):
73
@override
74
def _create_draft_vllm_config(self) -> VllmConfig:
75
base = super()._create_draft_vllm_config()
76
+ # The draft model is text-only — clear the target's multimodal
77
+ # flag so flash_attn is not rejected for mm_prefix support.
78
+ arch = base.model_config.model_arch_config
79
+ if arch.is_mm_prefix_lm:
80
+ base.model_config.model_arch_config = replace(arch, is_mm_prefix_lm=False)
81
return replace(
82
base,
83
attention_config=replace(
vllm/v1/worker/gpu_model_runner.py
+25
-70
@@ -2286,6 +2286,30 @@ class GPUModelRunner(
2286
seq_lens_cpu = None
2287
num_computed_tokens_cpu = None
2288
2289
+ # Compute mm_prefix bidirectional ranges before building
2290
+ # attention metadata so builders handle them during build().
2291
+ # Ranges exceeding sliding_window are skipped to prevent
2292
+ # early tokens from attending across the entire image span.
2293
+ req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None
2294
+ if self.is_mm_prefix_lm:
2295
+ req_doc_ranges = {}
2296
+ hf_text_config = self.model_config.hf_text_config
2297
+ _bidi_sw = getattr(hf_text_config, "sliding_window", None)
2298
+ for req_id in self.input_batch.req_ids:
2299
+ image_doc_ranges = []
2300
+ req_state = self.requests[req_id]
2301
+ for mm_feature in req_state.mm_features:
2302
+ if mm_feature.modality == "audio":
2303
+ continue
2304
+ pos_info = mm_feature.mm_position
2305
+ img_doc_range = pos_info.extract_embeds_range()
2306
+ for r in img_doc_range:
2307
+ if _bidi_sw is not None and (r[1] - r[0] + 1) > _bidi_sw:
2308
+ continue
2309
+ image_doc_ranges.append(r)
2310
+ req_idx = self.input_batch.req_id_to_index[req_id]
2311
+ req_doc_ranges[req_idx] = image_doc_ranges
2312
+
2313
cm_base = CommonAttentionMetadata(
2314
query_start_loc=self.query_start_loc.gpu[: num_reqs_padded + 1],
2315
query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs_padded + 1],
@@ -2302,6 +2326,7 @@ class GPUModelRunner(
2326
causal=True,
2327
is_prefilling=is_prefilling,
2328
positions=self.positions[:num_tokens_padded],
2329
+ mm_req_doc_ranges=req_doc_ranges,
2330
)
2331
2332
if self.dcp_world_size > 1:
@@ -2454,36 +2479,6 @@ class GPUModelRunner(
2479
else:
2480
_build_attn_group_metadata(kv_cache_gid, attn_gid, cm)
2481
2457
- if self.is_mm_prefix_lm:
2458
- req_doc_ranges = {}
2459
-
2460
- # Gemma4 bidi: skip ranges that exceed the sliding
2461
- # window. When image tokens > sliding_window, bidi causes
2462
- # early image tokens to attend to the entire image
2463
- # (e.g. 6 → 1092 targets), degrading spatial precision.
2464
- # Per-range filtering keeps bidi for small images/video
2465
- # frames while skipping oversized images.
2466
- hf_text_config = self.model_config.hf_text_config
2467
- _bidi_sw = getattr(hf_text_config, "sliding_window", None)
2468
-
2469
- for req_id in self.input_batch.req_ids:
2470
- image_doc_ranges = []
2471
- req_state = self.requests[req_id]
2472
- for mm_feature in req_state.mm_features:
2473
- if mm_feature.modality == "audio":
2474
- continue
2475
- pos_info = mm_feature.mm_position
2476
- img_doc_range = pos_info.extract_embeds_range()
2477
- for r in img_doc_range:
2478
- if _bidi_sw is not None and (r[1] - r[0] + 1) > _bidi_sw:
2479
- continue
2480
- image_doc_ranges.append(r)
2481
- req_idx = self.input_batch.req_id_to_index[req_id]
2482
- req_doc_ranges[req_idx] = image_doc_ranges
2483
-
2484
- # Set mm_prefix_range for all attention metadata
2485
- self._set_mm_prefix_range_for_metadata(attn_metadata, req_doc_ranges)
2486
-
2482
if spec_decode_common_attn_metadata is not None and (
2483
num_reqs != num_reqs_padded or num_tokens != num_tokens_padded
2484
):
@@ -6891,46 +6886,6 @@ class GPUModelRunner(
6886
return
6887
self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds) # type: ignore[assignment]
6888
6894
- def _set_mm_prefix_range_for_metadata(
6895
- self,
6896
- attn_metadata: Any,
6897
- req_doc_ranges: dict[int, list[tuple[int, int]]],
6898
- ) -> None:
6899
- """Set mm_prefix_range for all attention metadata objects.
6900
-
6901
- This method handles both list and non-list attention metadata,
6902
- computing mm_prefix_range_tensor once and sharing it across all
6903
- metadata objects to avoid redundant host-to-device transfers.
6904
- """
6905
- from vllm.v1.attention.backends.triton_attn import (
6906
- TritonAttentionMetadata,
6907
- )
6908
-
6909
- # Get all metadata objects from either list or dict structure
6910
- metadata_list = []
6911
- if isinstance(attn_metadata, list):
6912
- for ub_metadata in attn_metadata:
6913
- metadata_list.extend(ub_metadata.values())
6914
- else:
6915
- metadata_list.extend(attn_metadata.values())
6916
-
6917
- # Set mm_prefix_range for all metadata and compute tensor once
6918
- shared_tensor = None
6919
- for metadata in metadata_list:
6920
- metadata.mm_prefix_range = req_doc_ranges # type: ignore[attr-defined]
6921
-
6922
- # Only compute tensor for TritonAttentionMetadata
6923
- if isinstance(metadata, TritonAttentionMetadata):
6924
- if shared_tensor is None:
6925
- shared_tensor = (
6926
- TritonAttentionMetadata.compute_mm_prefix_range_tensor(
6927
- req_doc_ranges,
6928
- metadata.seq_lens.shape[0], # type: ignore[attr-defined]
6929
- metadata.seq_lens.device, # type: ignore[attr-defined]
6930
- )
6931
- )
6932
- metadata.mm_prefix_range_tensor = shared_tensor
6933
-
6889
def may_reinitialize_input_batch(
6890
self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
6891
) -> None:
vllm/vllm_flash_attn/flash_attn_interface.py
+13
@@ -206,6 +206,9 @@ def flash_attn_varlen_func(
206
cp_world_size=1,
207
cp_rank=0,
208
cp_tot_seqused_k=None,
209
+ # FA4 only
210
+ mask_mod=None,
211
+ aux_tensors=None,
212
):
213
"""dropout_p should be set to 0.0 during evaluation
214
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(
300
raise NotImplementedError("FA2 does not support s_aux")
301
if num_splits > 1:
302
raise NotImplementedError("FA2 does not support num_splits > 1")
303
+ if mask_mod is not None:
304
+ raise NotImplementedError("FA2 does not support mask_mod")
305
+ if aux_tensors is not None:
306
+ raise NotImplementedError("FA2 does not support aux_tensors")
307
out, softmax_lse = torch.ops._vllm_fa2_C.varlen_fwd(
308
q,
309
k,
@@ -325,6 +332,10 @@ def flash_attn_varlen_func(
332
)
333
elif fa_version == 3:
334
assert alibi_slopes is None, "Alibi is not supported in FA3"
335
+ if mask_mod is not None:
336
+ raise NotImplementedError("FA3 does not support mask_mod")
337
+ if aux_tensors is not None:
338
+ raise NotImplementedError("FA3 does not support aux_tensors")
339
out, softmax_lse, _, _ = torch.ops._vllm_fa3_C.fwd(
340
q,
341
k,
@@ -388,6 +399,8 @@ def flash_attn_varlen_func(
399
return_lse=return_softmax_lse,
400
out=out,
401
learnable_sink=s_aux,
402
+ mask_mod=mask_mod,
403
+ aux_tensors=aux_tensors,
404
)
405
else:
406
raise ValueError(f"Unsupported FA version: {fa_version}")