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, 21:30 UTC
[Fix] Gemma4 Mixed-Resolution Image Co-Batching Crash (#42217)
[Fix] Gemma4 Mixed-Resolution Image Co-Batching Crash (#42217) Signed-off-by: zitian.zhao <zitian.zhao@tencentmusic.com>
ZiTian Zhao committed
May 12, 2026 at 11:13 UTC
630492da308e2141396a41c120fae8802e9c5857
2 files changed
+45
-7
tests/models/multimodal/processing/test_gemma4.py
+33
@@ -4,9 +4,12 @@
4
from collections.abc import Mapping
5
6
import pytest
7
+import torch
8
from PIL import Image as PILImage
9
10
+from vllm.model_executor.models.gemma4_mm import Gemma4ImagePixelInputs
11
from vllm.multimodal import MULTIMODAL_REGISTRY
12
+from vllm.multimodal.inputs import MultiModalFieldConfig
13
14
from ....conftest import ImageTestAssets
15
from ...utils import build_model_context
@@ -15,6 +18,36 @@ from ...utils import build_model_context
18
GEMMA4_MODEL_ID = "google/gemma-4-E2B-it"
19
20
21
+def test_gemma4_image_schema_accepts_variable_patch_counts():
22
+ Gemma4ImagePixelInputs(
23
+ pixel_values=[
24
+ torch.randn(10080, 768),
25
+ torch.randn(2520, 768),
26
+ ],
27
+ pixel_position_ids=[
28
+ torch.zeros(10080, 2, dtype=torch.long),
29
+ torch.zeros(2520, 2, dtype=torch.long),
30
+ ],
31
+ )
32
+
33
+
34
+def test_gemma4_image_batching_keeps_variable_patch_counts_unstacked():
35
+ field = MultiModalFieldConfig.batched("image").field
36
+ elems = field.build_elems(
37
+ "image",
38
+ "pixel_values",
39
+ [torch.randn(10080, 768), torch.randn(2520, 768)],
40
+ )
41
+
42
+ reduced = field.reduce_data(list(elems))
43
+
44
+ assert isinstance(reduced, list)
45
+ assert [tensor.shape for tensor in reduced] == [
46
+ torch.Size([10080, 768]),
47
+ torch.Size([2520, 768]),
48
+ ]
49
+
50
+
51
@pytest.mark.parametrize(
52
"image_width,image_height,max_soft_tokens",
53
[
vllm/model_executor/models/gemma4_mm.py
+12
-7
@@ -124,12 +124,12 @@ class Gemma4ImagePixelInputs(TensorSchema):
124
125
type: Literal["pixel_values"] = "pixel_values"
126
pixel_values: Annotated[
127
- torch.Tensor,
128
- TensorShape("bn", "np", "pp"),
127
+ torch.Tensor | list[torch.Tensor],
128
+ TensorShape("bn", "np", "pp", dynamic_dims={"np"}),
129
]
130
pixel_position_ids: Annotated[
131
- torch.Tensor,
132
- TensorShape("bn", "np", 2),
131
+ torch.Tensor | list[torch.Tensor],
132
+ TensorShape("bn", "np", 2, dynamic_dims={"np"}),
133
]
134
135
@@ -1128,15 +1128,20 @@ class Gemma4ForConditionalGeneration(
1128
# metadata, and validating numerical equivalence with the
1129
# current per-image path.
1130
#
1131
+ # Concurrent requests with different image resolutions may
1132
+ # arrive as a list of per-image tensors, while same-resolution
1133
+ # batches may arrive as a stacked tensor. Both forms are
1134
+ # iterable over the per-image dimension.
1135
+
1136
# Process each image individually through the vision tower.
1137
# The vision tower's forward() strips padding and returns a
1138
# flat tensor of valid tokens. We process per-image to get
1139
# variable-length outputs matching the dynamic token count
1140
# from get_image_repl.
1141
per_image_features = []
1137
- for i in range(pixel_values.shape[0]):
1138
- pv = pixel_values[i].unsqueeze(0) # (1, max_patches, patch_pixels)
1139
- pp = pixel_position_ids[i].unsqueeze(0) # (1, max_patches, 2)
1142
+ for pv, pp in zip(pixel_values, pixel_position_ids, strict=True):
1143
+ pv = pv.unsqueeze(0) # (1, max_patches, patch_pixels)
1144
+ pp = pp.unsqueeze(0) # (1, max_patches, 2)
1145
1146
# Derive the pooler's output_length from the total patch
1147
# count (including padding). The vision tower encoder