@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 979f5511d7

[Bugfix][Gemma4] Keep image bidirectional attention within the sliding window (#47217)

Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com>

Luciano Martins committed Jul 2, 2026 at 23:57 UTC 979f5511d78b317760d45df9290233c27793a0af

8 files changed +89 -9

vllm/model_executor/layers/attention/attention.py

@@ -216,6 +216,7 @@ class Attention(nn.Module, AttentionLayerBase):

216 prefix: str = "",

217 attn_type: str = AttentionType.DECODER,

218 kv_sharing_target_layer_name: str | None = None,

219 + mm_prefix_clamp_sliding_window: bool = False,

220 attn_backend: type[AttentionBackend] | None = None,

221 head_size_v: int | None = None,

222 **extra_impl_args,

@@ -420,6 +421,9 @@ class Attention(nn.Module, AttentionLayerBase):

421 compilation_config.static_forward_context,

422 )

423 self.kv_sharing_target_layer_name = kv_sharing_target_layer_name

424 + # Gemma4: clamp mm_prefix bidirectional ranges by the sliding window

425 + # (read by the Triton backend impl). Default False for all other models.

426 + self.mm_prefix_clamp_sliding_window = mm_prefix_clamp_sliding_window

427

428 # use a placeholder kv cache tensor during init, which will be replaced

429 # by bind_kv_cache

vllm/model_executor/models/gemma4.py

@@ -499,6 +499,13 @@ class Gemma4Attention(nn.Module):

499 logits_soft_cap=attn_logits_soft_cap,

500 per_layer_sliding_window=sliding_window,

501 kv_sharing_target_layer_name=kv_sharing_target_layer_name,

502 + # Gemma4 vision bidi: on sliding layers the bidirectional image

503 + # block must stay within the sliding window, matching HF's

504 + # (causal OR blockwise) AND sliding_window. Without this the image

505 + # span (~1100 soft tokens at max_soft_tokens=1120) exceeds the 1024

506 + # window; the runner keeps the full range and the kernel bounds it

507 + # per-query here.

508 + mm_prefix_clamp_sliding_window=self.is_sliding,

509 prefix=f"{prefix}.attn",

510 )

511

vllm/model_executor/models/gemma4_mm.py

@@ -978,6 +978,13 @@ class Gemma4ForConditionalGeneration(

978 SupportsLoRA,

979 SupportsEagle3,

980 ):

981 + # Gemma4 clamps mm_prefix bidirectional ranges to the sliding window

982 + # in-kernel (HF's (causal OR blockwise) AND sliding_window). The model

983 + # runner reads this to keep image bidirectional ranges that exceed the

984 + # window instead of dropping them (which would make image attention

985 + # causal-only for images larger than the sliding window).

986 + mm_prefix_clamp_sliding_window: bool = True

987 +

