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 12, 04:45 UTC
[Multimodal][Core] Optimize multimodal preprocessing cache by hashing image bytes instead of pixel values (#29621)
Signed-off-by: Rahul Steiger <rasteiger@ethz.ch> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
ImaGoodFella committed
Dec 2, 2025 at 14:49 UTC
60c3d413afccab6a1f9a18cf3cd1fe11019c1040
8 files changed
+95
-19
tests/conftest.py
+6
-1
@@ -59,6 +59,7 @@ from vllm.distributed import (
59
)
60
from vllm.logger import init_logger
61
from vllm.logprobs import Logprob
62
+from vllm.multimodal.base import MediaWithBytes
63
from vllm.multimodal.utils import fetch_image
64
from vllm.outputs import RequestOutput
65
from vllm.sampling_params import BeamSearchParams
@@ -1389,7 +1390,11 @@ class LocalAssetServer:
1390
return f"{self.base_url}/{name}"
1391
1392
def get_image_asset(self, name: str) -> Image.Image:
1392
- return fetch_image(self.url_for(name))
1393
+ image = fetch_image(self.url_for(name))
1394
+ # Unwrap MediaWithBytes if present
1395
+ if isinstance(image, MediaWithBytes):
1396
+ image = image.media
1397
+ return image
1398
1399
1400
@pytest.fixture(scope="session")
tests/entrypoints/openai/test_vision.py
+6
-1
@@ -8,6 +8,7 @@ import pytest
8
import pytest_asyncio
9
from transformers import AutoProcessor
10
11
+from vllm.multimodal.base import MediaWithBytes
12
from vllm.multimodal.utils import encode_image_base64, fetch_image
13
14
from ...utils import RemoteOpenAIServer
@@ -111,7 +112,11 @@ def get_hf_prompt_tokens(model_name, content, image_url):
112
"content": f"{placeholder}{content}",
113
}
114
]
114
- images = [fetch_image(image_url)]
115
+ image = fetch_image(image_url)
116
+ # Unwrap MediaWithBytes if present
117
+ if isinstance(image, MediaWithBytes):
118
+ image = image.media
119
+ images = [image]
120
121
prompt = processor.tokenizer.apply_chat_template(
122
messages, tokenize=False, add_generation_prompt=True
tests/entrypoints/pooling/embed/test_online_vision.py
+6
-1
@@ -9,6 +9,7 @@ from transformers import AutoProcessor
9
10
from tests.utils import VLLM_PATH, RemoteOpenAIServer
11
from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
12
+from vllm.multimodal.base import MediaWithBytes
13
from vllm.multimodal.utils import encode_image_base64, fetch_image
14
15
MODEL_NAME = "TIGER-Lab/VLM2Vec-Full"
@@ -62,7 +63,11 @@ def get_hf_prompt_tokens(model_name, content, image_url):
63
64
placeholder = "<|image_1|> "
65
prompt = f"{placeholder}{content}"
65
- images = [fetch_image(image_url)]
66
+ image = fetch_image(image_url)
67
+ # Unwrap MediaWithBytes if present
68
+ if isinstance(image, MediaWithBytes):
69
+ image = image.media
70
+ images = [image]
71
inputs = processor(prompt, images, return_tensors="pt")
72
return inputs.input_ids.shape[1]
73
vllm/multimodal/base.py
+28
@@ -2,12 +2,40 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from abc import ABC, abstractmethod
5
+from dataclasses import dataclass
6
from pathlib import Path
7
from typing import Generic, TypeVar
8
9
+import numpy as np
10
+
11
_T = TypeVar("_T")
12
13
14
+@dataclass
15
+class MediaWithBytes(Generic[_T]):
16
+ """
17
+ Wrapper that couples a media object with its original encoded bytes.
18
+
19
+ This ensures the raw bytes and media object remain synchronized,
20
+ preventing cache corruption from in-place modifications.
21
+
22
+ The wrapper delegates attribute access to the underlying media object,
23
+ making it behave transparently like the wrapped type (e.g., PIL.Image).
24
+ """
25
+
26
+ media: _T
27
+ original_bytes: bytes
28
+
29
+ def __array__(self, *args, **kwargs) -> np.ndarray:
30
+ """Allow np.array(obj) to return np.array(obj.media)."""
31
+ return np.array(self.media, *args, **kwargs)
32
+
33
+ def __getattr__(self, name: str):
34
+ """Delegate attribute access to the underlying media object."""
35
+ # This is only called when the attribute is not found on self
36
+ return getattr(self.media, name)
37
+
38
+
39
class MediaIO(ABC, Generic[_T]):
40
@abstractmethod
41
def load_bytes(self, data: bytes) -> _T:
vllm/multimodal/hasher.py
+19
-5
@@ -12,6 +12,8 @@ from PIL import Image
12
13
from vllm.logger import init_logger
14
15
+from .base import MediaWithBytes
16
+
17
logger = init_logger(__name__)
18
19
@@ -31,14 +33,26 @@ class MultiModalHasher:
33
if Image.ExifTags.Base.ImageID in exif and isinstance(
34
exif[Image.ExifTags.Base.ImageID], uuid.UUID
35
):
34
- # If the image has exif ImageID tag, use that
36
return (exif[Image.ExifTags.Base.ImageID].bytes,)
37
+
38
data = {"mode": obj.mode, "data": np.asarray(obj)}
37
- if obj.palette is not None:
38
- data["palette"] = obj.palette.palette
39
- if obj.palette.rawmode is not None:
40
- data["palette_rawmode"] = obj.palette.rawmode
39
+ palette = obj.palette
40
+ if palette is not None:
41
+ data["palette"] = palette.palette
42
+ if palette.rawmode is not None:
43
+ data["palette_rawmode"] = palette.rawmode
44
+
45
return cls.iter_item_to_bytes("image", data)
46
+
47
+ if isinstance(obj, MediaWithBytes) and isinstance(obj.media, Image.Image):
48
+ exif = obj.media.getexif()
49
+ if Image.ExifTags.Base.ImageID in exif and isinstance(
50
+ exif[Image.ExifTags.Base.ImageID], uuid.UUID
51
+ ):
52
+ return (exif[Image.ExifTags.Base.ImageID].bytes,)
53
+
54
+ return cls.iter_item_to_bytes("image", obj.original_bytes)
55
+
56
if isinstance(obj, torch.Tensor):
57
tensor_obj: torch.Tensor = obj.cpu()
58
tensor_dtype = tensor_obj.dtype
vllm/multimodal/image.py
+14
-10
@@ -8,7 +8,7 @@ import pybase64
8
import torch
9
from PIL import Image
10
11
-from .base import MediaIO
11
+from .base import MediaIO, MediaWithBytes
12
13
14
def rescale_image_size(
@@ -74,8 +74,12 @@ class ImageMediaIO(MediaIO[Image.Image]):
74
)
75
self.rgba_background_color = rgba_bg
76
77
- def _convert_image_mode(self, image: Image.Image) -> Image.Image:
77
+ def _convert_image_mode(
78
+ self, image: Image.Image | MediaWithBytes[Image.Image]
79
+ ) -> Image.Image:
80
"""Convert image mode with custom background color."""
81
+ if isinstance(image, MediaWithBytes):
82
+ image = image.media
83
if image.mode == self.image_mode:
84
return image
85
elif image.mode == "RGBA" and self.image_mode == "RGB":
@@ -83,18 +87,18 @@ class ImageMediaIO(MediaIO[Image.Image]):
87
else:
88
return convert_image_mode(image, self.image_mode)
89
86
- def load_bytes(self, data: bytes) -> Image.Image:
90
+ def load_bytes(self, data: bytes) -> MediaWithBytes[Image.Image]:
91
image = Image.open(BytesIO(data))
88
- image.load()
89
- return self._convert_image_mode(image)
92
+ return MediaWithBytes(self._convert_image_mode(image), data)
93
91
- def load_base64(self, media_type: str, data: str) -> Image.Image:
94
+ def load_base64(self, media_type: str, data: str) -> MediaWithBytes[Image.Image]:
95
return self.load_bytes(pybase64.b64decode(data, validate=True))
96
94
- def load_file(self, filepath: Path) -> Image.Image:
95
- image = Image.open(filepath)
96
- image.load()
97
- return self._convert_image_mode(image)
97
+ def load_file(self, filepath: Path) -> MediaWithBytes[Image.Image]:
98
+ with open(filepath, "rb") as f:
99
+ data = f.read()
100
+ image = Image.open(BytesIO(data))
101
+ return MediaWithBytes(self._convert_image_mode(image), data)
102
103
def encode_base64(
104
self,
vllm/multimodal/parse.py
+15
@@ -23,6 +23,7 @@ from vllm.utils.collection_utils import is_list_of
23
from vllm.utils.import_utils import LazyLoader
24
25
from .audio import AudioResampler
26
+from .base import MediaWithBytes
27
from .inputs import (
28
AudioItem,
29
HfAudioItem,
@@ -84,6 +85,12 @@ class ModalityDataItems(ABC, Generic[_T, _I]):
85
"""Get all data items."""
86
return [self.get(idx) for idx in range(self.get_count())]
87
88
+ def get_item_for_hash(self, index: int) -> object:
89
+ return self.get(index)
90
+
91
+ def get_all_items_for_hash(self) -> list[object]:
92
+ return [self.get_item_for_hash(idx) for idx in range(self.get_count())]
93
+
94
@abstractmethod
95
def get_processor_data(self) -> Mapping[str, object]:
96
"""Get the data to pass to the HF processor."""
@@ -98,10 +105,18 @@ class ModalityDataItems(ABC, Generic[_T, _I]):
105
class ProcessorBatchItems(ModalityDataItems[Sequence[_T], _T]):
106
"""Base class for data items that are arranged in a list."""
107
108
+ def _unwrap(self, item: _T | MediaWithBytes[_T]) -> _T:
109
+ """Extract media from wrapper if present."""
110
+ return item.media if isinstance(item, MediaWithBytes) else item
111
+
112
def get_count(self) -> int:
113
return len(self.data)
114
115
def get(self, index: int) -> _T:
116
+ return self._unwrap(self.data[index])
117
+
118
+ def get_item_for_hash(self, index: int) -> _T | MediaWithBytes[_T]:
119
+ # Return raw item for hashing (preserves original_bytes if present)
120
return self.data[index]
121
122
def get_processor_data(self) -> Mapping[str, object]:
vllm/multimodal/processing.py
+1
-1
@@ -1684,7 +1684,7 @@ class BaseMultiModalProcessor(ABC, Generic[_I]):
1684
1685
# For None entries, compute a hash; otherwise, use provided ID.
1686
computed: list[str] = []
1687
- for i, item in enumerate(items):
1687
+ for i, item in enumerate(items.get_all_items_for_hash()):
1688
item_uuid = mm_uuids_per_modality[i]
1689
1690
# NOTE: Even if a item_uuid is provided, we still compute a