[CI] Fix "test_awq_load[gemma4-moe-*]" failure

tests/models/multimodal/processing/test_gemma4.py

@@ -7,9 +7,13 @@
 import torch
 from PIL import Image as PILImage
 
-from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs
+from vllm.model_executor.models.gemma4_mm import (
+    Gemma4ForConditionalGeneration,
+    Gemma4ImagePixelInputs,
+)
 from vllm.multimodal import MULTIMODAL_REGISTRY
 from vllm.multimodal.inputs import MultiModalFieldConfig
+from vllm.utils.mem_constants import GiB_bytes
 
 from ....conftest import ImageTestAssets
 from ...utils import build_model_context
@@ -224,3 +228,60 @@ def test_limit_mm_per_prompt(
             mm_items=processor.info.parse_mm_data(mm_data),
             hf_processor_mm_kwargs={},
         )
+
+
+# Regression guard for PR #43169 follow-up: the batched Gemma4 vision encoder
+# admitted ``chunk ~= 53`` on a 22 GiB L4 with a 26B AWQ model loaded,
+# allocating 2.43 GiB int64 inside
+# ``F.one_hot(num_classes=position_embedding_size)`` and OOMing because only
+# 2.41 GiB was actually free.  The fix sizes ``chunk`` from currently-free GPU
+# memory and counts the ``F.one_hot`` transient as the dominant cost.
+
+_encoder_chunk = Gemma4ForConditionalGeneration._encoder_chunk
+
+# Gemma4 vision config default (HF: configuration_gemma4.py).
+_POSITION_EMBEDDING_SIZE = 10240
+# Video frame: max_soft_tokens=70, pooling_kernel_size=2 -> 70 * 4 patches.
+_VIDEO_PATCHES_PER_FRAME = 280
+
+
+def test_encoder_chunk_tight_budget_fits_in_free():
+    free = 3 * GiB_bytes  # L4 22 GiB after 26B AWQ load.
+    total = 22 * GiB_bytes
+    chunk = _encoder_chunk(
+        _VIDEO_PATCHES_PER_FRAME, free, total, _POSITION_EMBEDDING_SIZE
+    )
+    one_hot_bytes = chunk * _VIDEO_PATCHES_PER_FRAME * 2 * _POSITION_EMBEDDING_SIZE * 8
+    assert one_hot_bytes <= free // 2
+
+
+def test_encoder_chunk_roomy_gpu_keeps_batching():
+    chunk = _encoder_chunk(
+        _VIDEO_PATCHES_PER_FRAME,
+        60 * GiB_bytes,
+        80 * GiB_bytes,
+        _POSITION_EMBEDDING_SIZE,
+    )
+    assert chunk > 8
+
+
+def test_encoder_chunk_zero_patches_is_safe():
+    assert (
+        _encoder_chunk(0, 60 * GiB_bytes, 80 * GiB_bytes, _POSITION_EMBEDDING_SIZE) == 1
+    )
+
+
+def test_encoder_chunk_zero_position_embedding_size_is_safe():
+    # Degenerate config: must not raise ZeroDivisionError.
+    assert (
+        _encoder_chunk(_VIDEO_PATCHES_PER_FRAME, 60 * GiB_bytes, 80 * GiB_bytes, 0) == 1
+    )
+
+
+def test_encoder_chunk_no_free_memory_falls_back_to_one():
+    assert (
+        _encoder_chunk(
+            _VIDEO_PATCHES_PER_FRAME, 0, 22 * GiB_bytes, _POSITION_EMBEDDING_SIZE
+        )
+        == 1
+    )

vllm/model_executor/models/gemma4_mm.py

@@ -961,9 +961,6 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
             self.embed_vision = Gemma4MultimodalEmbedder(
                 config.vision_config, config.text_config
             )
-        # Lazy-initialized on first encoder call (see _encoder_max_batch).
-        self._encoder_budget_bytes = 0
-        self._encoder_bytes_per_patch = 0
 
         # ---- Audio tower (variants with audio_config) ----
         if config.audio_config is not None:
@@ -1104,18 +1101,36 @@ def _parse_and_validate_multimodal_inputs(
                 )
         return mm_input_by_modality
 
