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 6, 00:45 UTC
[MultiModal] add support for numpy array embeddings (#38119)
Signed-off-by: guillaume_guy <guillaume.guy@airbnb.com> Signed-off-by: Guillaume Guy <guillaume.c.guy@gmail.com> Co-authored-by: guillaume_guy <guillaume.guy@airbnb.com> Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>
Guillaume Guy committed
Mar 25, 2026 at 15:13 UTC
70a2152830959d0d5e25d60869f2fb3e2fa3733c
2 files changed
+63
-3
tests/renderers/test_sparse_tensor_validation.py
+46
@@ -7,6 +7,7 @@ out-of-bounds memory writes during to_dense() operations.
7
8
import io
9
10
+import numpy as np
11
import pybase64 as base64
12
import pytest
13
import torch
@@ -190,6 +191,51 @@ class TestImageEmbedsValidation:
191
with pytest.raises((RuntimeError, ValueError)):
192
io_handler.load_bytes(buffer.read())
193
194
+ def test_valid_numpy_tensor_accepted(self):
195
+ """numpy .npy format should load and return correct tensor."""
196
+ io_handler = ImageEmbeddingMediaIO()
197
+
198
+ arr = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32)
199
+ buf = io.BytesIO()
200
+ np.save(buf, arr)
201
+ encoded = base64.b64encode(buf.getvalue()).decode("utf-8")
202
+
203
+ result = io_handler.load_base64("", encoded)
204
+ assert isinstance(result, torch.Tensor)
205
+ assert result.shape == torch.Size([2, 3])
206
+ assert result.dtype == torch.float32
207
+ assert torch.allclose(result, torch.from_numpy(arr))
208
+
209
+ def test_numpy_int32_tensor_accepted(self):
210
+ """numpy int32 arrays should round-trip correctly."""
211
+
212
+ io_handler = ImageEmbeddingMediaIO()
213
+
214
+ arr = np.arange(280, dtype=np.int32)
215
+ buf = io.BytesIO()
216
+ np.save(buf, arr)
217
+ encoded = base64.b64encode(buf.getvalue()).decode("utf-8")
218
+
219
+ result = io_handler.load_base64("", encoded)
220
+ assert result.dtype == torch.int32
221
+ assert result.shape == torch.Size([280])
222
+ assert (result == torch.from_numpy(arr)).all()
223
+
224
+ def test_load_file_numpy_tensor_accepted(self, tmp_path):
225
+ """numpy .npy files should load correctly via load_file."""
226
+
227
+ io_handler = ImageEmbeddingMediaIO()
228
+
229
+ arr = np.array([[1.5, 2.5], [3.5, 4.5]], dtype=np.float32)
230
+ npy_path = tmp_path / "image_embeds.npy"
231
+ np.save(npy_path, arr)
232
+
233
+ result = io_handler.load_file(npy_path)
234
+ assert isinstance(result, torch.Tensor)
235
+ assert result.shape == torch.Size([2, 2])
236
+ assert result.dtype == torch.float32
237
+ assert torch.allclose(result, torch.from_numpy(arr))
238
+
239
240
class TestAudioEmbedsValidation:
241
"""Test sparse tensor validation in audio embeddings (Chat API)."""
vllm/multimodal/media/image.py
+17
-3
@@ -4,6 +4,7 @@
4
from io import BytesIO
5
from pathlib import Path
6
7
+import numpy as np
8
import pybase64
9
import torch
10
from PIL import Image
@@ -13,6 +14,8 @@ from vllm.utils.serial_utils import tensor2base64
14
from ..image import convert_image_mode, rgba_to_rgb
15
from .base import MediaIO, MediaWithBytes
16
17
+MAGIC_NUMPY_PREFIX = b"\x93NUMPY" # https://numpy.org/devdocs/reference/generated/numpy.lib.format.html#format-version-1-0
18
+
19
20
class ImageMediaIO(MediaIO[Image.Image]):
21
"""Configuration values can be user-provided either by --media-io-kwargs or
@@ -104,7 +107,7 @@ class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]):
107
def __init__(self) -> None:
108
super().__init__()
109
107
- def load_bytes(self, data: bytes) -> torch.Tensor:
110
+ def _load_pickled_torch(self, data: bytes) -> torch.Tensor:
111
buffer = BytesIO(data)
112
# Enable sparse tensor integrity checks to prevent out-of-bounds
113
# writes from maliciously crafted tensors
@@ -112,12 +115,23 @@ class ImageEmbeddingMediaIO(MediaIO[torch.Tensor]):
115
tensor = torch.load(buffer, weights_only=True)
116
return tensor.to_dense()
117
118
+ def _load_numpy(self, data: bytes) -> torch.Tensor:
119
+ with BytesIO(data) as buffer:
120
+ return torch.from_numpy(np.load(buffer))
121
+
122
+ def load_bytes(self, data: bytes) -> torch.Tensor:
123
+ if data[:6] == MAGIC_NUMPY_PREFIX:
124
+ return self._load_numpy(data)
125
+
126
+ return self._load_pickled_torch(data)
127
+
128
def load_base64(self, media_type: str, data: str) -> torch.Tensor:
129
return self.load_bytes(pybase64.b64decode(data, validate=True))
130
131
def load_file(self, filepath: Path) -> torch.Tensor:
119
- # Enable sparse tensor integrity checks to prevent out-of-bounds
120
- # writes from maliciously crafted tensors
132
+ if filepath.suffix == ".npy":
133
+ return torch.from_numpy(np.load(filepath))
134
+
135
with torch.sparse.check_sparse_tensor_invariants():
136
tensor = torch.load(filepath, weights_only=True)
137
return tensor.to_dense()