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, 08:45 UTC
[Bugfix][Gemma 4] Clamp soft-token estimate to max_soft_tokens (#40796)
Signed-off-by: Hoang Nguyen <118159510+hnt2601@users.noreply.github.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Hoang Nguyen committed
May 2, 2026 at 13:34 UTC
c3ad791e1a9ad9a2bff082ebbdf2b03749deef3e
2 files changed
+62
-1
tests/models/multimodal/processing/test_gemma4.py
+54
@@ -12,6 +12,60 @@ from ...utils import build_model_context
12
GEMMA4_MODEL_ID = "google/gemma-4-E2B-it"
13
14
15
+@pytest.mark.parametrize(
16
+ "image_width,image_height,max_soft_tokens",
17
+ [
18
+ # Production repro: a 3x900 image (extreme aspect ratio) made the
19
+ # prompt-side estimator return 289 while the HF Gemma 4 image
20
+ # processor's vision tower output capped at 280, producing the
21
+ # "Attempted to assign 280 multimodal tokens to 289 placeholders"
22
+ # mismatch that crashed EngineCore.
23
+ (900, 3, 280),
24
+ (3, 900, 280),
25
+ # Same pathology should hold for the video-frame budget (70 tokens).
26
+ (900, 3, 70),
27
+ # And for any other supported budget.
28
+ (4000, 2, 1120),
29
+ ],
30
+)
31
+@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID])
32
+def test_compute_num_soft_tokens_does_not_exceed_max_soft_tokens(
33
+ model_id: str,
34
+ image_width: int,
35
+ image_height: int,
36
+ max_soft_tokens: int,
37
+):
38
+ """Regression for the Gemma 3/4 multimodal crash.
39
+
40
+ `_compute_num_soft_tokens` must never return a value larger than
41
+ `max_soft_tokens`. The HF Gemma 4 image processor clamps its vision
42
+ tower output to that value; if the prompt-side estimator returns more,
43
+ the prompt has more `image` placeholder tokens than the encoder will
44
+ fill, and `_merge_multimodal_embeddings` raises `ValueError` deep in
45
+ the model forward.
46
+ """
47
+ ctx = build_model_context(
48
+ model_id,
49
+ mm_processor_kwargs={"do_pan_and_scan": True},
50
+ limit_mm_per_prompt={"image": 1},
51
+ )
52
+ processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
53
+
54
+ num_soft_tokens = processor.info._compute_num_soft_tokens(
55
+ image_width=image_width,
56
+ image_height=image_height,
57
+ max_soft_tokens=max_soft_tokens,
58
+ )
59
+
60
+ assert num_soft_tokens <= max_soft_tokens, (
61
+ f"_compute_num_soft_tokens returned {num_soft_tokens} for "
62
+ f"image_width={image_width}, image_height={image_height}, "
63
+ f"max_soft_tokens={max_soft_tokens} — exceeds the cap that the HF "
64
+ f"image processor enforces on its vision tower output. This is "
65
+ f"the placeholder/encoder count mismatch that crashes EngineCore."
66
+ )
67
+
68
+
69
@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID])
70
def test_limit_mm_per_prompt(
71
image_assets: ImageTestAssets,
vllm/model_executor/models/gemma4_mm.py
+8
-1
@@ -265,7 +265,14 @@ class Gemma4ProcessingInfo(BaseProcessingInfo):
265
target_h = max(unit, int(math.floor(image_height * scale / unit)) * unit)
266
target_w = max(unit, int(math.floor(image_width * scale / unit)) * unit)
267
num_patches = (target_h // patch_size) * (target_w // patch_size)
268
- return num_patches // (pooling_kernel_size**2)
268
+ # Clamp to ``max_soft_tokens``: extreme aspect ratios (e.g. 3x900)
269
+ # cause the floor() above to round one dim up to ``unit`` while the
270
+ # other scales freely, which over-shoots ``max_patches``. The HF
271
+ # Gemma 4 image processor caps its vision-tower output at
272
+ # ``max_soft_tokens``, so without this clamp the prompt-side
273
+ # placeholder count exceeds the encoder output and
274
+ # ``_merge_multimodal_embeddings`` crashes.
275
+ return min(num_patches // (pooling_kernel_size**2), max_soft_tokens)
276
277
def get_image_repl(
278
self,