| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Mapping |
| 5 | |
| 6 | import numpy as np |
| 7 | import pytest |
| 8 | import torch |
| 9 | from PIL import Image as PILImage |
| 10 | |
| 11 | from vllm.exceptions import VLLMValidationError |
| 12 | from vllm.model_executor.models.gemma4_mm import ( |
| 13 | Gemma4ForConditionalGeneration, |
| 14 | Gemma4ImagePixelInputs, |
| 15 | ) |
| 16 | from vllm.multimodal import MULTIMODAL_REGISTRY |
| 17 | from vllm.multimodal.cache import MultiModalProcessorOnlyCache |
| 18 | from vllm.multimodal.inputs import MultiModalFieldConfig |
| 19 | from vllm.utils.mem_constants import GiB_bytes |
| 20 | |
| 21 | from ....conftest import ImageTestAssets |
| 22 | from ...utils import build_model_context |
| 23 | |
| 24 | # TODO: to be updated to "google/gemma-4-e2b-it" once the models are available |
| 25 | GEMMA4_MODEL_ID = "google/gemma-4-E2B-it" |
| 26 | |
| 27 | |
| 28 | def test_gemma4_image_schema_accepts_variable_patch_counts(): |
| 29 | Gemma4ImagePixelInputs( |
| 30 | pixel_values=[ |
| 31 | torch.randn(10080, 768), |
| 32 | torch.randn(2520, 768), |
| 33 | ], |
| 34 | pixel_position_ids=[ |
| 35 | torch.zeros(10080, 2, dtype=torch.long), |
| 36 | torch.zeros(2520, 2, dtype=torch.long), |
| 37 | ], |
| 38 | ) |
| 39 | |
| 40 | |
| 41 | def test_gemma4_image_batching_keeps_variable_patch_counts_unstacked(): |
| 42 | field = MultiModalFieldConfig.batched("image").field |
| 43 | elems = field.build_elems( |
| 44 | "image", |
| 45 | "pixel_values", |
| 46 | [torch.randn(10080, 768), torch.randn(2520, 768)], |
| 47 | ) |
| 48 | |
| 49 | reduced = field.reduce_data(list(elems)) |
| 50 | |
| 51 | assert isinstance(reduced, list) |
| 52 | assert [tensor.shape for tensor in reduced] == [ |
| 53 | torch.Size([10080, 768]), |
| 54 | torch.Size([2520, 768]), |
| 55 | ] |
| 56 | |
| 57 | |
| 58 | @pytest.mark.parametrize( |
| 59 | "image_width,image_height,max_soft_tokens", |
| 60 | [ |
| 61 | # Production repro: a 3x900 image (extreme aspect ratio) made the |
| 62 | # prompt-side estimator return 289 while the HF Gemma 4 image |
| 63 | # processor's vision tower output capped at 280, producing the |
| 64 | # "Attempted to assign 280 multimodal tokens to 289 placeholders" |
| 65 | # mismatch that crashed EngineCore. |
| 66 | (900, 3, 280), |
| 67 | (3, 900, 280), |
| 68 | # Same pathology should hold for the video-frame budget (70 tokens). |
| 69 | (900, 3, 70), |
| 70 | # And for any other supported budget. |
| 71 | (4000, 2, 1120), |
| 72 | ], |
| 73 | ) |
| 74 | @pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID]) |
| 75 | def test_compute_num_soft_tokens_does_not_exceed_max_soft_tokens( |
| 76 | model_id: str, |
| 77 | image_width: int, |
| 78 | image_height: int, |
| 79 | max_soft_tokens: int, |
| 80 | ): |
| 81 | """Regression for the Gemma 3/4 multimodal crash. |
| 82 | |
| 83 | `_compute_num_soft_tokens` must never return a value larger than |
| 84 | `max_soft_tokens`. The HF Gemma 4 image processor clamps its vision |
| 85 | tower output to that value; if the prompt-side estimator returns more, |
| 86 | the prompt has more `image` placeholder tokens than the encoder will |
| 87 | fill, and `_merge_multimodal_embeddings` raises `ValueError` deep in |
| 88 | the model forward. |
| 89 | """ |
| 90 | ctx = build_model_context( |
| 91 | model_id, |
| 92 | mm_processor_kwargs={"do_pan_and_scan": True}, |
| 93 | limit_mm_per_prompt={"image": 1}, |
| 94 | ) |
| 95 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 96 | |
| 97 | num_soft_tokens = processor.info._compute_num_soft_tokens( |
| 98 | image_width=image_width, |
| 99 | image_height=image_height, |
| 100 | max_soft_tokens=max_soft_tokens, |
| 101 | ) |
| 102 | |
| 103 | assert num_soft_tokens <= max_soft_tokens, ( |
| 104 | f"_compute_num_soft_tokens returned {num_soft_tokens} for " |
| 105 | f"image_width={image_width}, image_height={image_height}, " |
| 106 | f"max_soft_tokens={max_soft_tokens} — exceeds the cap that the HF " |
| 107 | f"image processor enforces on its vision tower output. This is " |
| 108 | f"the placeholder/encoder count mismatch that crashes EngineCore." |
| 109 | ) |
| 110 | |
| 111 | |
| 112 | @pytest.mark.parametrize( |
| 113 | ("mm_processor_kwargs", "expected_image_tokens"), |
| 114 | [ |
| 115 | ({}, 280), |
| 116 | ({"max_soft_tokens": 70}, 70), |
| 117 | ({"max_soft_tokens": 280}, 280), |
| 118 | ({"max_soft_tokens": 1120}, 1120), |
| 119 | ({"images_kwargs": {"max_soft_tokens": 560}}, 560), |
| 120 | ({"images_kwargs": None}, 280), |
| 121 | ({"images_kwargs": "not-a-dict"}, 280), |
| 122 | ], |
| 123 | ) |
| 124 | @pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID]) |
| 125 | def test_get_mm_max_tokens_per_item_respects_configured_max_soft_tokens( |
| 126 | model_id: str, |
| 127 | mm_processor_kwargs: dict[str, object], |
| 128 | expected_image_tokens: int, |
| 129 | ): |
| 130 | ctx = build_model_context( |
| 131 | model_id, |
| 132 | mm_processor_kwargs=mm_processor_kwargs, |
| 133 | limit_mm_per_prompt={"image": 1, "video": 1}, |
| 134 | ) |
| 135 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 136 | |
| 137 | tokens = processor.info.get_mm_max_tokens_per_item( |
| 138 | seq_len=ctx.model_config.max_model_len, |
| 139 | mm_counts={"image": 1, "video": 1}, |
| 140 | ) |
| 141 | |
| 142 | assert tokens is not None |
| 143 | assert tokens["image"] == expected_image_tokens |
| 144 | assert tokens["video"] == 32 * (70 + 2 + 6) |
| 145 | |
| 146 | |
| 147 | @pytest.mark.parametrize( |
| 148 | ("limit_mm_per_prompt", "expected_video_tokens"), |
| 149 | [ |
| 150 | ({"video": 1}, 32 * (70 + 2 + 6)), |
| 151 | ({"video": {"count": 1}}, 32 * (70 + 2 + 6)), |
| 152 | ({"video": {"count": 1, "num_frames": 1}}, 1 * (70 + 2 + 6)), |
| 153 | ({"video": {"count": 1, "num_frames": 8}}, 8 * (70 + 2 + 6)), |
| 154 | ({"video": {"count": 1, "num_frames": 32}}, 32 * (70 + 2 + 6)), |
| 155 | ({"video": {"count": 1, "num_frames": 40}}, 32 * (70 + 2 + 6)), |
| 156 | ], |
| 157 | ) |
| 158 | @pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID]) |
| 159 | def test_get_mm_max_tokens_per_item_respects_configured_video_num_frames( |
| 160 | model_id: str, |
| 161 | limit_mm_per_prompt: Mapping[str, int | Mapping[str, int]], |
| 162 | expected_video_tokens: int, |
| 163 | ): |
| 164 | ctx = build_model_context( |
| 165 | model_id, |
| 166 | limit_mm_per_prompt=limit_mm_per_prompt, |
| 167 | ) |
| 168 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 169 | |
| 170 | tokens = processor.info.get_mm_max_tokens_per_item( |
| 171 | seq_len=ctx.model_config.max_model_len, |
| 172 | mm_counts={"video": 1}, |
| 173 | ) |
| 174 | |
| 175 | assert tokens is not None |
| 176 | assert tokens["image"] == 280 |
| 177 | assert tokens["video"] == expected_video_tokens |
| 178 | |
| 179 | |
| 180 | @pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID]) |
| 181 | def test_get_prompt_updates_respects_nested_max_soft_tokens(model_id: str): |
| 182 | ctx = build_model_context( |
| 183 | model_id, |
| 184 | mm_processor_kwargs={"images_kwargs": {"max_soft_tokens": 560}}, |
| 185 | limit_mm_per_prompt={"image": 1}, |
| 186 | ) |
| 187 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 188 | image = PILImage.new("RGB", (1000, 1000), color="white") |
| 189 | image_size = image.size |
| 190 | mm_items = processor.info.parse_mm_data({"image": image}) |
| 191 | |
| 192 | prompt_update = processor._get_prompt_updates(mm_items, {}, {})[0] |
| 193 | replacement = prompt_update.resolve(0).content.full |
| 194 | expected = processor.info.get_image_repl( |
| 195 | image_width=image_size[0], |
| 196 | image_height=image_size[1], |
| 197 | processor=processor.info.get_hf_processor(), |
| 198 | max_soft_tokens=560, |
| 199 | ).full |
| 200 | |
| 201 | assert replacement == expected |
| 202 | |
| 203 | |
| 204 | @pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID]) |
| 205 | @pytest.mark.parametrize( |
| 206 | "image_kwargs", [{"rescale_factor": 1 / 127.5}, {"max_soft_tokens": 560}] |
| 207 | ) |
| 208 | @pytest.mark.parametrize("video_uuid", [None, "same-video"]) |
| 209 | def test_video_cache_is_independent_of_image_kwargs( |
| 210 | model_id: str, |
| 211 | image_kwargs: dict[str, object], |
| 212 | video_uuid: str | None, |
| 213 | ): |
| 214 | """Request image overrides must not change video frames or make results |
| 215 | depend on cache warmth. |
| 216 | """ |
| 217 | request_kwargs = {"images_kwargs": image_kwargs} |
| 218 | ctx = build_model_context( |
| 219 | model_id, |
| 220 | limit_mm_per_prompt={"image": 1, "video": 1}, |
| 221 | mm_processor_cache_gb=1, |
| 222 | ) |
| 223 | cache = MultiModalProcessorOnlyCache(ctx.model_config) |
| 224 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 225 | hf_processor = processor.info.get_hf_processor() |
| 226 | image = PILImage.new("RGB", (48, 48), color=(128, 128, 128)) |
| 227 | frames = np.stack([np.asarray(image)] * 2) |
| 228 | metadata = {"fps": 2.0, "frames_indices": [0, 1]} |
| 229 | mm_items = processor.info.parse_mm_data( |
| 230 | {"image": image, "video": [(frames, metadata)]} |
| 231 | ) |
| 232 | |
| 233 | def process(mm_kwargs, cache): |
| 234 | return processor( |
| 235 | hf_processor.image_token + hf_processor.video_token, |
| 236 | mm_items, |
| 237 | mm_uuid_items={"video": [video_uuid]}, |
| 238 | hf_processor_mm_kwargs=mm_kwargs, |
| 239 | cache=cache, |
| 240 | ) |
| 241 | |
| 242 | # Populate the video cache without image overrides, then change only |
| 243 | # images_kwargs on the next request. Since images_kwargs are excluded from |
| 244 | # the video hash, and must not affect the processed video, the cached result |
| 245 | # must match a fresh recomputation. |
| 246 | baseline = process({}, cache=cache) |
| 247 | cached = process(request_kwargs, cache=cache) |
| 248 | cache.clear_cache() |
| 249 | fresh = process(request_kwargs, cache=cache) |
| 250 | |
| 251 | def pixels(result, modality, field): |
| 252 | return result["mm_kwargs"][modality][0][field].data |
| 253 | |
| 254 | video_pixels = pixels(fresh, "video", "pixel_values_videos") |
| 255 | assert torch.equal(pixels(cached, "video", "pixel_values_videos"), video_pixels) |
| 256 | assert torch.equal(pixels(baseline, "video", "pixel_values_videos"), video_pixels) |
| 257 | assert baseline["mm_hashes"]["video"] == fresh["mm_hashes"]["video"] |
| 258 | assert not torch.equal( |
| 259 | pixels(baseline, "image", "pixel_values"), |
| 260 | pixels(fresh, "image", "pixel_values"), |
| 261 | ) |
| 262 | |
| 263 | |
| 264 | @pytest.mark.parametrize("model_id", [GEMMA4_MODEL_ID]) |
| 265 | def test_limit_mm_per_prompt( |
| 266 | image_assets: ImageTestAssets, |
| 267 | model_id: str, |
| 268 | ): |
| 269 | """Test that limit_mm_per_prompt accurately restricts multiple images.""" |
| 270 | # We only allow 1 image |
| 271 | ctx = build_model_context( |
| 272 | model_id, |
| 273 | mm_processor_kwargs={}, |
| 274 | limit_mm_per_prompt={"image": 1}, |
| 275 | ) |
| 276 | processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) |
| 277 | |
| 278 | # Provide 2 images in the prompt |
| 279 | prompt = "<image><image>" |
| 280 | # image_assets usually has multiple images |
| 281 | images = [asset.pil_image for asset in image_assets][:2] |
| 282 | if len(images) < 2: |
| 283 | images = [images[0], images[0]] |
| 284 | |
| 285 | mm_data = {"image": images} |
| 286 | |
| 287 | # Expect ValueError when exceeding limit |
| 288 | with pytest.raises(VLLMValidationError, match="At most 1 image"): |
| 289 | processor( |
| 290 | prompt, |
| 291 | mm_items=processor.info.parse_mm_data(mm_data), |
| 292 | hf_processor_mm_kwargs={}, |
| 293 | ) |
| 294 | |
| 295 | |
| 296 | # Regression guard for PR #43169 follow-up: the batched Gemma4 vision encoder |
| 297 | # admitted ``chunk ~= 53`` on a 22 GiB L4 with a 26B AWQ model loaded, |
| 298 | # allocating 2.43 GiB int64 inside |
| 299 | # ``F.one_hot(num_classes=position_embedding_size)`` and OOMing because only |
| 300 | # 2.41 GiB was actually free. The fix sizes ``chunk`` from currently-free GPU |
| 301 | # memory and counts the ``F.one_hot`` transient as the dominant cost. |
| 302 | |
| 303 | _encoder_chunk = Gemma4ForConditionalGeneration._encoder_chunk |
| 304 | |
| 305 | # Gemma4 vision config default (HF: configuration_gemma4.py). |
| 306 | _POSITION_EMBEDDING_SIZE = 10240 |
| 307 | # Video frame: max_soft_tokens=70, pooling_kernel_size=2 -> 70 * 4 patches. |
| 308 | _VIDEO_PATCHES_PER_FRAME = 280 |
| 309 | |
| 310 | |
| 311 | def test_encoder_chunk_tight_budget_fits_in_free(): |
| 312 | free = 3 * GiB_bytes # L4 22 GiB after 26B AWQ load. |
| 313 | total = 22 * GiB_bytes |
| 314 | chunk = _encoder_chunk( |
| 315 | _VIDEO_PATCHES_PER_FRAME, free, total, _POSITION_EMBEDDING_SIZE |
| 316 | ) |
| 317 | one_hot_bytes = chunk * _VIDEO_PATCHES_PER_FRAME * 2 * _POSITION_EMBEDDING_SIZE * 8 |
| 318 | assert one_hot_bytes <= free // 2 |
| 319 | |
| 320 | |
| 321 | def test_encoder_chunk_roomy_gpu_keeps_batching(): |
| 322 | chunk = _encoder_chunk( |
| 323 | _VIDEO_PATCHES_PER_FRAME, |
| 324 | 60 * GiB_bytes, |
| 325 | 80 * GiB_bytes, |
| 326 | _POSITION_EMBEDDING_SIZE, |
| 327 | ) |
| 328 | assert chunk > 8 |
| 329 | |
| 330 | |
| 331 | def test_encoder_chunk_zero_patches_is_safe(): |
| 332 | assert ( |
| 333 | _encoder_chunk(0, 60 * GiB_bytes, 80 * GiB_bytes, _POSITION_EMBEDDING_SIZE) == 1 |
| 334 | ) |
| 335 | |
| 336 | |
| 337 | def test_encoder_chunk_zero_position_embedding_size_is_safe(): |
| 338 | # Degenerate config: must not raise ZeroDivisionError. |
| 339 | assert ( |
| 340 | _encoder_chunk(_VIDEO_PATCHES_PER_FRAME, 60 * GiB_bytes, 80 * GiB_bytes, 0) == 1 |
| 341 | ) |
| 342 | |
| 343 | |
| 344 | def test_encoder_chunk_no_free_memory_falls_back_to_one(): |
| 345 | assert ( |
| 346 | _encoder_chunk( |
| 347 | _VIDEO_PATCHES_PER_FRAME, 0, 22 * GiB_bytes, _POSITION_EMBEDDING_SIZE |
| 348 | ) |
| 349 | == 1 |
| 350 | ) |