988 packed_modules_mapping = {

989 "qkv_proj": [

990 "q_proj",

vllm/v1/attention/backends/flash_attn.py

+32 -4

@@ -892,7 +892,22 @@ class FlashAttentionImpl(AttentionImpl):

892 and self.vllm_flash_attn_version == 4

893 ):

894 max_ranges = mm_prefix_ranges.shape[1]

895 - mm_mask_mod = _make_mm_prefix_mask_mod(max_ranges)

895 + # Gemma4: clamp the bidirectional block to the sliding

896 + # window (HF (causal OR blockwise) AND sliding_window on

897 + # local layers). Other mm-prefix models pass sliding_window=0

898 + # -> unclamped, behavior unchanged. sliding_window_size[0]+1

899 + # is the window (self.sliding_window = (sw-1, 0) on sliding

900 + # layers); causal mm_prefix never hits the symmetric path.

901 + mm_clamp_sw = 0

902 + if (

903 + getattr(layer, "mm_prefix_clamp_sliding_window", False)

904 + and sliding_window_size is not None

905 + and sliding_window_size[0] >= 0

906 + ):

907 + mm_clamp_sw = sliding_window_size[0] + 1

908 + mm_mask_mod = _make_mm_prefix_mask_mod(

909 + max_ranges, sliding_window=mm_clamp_sw

910 + )

911 mm_aux = [mm_prefix_ranges]

912

913 # R-SWA: use CuTE-DSL mask_mod on FA4 for exact token-level

@@ -1189,13 +1204,23 @@ class FlashAttentionImpl(AttentionImpl):

1204 return output

1205

1206

1192 -def _make_mm_prefix_mask_mod(max_ranges: int):

1207 +def _make_mm_prefix_mask_mod(max_ranges: int, sliding_window: int = 0):

1208 """Build a CuTE-DSL mask_mod implementing (causal OR mm_prefix).

1209

1210 Returns a @cute.jit callable that evaluates:

1211 keep = (kv_idx <= q_idx) OR

1197 - (q_idx in [r_start,r_end] AND kv_idx in [r_start,r_end])

1212 + (q_idx in [r_start,r_end] AND kv_idx in [r_start,r_end]

1213 + [AND (q_idx - kv_idx) < sliding_window])

1214 for each mm_prefix range stored in aux_tensors[0].

1215 +

1216 + When sliding_window > 0 (Gemma4 sliding layers), the bidirectional block is

1217 + clamped to the sliding window, matching HF's

1218 + (causal OR blockwise) AND sliding_window (== sliding_window_overlay

1219 + kv > q - sw; future kv passes since q - kv < 0). sliding_window <= 0

1220 + (default) leaves the block unclamped, so other mm-prefix models are

1221 + unchanged. q_idx/kv_idx are local offsets, but their difference is

1222 + frame-invariant and image bidirectional attention only matters during

1223 + prefill (where local == absolute).

1224 """

1225 import cutlass

1226 import cutlass.cute as cute

@@ -1223,7 +1248,10 @@ def _make_mm_prefix_mask_mod(max_ranges: int):

1248 valid = r_start < r_end

1249 q_in = (q_idx >= r_start) & (q_idx <= r_end) & valid

1250 k_in = (kv_idx >= r_start) & (kv_idx <= r_end) & valid

1226 - keep = keep | (q_in & k_in)

1251 + mm = q_in & k_in

1252 + if sliding_window > 0:

1253 + mm = mm & ((q_idx - kv_idx) < sliding_window)

1254 + keep = keep | mm

1255 return keep

1256

1257 mm_prefix_mask_mod.use_fast_sampling = True

vllm/v1/attention/backends/triton_attn.py

@@ -676,6 +676,9 @@ class TritonAttentionImpl(AttentionImpl):

676 v_scale_cache=v_scale_cache,

677 chunk_lookback=self.chunk_lookback,

678 use_td=self.use_td,

679 + mm_prefix_clamp_sliding_window=getattr(

680 + layer, "mm_prefix_clamp_sliding_window", False

681 + ),

682 )

683

684 return output

vllm/v1/attention/ops/triton_attention_helpers.py

+13 -2

@@ -283,6 +283,7 @@ def compute_kv_seq_mask(

283 per_seq_causal_ptr=None,

284 CHUNK_LOOKBACK: tl.constexpr = -1,

285 CHUNK_SIZE: tl.constexpr = -1,

286 + MM_PREFIX_CLAMP_SW: tl.constexpr = False,

287 ):

