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 8, 20:15 UTC
[Bugfix] Make MM batching more robust (#33817)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Feb 6, 2026 at 04:40 UTC
116880a5a0af3a226e29f4716c484a4cc8422fc1
13 files changed
+625
-428
.buildkite/test-amd.yaml
+2
-1
@@ -863,10 +863,11 @@ steps:
863
torch_nightly: true
864
source_file_dependencies:
865
- vllm/
866
+ - tests/models/test_terratorch.py
867
- tests/models/test_transformers.py
868
- tests/models/test_registry.py
869
commands:
869
- - pytest -v -s models/test_transformers.py models/test_registry.py
870
+ - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
871
872
- label: Basic Models Test (Other CPU) # 5min
873
mirror_hardwares: [amdexperimental, amdproduction]
.buildkite/test-pipeline.yaml
+2
-1
@@ -804,10 +804,11 @@ steps:
804
torch_nightly: true
805
source_file_dependencies:
806
- vllm/
807
+ - tests/models/test_terratorch.py
808
- tests/models/test_transformers.py
809
- tests/models/test_registry.py
810
commands:
810
- - pytest -v -s models/test_transformers.py models/test_registry.py
811
+ - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
812
813
- label: Basic Models Test (Other CPU) # 5min
814
timeout_in_minutes: 10
.buildkite/test_areas/models_basic.yaml
+2
-1
@@ -33,10 +33,11 @@ steps:
33
timeout_in_minutes: 45
34
source_file_dependencies:
35
- vllm/
36
+ - tests/models/test_terratorch.py
37
- tests/models/test_transformers.py
38
- tests/models/test_registry.py
39
commands:
39
- - pytest -v -s models/test_transformers.py models/test_registry.py
40
+ - pytest -v -s models/test_terratorch.py models/test_transformers.py models/test_registry.py
41
42
- label: Basic Models Test (Other CPU) # 5min
43
depends_on:
tests/models/test_terratorch.py
+9
-3
@@ -5,8 +5,10 @@ import pytest
5
import torch
6
7
from tests.conftest import VllmRunner
8
+from tests.utils import create_new_process_for_each_test
9
10
11
+@create_new_process_for_each_test() # Memory is not cleaned up properly otherwise
12
@pytest.mark.parametrize(
13
"model",
14
[
@@ -22,10 +24,14 @@ def test_inference(
24
location_coords = torch.full((1, 2), 1.0, dtype=torch.float16)
25
prompt = dict(
26
prompt_token_ids=[1],
25
- multi_modal_data=dict(
26
- pixel_values=pixel_values, location_coords=location_coords
27
- ),
27
+ multi_modal_data={
28
+ "image": {
29
+ "pixel_values": pixel_values,
30
+ "location_coords": location_coords,
31
+ }
32
+ },
33
)
34
+
35
with vllm_runner(
36
model,
37
runner="pooling",
tests/multimodal/media/test_connector.py
new
+320
@@ -0,0 +1,320 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import asyncio
5
+import base64
6
+import mimetypes
7
+import os
8
+from tempfile import NamedTemporaryFile, TemporaryDirectory
9
+
10
+import numpy as np
11
+import pytest
12
+import torch
13
+from PIL import Image, ImageChops
14
+
15
+from vllm.multimodal.image import convert_image_mode
16
+from vllm.multimodal.inputs import PlaceholderRange
17
+from vllm.multimodal.media import MediaConnector
18
+
19
+# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
20
+TEST_IMAGE_ASSETS = [
21
+ "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
22
+ "Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
23
+ "1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
24
+ "RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
25
+]
26
+
27
+TEST_VIDEO_URLS = [
28
+ "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4",
29
+ "https://github.com/opencv/opencv/raw/refs/tags/4.12.0/samples/data/vtest.avi",
30
+]
31
+
32
+
33
+@pytest.fixture(scope="module")
34
+def url_images(local_asset_server) -> dict[str, Image.Image]:
35
+ return {
36
+ image_url: local_asset_server.get_image_asset(image_url)
37
+ for image_url in TEST_IMAGE_ASSETS
38
+ }
39
+
40
+
41
+def get_supported_suffixes() -> tuple[str, ...]:
42
+ # We should at least test the file types mentioned in GPT-4 with Vision
43
+ OPENAI_SUPPORTED_SUFFIXES = (".png", ".jpeg", ".jpg", ".webp", ".gif")
44
+
45
+ # Additional file types that are supported by us
46
+ EXTRA_SUPPORTED_SUFFIXES = (".bmp", ".tiff")
47
+
48
+ return OPENAI_SUPPORTED_SUFFIXES + EXTRA_SUPPORTED_SUFFIXES
49
+
50
+
51
+def _image_equals(a: Image.Image, b: Image.Image) -> bool:
52
+ return (np.asarray(a) == np.asarray(convert_image_mode(b, a.mode))).all()
53
+
54
+
55
+@pytest.mark.asyncio
56
+@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
57
+async def test_fetch_image_http(image_url: str):
58
+ connector = MediaConnector()
59
+
60
+ image_sync = connector.fetch_image(image_url)
61
+ image_async = await connector.fetch_image_async(image_url)
62
+ assert _image_equals(image_sync, image_async)
63
+
64
+
65
+@pytest.mark.asyncio
66
+@pytest.mark.parametrize("raw_image_url", TEST_IMAGE_ASSETS)
67
+@pytest.mark.parametrize("suffix", get_supported_suffixes())
68
+async def test_fetch_image_base64(
69
+ url_images: dict[str, Image.Image], raw_image_url: str, suffix: str
70
+):
71
+ connector = MediaConnector(
72
+ # Domain restriction should not apply to data URLs.
73
+ allowed_media_domains=[
74
+ "www.bogotobogo.com",
75
+ "github.com",
76
+ ]
77
+ )
78
+ url_image = url_images[raw_image_url]
79
+
80
+ try:
81
+ mime_type = Image.MIME[Image.registered_extensions()[suffix]]
82
+ except KeyError:
83
+ try:
84
+ mime_type = mimetypes.types_map[suffix]
85
+ except KeyError:
86
+ pytest.skip("No MIME type")
87
+
88
+ with NamedTemporaryFile(suffix=suffix) as f:
89
+ try:
90
+ url_image.save(f.name)
91
+ except Exception as e:
92
+ if e.args[0] == "cannot write mode RGBA as JPEG":
93
+ pytest.skip("Conversion not supported")
94
+
95
+ raise
96
+
97
+ base64_image = base64.b64encode(f.read()).decode("utf-8")
98
+ data_url = f"data:{mime_type};base64,{base64_image}"
99
+
100
+ data_image_sync = connector.fetch_image(data_url)
101
+ if _image_equals(url_image, Image.open(f)):
102
+ assert _image_equals(url_image, data_image_sync)
103
+ else:
104
+ pass # Lossy format; only check that image can be opened
105
+
106
+ data_image_async = await connector.fetch_image_async(data_url)
107
+ assert _image_equals(data_image_sync, data_image_async)
108
+
109
+
110
+@pytest.mark.asyncio
111
+@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
112
+async def test_fetch_image_local_files(image_url: str):
113
+ connector = MediaConnector()
114
+
115
+ with TemporaryDirectory() as temp_dir:
116
+ local_connector = MediaConnector(allowed_local_media_path=temp_dir)
117
+
118
+ origin_image = connector.fetch_image(image_url)
119
+ origin_image.save(
120
+ os.path.join(temp_dir, os.path.basename(image_url)),
121
+ quality=100,
122
+ icc_profile=origin_image.info.get("icc_profile"),
123
+ )
124
+
125
+ image_async = await local_connector.fetch_image_async(
126
+ f"file://{temp_dir}/{os.path.basename(image_url)}"
127
+ )
128
+ image_sync = local_connector.fetch_image(
129
+ f"file://{temp_dir}/{os.path.basename(image_url)}"
130
+ )
131
+ # Check that the images are equal
132
+ assert not ImageChops.difference(image_sync, image_async).getbbox()
133
+
134
+ with pytest.raises(ValueError, match="must be a subpath"):
135
+ await local_connector.fetch_image_async(
136
+ f"file://{temp_dir}/../{os.path.basename(image_url)}"
137
+ )
138
+ with pytest.raises(RuntimeError, match="Cannot load local files"):
139
+ await connector.fetch_image_async(
140
+ f"file://{temp_dir}/../{os.path.basename(image_url)}"
141
+ )
142
+
143
+ with pytest.raises(ValueError, match="must be a subpath"):
144
+ local_connector.fetch_image(
145
+ f"file://{temp_dir}/../{os.path.basename(image_url)}"
146
+ )
147
+ with pytest.raises(RuntimeError, match="Cannot load local files"):
148
+ connector.fetch_image(f"file://{temp_dir}/../{os.path.basename(image_url)}")
149
+
150
+
151
+@pytest.mark.asyncio
152
+@pytest.mark.parametrize("image_url", [TEST_IMAGE_ASSETS[0]], indirect=True)
153
+async def test_fetch_image_local_files_with_space_in_name(image_url: str):
154
+ connector = MediaConnector()
155
+
156
+ with TemporaryDirectory() as temp_dir:
157
+ local_connector = MediaConnector(allowed_local_media_path=temp_dir)
158
+
159
+ origin_image = connector.fetch_image(image_url)
160
+ filename = "file name with space.jpg"
161
+ origin_image.save(
162
+ os.path.join(temp_dir, filename),
163
+ quality=100,
164
+ icc_profile=origin_image.info.get("icc_profile"),
165
+ )
166
+
167
+ try:
168
+ image_async = await local_connector.fetch_image_async(
169
+ f"file://{temp_dir}/{filename}"
170
+ )
171
+ image_sync = local_connector.fetch_image(f"file://{temp_dir}/{filename}")
172
+ except FileNotFoundError as e:
173
+ pytest.fail("Failed to fetch image with space in name: {}".format(e))
174
+ # Check that the images are equal
175
+ assert not ImageChops.difference(image_sync, image_async).getbbox()
176
+
177
+
178
+@pytest.mark.asyncio
179
+async def test_fetch_image_error_conversion():
180
+ connector = MediaConnector()
181
+ broken_img = "data:image/png;base64,aGVsbG9fdmxsbV9jb21tdW5pdHkK"
182
+
183
+ # PIL.UnidentifiedImageError should be converted to ValueError
184
+ with pytest.raises(ValueError):
185
+ await connector.fetch_image_async(broken_img)
186
+
187
+ with pytest.raises(ValueError):
188
+ connector.fetch_image(broken_img)
189
+
190
+
191
+@pytest.mark.flaky(reruns=3, reruns_delay=5)
192
+@pytest.mark.asyncio
193
+@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
194
+@pytest.mark.parametrize("num_frames", [-1, 32, 1800])
195
+async def test_fetch_video_http(video_url: str, num_frames: int):
196
+ connector = MediaConnector(
197
+ media_io_kwargs={
198
+ "video": {
199
+ "num_frames": num_frames,
200
+ }
201
+ }
202
+ )
203
+
204
+ try:
205
+ video_sync, metadata_sync = connector.fetch_video(video_url)
206
+ video_async, metadata_async = await connector.fetch_video_async(video_url)
207
+ except (TimeoutError, asyncio.TimeoutError) as e:
208
+ pytest.skip(f"Timeout fetching video (CI network flakiness): {e}")
209
+
210
+ assert np.array_equal(video_sync, video_async)
211
+ assert metadata_sync == metadata_async
212
+
213
+
214
+@pytest.mark.asyncio
215
+@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
216
+@pytest.mark.parametrize("max_duration", [1, 60, 1800])
217
+@pytest.mark.parametrize("requested_fps", [2, 24])
218
+async def test_fetch_video_http_with_dynamic_loader(
219
+ video_url: str,
220
+ max_duration: int,
221
+ requested_fps: int,
222
+ monkeypatch: pytest.MonkeyPatch,
223
+):
224
+ with monkeypatch.context() as m:
225
+ m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv_dynamic")
226
+ connector = MediaConnector(
227
+ media_io_kwargs={
228
+ "video": {
229
+ "max_duration": max_duration,
230
+ "requested_fps": requested_fps,
231
+ }
232
+ }
233
+ )
234
+
235
+ video_sync, metadata_sync = connector.fetch_video(video_url)
236
+ video_async, metadata_async = await connector.fetch_video_async(video_url)
237
+
238
+ assert np.array_equal(video_sync, video_async)
239
+ assert metadata_sync == metadata_async
240
+ assert metadata_sync["video_backend"] == "opencv_dynamic"
241
+
242
+
243
+@pytest.mark.parametrize(
244
+ "is_embed,start_idx,end_idx,expected",
245
+ [
246
+ (None, 2, 4, (2, 4)),
247
+ (
248
+ torch.tensor([False, True, False, True, True]),
249
+ 3,
250
+ 5,
251
+ (1, 3),
252
+ ),
253
+ (
254
+ torch.tensor([False, True, False, True, True]),
255
+ 0,
256
+ 2,
257
+ (0, 1),
258
+ ),
259
+ (
260
+ torch.tensor([True, False, True, False]),
261
+ 2,
262
+ 2,
263
+ (1, 1),
264
+ ),
265
+ ],
266
+)
267
+def test_placeholder_range_get_embeds_indices_in_range(
268
+ is_embed, start_idx, end_idx, expected
269
+):
270
+ length = len(is_embed) if is_embed is not None else 5
271
+ pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
272
+ assert pr.get_embeds_indices_in_range(start_idx, end_idx) == expected
273
+
274
+
275
+@pytest.mark.parametrize(
276
+ "offset,is_embed,expected",
277
+ [
278
+ (0, None, [(0, 4)]),
279
+ (
280
+ 2,
281
+ torch.tensor([False, True, False, True, True]),
282
+ [(3, 3), (5, 6)],
283
+ ),
284
+ (0, torch.tensor([True, True, True, True]), [(0, 3)]),
285
+ (0, torch.tensor([False, False, False, False]), []),
286
+ ],
287
+)
288
+def test_placeholder_range_extract_embeds_range(offset, is_embed, expected):
289
+ length = len(is_embed) if is_embed is not None else 5
290
+ pr = PlaceholderRange(offset=offset, length=length, is_embed=is_embed)
291
+ assert pr.extract_embeds_range() == expected
292
+
293
+
294
+@pytest.mark.asyncio
295
+@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
296
+@pytest.mark.parametrize("num_frames", [-1, 32, 1800])
297
+async def test_allowed_media_domains(video_url: str, num_frames: int):
298
+ connector = MediaConnector(
299
+ media_io_kwargs={
300
+ "video": {
301
+ "num_frames": num_frames,
302
+ }
303
+ },
304
+ allowed_media_domains=[
305
+ "www.bogotobogo.com",
306
+ "github.com",
307
+ ],
308
+ )
309
+
310
+ video_sync, metadata_sync = connector.fetch_video(video_url)
311
+ video_async, metadata_async = await connector.fetch_video_async(video_url)
312
+ assert np.array_equal(video_sync, video_async)
313
+ assert metadata_sync == metadata_async
314
+
315
+ disallowed_url = "https://upload.wikimedia.org/wikipedia/commons/4/47/PNG_transparency_demonstration_1.png"
316
+ with pytest.raises(ValueError):
317
+ _, _ = connector.fetch_video(disallowed_url)
318
+
319
+ with pytest.raises(ValueError):
320
+ _, _ = await connector.fetch_video_async(disallowed_url)
tests/multimodal/test_hasher.py
+16
@@ -16,6 +16,22 @@ ASSETS_DIR = Path(__file__).parent / "assets"
16
assert ASSETS_DIR.exists()
17
18
19
+def test_hash_single_item_different_shape():
20
+ x1 = torch.zeros(())
21
+ x2 = torch.zeros((1,))
22
+
23
+ hasher = MultiModalHasher
24
+ assert hasher.hash_kwargs(x=x1) != hasher.hash_kwargs(x=x2)
25
+
26
+
27
+def test_hash_key_order_invariant():
28
+ x = torch.zeros((5, 10))
29
+ y = torch.ones((5, 10))
30
+
31
+ hasher = MultiModalHasher
32
+ assert hasher.hash_kwargs(x=x, y=y) == hasher.hash_kwargs(y=y, x=x)
33
+
34
+
35
# NOTE: Images that are the same visually are allowed to have the same hash
36
@pytest.mark.parametrize("mode_pair", [("1", "L"), ("RGBA", "CMYK")])
37
def test_hash_collision_image_mode(mode_pair):
tests/multimodal/test_inputs.py
new
+46
@@ -0,0 +1,46 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+
6
+from vllm.multimodal.inputs import PlaceholderRange
7
+
8
+
9
+@pytest.mark.parametrize(
10
+ "is_embed,expected",
11
+ [
12
+ (None, 5),
13
+ (torch.tensor([True, True, True, True, True]), 5),
14
+ (torch.tensor([False, False, False, False, False]), 0),
15
+ (torch.tensor([True, False, True, False, True]), 3),
16
+ (torch.tensor([True]), 1),
17
+ ],
18
+)
19
+def test_placeholder_range_get_num_embeds(is_embed, expected):
20
+ length = len(is_embed) if is_embed is not None else 5
21
+ pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
22
+ assert pr.get_num_embeds == expected
23
+
24
+
25
+@pytest.mark.parametrize(
26
+ "is_embed,expected",
27
+ [
28
+ (None, None),
29
+ (
30
+ torch.tensor([False, True, False, True, True]),
31
+ torch.tensor([0, 1, 1, 2, 3]),
32
+ ),
33
+ (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])),
34
+ ],
35
+)
36
+def test_placeholder_range_embeds_cumsum(is_embed, expected):
37
+ length = len(is_embed) if is_embed is not None else 5
38
+ pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
39
+
40
+ if expected is None:
41
+ assert pr.embeds_cumsum is None
42
+ return
43
+
44
+ assert torch.equal(pr.embeds_cumsum, expected)
45
+ # cached_property should return the same object on repeated access
46
+ assert pr.embeds_cumsum is pr.embeds_cumsum
tests/multimodal/test_utils.py
+36
-350
@@ -1,244 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import asyncio
5
-import base64
6
-import mimetypes
7
-import os
8
-from tempfile import NamedTemporaryFile, TemporaryDirectory
9
-
10
-import numpy as np
3
import pytest
4
import torch
13
-from PIL import Image, ImageChops
14
-
15
-from vllm.multimodal.image import convert_image_mode
16
-from vllm.multimodal.inputs import PlaceholderRange
17
-from vllm.multimodal.media import MediaConnector
18
-from vllm.multimodal.utils import argsort_mm_positions
19
-
20
-# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
21
-TEST_IMAGE_ASSETS = [
22
- "2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
23
- "Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
24
- "1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
25
- "RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
26
-]
27
-
28
-TEST_VIDEO_URLS = [
29
- "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4",
30
- "https://github.com/opencv/opencv/raw/refs/tags/4.12.0/samples/data/vtest.avi",
31
-]
32
-
33
-
34
-@pytest.fixture(scope="module")
35
-def url_images(local_asset_server) -> dict[str, Image.Image]:
36
- return {
37
- image_url: local_asset_server.get_image_asset(image_url)
38
- for image_url in TEST_IMAGE_ASSETS
39
- }
40
-
41
-
42
-def get_supported_suffixes() -> tuple[str, ...]:
43
- # We should at least test the file types mentioned in GPT-4 with Vision
44
- OPENAI_SUPPORTED_SUFFIXES = (".png", ".jpeg", ".jpg", ".webp", ".gif")
45
-
46
- # Additional file types that are supported by us
47
- EXTRA_SUPPORTED_SUFFIXES = (".bmp", ".tiff")
48
-
49
- return OPENAI_SUPPORTED_SUFFIXES + EXTRA_SUPPORTED_SUFFIXES
50
-
51
-
52
-def _image_equals(a: Image.Image, b: Image.Image) -> bool:
53
- return (np.asarray(a) == np.asarray(convert_image_mode(b, a.mode))).all()
54
-
55
-
56
-@pytest.mark.asyncio
57
-@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
58
-async def test_fetch_image_http(image_url: str):
59
- connector = MediaConnector()
60
-
61
- image_sync = connector.fetch_image(image_url)
62
- image_async = await connector.fetch_image_async(image_url)
63
- assert _image_equals(image_sync, image_async)
64
-
65
-
66
-@pytest.mark.asyncio
67
-@pytest.mark.parametrize("raw_image_url", TEST_IMAGE_ASSETS)
68
-@pytest.mark.parametrize("suffix", get_supported_suffixes())
69
-async def test_fetch_image_base64(
70
- url_images: dict[str, Image.Image], raw_image_url: str, suffix: str
71
-):
72
- connector = MediaConnector(
73
- # Domain restriction should not apply to data URLs.
74
- allowed_media_domains=[
75
- "www.bogotobogo.com",
76
- "github.com",
77
- ]
78
- )
79
- url_image = url_images[raw_image_url]
80
-
81
- try:
82
- mime_type = Image.MIME[Image.registered_extensions()[suffix]]
83
- except KeyError:
84
- try:
85
- mime_type = mimetypes.types_map[suffix]
86
- except KeyError:
87
- pytest.skip("No MIME type")
88
-
89
- with NamedTemporaryFile(suffix=suffix) as f:
90
- try:
91
- url_image.save(f.name)
92
- except Exception as e:
93
- if e.args[0] == "cannot write mode RGBA as JPEG":
94
- pytest.skip("Conversion not supported")
95
-
96
- raise
97
-
98
- base64_image = base64.b64encode(f.read()).decode("utf-8")
99
- data_url = f"data:{mime_type};base64,{base64_image}"
100
-
101
- data_image_sync = connector.fetch_image(data_url)
102
- if _image_equals(url_image, Image.open(f)):
103
- assert _image_equals(url_image, data_image_sync)
104
- else:
105
- pass # Lossy format; only check that image can be opened
106
-
107
- data_image_async = await connector.fetch_image_async(data_url)
108
- assert _image_equals(data_image_sync, data_image_async)
109
-
110
-
111
-@pytest.mark.asyncio
112
-@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
113
-async def test_fetch_image_local_files(image_url: str):
114
- connector = MediaConnector()
115
-
116
- with TemporaryDirectory() as temp_dir:
117
- local_connector = MediaConnector(allowed_local_media_path=temp_dir)
118
-
119
- origin_image = connector.fetch_image(image_url)
120
- origin_image.save(
121
- os.path.join(temp_dir, os.path.basename(image_url)),
122
- quality=100,
123
- icc_profile=origin_image.info.get("icc_profile"),
124
- )
125
-
126
- image_async = await local_connector.fetch_image_async(
127
- f"file://{temp_dir}/{os.path.basename(image_url)}"
128
- )
129
- image_sync = local_connector.fetch_image(
130
- f"file://{temp_dir}/{os.path.basename(image_url)}"
131
- )
132
- # Check that the images are equal
133
- assert not ImageChops.difference(image_sync, image_async).getbbox()
134
-
135
- with pytest.raises(ValueError, match="must be a subpath"):
136
- await local_connector.fetch_image_async(
137
- f"file://{temp_dir}/../{os.path.basename(image_url)}"
138
- )
139
- with pytest.raises(RuntimeError, match="Cannot load local files"):
140
- await connector.fetch_image_async(
141
- f"file://{temp_dir}/../{os.path.basename(image_url)}"
142
- )
143
-
144
- with pytest.raises(ValueError, match="must be a subpath"):
145
- local_connector.fetch_image(
146
- f"file://{temp_dir}/../{os.path.basename(image_url)}"
147
- )
148
- with pytest.raises(RuntimeError, match="Cannot load local files"):
149
- connector.fetch_image(f"file://{temp_dir}/../{os.path.basename(image_url)}")
150
-
151
-
152
-@pytest.mark.asyncio
153
-@pytest.mark.parametrize("image_url", [TEST_IMAGE_ASSETS[0]], indirect=True)
154
-async def test_fetch_image_local_files_with_space_in_name(image_url: str):
155
- connector = MediaConnector()
156
-
157
- with TemporaryDirectory() as temp_dir:
158
- local_connector = MediaConnector(allowed_local_media_path=temp_dir)
159
-
160
- origin_image = connector.fetch_image(image_url)
161
- filename = "file name with space.jpg"
162
- origin_image.save(
163
- os.path.join(temp_dir, filename),
164
- quality=100,
165
- icc_profile=origin_image.info.get("icc_profile"),
166
- )
167
-
168
- try:
169
- image_async = await local_connector.fetch_image_async(
170
- f"file://{temp_dir}/{filename}"
171
- )
172
- image_sync = local_connector.fetch_image(f"file://{temp_dir}/{filename}")
173
- except FileNotFoundError as e:
174
- pytest.fail("Failed to fetch image with space in name: {}".format(e))
175
- # Check that the images are equal
176
- assert not ImageChops.difference(image_sync, image_async).getbbox()
177
-
178
-
179
-@pytest.mark.asyncio
180
-async def test_fetch_image_error_conversion():
181
- connector = MediaConnector()
182
- broken_img = "data:image/png;base64,aGVsbG9fdmxsbV9jb21tdW5pdHkK"
183
-
184
- # PIL.UnidentifiedImageError should be converted to ValueError
185
- with pytest.raises(ValueError):
186
- await connector.fetch_image_async(broken_img)
187
-
188
- with pytest.raises(ValueError):
189
- connector.fetch_image(broken_img)
190
-
5
192
-@pytest.mark.flaky(reruns=3, reruns_delay=5)
193
-@pytest.mark.asyncio
194
-@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
195
-@pytest.mark.parametrize("num_frames", [-1, 32, 1800])
196
-async def test_fetch_video_http(video_url: str, num_frames: int):
197
- connector = MediaConnector(
198
- media_io_kwargs={
199
- "video": {
200
- "num_frames": num_frames,
201
- }
202
- }
203
- )
204
-
205
- try:
206
- video_sync, metadata_sync = connector.fetch_video(video_url)
207
- video_async, metadata_async = await connector.fetch_video_async(video_url)
208
- except (TimeoutError, asyncio.TimeoutError) as e:
209
- pytest.skip(f"Timeout fetching video (CI network flakiness): {e}")
210
-
211
- assert np.array_equal(video_sync, video_async)
212
- assert metadata_sync == metadata_async
213
-
214
-
215
-@pytest.mark.asyncio
216
-@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
217
-@pytest.mark.parametrize("max_duration", [1, 60, 1800])
218
-@pytest.mark.parametrize("requested_fps", [2, 24])
219
-async def test_fetch_video_http_with_dynamic_loader(
220
- video_url: str,
221
- max_duration: int,
222
- requested_fps: int,
223
- monkeypatch: pytest.MonkeyPatch,
224
-):
225
- with monkeypatch.context() as m:
226
- m.setenv("VLLM_VIDEO_LOADER_BACKEND", "opencv_dynamic")
227
- connector = MediaConnector(
228
- media_io_kwargs={
229
- "video": {
230
- "max_duration": max_duration,
231
- "requested_fps": requested_fps,
232
- }
233
- }
234
- )
235
-
236
- video_sync, metadata_sync = connector.fetch_video(video_url)
237
- video_async, metadata_async = await connector.fetch_video_async(video_url)
238
-
239
- assert np.array_equal(video_sync, video_async)
240
- assert metadata_sync == metadata_async
241
- assert metadata_sync["video_backend"] == "opencv_dynamic"
6
+from vllm.multimodal.inputs import (
7
+ MultiModalBatchedField,
8
+ MultiModalFieldElem,
9
+ MultiModalKwargsItem,
10
+ MultiModalSharedField,
11
+ PlaceholderRange,
12
+)
13
+from vllm.multimodal.utils import argsort_mm_positions, group_and_batch_mm_items
14
15
16
@pytest.mark.parametrize(
@@ -412,121 +184,35 @@ def test_argsort_mm_positions(case):
184
assert modality_idxs == expected_modality_idxs
185
186
415
-@pytest.mark.parametrize(
416
- "is_embed,expected",
417
- [
418
- (None, 5),
419
- (torch.tensor([True, True, True, True, True]), 5),
420
- (torch.tensor([False, False, False, False, False]), 0),
421
- (torch.tensor([True, False, True, False, True]), 3),
422
- (torch.tensor([True]), 1),
423
- ],
424
-)
425
-def test_placeholder_range_get_num_embeds(is_embed, expected):
426
- length = len(is_embed) if is_embed is not None else 5
427
- pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
428
- assert pr.get_num_embeds == expected
429
-
430
-
431
-@pytest.mark.parametrize(
432
- "is_embed,expected",
433
- [
434
- (None, None),
435
- (
436
- torch.tensor([False, True, False, True, True]),
437
- torch.tensor([0, 1, 1, 2, 3]),
438
- ),
439
- (torch.tensor([True, True, True]), torch.tensor([1, 2, 3])),
440
- ],
441
-)
442
-def test_placeholder_range_embeds_cumsum(is_embed, expected):
443
- length = len(is_embed) if is_embed is not None else 5
444
- pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
445
-
446
- if expected is None:
447
- assert pr.embeds_cumsum is None
448
- return
449
-
450
- assert torch.equal(pr.embeds_cumsum, expected)
451
- # cached_property should return the same object on repeated access
452
- assert pr.embeds_cumsum is pr.embeds_cumsum
453
-
454
-
455
-@pytest.mark.parametrize(
456
- "is_embed,start_idx,end_idx,expected",
457
- [
458
- (None, 2, 4, (2, 4)),
459
- (
460
- torch.tensor([False, True, False, True, True]),
461
- 3,
462
- 5,
463
- (1, 3),
464
- ),
465
- (
466
- torch.tensor([False, True, False, True, True]),
467
- 0,
468
- 2,
469
- (0, 1),
470
- ),
471
- (
472
- torch.tensor([True, False, True, False]),
473
- 2,
474
- 2,
475
- (1, 1),
476
- ),
477
- ],
478
-)
479
-def test_placeholder_range_get_embeds_indices_in_range(
480
- is_embed, start_idx, end_idx, expected
481
-):
482
- length = len(is_embed) if is_embed is not None else 5
483
- pr = PlaceholderRange(offset=0, length=length, is_embed=is_embed)
484
- assert pr.get_embeds_indices_in_range(start_idx, end_idx) == expected
485
-
486
-
487
-@pytest.mark.parametrize(
488
- "offset,is_embed,expected",
489
- [
490
- (0, None, [(0, 4)]),
491
- (
492
- 2,
493
- torch.tensor([False, True, False, True, True]),
494
- [(3, 3), (5, 6)],
495
- ),
496
- (0, torch.tensor([True, True, True, True]), [(0, 3)]),
497
- (0, torch.tensor([False, False, False, False]), []),
498
- ],
499
-)
500
-def test_placeholder_range_extract_embeds_range(offset, is_embed, expected):
501
- length = len(is_embed) if is_embed is not None else 5
502
- pr = PlaceholderRange(offset=offset, length=length, is_embed=is_embed)
503
- assert pr.extract_embeds_range() == expected
504
-
505
-
506
-@pytest.mark.asyncio
507
-@pytest.mark.parametrize("video_url", TEST_VIDEO_URLS)
508
-@pytest.mark.parametrize("num_frames", [-1, 32, 1800])
509
-async def test_allowed_media_domains(video_url: str, num_frames: int):
510
- connector = MediaConnector(
511
- media_io_kwargs={
512
- "video": {
513
- "num_frames": num_frames,
514
- }
515
- },
516
- allowed_media_domains=[
517
- "www.bogotobogo.com",
518
- "github.com",
519
- ],
187
+def test_group_and_batch_mm_items_split_by_fieldset():
188
+ elem = MultiModalFieldElem(
189
+ data=torch.empty(1, dtype=torch.uint8),
190
+ field=MultiModalBatchedField(),
191
)
192
+ item1 = MultiModalKwargsItem({"x": elem, "y": elem})
193
+ item2 = MultiModalKwargsItem({"y": elem, "x": elem})
194
+ item3 = MultiModalKwargsItem({"x": elem, "y": elem, "z": elem})
195
+ item4 = MultiModalKwargsItem({"x": elem})
196
+ item5 = MultiModalKwargsItem({"x": elem, "y": elem})
197
522
- video_sync, metadata_sync = connector.fetch_video(video_url)
523
- video_async, metadata_async = await connector.fetch_video_async(video_url)
524
- assert np.array_equal(video_sync, video_async)
525
- assert metadata_sync == metadata_async
198
+ res = group_and_batch_mm_items([item1, item2, item3, item4, item5])
199
+ assert [num_items for num_items, _ in res] == [2, 1, 1, 1]
200
527
- disallowed_url = "https://upload.wikimedia.org/wikipedia/commons/4/47/PNG_transparency_demonstration_1.png"
528
- with pytest.raises(ValueError):
529
- _, _ = connector.fetch_video(disallowed_url)
201
531
- with pytest.raises(ValueError):
532
- _, _ = await connector.fetch_video_async(disallowed_url)
202
+def test_group_and_batch_mm_items_split_by_shared_data():
203
+ elem1 = MultiModalFieldElem(
204
+ data=torch.zeros(1, dtype=torch.uint8),
205
+ field=MultiModalSharedField(batch_size=1),
206
+ )
207
+ elem2 = MultiModalFieldElem(
208
+ data=torch.zeros(2, dtype=torch.uint8),
209
+ field=MultiModalSharedField(batch_size=1),
210
+ )
211
+ item1 = MultiModalKwargsItem({"x": elem1})
212
+ item2 = MultiModalKwargsItem({"x": elem1})
213
+ item3 = MultiModalKwargsItem({"x": elem2})
214
+ item4 = MultiModalKwargsItem({"x": elem1})
215
+ item5 = MultiModalKwargsItem({"x": elem2})
216
+
217
+ res = group_and_batch_mm_items([item1, item2, item3, item4, item5])
218
+ assert [num_items for num_items, _ in res] == [2, 1, 1, 1]
vllm/model_executor/models/step3_vl.py
+18
-18
@@ -71,9 +71,7 @@ class Step3VLImagePixelInputs(TensorSchema):
71
72
type: Literal["pixel_values"]
73
pixel_values: Annotated[torch.Tensor, TensorShape("bn", 3, "h", "w")]
74
- patch_pixel_values: Annotated[
75
- torch.Tensor | None, TensorShape("bnp", 3, "hp", "wp")
76
- ]
74
+ patch_pixel_values: Annotated[torch.Tensor, TensorShape("bnp", 3, "hp", "wp")]
75
num_patches: Annotated[torch.Tensor, TensorShape("bn")]
76
77
@@ -91,7 +89,7 @@ class Step3VLImageEmbeddingInputs(TensorSchema):
89
90
Step3VLImageInputs: TypeAlias = Step3VLImagePixelInputs | Step3VLImageEmbeddingInputs
91
94
-ImageWithPatches = tuple[Image.Image, list[Image.Image], list[int] | None]
92
+ImageWithPatches = tuple[Image.Image, list[Image.Image], list[bool] | None]
93
94
MAX_IMAGE_SIZE: int = 3024
95
@@ -432,7 +430,7 @@ class Step3VLProcessor:
430
431
if len(parts) - 1 != len(repls):
432
raise ValueError(
435
- "The number of placeholders does not match the number of replacements." # noqa: E501
433
+ "The number of placeholders does not match the number of replacements."
434
)
435
436
result = [parts[0]]
@@ -468,7 +466,7 @@ class Step3VLProcessor:
466
image_repl_str_lst = []
467
image_repl_ids_lst = []
468
num_patches = []
471
- for raw_img, img_patches, patch_newline_mask in splitted_images_data: # noqa: E501
469
+ for raw_img, img_patches, patch_newline_mask in splitted_images_data:
470
pixel_values_lst.extend(self._convert_images_to_pixel_values([raw_img]))
471
472
if len(img_patches) > 0:
@@ -486,16 +484,20 @@ class Step3VLProcessor:
484
if patch_newline_mask is not None:
485
patch_newline_mask_lst.extend(patch_newline_mask)
486
487
+ pixel_values = torch.cat(pixel_values_lst)
488
+ patch_size = self.patch_size
489
image_inputs = {
490
- "pixel_values": torch.cat(pixel_values_lst),
490
+ "pixel_values": pixel_values,
491
"num_patches": num_patches,
492
- }
493
- if patch_pixel_values_lst:
494
- image_inputs["patch_pixel_values"] = torch.cat(patch_pixel_values_lst)
495
- if patch_newline_mask_lst:
496
- image_inputs["patch_newline_mask"] = torch.tensor(
492
+ "patch_pixel_values": (
493
+ torch.cat(patch_pixel_values_lst)
494
+ if patch_pixel_values_lst
495
+ else pixel_values.new_empty((0, 3, patch_size, patch_size))
496
+ ),
497
+ "patch_newline_mask": torch.tensor(
498
patch_newline_mask_lst, dtype=torch.bool
498
- )
499
+ ),
500
+ }
501
502
text = [
503
self.replace_placeholder(t, self.image_token, image_repl_str_lst)
@@ -998,13 +1000,11 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
1000
if pixel_values is None and image_embeds is None:
1001
return None
1002
1001
- if pixel_values is not None:
1003
+ if pixel_values is not None and patch_pixel_values is not None:
1004
return Step3VLImagePixelInputs(
1005
type="pixel_values",
1006
pixel_values=pixel_values.to(self.dtype),
1005
- patch_pixel_values=patch_pixel_values.to(self.dtype)
1006
- if patch_pixel_values is not None
1007
- else None,
1007
+ patch_pixel_values=patch_pixel_values.to(self.dtype),
1008
num_patches=num_patches,
1009
)
1010
@@ -1039,7 +1039,7 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
1039
image_features = self._get_vision_model_output(image_input["pixel_values"])
1040
patch_image_features = (
1041
self._get_vision_model_output(image_input["patch_pixel_values"])
1042
- if image_input["patch_pixel_values"] is not None
1042
+ if len(image_input["patch_pixel_values"]) > 0
1043
else None
1044
)
1045
num_patches = image_input["num_patches"]
vllm/model_executor/models/terratorch.py
+28
-18
@@ -62,7 +62,6 @@ from vllm.multimodal.processing import (
62
PromptUpdate,
63
)
64
from vllm.sequence import IntermediateTensors
65
-from vllm.utils import length_from_prompt_token_ids_or_embeds
65
66
from .interfaces import IsAttentionFree, MultiModalEmbeddings, SupportsMultiModal
67
from .interfaces_base import attn_type
@@ -74,7 +73,11 @@ def _terratorch_field_names(input_definition: InputDefinition):
73
return set(input_definition.data.keys())
74
75
77
-def _terratorch_field_factory(input_definition: InputDefinition):
76
+def _terratorch_field_factory(
77
+ input_definition: InputDefinition,
78
+ *,
79
+ is_shared: bool = True, # True for unprocessed data, False for processed data
80
+):
81
def _terratorch_field_config(
82
hf_inputs: Mapping[str, torch.Tensor],
83
) -> Mapping[str, MultiModalFieldConfig]:
@@ -82,7 +85,11 @@ def _terratorch_field_factory(input_definition: InputDefinition):
85
for name, input in input_definition.data.items():
86
modality = "image"
87
if input.type == InputTypeEnum.tensor:
85
- fields[name] = MultiModalFieldConfig.shared(modality, batch_size=1)
88
+ fields[name] = (
89
+ MultiModalFieldConfig.shared(modality, batch_size=1)
90
+ if is_shared
91
+ else MultiModalFieldConfig.batched(modality)
92
+ )
93
94
return fields
95
@@ -166,8 +173,14 @@ class TerratorchMultiModalProcessor(BaseMultiModalProcessor[TerratorchProcessing
173
self,
174
hf_inputs: BatchFeature,
175
hf_processor_mm_kwargs: Mapping[str, object],
176
+ *,
177
+ is_shared: bool = True,
178
) -> Mapping[str, MultiModalFieldConfig]:
170
- return _terratorch_field_factory(self.info.input_definition)(hf_inputs)
179
+ factory = _terratorch_field_factory(
180
+ self.info.input_definition,
181
+ is_shared=is_shared,
182
+ )
183
+ return factory(hf_inputs)
184
185
def _get_prompt_updates(
186
self,
@@ -193,12 +206,19 @@ class TerratorchMultiModalProcessor(BaseMultiModalProcessor[TerratorchProcessing
206
)
207
208
_, passthrough_data = self._get_hf_mm_data(mm_items)
196
- mm_processed_data = BatchFeature(dict(passthrough_data), tensor_type="pt")
209
+ mm_processed_data = BatchFeature(
210
+ {k: torch.tensor(v).unsqueeze(0) for k, v in passthrough_data.items()},
211
+ tensor_type="pt",
212
+ )
213
mm_placeholders = {"image": [PlaceholderRange(offset=0, length=0)]}
214
215
mm_kwargs = MultiModalKwargsItems.from_hf_inputs(
216
mm_processed_data,
201
- self._get_mm_fields_config(mm_processed_data, hf_processor_mm_kwargs),
217
+ self._get_mm_fields_config(
218
+ mm_processed_data,
219
+ hf_processor_mm_kwargs,
220
+ is_shared=False,
221
+ ),
222
)
223
224
return MultiModalInputs(
@@ -235,9 +255,6 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):
255
self.inference_runner = InferenceRunner(config)
256
self.model = self.inference_runner.model
257
238
- pooler_config = vllm_config.model_config.pooler_config
239
- assert pooler_config is not None
240
-
258
self.pooler = IdentityPooler()
259
260
def embed_input_ids(
@@ -262,15 +279,8 @@ class Terratorch(nn.Module, IsAttentionFree, SupportsMultiModal):
279
inputs_embeds: torch.Tensor | None = None,
280
**kwargs: object,
281
):
265
- input_len = length_from_prompt_token_ids_or_embeds(input_ids, inputs_embeds)
266
-
267
- batched_kwargs = {k: v.unsqueeze(0) for k, v in kwargs.items()}
268
- model_output = self.inference_runner.forward(**batched_kwargs).output
269
-
270
- # The leading dimension of hidden states needs to equal input length
271
- return model_output.expand(
272
- input_len, *(-1 for _ in range(model_output.ndim - 1))
273
- )
282
+ model_output = self.inference_runner.forward(**kwargs)
283
+ return model_output.output
284
285
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
286
params_list = []
vllm/multimodal/hasher.py
+13
-5
@@ -102,12 +102,19 @@ class MultiModalHasher:
102
"data": tensor_obj.numpy(),
103
},
104
)
105
+
106
return cls.iter_item_to_bytes("tensor", tensor_obj.numpy())
107
+
108
if isinstance(obj, np.ndarray):
107
- # If the array is non-contiguous, we need to copy it first
108
- arr_data = (
109
- obj.view(np.uint8).data if obj.flags.c_contiguous else obj.tobytes()
110
- )
109
+ if obj.ndim == 0:
110
+ arr_data = obj.item()
111
+ elif obj.flags.c_contiguous:
112
+ # Not valid for 0-D arrays
113
+ arr_data = obj.view(np.uint8).data
114
+ else:
115
+ # If the array is non-contiguous, we need to copy it first
116
+ arr_data = obj.tobytes()
117
+
118
return cls.iter_item_to_bytes(
119
"ndarray",
120
{
@@ -116,6 +123,7 @@ class MultiModalHasher:
123
"data": arr_data,
124
},
125
)
126
+
127
logger.warning(
128
"No serialization method found for %s. Falling back to pickle.", type(obj)
129
)
@@ -147,7 +155,7 @@ class MultiModalHasher:
155
hasher_factory = _get_hasher_factory(envs.VLLM_MM_HASHER_ALGORITHM)
156
hasher = hasher_factory()
157
150
- for k, v in kwargs.items():
158
+ for k, v in sorted(kwargs.items(), key=lambda kv: kv[0]):
159
for bytes_ in cls.iter_item_to_bytes(k, v):
160
hasher.update(bytes_)
161
vllm/multimodal/inputs.py
+32
-20
@@ -424,8 +424,9 @@ class BaseMultiModalField(ABC):
424
425
keep_on_cpu: bool = False
426
"""
427
- If `True`, then this field is excluded from being moved to the accelerator
428
- when `MultiModalKwargsItems.get_data()` is called to batch the data.
427
+ If `True`, then this field is excluded from being moved to the accelerator when
428
+ [`group_and_batch_mm_items`][vllm.multimodal.utils.group_and_batch_mm_items]
429
+ is called to batch the data.
430
"""
431
432
def _field_factory(self):
@@ -1006,27 +1007,38 @@ class MultiModalKwargsItems(UserDict[str, Sequence[_I]]):
1007
pin_memory: bool = False,
1008
) -> BatchedTensorInputs:
1009
"""Construct a dictionary of keyword arguments to pass to the model."""
1009
- elems_by_key = defaultdict[str, list[MultiModalFieldElem]](list)
1010
- for modality, items in self.items():
1011
- for i, item in enumerate(items):
1012
- if item is None:
1013
- raise RuntimeError(
1014
- f"Cannot build data from empty mm_items[{modality}][{i}]"
1015
- )
1010
+ from .utils import group_and_batch_mm_items
1011
+
1012
+ items_by_modality = self.require_data()
1013
+ batches_by_modality = {
1014
+ modality: [
1015
+ data
1016
+ for _, data in group_and_batch_mm_items(
1017
+ items,
1018
+ device=device,
1019
+ pin_memory=pin_memory,
1020
+ )
1021
+ ]
1022
+ for modality, items in items_by_modality.items()
1023
+ if len(items) > 0
1024
+ }
1025
1017
- for key, elem in item.items():
1018
- elems_by_key[key].append(elem)
1026
+ out_data: BatchedTensorInputs = {}
1027
+ for _, batches in batches_by_modality.items():
1028
+ if len(batches) != 1:
1029
+ num_batches_by_modality = {
1030
+ modality: len(batches)
1031
+ for modality, batches in batches_by_modality.items()
1032
+ }
1033
+
1034
+ raise RuntimeError(
1035
+ f"Some modalities cannot be merged into a single batch "
1036
+ f"({num_batches_by_modality=})"
1037
+ )
1038
1020
- data = {
1021
- key: elems[0].field.reduce_data(
1022
- elems,
1023
- device=device,
1024
- pin_memory=pin_memory,
1025
- )
1026
- for key, elems in elems_by_key.items()
1027
- }
1039
+ out_data.update(batches[0])
1040
1029
- return data
1041
+ return out_data
1042
1043
1044
MultiModalKwargsOptionalItems: TypeAlias = (
vllm/multimodal/utils.py
+101
-11
@@ -3,7 +3,8 @@
3
4
import mimetypes
5
import warnings
6
-from collections.abc import Generator
6
+from collections import defaultdict
7
+from collections.abc import Generator, Sequence
8
from itertools import groupby
9
from typing import TYPE_CHECKING, Any
10
@@ -13,11 +14,13 @@ from PIL import Image
14
15
from vllm.utils.import_utils import LazyLoader
16
17
+from .hasher import MultiModalHasher
18
from .inputs import (
19
BatchedTensorInputs,
20
+ MultiModalFieldElem,
21
MultiModalKwargsItem,
19
- MultiModalKwargsItems,
22
MultiModalPlaceholderDict,
23
+ MultiModalSharedField,
24
)
25
from .media import AudioMediaIO, ImageMediaIO, MediaConnector, VideoMediaIO
26
@@ -146,32 +149,119 @@ def argsort_mm_positions(
149
return [(modality, idx) for modality, idx, _ in sorted_flat_items]
150
151
152
+def _get_group_hash(elem: MultiModalFieldElem):
153
+ if not isinstance(elem.field, MultiModalSharedField):
154
+ return None
155
+
156
+ return MultiModalHasher.hash_kwargs(data=elem.data)
157
+
158
+
159
+def _batch_mm_items(
160
+ items: Sequence[MultiModalKwargsItem],
161
+ *,
162
+ device: torch.types.Device = None,
163
+ pin_memory: bool = False,
164
+):
165
+ elems = defaultdict[str, list[MultiModalFieldElem]](list)
166
+ for item in items:
167
+ for key, elem in item.items():
168
+ elems[key].append(elem)
169
+
170
+ return {
171
+ key: elems[0].field.reduce_data(
172
+ elems,
173
+ device=device,
174
+ pin_memory=pin_memory,
175
+ )
176
+ for key, elems in elems.items()
177
+ }
178
+
179
+
180
+def group_and_batch_mm_items(
181
+ items: Sequence[MultiModalKwargsItem],
182
+ *,
183
+ device: torch.types.Device = None,
184
+ pin_memory: bool = False,
185
+) -> Generator[tuple[int, BatchedTensorInputs]]:
186
+ """
187
+ Group consecutive items (possibly from different requests) into batches.
188
+
189
+ Items must be split across groups if any of the following occurs,
190
+ as the batch would otherwise be invalid:
191
+ - They have different fields (e.g. mixed image and embedding inputs).
192
+ - They have different values in `MultiModalSharedField`.
193
+
194
+ Args:
195
+ items: List of `MultiModalKwargsItem`.
196
+ device: The device to place the grouped tensors on.
197
+ pin_memory: Whether to pin memory for faster host-to-device transfer.
198
+
199
+ Yields:
200
+ A tuple `(num_items, grouped_kwargs)`, where:
201
+ - `kwargs` is a dictionary of keyword arguments to pass to the model;
202
+ - `num_items` is the corresponding number of items.
203
+ """
204
+ group_ids = [
205
+ tuple(
206
+ (key, _get_group_hash(elem))
207
+ for key, elem in sorted(item.items(), key=lambda kv: kv[0])
208
+ )
209
+ for item in items
210
+ ]
211
+ group_sizes = [sum(1 for _ in group) for _, group in groupby(group_ids)]
212
+
213
+ start_idx = 0
214
+ for group_size in group_sizes:
215
+ group_data = _batch_mm_items(
216
+ items[start_idx : start_idx + group_size],
217
+ device=device,
218
+ pin_memory=pin_memory,
219
+ )
220
+
221
+ yield group_size, group_data
222
+
223
+ start_idx += group_size
224
+
225
+ assert start_idx == len(items)
226
+
227
+
228
def group_mm_kwargs_by_modality(
229
mm_kwargs: list[tuple[str, MultiModalKwargsItem]],
230
*,
231
device: torch.types.Device = None,
232
pin_memory: bool = False,
233
) -> Generator[tuple[str, int, BatchedTensorInputs], None, None]:
155
- """Group consecutive `MultiModalKwargsItem`s from `mm_kwargs` with the same
156
- modality together into the same `MultiModalKwargs` instance.
234
+ """
235
+ Group consecutive items (possibly from different requests) into batches.
236
+
237
+ Items must be split across groups if any of the following occurs,
238
+ as the batch would otherwise be invalid:
239
+ - They have different fields (e.g. mixed image and embedding inputs).
240
+ - They have different values in `MultiModalSharedField`.
241
+
242
+ To simplify the implementation of `embed_multimodal`, we add another
243
+ restriction that the items in a batch must belong to the same modality.
244
245
Args:
159
- mm_kwargs: List of `MultiModalKwargsItem`.
246
+ mm_kwargs: List of `(modality, item)`.
247
device: The device to place the grouped tensors on.
248
pin_memory: Whether to pin memory for faster host-to-device transfer.
249
250
Yields:
164
- A tuple `(modality, num_items, grouped_kwargs)`.
251
+ A tuple `(modality, num_items, grouped_kwargs)`, where:
252
+ - `modality` is the modality of the batch;
253
+ - `kwargs` is a dictionary of keyword arguments to pass to the model;
254
+ - `num_items` is the corresponding number of items.
255
"""
256
for modality, group in groupby(mm_kwargs, key=lambda x: x[0]):
257
items_lst = [item for _, item in group]
168
- mm_kwargs_items = MultiModalKwargsItems({modality: items_lst})
169
- mm_kwargs_data = mm_kwargs_items.get_data(
258
+
259
+ for num_items, mm_kwargs_batch in group_and_batch_mm_items(
260
+ items_lst,
261
device=device,
262
pin_memory=pin_memory,
172
- )
173
-
174
- yield modality, len(items_lst), mm_kwargs_data
263
+ ):
264
+ yield modality, num_items, mm_kwargs_batch
265
266
267
def fetch_audio(