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 3, 01:45 UTC
[CI] Fix "test_awq_load[gemma4-moe-*]" failure (#43296)
Signed-off-by: haosdent <haosdent@gmail.com>
haosdent committed
May 22, 2026 at 15:13 UTC
025d4f5cd2617bb767663f9e7d62354039887757
2 files changed
+116
-31
tests/models/multimodal/processing/test_gemma4.py
+62
-1
@@ -7,9 +7,13 @@ import pytest
7
import torch
8
from PIL import Image as PILImage
9
10
-from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs
10
+from vllm.model_executor.models.gemma4_mm import (
11
+ Gemma4ForConditionalGeneration,
12
+ Gemma4ImagePixelInputs,
13
+)
14
from vllm.multimodal import MULTIMODAL_REGISTRY
15
from vllm.multimodal.inputs import MultiModalFieldConfig
16
+from vllm.utils.mem_constants import GiB_bytes
17
18
from ....conftest import ImageTestAssets
19
from ...utils import build_model_context
@@ -224,3 +228,60 @@ def test_limit_mm_per_prompt(
228
mm_items=processor.info.parse_mm_data(mm_data),
229
hf_processor_mm_kwargs={},
230
)
231
+
232
+
233
+# Regression guard for PR #43169 follow-up: the batched Gemma4 vision encoder
234
+# admitted ``chunk ~= 53`` on a 22 GiB L4 with a 26B AWQ model loaded,
235
+# allocating 2.43 GiB int64 inside
236
+# ``F.one_hot(num_classes=position_embedding_size)`` and OOMing because only
237
+# 2.41 GiB was actually free. The fix sizes ``chunk`` from currently-free GPU
238
+# memory and counts the ``F.one_hot`` transient as the dominant cost.
239
+
240
+_encoder_chunk = Gemma4ForConditionalGeneration._encoder_chunk
241
+
242
+# Gemma4 vision config default (HF: configuration_gemma4.py).
243
+_POSITION_EMBEDDING_SIZE = 10240
244
+# Video frame: max_soft_tokens=70, pooling_kernel_size=2 -> 70 * 4 patches.
245
+_VIDEO_PATCHES_PER_FRAME = 280
246
+
247
+
248
+def test_encoder_chunk_tight_budget_fits_in_free():
249
+ free = 3 * GiB_bytes # L4 22 GiB after 26B AWQ load.
250
+ total = 22 * GiB_bytes
251
+ chunk = _encoder_chunk(
252
+ _VIDEO_PATCHES_PER_FRAME, free, total, _POSITION_EMBEDDING_SIZE
253
+ )
254
+ one_hot_bytes = chunk * _VIDEO_PATCHES_PER_FRAME * 2 * _POSITION_EMBEDDING_SIZE * 8
255
+ assert one_hot_bytes <= free // 2
256
+
257
+
258
+def test_encoder_chunk_roomy_gpu_keeps_batching():
259
+ chunk = _encoder_chunk(
260
+ _VIDEO_PATCHES_PER_FRAME,
261
+ 60 * GiB_bytes,
262
+ 80 * GiB_bytes,
263
+ _POSITION_EMBEDDING_SIZE,
264
+ )
265
+ assert chunk > 8
266
+
267
+
268
+def test_encoder_chunk_zero_patches_is_safe():
269
+ assert (
270
+ _encoder_chunk(0, 60 * GiB_bytes, 80 * GiB_bytes, _POSITION_EMBEDDING_SIZE) == 1
271
+ )
272
+
273
+
274
+def test_encoder_chunk_zero_position_embedding_size_is_safe():
275
+ # Degenerate config: must not raise ZeroDivisionError.
276
+ assert (
277
+ _encoder_chunk(_VIDEO_PATCHES_PER_FRAME, 60 * GiB_bytes, 80 * GiB_bytes, 0) == 1
278
+ )
279
+
280
+
281
+def test_encoder_chunk_no_free_memory_falls_back_to_one():
282
+ assert (
283
+ _encoder_chunk(
284
+ _VIDEO_PATCHES_PER_FRAME, 0, 22 * GiB_bytes, _POSITION_EMBEDDING_SIZE
285
+ )
286
+ == 1
287
+ )
vllm/model_executor/models/gemma4_mm.py
+54
-30
@@ -961,9 +961,6 @@ class Gemma4ForConditionalGeneration(
961
self.embed_vision = Gemma4MultimodalEmbedder(
962
config.vision_config, config.text_config
963
)
964
- # Lazy-initialized on first encoder call (see _encoder_max_batch).
965
- self._encoder_budget_bytes = 0
966
- self._encoder_bytes_per_patch = 0
964
965
# ---- Audio tower (variants with audio_config) ----
966
if config.audio_config is not None:
@@ -1104,18 +1101,36 @@ class Gemma4ForConditionalGeneration(
1101
)
1102
return mm_input_by_modality
1103
1107
- def _encoder_max_batch(self, patches_per_item: int) -> int:
1108
- """Max items per encoder call given per-item patch count."""
1109
- if self._encoder_budget_bytes == 0:
1110
- total_mem = current_platform.get_device_total_memory()
1111
- self._encoder_budget_bytes = int(total_mem * 0.05)
1112
- logger.info(
1113
- "Encoder memory budget: %.1fGB (total=%.1fGB)",
1114
- self._encoder_budget_bytes / 1024**3,
1115
- total_mem / 1024**3,
1116
- )
1117
- cost = patches_per_item * self._encoder_bytes_per_patch
1118
- return max(1, self._encoder_budget_bytes // cost) if cost > 0 else 1
1104
+ @staticmethod
1105
+ def _encoder_chunk(
1106
+ patches_per_item: int,
1107
+ free_bytes: int,
1108
+ total_bytes: int,
1109
+ position_embedding_size: int,
1110
+ ) -> int:
1111
+ """Max chunk size whose F.one_hot transient fits in the budget.
1112
+
1113
+ The dominant transient inside HF's ``Gemma4VisionPatchEmbedder.
1114
+ _position_embeddings`` is
1115
+ ``F.one_hot(clamped_positions, num_classes=position_embedding_size)``
1116
+ with shape ``(chunk, patches, 2, position_embedding_size)``,
1117
+ int64, plus its simultaneous cast to the position embedding
1118
+ table dtype. That, not the encoder residual stream, sets peak
1119
+ memory.
1120
+ """
1121
+ if patches_per_item <= 0:
1122
+ return 1
1123
+ # Half of currently-free, capped at 10% of total so we leave room
1124
+ # for the rest of profile_run / the subsequent encoder + pooler.
1125
+ budget = min(free_bytes // 2, total_bytes // 10)
1126
+ if budget <= 0:
1127
+ return 1
1128
+ # F.one_hot allocates (chunk, patches, 2, pos_emb_size) int64
1129
+ # (the inner 2 is the (x, y) coordinate axis, 8 is sizeof(int64)).
1130
+ # Outer 2x covers the int64 buffer and its concurrent bf16 cast
1131
+ # plus the matmul output that live alongside it at peak.
1132
+ cost = patches_per_item * 4 * position_embedding_size * 8
1133
+ return max(1, budget // cost) if cost > 0 else 1
1134
1135
# ------------------------------------------------------------------ #
1136
# Image processing
@@ -1136,7 +1151,8 @@ class Gemma4ForConditionalGeneration(
1151
pixel_position_ids = image_input["pixel_position_ids"]
1152
1153
vt = self.vision_tower
1139
- pooling_k2 = self.config.vision_config.pooling_kernel_size**2
1154
+ vision_cfg = self.config.vision_config
1155
+ pooling_k2 = vision_cfg.pooling_kernel_size**2
1156
1157
# Concurrent requests with different image resolutions may
1158
# arrive as a list of per-image tensors, while same-resolution
@@ -1153,10 +1169,18 @@ class Gemma4ForConditionalGeneration(
1169
pp = pixel_position_ids[idx]
1170
buckets.setdefault(pv.shape[0], []).append((idx, pv, pp))
1171
1156
- # Encode each resolution bucket in memory-safe chunks.
1172
+ # Encode each resolution bucket in memory-safe chunks. Re-read
1173
+ # free memory per bucket because the previous bucket's encoder
1174
+ # pass has already allocated activations we should account for.
1175
last_hidden_states_map: dict[int, torch.Tensor] = {}
1176
for patches, items in buckets.items():
1159
- max_batch_size = min(len(items), self._encoder_max_batch(patches))
1177
+ free, total = current_platform.mem_get_info()
1178
+ max_batch_size = min(
1179
+ len(items),
1180
+ self._encoder_chunk(
1181
+ patches, free, total, vision_cfg.position_embedding_size
1182
+ ),
1183
+ )
1184
1185
for chunk_idx in range(0, len(items), max_batch_size):
1186
chunk_items = items[chunk_idx : chunk_idx + max_batch_size]
@@ -1247,7 +1271,8 @@ class Gemma4ForConditionalGeneration(
1271
frame_counts = video_input["video_frame_counts"]
1272
1273
vt = self.vision_tower
1250
- pooling_k2 = self.config.vision_config.pooling_kernel_size**2
1274
+ vision_cfg = self.config.vision_config
1275
+ pooling_k2 = vision_cfg.pooling_kernel_size**2
1276
target_dtype = self.embed_vision.embedding_projection.weight.dtype
1277
1278
if isinstance(frame_counts, torch.Tensor):
@@ -1256,13 +1281,20 @@ class Gemma4ForConditionalGeneration(
1281
fc_list = list(frame_counts)
1282
1283
total_frames = pixel_values.shape[0]
1284
+ free, total = current_platform.mem_get_info()
1285
max_batch_size = min(
1260
- total_frames, self._encoder_max_batch(pixel_values.shape[1])
1286
+ total_frames,
1287
+ self._encoder_chunk(
1288
+ pixel_values.shape[1],
1289
+ free,
1290
+ total,
1291
+ vision_cfg.position_embedding_size,
1292
+ ),
1293
)
1294
1295
padding_positions = (pixel_position_ids == -1).all(dim=-1)
1296
1265
- # Encode frames in chunks bounded by _encoder_max_batch.
1297
+ # Encode frames in chunks bounded by _encoder_chunk.
1298
last_hidden_states_list: list[torch.Tensor] = []
1299
for i in range(0, total_frames, max_batch_size):
1300
pv_chunk = pixel_values[i : i + max_batch_size]
@@ -1538,15 +1570,7 @@ class Gemma4ForConditionalGeneration(
1570
self,
1571
ignore_unexpected_prefixes=ignore_prefixes,
1572
)
1541
- loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1542
-
1543
- # Per-patch activation cost for dynamic encoder batch sizing.
1544
- vis_cfg = self.config.vision_config
1545
- self._encoder_bytes_per_patch = (
1546
- vis_cfg.hidden_size * 2 * vis_cfg.num_hidden_layers
1547
- )
1548
-
1549
- return loaded
1573
+ return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1574
1575
# ------------------------------------------------------------------ #
1576
# LoRA / multimodal mapping