288 """Build the KV mask for one tile.

289

@@ -333,7 +334,14 @@ def compute_kv_seq_mask(

334 seq_mask = seq_mask & sw_left

335

336 # PrefixLM: extend mask with bidirectional ranges for multimodal tokens.

336 - # Applied AFTER sliding window so mm_prefix ranges override SW restriction.

337 + # Default (MM_PREFIX_CLAMP_SW=False): applied AFTER sliding window so

338 + # mm_prefix ranges override the SW restriction -> (causal AND SW) OR mm.

339 + # Gemma4 (MM_PREFIX_CLAMP_SW=True): the bidirectional image block must stay

340 + # within the sliding window, matching HF's (causal OR blockwise) AND

341 + # sliding_window. We AND each range with the SW past-bound

342 + # (query_abs_pos - seq_offset) < SLIDING_WINDOW (== HF sliding_window_overlay

343 + # kv > q - sw; future kv passes trivially). Inert for full-attention layers

344 + # (SLIDING_WINDOW <= 0).

345 if USE_MM_PREFIX:

346 for i in range(MAX_MM_RANGES):

347 range_start = tl.load(

@@ -351,7 +359,10 @@ def compute_kv_seq_mask(

359 & (seq_offset[None, :] <= range_end)

360 & is_valid

361 )

354 - seq_mask |= q_in_range & k_in_range

362 + mm_mask = q_in_range & k_in_range

363 + if MM_PREFIX_CLAMP_SW and SLIDING_WINDOW > 0:

364 + mm_mask = mm_mask & ((query_abs_pos - seq_offset) < SLIDING_WINDOW)

365 + seq_mask |= mm_mask

366 return seq_mask

367

368

vllm/v1/attention/ops/triton_unified_attention.py

@@ -275,6 +275,10 @@ def kernel_unified_attention(

275 USE_TD: tl.constexpr = False,

276 USE_TD_QO: tl.constexpr = False,

277 Q_IS_FP8: tl.constexpr = False,

278 + # Gemma4: clamp mm_prefix bidirectional ranges by the sliding window

279 + # instead of letting them override it. Default False preserves the

280 + # original (causal AND SW) OR mm_prefix behavior for all other models.

281 + MM_PREFIX_CLAMP_SW: tl.constexpr = False,

282 ):

283 USE_PER_TOKEN_HEAD_SCALES: tl.constexpr = KV_QUANT_MODE >= 2

284 USE_FP8_Q_DESCALE: tl.constexpr = KV_QUANT_MODE == 1 and Q_IS_FP8

@@ -509,6 +513,7 @@ def kernel_unified_attention(

513 per_seq_causal_ptr,

514 CHUNK_LOOKBACK,

515 CHUNK_SIZE,

516 + MM_PREFIX_CLAMP_SW,

517 )

518

519 # S : (BLOCK_M, TILE_SIZE)

@@ -819,6 +824,9 @@ def unified_attention(

824 # The non-TD branch is dead-code-eliminated at Triton compile time so

825 # disabling this flag costs nothing.

826 use_td: bool = False,

827 + # Gemma4: clamp mm_prefix bidirectional ranges by the sliding window.

828 + # Default False keeps the original behavior for every other model.

829 + mm_prefix_clamp_sliding_window: bool = False,

830 ):

831 # Resolve causal: bool or per-seq tensor.

832 use_per_seq_causal = isinstance(causal, torch.Tensor)

@@ -1130,6 +1138,7 @@ def unified_attention(

1138 CHUNK_SIZE=chunk_size,

1139 USE_TD=use_td,

1140 USE_TD_QO=use_td_qo,

1141 + MM_PREFIX_CLAMP_SW=mm_prefix_clamp_sliding_window,

1142 **launch_kwargs,

1143 )

1144

vllm/v1/worker/gpu_model_runner.py

+14 -3

@@ -2330,13 +2330,20 @@ class GPUModelRunner(

2330

2331 # Compute mm_prefix bidirectional ranges before building

2332 # attention metadata so builders handle them during build().

2333 - # Ranges exceeding sliding_window are skipped to prevent

2334 - # early tokens from attending across the entire image span.

2333 + # By default, ranges exceeding sliding_window are skipped to prevent

2334 + # early tokens from attending across the entire image span. Models that

2335 + # clamp mm_prefix to the sliding window *in-kernel* (e.g. Gemma4, which

2336 + # needs HF's (causal OR blockwise) AND sliding_window on sliding layers)

2337 + # opt out of the skip so the bidirectional range survives for images

2338 + # larger than the window; the kernel then bounds it per-query.

2339 req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None

2340 if self.is_mm_prefix_lm:

2341 req_doc_ranges = {}

2342 hf_text_config = self.model_config.hf_text_config

2343 _bidi_sw = getattr(hf_text_config, "sliding_window", None)

2344 + _clamps_in_kernel = getattr(

2345 + self.model, "mm_prefix_clamp_sliding_window", False

2346 + )

2347 for req_id in self.input_batch.req_ids:

2348 image_doc_ranges = []

2349 req_state = self.requests[req_id]

@@ -2346,7 +2353,11 @@ class GPUModelRunner(

2353 pos_info = mm_feature.mm_position

2354 img_doc_range = pos_info.extract_embeds_range()

2355 for r in img_doc_range:

2349 - if _bidi_sw is not None and (r[1] - r[0] + 1) > _bidi_sw:

2356 + if (

2357 + not _clamps_in_kernel

2358 + and _bidi_sw is not None

2359 + and (r[1] - r[0] + 1) > _bidi_sw

2360 + ):

2361 continue

2362 image_doc_ranges.append(r)

2363 req_idx = self.input_batch.req_id_to_index[req_id]