A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 17:30 UTC
[Model] Gemma4: add bidirectional vision attention for sliding layers with window guard (#40534)
Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Signed-off-by: Luciano Martins <lucianomartins@google.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Luciano Martins committed
Apr 24, 2026 at 01:27 UTC
512f52219240b0aa1be687955ab52fcdd0c5a40e
2 files changed
+73
-1
vllm/model_executor/models/gemma4_mm.py
+59
@@ -969,6 +969,16 @@ class Gemma4ForConditionalGeneration(
969
self.language_model.make_empty_intermediate_tensors
970
)
971
972
+ # --- Precompute full-attention layer indices for bidi clearing ---
973
+ self._full_attn_layer_idxs: frozenset[int] = frozenset()
974
+ text_config = config.text_config
975
+ if getattr(text_config, "use_bidirectional_attention", None) == "vision":
976
+ layer_types = getattr(text_config, "layer_types", None)
977
+ if layer_types:
978
+ self._full_attn_layer_idxs = frozenset(
979
+ i for i, lt in enumerate(layer_types) if lt != "sliding_attention"
980
+ )
981
+
982
# --- MixtureOfExperts delegation to language_model ---
983
self.expert_weights = self.language_model.expert_weights
984
self.moe_layers = self.language_model.moe_layers
@@ -1310,6 +1320,12 @@ class Gemma4ForConditionalGeneration(
1320
else None
1321
)
1322
1323
+ # Gemma4 bidi: clear mm_prefix_range for full_attention layers.
1324
+ # Must run here (outside @support_torch_compile boundary) because
1325
+ # _run_decoder_layers is inside a compiled graph where Python
1326
+ # side effects are eliminated.
1327
+ self._clear_mm_prefix_for_full_attn_layers()
1328
+
1329
hidden_states = self.language_model.model(
1330
input_ids,
1331
positions,
@@ -1327,6 +1343,49 @@ class Gemma4ForConditionalGeneration(
1343
) -> torch.Tensor | None:
1344
return self.language_model.compute_logits(hidden_states)
1345
1346
+ # ------------------------------------------------------------------ #
1347
+ # Bidirectional attention helpers
1348
+ # ------------------------------------------------------------------ #
1349
+
1350
+ def _clear_mm_prefix_for_full_attn_layers(self) -> None:
1351
+ """Clear mm_prefix_range for non-sliding layers.
1352
+
1353
+ Gemma4 with use_bidirectional_attention='vision' applies
1354
+ bidirectional attention only to sliding_attention layers.
1355
+ Full attention layers use plain causal masking.
1356
+
1357
+ Uses _full_attn_layer_idxs (precomputed in __init__) for O(1)
1358
+ lookup instead of per-call regex parsing.
1359
+ """
1360
+ if not self._full_attn_layer_idxs:
1361
+ return
1362
+
1363
+ from vllm.forward_context import get_forward_context
1364
+
1365
+ attn_metadata = get_forward_context().attn_metadata
1366
+ if attn_metadata is None:
1367
+ return
1368
+
1369
+ def _process(metadata_dict: dict) -> None:
1370
+ for layer_name, metadata in metadata_dict.items():
1371
+ if ".layers." not in layer_name:
1372
+ continue
1373
+ try:
1374
+ layer_idx = int(layer_name.split(".layers.")[1].split(".")[0])
1375
+ except (ValueError, IndexError):
1376
+ continue
1377
+ if layer_idx in self._full_attn_layer_idxs:
1378
+ if hasattr(metadata, "mm_prefix_range"):
1379
+ metadata.mm_prefix_range = None
1380
+ if hasattr(metadata, "mm_prefix_range_tensor"):
1381
+ metadata.mm_prefix_range_tensor = None
1382
+
1383
+ if isinstance(attn_metadata, list):
1384
+ for ub_metadata in attn_metadata:
1385
+ _process(ub_metadata)
1386
+ elif isinstance(attn_metadata, dict):
1387
+ _process(attn_metadata)
1388
+
1389
# ------------------------------------------------------------------ #
1390
# Weight loading
1391
# ------------------------------------------------------------------ #
vllm/v1/worker/gpu_model_runner.py
+14
-1
@@ -2314,13 +2314,26 @@ class GPUModelRunner(
2314
2315
if self.is_mm_prefix_lm:
2316
req_doc_ranges = {}
2317
+
2318
+ # Gemma4 bidi: skip ranges that exceed the sliding
2319
+ # window. When image tokens > sliding_window, bidi causes
2320
+ # early image tokens to attend to the entire image
2321
+ # (e.g. 6 → 1092 targets), degrading spatial precision.
2322
+ # Per-range filtering keeps bidi for small images/video
2323
+ # frames while skipping oversized images.
2324
+ hf_text_config = self.model_config.hf_text_config
2325
+ _bidi_sw = getattr(hf_text_config, "sliding_window", None)
2326
+
2327
for req_id in self.input_batch.req_ids:
2328
image_doc_ranges = []
2329
req_state = self.requests[req_id]
2330
for mm_feature in req_state.mm_features:
2331
pos_info = mm_feature.mm_position
2332
img_doc_range = pos_info.extract_embeds_range()
2323
- image_doc_ranges.extend(img_doc_range)
2333
+ for r in img_doc_range:
2334
+ if _bidi_sw is not None and (r[1] - r[0] + 1) > _bidi_sw:
2335
+ continue
2336
+ image_doc_ranges.append(r)
2337
req_idx = self.input_batch.req_id_to_index[req_id]
2338
req_doc_ranges[req_idx] = image_doc_ranges
2339