-    def _encoder_max_batch(self, patches_per_item: int) -> int:
-        """Max items per encoder call given per-item patch count."""
-        if self._encoder_budget_bytes == 0:
-            total_mem = current_platform.get_device_total_memory()
-            self._encoder_budget_bytes = int(total_mem * 0.05)
-            logger.info(
-                "Encoder memory budget: %.1fGB (total=%.1fGB)",
-                self._encoder_budget_bytes / 1024**3,
-                total_mem / 1024**3,
-            )
-        cost = patches_per_item * self._encoder_bytes_per_patch
-        return max(1, self._encoder_budget_bytes // cost) if cost > 0 else 1
+    @staticmethod
+    def _encoder_chunk(
+        patches_per_item: int,
+        free_bytes: int,
+        total_bytes: int,
+        position_embedding_size: int,
+    ) -> int:
+        """Max chunk size whose F.one_hot transient fits in the budget.
+
+        The dominant transient inside HF's ``Gemma4VisionPatchEmbedder.
+        _position_embeddings`` is
+        ``F.one_hot(clamped_positions, num_classes=position_embedding_size)``
+        with shape ``(chunk, patches, 2, position_embedding_size)``,
+        int64, plus its simultaneous cast to the position embedding
+        table dtype. That, not the encoder residual stream, sets peak
+        memory.
+        """
+        if patches_per_item <= 0:
+            return 1
+        # Half of currently-free, capped at 10% of total so we leave room
+        # for the rest of profile_run / the subsequent encoder + pooler.
+        budget = min(free_bytes // 2, total_bytes // 10)
+        if budget <= 0:
+            return 1
+        # F.one_hot allocates (chunk, patches, 2, pos_emb_size) int64
+        # (the inner 2 is the (x, y) coordinate axis, 8 is sizeof(int64)).
+        # Outer 2x covers the int64 buffer and its concurrent bf16 cast
+        # plus the matmul output that live alongside it at peak.
+        cost = patches_per_item * 4 * position_embedding_size * 8
+        return max(1, budget // cost) if cost > 0 else 1
 
     # ------------------------------------------------------------------ #
     # Image processing
@@ -1136,7 +1151,8 @@ def _process_image_input(
         pixel_position_ids = image_input["pixel_position_ids"]
 
         vt = self.vision_tower
-        pooling_k2 = self.config.vision_config.pooling_kernel_size**2
+        vision_cfg = self.config.vision_config
+        pooling_k2 = vision_cfg.pooling_kernel_size**2
 
         # Concurrent requests with different image resolutions may
         # arrive as a list of per-image tensors, while same-resolution
@@ -1153,10 +1169,18 @@ def _process_image_input(
             pp = pixel_position_ids[idx]
             buckets.setdefault(pv.shape[0], []).append((idx, pv, pp))
 
-        # Encode each resolution bucket in memory-safe chunks.
+        # Encode each resolution bucket in memory-safe chunks. Re-read
+        # free memory per bucket because the previous bucket's encoder
+        # pass has already allocated activations we should account for.
         last_hidden_states_map: dict[int, torch.Tensor] = {}
         for patches, items in buckets.items():
-            max_batch_size = min(len(items), self._encoder_max_batch(patches))
+            free, total = current_platform.mem_get_info()
+            max_batch_size = min(
+                len(items),
+                self._encoder_chunk(
+                    patches, free, total, vision_cfg.position_embedding_size
+                ),
+            )
 
             for chunk_idx in range(0, len(items), max_batch_size):
                 chunk_items = items[chunk_idx : chunk_idx + max_batch_size]
@@ -1247,7 +1271,8 @@ def _process_video_input(
         frame_counts = video_input["video_frame_counts"]
 
         vt = self.vision_tower
-        pooling_k2 = self.config.vision_config.pooling_kernel_size**2
+        vision_cfg = self.config.vision_config
+        pooling_k2 = vision_cfg.pooling_kernel_size**2
         target_dtype = self.embed_vision.embedding_projection.weight.dtype
 
         if isinstance(frame_counts, torch.Tensor):
@@ -1256,13 +1281,20 @@ def _process_video_input(
             fc_list = list(frame_counts)
 
         total_frames = pixel_values.shape[0]
+        free, total = current_platform.mem_get_info()
         max_batch_size = min(
-            total_frames, self._encoder_max_batch(pixel_values.shape[1])
+            total_frames,
+            self._encoder_chunk(
+                pixel_values.shape[1],
+                free,
+                total,
+                vision_cfg.position_embedding_size,
+            ),
         )
 
         padding_positions = (pixel_position_ids == -1).all(dim=-1)
 
-        # Encode frames in chunks bounded by _encoder_max_batch.
+        # Encode frames in chunks bounded by _encoder_chunk.
         last_hidden_states_list: list[torch.Tensor] = []
         for i in range(0, total_frames, max_batch_size):
             pv_chunk = pixel_values[i : i + max_batch_size]
@@ -1538,15 +1570,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
             self,
             ignore_unexpected_prefixes=ignore_prefixes,
         )
-        loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
-
-        # Per-patch activation cost for dynamic encoder batch sizing.
-        vis_cfg = self.config.vision_config
-        self._encoder_bytes_per_patch = (
-            vis_cfg.hidden_size * 2 * vis_cfg.num_hidden_layers
-        )
-
-        return loaded
+        return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     # ------------------------------------------------------------------ #
     # LoRA / multimodal mapping