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, 06:15 UTC
[MM][Gemma4] Respect max_soft_tokens in encoder budget (#41799)
Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com> Co-authored-by: lesj0610 <lesj0610@users.noreply.github.com> Co-authored-by: gemini-code-assist <gemini-code-assist@google.com>
SeongJun Lee committed
May 6, 2026 at 21:54 UTC
242afc6bf40d6d088b7b97eda1dd7f00ddcdfa21
2 files changed
+91
-11
tests/models/multimodal/processing/test_gemma4.py
+60
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import pytest
5
+from PIL import Image as PILImage
6
7
from vllm.multimodal import MULTIMODAL_REGISTRY
8
@@ -66,6 +67,65 @@ def test_compute_num_soft_tokens_does_not_exceed_max_soft_tokens(
67
)
68
69
70
+@pytest.mark.parametrize(
71
+ ("mm_processor_kwargs", "expected_image_tokens"),
72
+ [
73
+ ({}, 280),
74
+ ({"max_soft_tokens": 70}, 70),
75
+ ({"max_soft_tokens": 280}, 280),
76
+ ({"max_soft_tokens": 1120}, 1120),
77
+ ({"images_kwargs": {"max_soft_tokens": 560}}, 560),
78
+ ({"images_kwargs": None}, 280),
79
+ ({"images_kwargs": "not-a-dict"}, 280),
80
+ ],
81
+)
82
+@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID])
83
+def test_get_mm_max_tokens_per_item_respects_configured_max_soft_tokens(
84
+ model_id: str,
85
+ mm_processor_kwargs: dict[str, object],
86
+ expected_image_tokens: int,
87
+):
88
+ ctx = build_model_context(
89
+ model_id,
90
+ mm_processor_kwargs=mm_processor_kwargs,
91
+ limit_mm_per_prompt={"image": 1, "video": 1},
92
+ )
93
+ processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
94
+
95
+ tokens = processor.info.get_mm_max_tokens_per_item(
96
+ seq_len=ctx.model_config.max_model_len,
97
+ mm_counts={"image": 1, "video": 1},
98
+ )
99
+
100
+ assert tokens is not None
101
+ assert tokens["image"] == expected_image_tokens
102
+ assert tokens["video"] == 32 * (70 + 2 + 6)
103
+
104
+
105
+@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID])
106
+def test_get_prompt_updates_respects_nested_max_soft_tokens(model_id: str):
107
+ ctx = build_model_context(
108
+ model_id,
109
+ mm_processor_kwargs={"images_kwargs": {"max_soft_tokens": 560}},
110
+ limit_mm_per_prompt={"image": 1},
111
+ )
112
+ processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
113
+ image = PILImage.new("RGB", (1000, 1000), color="white")
114
+ image_size = image.size
115
+ mm_items = processor.info.parse_mm_data({"image": image})
116
+
117
+ prompt_update = processor._get_prompt_updates(mm_items, {}, {})[0]
118
+ replacement = prompt_update.resolve(0).content.full
119
+ expected = processor.info.get_image_repl(
120
+ image_width=image_size[0],
121
+ image_height=image_size[1],
122
+ processor=processor.info.get_hf_processor(),
123
+ max_soft_tokens=560,
124
+ ).full
125
+
126
+ assert replacement == expected
127
+
128
+
129
@pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID])
130
def test_limit_mm_per_prompt(
131
image_assets: ImageTestAssets,
vllm/model_executor/models/gemma4_mm.py
+31
-11
@@ -81,10 +81,26 @@ from .utils import (
81
logger = init_logger(__name__)
82
83
# Video constants — match transformers Gemma4VideoProcessor defaults.
84
+_SUPPORTED_SOFT_TOKENS = (70, 140, 280, 560, 1120)
85
_VIDEO_MAX_SOFT_TOKENS = 70 # soft tokens per video frame (vs 280 for images)
86
_VIDEO_MAX_FRAMES = 32 # max sampled frames per video
87
88
89
+def _get_max_soft_tokens(
90
+ merged_kwargs: Mapping[str, object],
91
+) -> tuple[object | None, bool]:
92
+ """Return configured image max_soft_tokens and whether it is top-level."""
93
+ val = merged_kwargs.get("max_soft_tokens")
94
+ if val is not None:
95
+ return val, True
96
+
97
+ images_kwargs = merged_kwargs.get("images_kwargs")
98
+ if isinstance(images_kwargs, Mapping):
99
+ return images_kwargs.get("max_soft_tokens"), False
100
+
101
+ return None, False
102
+
103
+
104
# ---------------------------------------------------------------------------
105
# Input schema
106
# ---------------------------------------------------------------------------
@@ -216,10 +232,14 @@ class Gemma4ProcessingInfo(BaseProcessingInfo):
232
self, seq_len: int, mm_counts: Mapping[str, int]
233
) -> Mapping[str, int] | None:
234
config = self.get_hf_config()
219
- # Upper bound: the pooler outputs default_output_length slots
220
- # per image (280). After padding is stripped the actual count
221
- # is ≤ this value, but vLLM needs the max for memory planning.
235
+ # Upper bound: the pooler outputs max_soft_tokens slots per image.
236
+ # After padding is stripped the actual count is ≤ this value, but
237
+ # vLLM needs the max for memory planning.
238
tokens_per_image = config.vision_config.default_output_length
239
+ merged_kwargs = self.ctx.get_merged_mm_kwargs({})
240
+ val, _ = _get_max_soft_tokens(merged_kwargs)
241
+ if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS:
242
+ tokens_per_image = val
243
tokens: dict[str, int] = {"image": tokens_per_image}
244
if config.audio_config is not None:
245
# Audio max tokens from the processor's audio_seq_length.
@@ -492,13 +512,8 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):
512
mm_kwargs: Mapping[str, object],
513
tok_kwargs: Mapping[str, object],
514
) -> BatchFeature:
495
- # Validate max_soft_tokens early and exit cleanly on bad values.
496
- _SUPPORTED_SOFT_TOKENS = (70, 140, 280, 560, 1120)
497
-
515
merged_kwargs = self.info.ctx.get_merged_mm_kwargs(mm_kwargs)
499
- val = merged_kwargs.get("max_soft_tokens")
500
- if val is None:
501
- val = merged_kwargs.get("images_kwargs", {}).get("max_soft_tokens")
516
+ val, is_top_level_max_soft_tokens = _get_max_soft_tokens(merged_kwargs)
517
518
if val is not None and val not in _SUPPORTED_SOFT_TOKENS:
519
raise ValueError(
@@ -645,7 +660,7 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):
660
# HF side (Gemma4ProcessorKwargs.images_kwargs) so that
661
# _merge_kwargs routes max_soft_tokens into images_kwargs.
662
patched_mm_kwargs = dict(mm_kwargs)
648
- if val is not None:
663
+ if val is not None and is_top_level_max_soft_tokens:
664
patched_mm_kwargs["max_soft_tokens"] = val
665
666
processed_outputs = super()._call_hf_processor(
@@ -755,7 +770,12 @@ class Gemma4MultiModalProcessor(BaseMultiModalProcessor[Gemma4ProcessingInfo]):
770
merged_kwargs = self.info.ctx.get_merged_mm_kwargs(
771
hf_processor_mm_kwargs,
772
)
758
- max_soft_tokens = merged_kwargs.get("max_soft_tokens")
773
+ val, _ = _get_max_soft_tokens(merged_kwargs)
774
+ max_soft_tokens = (
775
+ val
776
+ if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS
777
+ else None
778
+ )
779
return self.info.get_image_repl(
780
image_width=image_size.width,
781
image_height=image_size.height,