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 14, 05:30 UTC
[Multimodal] Update Tensor schema test to cover arbitrary shape mm inputs (#22867)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Aug 16, 2025 at 15:44 UTC
cc826a202b7b66af222374129573763237db3c1c
2 files changed
+138
-27
tests/models/multimodal/test_tensor_schema.py
+124
-19
@@ -1,17 +1,26 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterable
4
from functools import partial
5
+from typing import Any, Union
6
from unittest.mock import patch
7
8
+import numpy as np
9
import pytest
10
+from mistral_common.protocol.instruct.messages import (ImageChunk, TextChunk,
11
+ UserMessage)
12
+from mistral_common.protocol.instruct.request import ChatCompletionRequest
13
+from PIL import Image
14
15
from vllm.config import ModelConfig
16
from vllm.engine.llm_engine import LLMEngine as V0LLMEngine
17
from vllm.inputs import InputProcessingContext
11
-from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalKwargs
18
+from vllm.multimodal import (MULTIMODAL_REGISTRY, BatchedTensorInputs,
19
+ MultiModalKwargs)
20
from vllm.multimodal.processing import BaseMultiModalProcessor
21
+from vllm.multimodal.utils import group_mm_kwargs_by_modality
22
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
14
-from vllm.utils import GiB_bytes, set_default_torch_num_threads
23
+from vllm.utils import GiB_bytes, is_list_of, set_default_torch_num_threads
24
from vllm.v1.core.kv_cache_utils import get_kv_cache_config
25
from vllm.v1.engine.core import EngineCore as V1EngineCore
26
@@ -23,12 +32,64 @@ ARCH_TO_SKIP = {
32
"MolmoForCausalLM": "incompatible requirements",
33
"MiniMaxVL01ForConditionalGeneration": "broken model",
34
}
35
+ARCH_NEEDS_EXTRAS = [
36
+ "InternVLChatModel",
37
+ "Idefics3ForConditionalGeneration",
38
+ "LlavaForConditionalGeneration",
39
+ "MiniCPMV",
40
+ "PaliGemmaForConditionalGeneration",
41
+]
42
+REPO_ID_TO_SKIP = {"nm-testing/pixtral-12b-FP8-dynamic": "duplicated test"}
43
+
44
+ImageInput = list[Image.Image]
45
+VideoInput = Union[list[Image.Image], list[np.ndarray],
46
+ list[tuple[np.ndarray, dict[str, Any]]]]
47
+AudioInput = list[tuple[np.ndarray, int]]
48
+
49
+
50
+def _resize_data(_data: Union[Image.Image, np.ndarray],
51
+ size_factor: float) -> Union[Image.Image, np.ndarray]:
52
+ assert size_factor <= 1, "Size factor must be less than 1"
53
+ # Image input
54
+ if isinstance(_data, Image.Image):
55
+ W, H = _data.width, _data.height
56
+ W, H = map(lambda x: int(x * size_factor), (W, H))
57
+ return _data.resize((W, H))
58
+ # Video input with PIL Images
59
+ elif is_list_of(_data, Image.Image):
60
+ W, H = next(iter(_data)).width, next(iter(_data)).height
61
+ T = len(_data)
62
+ T, W, H = map(lambda x: max(int(x * size_factor), 1), (T, W, H))
63
+ return [d.resize((W, H)) for d in _data[:T]]
64
+ # Video input with numpy arrays
65
+ elif isinstance(_data, np.ndarray) and _data.ndim >= 4:
66
+ T, H, W, C = _data.shape[-4:]
67
+ T, H, W = map(lambda x: max(int(x * size_factor), 1), (T, H, W))
68
+ return _data[..., :T, :H, :W, :C]
69
+ # Audio input
70
+ elif isinstance(_data, np.ndarray) and _data.ndim == 1:
71
+ return _data[:int(len(_data) * size_factor)]
72
+ raise AssertionError("This line should be unreachable.")
73
+
74
+
75
+def resize_mm_data(
76
+ data: Union[ImageInput, VideoInput, AudioInput],
77
+ size_factors: tuple[float,
78
+ ...]) -> Union[ImageInput, VideoInput, AudioInput]:
79
+ size_factors = size_factors[:len(data)]
80
+ if is_list_of(data, (Image.Image, np.ndarray, list)):
81
+ return [_resize_data(d, s) for d, s in zip(data, size_factors)]
82
+ elif is_list_of(data, tuple):
83
+ return [(_resize_data(d, s), meta)
84
+ for (d, meta), s in zip(data, size_factors)]
85
+ raise ValueError("Unsupported multimodal data type.")
86
87
88
def create_batched_mm_kwargs(
89
model_config: ModelConfig,
90
processor: BaseMultiModalProcessor,
31
-) -> MultiModalKwargs:
91
+ size_factors: tuple[float, ...] = (1.0, 0.5, 0.25),
92
+) -> Iterable[tuple[str, int, BatchedTensorInputs]]:
93
processing_info = processor.info
94
dummy_inputs = processor.dummy_inputs
95
supported_mm_limits = processing_info.get_supported_mm_limits()
@@ -40,30 +101,69 @@ def create_batched_mm_kwargs(
101
seq_len=model_config.max_model_len,
102
mm_counts=mm_counts,
103
)
104
+ mm_data = processor_inputs.mm_data
105
+ resized_mm_data = {
106
+ modality: resize_mm_data(data, size_factors)
107
+ for modality, data in mm_data.items()
108
+ }
109
+ # Mistral chat outputs tokens directly, rather than text prompts
110
+ if model_config.tokenizer_mode == "mistral":
111
+ images = resized_mm_data.get("image", [])
112
+ request = ChatCompletionRequest(messages=[
113
+ UserMessage(content=[
114
+ TextChunk(text=""),
115
+ *(ImageChunk(image=image) for image in images),
116
+ ]),
117
+ ])
118
+ tokenizer = processing_info.get_tokenizer()
119
+ res = tokenizer.mistral.encode_chat_completion(request)
120
+ prompt = res.tokens
121
+ else:
122
+ prompt = processor_inputs.prompt
123
mm_kwargs = processor.apply(
44
- prompt=processor_inputs.prompt,
45
- mm_data=processor_inputs.mm_data,
124
+ prompt=prompt,
125
+ mm_data=resized_mm_data,
126
hf_processor_mm_kwargs=processor_inputs.hf_processor_mm_kwargs,
127
tokenization_kwargs=processor_inputs.tokenization_kwargs,
128
)["mm_kwargs"]
49
- mm_kwargs = MultiModalKwargs.batch([mm_kwargs])
50
- return mm_kwargs
129
+ items = [
130
+ item for modality in supported_mm_limits
131
+ for item in mm_kwargs.get_items(modality)
132
+ ]
133
+ return group_mm_kwargs_by_modality(items)
134
+
135
+
136
+def get_model_id_to_test(
137
+ model_arch_list: Iterable[str]) -> list[tuple[str, str]]:
138
+ filtered_results = []
139
+ for model_arch in model_arch_list:
140
+ model_info = HF_EXAMPLE_MODELS.get_hf_info(model_arch)
141
+ if model_info.extras and model_arch in ARCH_NEEDS_EXTRAS:
142
+ available_repos = list(
143
+ map(lambda model_id: (model_arch, model_id),
144
+ [model_info.default, *model_info.extras.values()]))
145
+ filtered_results.extend(available_repos)
146
+ else:
147
+ filtered_results.append((model_arch, model_info.default))
148
+ return filtered_results
149
150
151
@pytest.mark.core_model
54
-@pytest.mark.parametrize("model_arch", list(_MULTIMODAL_EXAMPLE_MODELS.keys()))
55
-def test_model_tensor_schema(model_arch: str, vllm_runner: type[VllmRunner],
56
- monkeypatch):
152
+@pytest.mark.parametrize(
153
+ "model_arch, model_id",
154
+ get_model_id_to_test(_MULTIMODAL_EXAMPLE_MODELS.keys()))
155
+def test_model_tensor_schema(model_arch: str, model_id: str,
156
+ vllm_runner: type[VllmRunner], monkeypatch):
157
if model_arch in ARCH_TO_SKIP:
158
pytest.skip(f"Skipping {model_arch} due to {ARCH_TO_SKIP[model_arch]}")
159
+ if model_id in REPO_ID_TO_SKIP:
160
+ pytest.skip(f"Skipping {model_id} due to {REPO_ID_TO_SKIP[model_id]}")
161
162
model_info = HF_EXAMPLE_MODELS.get_hf_info(model_arch)
163
model_info.check_available_online(on_fail="skip")
164
model_info.check_transformers_version(on_fail="skip",
165
check_max_version=False)
166
65
- model_id = model_info.default
66
-
167
hf_overrides_fn = partial(dummy_hf_overrides,
168
model_arch=model_arch,
169
exist_overrides=model_info.hf_overrides)
@@ -119,6 +219,7 @@ def test_model_tensor_schema(model_arch: str, vllm_runner: type[VllmRunner],
219
if model_info.v0_only:
220
m.setenv("VLLM_USE_V1", "0")
221
222
+ # TODO(Isotr0py): Can we avoid initializing engine?
223
with (
224
set_default_torch_num_threads(1),
225
vllm_runner(
@@ -145,12 +246,16 @@ def test_model_tensor_schema(model_arch: str, vllm_runner: type[VllmRunner],
246
mm_registry = llm_engine.input_preprocessor.mm_registry
247
248
processor = mm_registry.create_processor(model_config)
148
- mm_kwargs = create_batched_mm_kwargs(model_config, processor)
249
150
- def validate_model_input(model):
151
- for modality in ("audio", "image", "video"):
152
- method_name = f"_parse_and_validate_{modality}_input"
153
- if hasattr(model, method_name):
154
- getattr(model, method_name)(**mm_kwargs)
250
+ def validate_model_input(model, modality: str,
251
+ mm_kwargs: MultiModalKwargs):
252
+ method_name = f"_parse_and_validate_{modality}_input"
253
+ if hasattr(model, method_name):
254
+ getattr(model, method_name)(**mm_kwargs)
255
156
- vllm_model.apply_model(validate_model_input)
\ No newline at end of file
256
+ for modality, _, mm_kwargs in create_batched_mm_kwargs(
257
+ model_config, processor):
258
+ valid_func = partial(validate_model_input,
259
+ modality=modality,
260
+ mm_kwargs=mm_kwargs)
261
+ vllm_model.apply_model(valid_func)
vllm/model_executor/models/keye.py
+14
-8
@@ -30,7 +30,7 @@ from vllm.model_executor.layers.quantization.gptq_marlin import (
30
from vllm.model_executor.model_loader.weight_utils import (
31
default_weight_loader, maybe_remap_kv_scale_name)
32
from vllm.model_executor.models.module_mapping import MultiModelKeys
33
-from vllm.multimodal import MULTIMODAL_REGISTRY
33
+from vllm.multimodal import MULTIMODAL_REGISTRY, NestedTensors
34
from vllm.multimodal.inputs import (ImageItem, ModalityData,
35
MultiModalDataDict, MultiModalFieldConfig,
36
MultiModalKwargs, VideoItem)
@@ -44,6 +44,7 @@ from vllm.multimodal.profiling import BaseDummyInputsBuilder
44
from vllm.platforms import _Backend
45
from vllm.sequence import IntermediateTensors
46
from vllm.transformers_utils.config import uses_mrope
47
+from vllm.utils import is_list_of
48
from vllm.utils.tensor_schema import TensorSchema, TensorShape
49
50
from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
@@ -112,8 +113,9 @@ class KeyeImagePixelInputs(TensorSchema):
113
- g: Grid dimensions (3 for t, h, w)
114
"""
115
type: Literal["pixel_values"]
115
- pixel_values: Annotated[torch.Tensor,
116
- TensorShape("b", "np", 3, "ps", "ps")]
116
+ pixel_values: Annotated[
117
+ torch.Tensor,
118
+ TensorShape("b", "np", 3, "ps", "ps", dynamic_dims={"np"})]
119
image_grid_thw: Annotated[torch.Tensor, TensorShape("ni", 3)]
120
121
@@ -145,8 +147,9 @@ class KeyeVideoPixelInputs(TensorSchema):
147
- g: Grid dimensions (3 for t, h, w)
148
"""
149
type: Literal["pixel_values_videos"]
148
- pixel_values_videos: Annotated[torch.Tensor,
149
- TensorShape("b", "np", 3, "ps", "ps")]
150
+ pixel_values_videos: Annotated[
151
+ torch.Tensor,
152
+ TensorShape("b", "np", 3, "ps", "ps", dynamic_dims={"np"})]
153
video_grid_thw: Annotated[torch.Tensor, TensorShape("nv", 3)]
154
155
@@ -1295,7 +1298,7 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,
1298
return None
1299
return quant_config
1300
1298
- def _validate_and_reshape_mm_tensor(self, mm_input: object,
1301
+ def _validate_and_reshape_mm_tensor(self, mm_input: NestedTensors,
1302
name: str) -> torch.Tensor:
1303
if not isinstance(mm_input, (torch.Tensor, list)):
1304
raise ValueError(f"Incorrect type of {name}. "
@@ -1310,8 +1313,11 @@ class KeyeForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsLoRA,
1313
f"Got ndim: {mm_input.ndim} "
1314
f"(shape={mm_input.shape})")
1315
return torch.concat(list(mm_input))
1313
- else:
1314
- return torch.concat(mm_input)
1316
+ elif is_list_of(mm_input, torch.Tensor):
1317
+ if all(p.dim() == 4 for p in mm_input) or all(p.dim() == 2
1318
+ for p in mm_input):
1319
+ return mm_input
1320
+ return torch.concat(list(mm_input))
1321
1322
def _parse_and_validate_image_input(
1323
self, **kwargs: object) -> Optional[KeyeImageInputs]: