[Security] Authenticate shared-memory multimodal cache handles (#59357)
Signed-off-by: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Clinton Thomas committed
Sep 30, 2026 at 08:16 UTC
3d5625cc3c2474994107047d7a6ecb8e4008f9b0
8 files changed
+574
-41
tests/distributed/test_shm_buffer.py
+3
-3
@@ -108,10 +108,10 @@ class TestSingleWriterShmRingBuffer(unittest.TestCase):
108
# Clear the buffer
109
self.ring_buffer.clear()
110
111
- # Check that metadata is empty and IDs reset
111
+ # Check that metadata is empty and IDs keep increasing
112
self.assertEqual(len(self.ring_buffer.metadata), 0)
113
- self.assertEqual(self.ring_buffer.monotonic_id_start, 0)
114
- self.assertEqual(self.ring_buffer.monotonic_id_end, 0)
113
+ self.assertEqual(self.ring_buffer.monotonic_id_start, 3)
114
+ self.assertEqual(self.ring_buffer.monotonic_id_end, 3)
115
self.assertEqual(self.ring_buffer.data_buffer_start, 0)
116
self.assertEqual(self.ring_buffer.data_buffer_end, 0)
117
tests/distributed/test_shm_storage.py
+64
-5
@@ -227,10 +227,68 @@ class TestSingleWriterShmObjectStorage(unittest.TestCase):
227
self.assertEqual(len(self.storage.id_index), 0)
228
self.assertEqual(len(self.storage.ring_buffer.metadata), 0)
229
230
- # Verify that new items can be added after clearing
230
+ # Verify that new items can be added after clearing. IDs keep
231
+ # increasing so handles issued before the clear cannot alias new data.
232
address, monotonic_id = self.storage.put("new_item", "new_value")
233
self.assertIn("new_item", self.storage.key_index)
233
- self.assertEqual((address, monotonic_id), (0, 0))
234
+ self.assertEqual((address, monotonic_id), (0, 5))
235
+
236
+ def _open_reader(self) -> SingleWriterShmObjectStorage:
237
+ reader = SingleWriterShmObjectStorage.create_from_handle(self.storage.handle())
238
+ self.addCleanup(reader.close)
239
+ return reader
240
+
241
+ def test_reader_get_and_touch_each_require_valid_signature(self):
242
+ address, monotonic_id = self.storage.put("key", "value")
243
+ signature = self.storage.get_signature("key")
244
+ reader = self._open_reader()
245
+
246
+ reader.touch("key", address, monotonic_id, signature)
247
+ self.assertEqual(reader.get(address, monotonic_id, signature, "key"), "value")
248
+
249
+ flipped = [signature[0] ^ 1, *signature[1:]]
250
+ for bad in (None, flipped, [signature], signature[:-1], [256] * 32):
251
+ with self.assertRaisesRegex(ValueError, "SHM handle signature"):
252
+ reader.get(address, monotonic_id, bad, "key")
253
+ with self.assertRaisesRegex(ValueError, "SHM handle signature"):
254
+ reader.touch("key", address, monotonic_id, bad)
255
+
256
+ def test_signature_is_bound_to_key_and_address(self):
257
+ address, monotonic_id = self.storage.put("key", "value")
258
+ signature = self.storage.get_signature("key")
259
+ reader = self._open_reader()
260
+
261
+ with self.assertRaisesRegex(ValueError, "SHM handle signature"):
262
+ reader.get(address, monotonic_id, signature, "other_key")
263
+ with self.assertRaisesRegex(ValueError, "SHM handle signature"):
264
+ reader.get(address + 1, monotonic_id, signature, "key")
265
+
266
+ def test_writer_rejects_handle_no_longer_in_key_index(self):
267
+ address, monotonic_id = self.storage.put("key", "value")
268
+ signature = self.storage.get_signature("key")
269
+ self.storage.verify_signature("key", address, monotonic_id, signature)
270
+
271
+ self.storage.clear()
272
+
273
+ with self.assertRaisesRegex(ValueError, "SHM handle signature"):
274
+ self.storage.verify_signature("key", address, monotonic_id, signature)
275
+
276
+ def test_handle_issued_before_clear_reads_own_data_until_reused(self):
277
+ address, monotonic_id = self.storage.put("key", "value")
278
+ signature = self.storage.get_signature("key")
279
+ reader = self._open_reader()
280
+
281
+ self.storage.clear()
282
+
283
+ # Requests still draining after a clear keep reading their own data.
284
+ self.assertEqual(reader.get(address, monotonic_id, signature, "key"), "value")
285
+
286
+ self.storage.put("key", "new_value")
287
+
288
+ with self.assertRaisesRegex(ValueError, "modified or is invalid"):
289
+ reader.get(address, monotonic_id, signature, "key")
290
+ with self.assertRaisesRegex(ValueError, "modified or is invalid"):
291
+ reader.touch("key", address, monotonic_id, signature)
292
293
294
# Reader process function
@@ -242,11 +300,11 @@ def reader_process(process_id, storage_handle, items_to_read):
300
301
errors = []
302
245
- for key, original_value, address, monotonic_id in items_to_read:
303
+ for key, original_value, address, monotonic_id, signature in items_to_read:
304
time.sleep(random.random() / 100)
305
try:
306
# Read data from shared memory
249
- retrieved_value = reader_storage.get(address, monotonic_id)
307
+ retrieved_value = reader_storage.get(address, monotonic_id, signature, key)
308
309
# Verify data integrity
310
assert retrieved_value == original_value
@@ -289,7 +347,8 @@ def run_multiprocess_example():
347
for key, value in test_data:
348
print(f"Storing {key}: {value}")
349
address, monotonic_id = storage.put(key, value)
292
- stored_items.append((key, value, address, monotonic_id))
350
+ signature = storage.get_signature(key)
351
+ stored_items.append((key, value, address, monotonic_id, signature))
352
print(f" -> Stored at address {address}, ID {monotonic_id}")
353
354
print("\n--- Retrieving Data ---")
tests/entrypoints/scale_out/token_in_token_out/test_generate_stream.py
+52
-1
@@ -15,14 +15,25 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels
15
from vllm.entrypoints.scale_out.token_in_token_out.mm_features import (
16
placeholder_ranges_from_engine_input,
17
)
18
+from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import (
19
+ encode_mm_kwargs_item,
20
+)
21
from vllm.entrypoints.scale_out.token_in_token_out.protocol import (
22
GenerateRequest,
23
GenerateResponse,
24
+ MultiModalFeatures,
25
+ PlaceholderRangeInfo,
26
)
27
from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens
28
+from vllm.entrypoints.serve.engine.protocol import ErrorResponse
29
from vllm.exceptions import GenerationError
30
from vllm.logprobs import Logprob
25
-from vllm.multimodal.inputs import PlaceholderRange
31
+from vllm.multimodal.inputs import (
32
+ MultiModalBatchedField,
33
+ MultiModalFieldElem,
34
+ MultiModalKwargsItem,
35
+ PlaceholderRange,
36
+)
37
from vllm.outputs import CompletionOutput, RequestOutput
38
from vllm.renderers import renderer_from_config
39
from vllm.renderers.online_renderer import OnlineRenderer
@@ -285,6 +296,46 @@ async def test_non_stream_error():
296
await serving.serve_tokens(request)
297
298
299
+@pytest.mark.asyncio
300
+async def test_serve_tokens_rejects_invalid_shm_handle_before_engine_handoff():
301
+ engine = _mock_engine()
302
+ engine.generate = MagicMock()
303
+ serving = _build_serving_tokens(engine)
304
+ cache = MagicMock()
305
+ cache.validate_input_item.side_effect = ValueError(
306
+ "Invalid SHM handle signature for cache key."
307
+ )
308
+ serving.online_renderer.renderer._mm_processor_cache = cache
309
+ forged_item = MultiModalKwargsItem(
310
+ {
311
+ "address": MultiModalFieldElem(data=224, field=MultiModalBatchedField()),
312
+ "monotonic_id": MultiModalFieldElem(
313
+ data=4242,
314
+ field=MultiModalBatchedField(),
315
+ ),
316
+ }
317
+ )
318
+ request = GenerateRequest(
319
+ token_ids=[1, 2, 3],
320
+ sampling_params=SamplingParams(max_tokens=1),
321
+ model=MODEL_NAME,
322
+ features=MultiModalFeatures(
323
+ mm_hashes={"image": ["image_A"]},
324
+ mm_placeholders={"image": [PlaceholderRangeInfo(offset=0, length=1)]},
325
+ kwargs_data={"image": [encode_mm_kwargs_item(forged_item)]},
326
+ ),
327
+ )
328
+
329
+ response = await serving.serve_tokens(request)
330
+
331
+ assert isinstance(response, ErrorResponse)
332
+ cache.validate_input_item.assert_called_once()
333
+ validated_item, validated_hash = cache.validate_input_item.call_args.args
334
+ assert validated_hash == "image_A"
335
+ assert validated_item["address"].data == 224
336
+ engine.generate.assert_not_called()
337
+
338
+
339
@pytest.mark.asyncio
340
async def test_stream_basic():
341
"""Streaming returns SSE chunks with correct token_ids and ends with [DONE]."""
tests/multimodal/test_cache.py
+208
@@ -29,6 +29,7 @@ from vllm.multimodal.cache.base import (
29
)
30
from vllm.multimodal.hasher import MultiModalHasher
31
from vllm.multimodal.inputs import (
32
+ MultiModalBatchedField,
33
MultiModalFeatureSpec,
34
MultiModalFieldElem,
35
MultiModalKwargsItem,
@@ -42,6 +43,7 @@ from vllm.utils.async_utils import make_async
43
from vllm.utils.mem_constants import GiB_bytes, MiB_bytes
44
from vllm.v1.engine.async_llm import AsyncLLM
45
from vllm.v1.engine.llm_engine import LLMEngine
46
+from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
47
48
from ..models.utils import build_model_context
49
from ..utils import create_new_process_for_each_test
@@ -778,6 +780,212 @@ def test_cache_eviction_shm_cache():
780
_run_test_cache_eviction_shm(sender_cache, receiver_cache, base_item_size=MiB_bytes)
781
782
783
+def test_shm_receiver_requires_sender_handle_for_current_cache_key():
784
+ vllm_config = VllmConfig(
785
+ model_config=ModelConfig(
786
+ model="llava-hf/llava-onevision-qwen2-0.5b-ov-hf",
787
+ mm_processor_cache_type="shm",
788
+ mm_shm_cache_max_object_size_mb=1,
789
+ mm_processor_cache_gb=2 * MiB_bytes / GiB_bytes,
790
+ ),
791
+ )
792
+ sender_cache = ShmObjectStoreSenderCache(vllm_config)
793
+ receiver_cache = ShmObjectStoreReceiverCache(vllm_config, mp.Lock())
794
+ mm_hash = "image_A"
795
+ item = MultiModalKwargsItem.dummy(1024)
796
+
797
+ try:
798
+ sender_item, _ = sender_cache.get_and_update_item((item, []), mm_hash)
799
+
800
+ valid_feature = MultiModalFeatureSpec(
801
+ data=sender_item,
802
+ modality="image",
803
+ identifier=mm_hash,
804
+ mm_position=PlaceholderRange(offset=0, length=1),
805
+ mm_hash=mm_hash,
806
+ )
807
+ receiver_cache.get_and_update_features([valid_feature])
808
+ assert valid_feature.data == item
809
+
810
+ forged_item = MultiModalKwargsItem(
811
+ {
812
+ "address": sender_item["address"],
813
+ "monotonic_id": sender_item["monotonic_id"],
814
+ }
815
+ )
816
+ forged_feature = MultiModalFeatureSpec(
817
+ data=forged_item,
818
+ modality="image",
819
+ identifier=mm_hash,
820
+ mm_position=PlaceholderRange(offset=0, length=1),
821
+ mm_hash=mm_hash,
822
+ )
823
+ with pytest.raises(ValueError, match="SHM handle signature"):
824
+ receiver_cache.get_and_update_features([forged_feature])
825
+
826
+ wrong_key = "image_B"
827
+ wrong_key_feature = MultiModalFeatureSpec(
828
+ data=sender_item,
829
+ modality="image",
830
+ identifier=wrong_key,
831
+ mm_position=PlaceholderRange(offset=0, length=1),
832
+ mm_hash=wrong_key,
833
+ )
834
+ with pytest.raises(ValueError, match="SHM handle signature"):
835
+ receiver_cache.get_and_update_features([wrong_key_feature])
836
+ finally:
837
+ receiver_cache._shm_cache.close()
838
+ sender_cache.close()
839
+
840
+
841
+def test_shm_sender_handle_survives_engine_msgpack_round_trip():
842
+ vllm_config = VllmConfig(
843
+ model_config=ModelConfig(
844
+ model="llava-hf/llava-onevision-qwen2-0.5b-ov-hf",
845
+ mm_processor_cache_type="shm",
846
+ mm_shm_cache_max_object_size_mb=1,
847
+ mm_processor_cache_gb=2 * MiB_bytes / GiB_bytes,
848
+ ),
849
+ )
850
+ sender_cache = ShmObjectStoreSenderCache(vllm_config)
851
+ receiver_cache = ShmObjectStoreReceiverCache(vllm_config, mp.Lock())
852
+ mm_hash = "image_A"
853
+ item = MultiModalKwargsItem.dummy(1024)
854
+
855
+ try:
856
+ sender_item, _ = sender_cache.get_and_update_item((item, []), mm_hash)
857
+ encoded = MsgpackEncoder().encode(sender_item)
858
+ decoded = MsgpackDecoder(MultiModalKwargsItem).decode(encoded)
859
+ feature = MultiModalFeatureSpec(
860
+ data=decoded,
861
+ modality="image",
862
+ identifier=mm_hash,
863
+ mm_position=PlaceholderRange(offset=0, length=1),
864
+ mm_hash=mm_hash,
865
+ )
866
+
867
+ receiver_cache.get_and_update_features([feature])
868
+
869
+ assert feature.data == item
870
+ finally:
871
+ receiver_cache._shm_cache.close()
872
+ sender_cache.close()
873
+
874
+
875
+def _small_shm_vllm_config() -> VllmConfig:
876
+ return VllmConfig(
877
+ model_config=ModelConfig(
878
+ model="llava-hf/llava-onevision-qwen2-0.5b-ov-hf",
879
+ mm_processor_cache_type="shm",
880
+ mm_shm_cache_max_object_size_mb=1,
881
+ mm_processor_cache_gb=2 * MiB_bytes / GiB_bytes,
882
+ ),
883
+ )
884
+
885
+
886
+def _with_elem(item: MultiModalKwargsItem, key: str, data) -> MultiModalKwargsItem:
887
+ elems = dict(item)
888
+ elems[key] = MultiModalFieldElem(data=data, field=MultiModalBatchedField())
889
+ return MultiModalKwargsItem(elems)
890
+
891
+
892
+def test_shm_sender_validates_client_supplied_handles():
893
+ sender_cache = ShmObjectStoreSenderCache(_small_shm_vllm_config())
894
+ mm_hash = "image_A"
895
+
896
+ try:
897
+ handle, _ = sender_cache.get_and_update_item(
898
+ (MultiModalKwargsItem.dummy(1024), []), mm_hash
899
+ )
900
+ sender_cache.validate_input_item(handle, mm_hash)
901
+ sender_cache.validate_input_item(MultiModalKwargsItem.dummy(16), mm_hash)
902
+
903
+ signature = handle["signature"].data
904
+ forged = [
905
+ _with_elem(handle, "signature", [signature[0] ^ 1, *signature[1:]]),
906
+ _with_elem(handle, "signature", [signature]),
907
+ _with_elem(handle, "address", True),
908
+ _with_elem(handle, "address", -1),
909
+ _with_elem(handle, "monotonic_id", 2**32),
910
+ MultiModalKwargsItem(
911
+ {"address": handle["address"], "monotonic_id": handle["monotonic_id"]}
912
+ ),
913
+ ]
914
+ for item in forged:
915
+ with pytest.raises(ValueError, match="SHM handle signature"):
916
+ sender_cache.validate_input_item(item, mm_hash)
917
+ with pytest.raises(ValueError, match="SHM handle signature"):
918
+ sender_cache.validate_input_item(handle, "image_B")
919
+
920
+ sender_cache.clear_cache()
921
+ with pytest.raises(ValueError, match="SHM handle signature"):
922
+ sender_cache.validate_input_item(handle, mm_hash)
923
+ finally:
924
+ sender_cache.close()
925
+
926
+
927
+def test_shm_receiver_rejects_out_of_range_handle_fields():
928
+ vllm_config = _small_shm_vllm_config()
929
+ sender_cache = ShmObjectStoreSenderCache(vllm_config)
930
+ receiver_cache = ShmObjectStoreReceiverCache(vllm_config, mp.Lock())
931
+ mm_hash = "image_A"
932
+
933
+ try:
934
+ handle, _ = sender_cache.get_and_update_item(
935
+ (MultiModalKwargsItem.dummy(1024), []), mm_hash
936
+ )
937
+ for key, value in [
938
+ ("address", -1),
939
+ ("address", 2**64),
940
+ ("monotonic_id", -1),
941
+ ("monotonic_id", 2**32),
942
+ ]:
943
+ with pytest.raises(ValueError, match="SHM handle signature"):
944
+ receiver_cache.get_and_update_item(
945
+ _with_elem(handle, key, value), mm_hash
946
+ )
947
+ finally:
948
+ receiver_cache._shm_cache.close()
949
+ sender_cache.close()
950
+
951
+
952
+def test_shm_handle_issued_before_cache_clear():
953
+ vllm_config = _small_shm_vllm_config()
954
+ sender_cache = ShmObjectStoreSenderCache(vllm_config)
955
+ receiver_cache = ShmObjectStoreReceiverCache(vllm_config, mp.Lock())
956
+ mm_hash = "image_A"
957
+ item = MultiModalKwargsItem.dummy(1024)
958
+
959
+ def feature(data: MultiModalKwargsItem) -> MultiModalFeatureSpec:
960
+ return MultiModalFeatureSpec(
961
+ data=data,
962
+ modality="image",
963
+ identifier=mm_hash,
964
+ mm_position=PlaceholderRange(offset=0, length=1),
965
+ mm_hash=mm_hash,
966
+ )
967
+
968
+ try:
969
+ stale_item, _ = sender_cache.get_and_update_item((item, []), mm_hash)
970
+ sender_cache.clear_cache()
971
+
972
+ # A request still draining after the clear resolves to its own data.
973
+ draining = feature(stale_item)
974
+ receiver_cache.get_and_update_features([draining])
975
+ assert draining.data == item
976
+
977
+ sender_cache.get_and_update_item(
978
+ (MultiModalKwargsItem.dummy(2048), []),
979
+ mm_hash,
980
+ )
981
+
982
+ with pytest.raises(ValueError, match="modified or is invalid"):
983
+ receiver_cache.get_and_update_features([feature(stale_item)])
984
+ finally:
985
+ receiver_cache._shm_cache.close()
986
+ sender_cache.close()
987
+
988
+
989
def test_processor_cache_shared_across_loras():
990
"""Test that processor cache uses mm_hash to share data across LoRAs."""
991
model_config = ModelConfig(
vllm/distributed/device_communicators/shm_object_storage.py
+146
-20
@@ -1,7 +1,10 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import hashlib
5
+import hmac
6
import pickle
7
+import secrets
8
from abc import ABC, abstractmethod
9
from collections.abc import Callable, Iterable
10
from contextlib import contextmanager, suppress
@@ -97,6 +100,9 @@ class SingleWriterShmRingBuffer:
100
- Reader synchronization handled by is_free_fn callback
101
- Writer handles garbage collection (free_buf) based on reader feedback
102
103
+ Shared Memory Layout:
104
+ `[32-byte handle-signing secret][ring buffer chunks...]`
105
+
106
Memory Layout per Buffer Chunk:
107
`[4-byte monotonic_id][4-byte chunk_size][actual_data...]`
108
^metadata_start ^data_start
@@ -119,6 +125,8 @@ class SingleWriterShmRingBuffer:
125
self.SIZE_NBYTES = 4
126
# 4 bytes for id, 4 bytes for buffer size
127
self.MD_SIZE = self.ID_NBYTES + self.SIZE_NBYTES
128
+ self.AUTH_SECRET_BYTES = 32
129
+ self.DATA_OFFSET = self.AUTH_SECRET_BYTES
130
self.monotonic_id_end = 0
131
self.monotonic_id_start = 0
132
self.data_buffer_start = 0
@@ -129,8 +137,11 @@ class SingleWriterShmRingBuffer:
137
# we are creating a buffer
138
self.metadata: dict[int, int] = {} # monotonic_id -> start address
139
self.shared_memory = shared_memory.SharedMemory(
132
- create=True, size=self.data_buffer_size, name=name
140
+ create=True,
141
+ size=self.data_buffer_size + self.DATA_OFFSET,
142
+ name=name,
143
)
144
+ self._init_auth_secret()
145
else:
146
# we are opening an existing buffer
147
# fix to https://stackoverflow.com/q/62748654/9191338
@@ -146,7 +157,9 @@ class SingleWriterShmRingBuffer:
157
# so the shared memory block size may be larger or equal
158
# to the requested size. The size parameter is ignored
159
# when attaching to an existing block.
149
- assert self.shared_memory.size >= self.data_buffer_size
160
+ assert (
161
+ self.shared_memory.size >= self.data_buffer_size + self.DATA_OFFSET
162
+ )
163
164
logger.debug(
165
"Shared memory created/opened with name: %s, size: %d",
@@ -164,11 +177,25 @@ class SingleWriterShmRingBuffer:
177
"""Clear the ring buffer."""
178
assert self.is_writer, "Only the writer can clear the buffer."
179
self.metadata.clear()
167
- self.monotonic_id_end = 0
168
- self.monotonic_id_start = 0
180
+ # Keep IDs increasing so handles issued before the clear (e.g. for
181
+ # requests still draining) fail the ID check once their chunk is
182
+ # reused, instead of resolving to a different object.
183
+ self.monotonic_id_start = self.monotonic_id_end
184
self.data_buffer_start = 0
185
self.data_buffer_end = 0
186
187
+ def get_auth_secret(self) -> bytes:
188
+ """Return the current handle-signing secret."""
189
+ assert self.shared_memory.buf is not None, "Buffer has been closed"
190
+ return bytes(self.shared_memory.buf[: self.AUTH_SECRET_BYTES])
191
+
192
+ def _init_auth_secret(self) -> None:
193
+ assert self.is_writer, "Only the writer can set the auth secret."
194
+ assert self.shared_memory.buf is not None, "Buffer has been closed"
195
+ self.shared_memory.buf[: self.AUTH_SECRET_BYTES] = secrets.token_bytes(
196
+ self.AUTH_SECRET_BYTES
197
+ )
198
+
199
def close(self) -> None:
200
"""Close the shared memory."""
201
if hasattr(self, "shared_memory"):
@@ -219,13 +246,14 @@ class SingleWriterShmRingBuffer:
246
247
# first 4 bytes as the monotonic id
248
buf_idx = self.data_buffer_end % self.data_buffer_size
222
- self.shared_memory.buf[buf_idx : buf_idx + self.ID_NBYTES] = self.int2byte(
223
- self.monotonic_id_end
249
+ physical_idx = self.DATA_OFFSET + buf_idx
250
+ self.shared_memory.buf[physical_idx : physical_idx + self.ID_NBYTES] = (
251
+ self.int2byte(self.monotonic_id_end)
252
)
253
# next 4 bytes as the size of the data buffer
226
- self.shared_memory.buf[buf_idx + self.ID_NBYTES : buf_idx + self.MD_SIZE] = (
227
- self.int2byte(size)
228
- )
254
+ self.shared_memory.buf[
255
+ physical_idx + self.ID_NBYTES : physical_idx + self.MD_SIZE
256
+ ] = self.int2byte(size)
257
258
# record metadata
259
self.metadata[self.monotonic_id_end % self.ID_MAX] = self.data_buffer_end
@@ -240,14 +268,19 @@ class SingleWriterShmRingBuffer:
268
def access_buf(self, address: int):
269
assert self.shared_memory.buf is not None, "Buffer has been closed"
270
buf_idx = address % self.data_buffer_size
271
+ physical_idx = self.DATA_OFFSET + buf_idx
272
273
# read metadata
245
- metadata_buff = self.shared_memory.buf[buf_idx : buf_idx + self.MD_SIZE]
274
+ metadata_buff = self.shared_memory.buf[
275
+ physical_idx : physical_idx + self.MD_SIZE
276
+ ]
277
id = self.byte2int(metadata_buff[: self.ID_NBYTES])
278
size = self.byte2int(metadata_buff[self.ID_NBYTES : self.MD_SIZE])
279
280
# yield the data buffer and metadata
250
- data_buff = self.shared_memory.buf[buf_idx + self.MD_SIZE : buf_idx + size]
281
+ data_buff = self.shared_memory.buf[
282
+ physical_idx + self.MD_SIZE : physical_idx + size
283
+ ]
284
with (
285
memoryview(data_buff) as data_view,
286
):
@@ -423,7 +456,8 @@ class SingleWriterShmObjectStorage:
456
457
Architecture:
458
- Single writer process can put(key, value) objects
426
- - Multiple reader processes can get(address, monotonic_id) objects
459
+ - Multiple reader processes can get(address, monotonic_id, signature, key)
460
+ objects
461
- Built on SingleWriterShmRingBuffer for efficient shared memory management
462
- Thread-safe operations with reader synchronization via locks
463
@@ -439,6 +473,13 @@ class SingleWriterShmObjectStorage:
473
- Automatic Cleanup: Garbage collection happens transparently during
474
allocation
475
476
+ Handle Signatures:
477
+ - Handles are signed with an HMAC of (key, address, monotonic_id), keyed by
478
+ a random secret stored at the start of the shared memory
479
+ - The writer signs the handles it issues and rejects stale or forged
480
+ handles supplied by requests
481
+ - Readers verify the signature in get/touch before reading shared memory
482
+
483
Memory Layout per Object:
484
`[4-byte reference_count][metadata_size][serialized_object_data]`
485
@@ -521,6 +562,76 @@ class SingleWriterShmObjectStorage:
562
else:
563
raise ValueError(f"Unsupported data type for serialization: {type(data)}")
564
565
+ def _validate_monotonic_id(
566
+ self,
567
+ address: int,
568
+ monotonic_id: int,
569
+ buf_metadata: tuple[int, int],
570
+ ) -> None:
571
+ if buf_metadata[0] != monotonic_id:
572
+ raise ValueError(
573
+ f"Data for address:id '{address}:{monotonic_id}'"
574
+ " has been modified or is invalid."
575
+ )
576
+
577
+ def _verify_reader_signature(
578
+ self,
579
+ address: int,
580
+ monotonic_id: int,
581
+ signature: list[int] | None,
582
+ key: str | None,
583
+ ) -> None:
584
+ if signature is None or key is None:
585
+ raise ValueError("Missing SHM handle signature for cache key.")
586
+ if (
587
+ not isinstance(signature, list)
588
+ or len(signature) != hashlib.sha256().digest_size
589
+ or any(type(byte) is not int or not 0 <= byte <= 255 for byte in signature)
590
+ or not isinstance(key, str)
591
+ ):
592
+ raise ValueError("Invalid SHM handle signature for cache key.")
593
+
594
+ expected_signature = self._make_signature(key, address, monotonic_id)
595
+ if not hmac.compare_digest(bytes(signature), expected_signature):
596
+ raise ValueError("Invalid SHM handle signature for cache key.")
597
+
598
+ def verify_signature(
599
+ self,
600
+ key: str | None,
601
+ address: int,
602
+ monotonic_id: int,
603
+ signature: list[int] | None,
604
+ ) -> None:
605
+ """Verify that a handle was issued for the given key.
606
+
607
+ For writers: the handle must also be the key's current entry
608
+ For readers: only the signature is checked
609
+
610
+ Args:
611
+ key: String key the handle belongs to
612
+ address: Address of the object
613
+ monotonic_id: Monotonic ID of the object
614
+ signature: Signature issued with the handle
615
+
616
+ """
617
+ if self.is_writer and (
618
+ key is None or self.key_index.get(key) != (address, monotonic_id)
619
+ ):
620
+ raise ValueError("Invalid SHM handle signature for cache key.")
621
+ self._verify_reader_signature(address, monotonic_id, signature, key)
622
+
623
+ def _make_signature(self, key: str, address: int, monotonic_id: int) -> bytes:
624
+ payload = (
625
+ hashlib.sha256(key.encode("utf-8")).digest()
626
+ + address.to_bytes(8, "little", signed=False)
627
+ + monotonic_id.to_bytes(4, "little", signed=False)
628
+ )
629
+ return hmac.new(
630
+ self.ring_buffer.get_auth_secret(),
631
+ payload,
632
+ hashlib.sha256,
633
+ ).digest()
634
+
635
def increment_writer_flag(self, id: int) -> None:
636
"""Set the in-use flag for the writer."""
637
self.writer_flag[id] = self.writer_flag.get(id, 0) + 1
@@ -555,6 +666,11 @@ class SingleWriterShmObjectStorage:
666
self.increment_writer_flag(monotonic_id)
667
return address, monotonic_id
668
669
+ def get_signature(self, key: str) -> list[int]:
670
+ """Sign the handle of a cached object so readers can verify it."""
671
+ address, monotonic_id = self.key_index[key]
672
+ return list(self._make_signature(key, address, monotonic_id))
673
+
674
def put(self, key: str, value: Any) -> tuple[int, int]:
675
"""Store a key-value pair in the object storage.
676
Attempts to free max_object_size bytes using FIFO order
@@ -605,15 +721,20 @@ class SingleWriterShmObjectStorage:
721
self.id_index[monotonic_id] = key
722
return address, monotonic_id
723
608
- def get(self, address: int, monotonic_id: int) -> Any:
724
+ def get(
725
+ self,
726
+ address: int,
727
+ monotonic_id: int,
728
+ signature: list[int] | None = None,
729
+ key: str | None = None,
730
+ ) -> Any:
731
+ if not self.is_writer:
732
+ # Reject forged handles before dereferencing the supplied address.
733
+ self.verify_signature(key, address, monotonic_id, signature)
734
+
735
# Read data from buffer
736
with self.ring_buffer.access_buf(address) as (data_view, buf_metadata):
611
- # check id from metadata
612
- if buf_metadata[0] != monotonic_id:
613
- raise ValueError(
614
- f"Data for address:id '{address}:{monotonic_id}'"
615
- " has been modified or is invalid."
616
- )
737
+ self._validate_monotonic_id(address, monotonic_id, buf_metadata)
738
739
obj = self.ser_de.deserialize(data_view[self.flag_bytes :])
740
@@ -633,6 +754,7 @@ class SingleWriterShmObjectStorage:
754
key: str,
755
address: int = 0,
756
monotonic_id: int = 0,
757
+ signature: list[int] | None = None,
758
) -> None:
759
"""Touch an existing cached item to update its eviction status.
760
@@ -643,6 +765,7 @@ class SingleWriterShmObjectStorage:
765
key: String key of the object to touch
766
address: Address of the object (only for readers)
767
monotonic_id: Monotonic ID of the object (only for readers)
768
+ signature: Server-issued handle signature (only for readers)
769
770
"""
771
if self._reader_lock is None:
@@ -652,10 +775,13 @@ class SingleWriterShmObjectStorage:
775
# Writer side: increment writer_flag to raise eviction threshold
776
self.increment_writer_flag(monotonic_id)
777
else:
778
+ # Reject forged handles before dereferencing the supplied address.
779
+ self.verify_signature(key, address, monotonic_id, signature)
780
with (
781
self._reader_lock,
657
- self.ring_buffer.access_buf(address) as (data_view, _),
782
+ self.ring_buffer.access_buf(address) as (data_view, buf_metadata),
783
):
784
+ self._validate_monotonic_id(address, monotonic_id, buf_metadata)
785
reader_count = self.ring_buffer.byte2int(data_view[: self.flag_bytes])
786
787
# NOTE(Long):
vllm/entrypoints/scale_out/token_in_token_out/serving.py
+20
@@ -39,6 +39,7 @@ from vllm.inputs import EngineInput, TokensPrompt, mm_input
39
from vllm.logger import init_logger
40
from vllm.logprobs import Logprob
41
from vllm.multimodal.inputs import (
42
+ MultiModalKwargsItem,
43
MultiModalKwargsItems,
44
PlaceholderRange,
45
)
@@ -106,6 +107,23 @@ class ServingTokens(GenerateBaseServing):
107
else getattr(mc, "override_generation_config", {}).get("max_new_tokens")
108
)
109
110
+ def _validate_mm_cache_handles(
111
+ self,
112
+ mm_kwargs: dict[str, list[MultiModalKwargsItem | None]],
113
+ mm_hashes: dict[str, list[str]],
114
+ ) -> ErrorResponse | None:
115
+ cache = self.online_renderer.renderer.mm_processor_cache
116
+ if cache is None:
117
+ return None
118
+ try:
119
+ for modality, items in mm_kwargs.items():
120
+ for mm_hash, item in zip(mm_hashes[modality], items, strict=True):
121
+ if item is not None:
122
+ cache.validate_input_item(item, mm_hash)
123
+ except ValueError as error:
124
+ return self.create_error_response(error)
125
+ return None
126
+
127
async def serve_tokens(
128
self,
129
request: GenerateRequest,
@@ -201,6 +219,8 @@ class ServingTokens(GenerateBaseServing):
219
# Deserialize full tensor data and optional metadata-only data.
220
# Metadata-only items are valid when ec_transfer_params is set.
221
mm_kwargs = mm_kwargs_from_features(features)
222
+ if error := self._validate_mm_cache_handles(mm_kwargs, features.mm_hashes):
223
+ return error
224
225
engine_input = mm_input(
226
prompt_token_ids=request.token_ids,
vllm/multimodal/cache/base.py
+8
@@ -356,6 +356,14 @@ class BaseMultiModalProcessorCache(
356
"""Close the underlying cache, if needed."""
357
pass
358
359
+ def validate_input_item(
360
+ self,
361
+ mm_item: MultiModalKwargsItem,
362
+ mm_hash: str,
363
+ ) -> None:
364
+ """Validate externally supplied cache metadata before engine handoff."""
365
+ return None
366
+
367
@abstractmethod
368
def touch_sender_cache_item(self, mm_hash: str) -> None:
369
"""Update the cache eviction order for a multi-modal item.
vllm/multimodal/cache/shm.py
+73
-12
@@ -4,7 +4,7 @@
4
5
from collections.abc import Sequence
6
from multiprocessing.synchronize import Lock as LockType
7
-from typing import TYPE_CHECKING, cast
7
+from typing import TYPE_CHECKING, TypeAlias
8
9
from typing_extensions import override
10
@@ -37,6 +37,29 @@ if TYPE_CHECKING:
37
38
logger = init_logger(__name__)
39
40
+ShmHandle: TypeAlias = tuple[int, int, list[int]]
41
+
42
+
43
+def _get_shm_handle(mm_item: MultiModalKwargsItem) -> ShmHandle | None:
44
+ if "address" not in mm_item:
45
+ return None
46
+ if "monotonic_id" not in mm_item or "signature" not in mm_item:
47
+ raise ValueError("Missing SHM handle signature for cache key.")
48
+
49
+ address = mm_item["address"].data
50
+ monotonic_id = mm_item["monotonic_id"].data
51
+ signature = mm_item["signature"].data
52
+ if (
53
+ type(address) is not int
54
+ or not 0 <= address < 2**64
55
+ or type(monotonic_id) is not int
56
+ or not 0 <= monotonic_id < 2**32
57
+ or not isinstance(signature, list)
58
+ ):
59
+ raise ValueError("Invalid SHM handle signature for cache key.")
60
+
61
+ return address, monotonic_id, signature
62
+
63
64
class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
65
"""The cache which is used on P0 when SHM caching is enabled.
@@ -98,8 +121,12 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
121
self._total += 1
122
123
address, monotonic_id = self._shm_cache.get_cached(mm_hash)
124
+ signature = self._shm_cache.get_signature(mm_hash)
125
prompt_updates = self._p0_cache[mm_hash]
102
- return self.address_as_item(address, monotonic_id), prompt_updates
126
+ return (
127
+ self.address_as_item(address, monotonic_id, signature),
128
+ prompt_updates,
129
+ )
130
131
assert mm_item is not None, f"Expected a cached item for {mm_hash=}"
132
item, prompt_updates = mm_item
@@ -108,12 +135,16 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
135
136
try:
137
address, monotonic_id = self._shm_cache.put(mm_hash, item)
138
+ signature = self._shm_cache.get_signature(mm_hash)
139
# Try to remove dangling items if p0 cache is too large.
140
if len(self._p0_cache) >= 2 * len(self._shm_cache.key_index):
141
self.remove_dangling_items()
142
143
self._p0_cache[mm_hash] = prompt_updates
116
- return self.address_as_item(address, monotonic_id), prompt_updates
144
+ return (
145
+ self.address_as_item(address, monotonic_id, signature),
146
+ prompt_updates,
147
+ )
148
except ValueError as e:
149
# `put` raises ValueError either for an oversize item or for a
150
# duplicate key (concurrent insert); the latter is benign so we
@@ -143,6 +174,21 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
174
Increments writer_flag on sender side."""
175
self._shm_cache.touch(mm_hash)
176
177
+ @override
178
+ def validate_input_item(
179
+ self,
180
+ mm_item: MultiModalKwargsItem,
181
+ mm_hash: str,
182
+ ) -> None:
183
+ if (handle := _get_shm_handle(mm_item)) is not None:
184
+ address, monotonic_id, signature = handle
185
+ self._shm_cache.verify_signature(
186
+ mm_hash,
187
+ address,
188
+ monotonic_id,
189
+ signature,
190
+ )
191
+
192
@override
193
def clear_cache(self) -> None:
194
self._shm_cache.clear()
@@ -171,6 +217,7 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
217
self,
218
address: int,
219
monotonic_id: int,
220
+ signature: list[int],
221
) -> MultiModalKwargsItem:
222
addr_elem = MultiModalFieldElem(
223
data=address,
@@ -181,7 +228,18 @@ class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
228
field=MultiModalBatchedField(),
229
)
230
184
- return MultiModalKwargsItem({"address": addr_elem, "monotonic_id": id_elem})
231
+ signature_elem = MultiModalFieldElem(
232
+ data=signature,
233
+ field=MultiModalBatchedField(),
234
+ )
235
+
236
+ return MultiModalKwargsItem(
237
+ {
238
+ "address": addr_elem,
239
+ "monotonic_id": id_elem,
240
+ "signature": signature_elem,
241
+ }
242
+ )
243
244
245
class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache):
@@ -235,10 +293,9 @@ class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache):
293
mm_hash: str,
294
) -> MultiModalKwargsItem:
295
assert mm_item is not None, f"Expected an address item for {mm_hash=}"
238
- if "address" in mm_item:
239
- address = cast(int, mm_item["address"].data)
240
- monotonic_id = cast(int, mm_item["monotonic_id"].data)
241
- return self._shm_cache.get(address, monotonic_id)
296
+ if (handle := _get_shm_handle(mm_item)) is not None:
297
+ address, monotonic_id, signature = handle
298
+ return self._shm_cache.get(address, monotonic_id, signature, mm_hash)
299
300
return mm_item
301
@@ -251,10 +308,14 @@ class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache):
308
"""Touch the item in shared memory cache to prevent eviction.
309
Increments reader_count on receiver side."""
310
assert mm_item is not None
254
- if "address" in mm_item:
255
- address = cast(int, mm_item["address"].data)
256
- monotonic_id = cast(int, mm_item["monotonic_id"].data)
257
- self._shm_cache.touch(mm_hash, address=address, monotonic_id=monotonic_id)
311
+ if (handle := _get_shm_handle(mm_item)) is not None:
312
+ address, monotonic_id, signature = handle
313
+ self._shm_cache.touch(
314
+ mm_hash,
315
+ address=address,
316
+ monotonic_id=monotonic_id,
317
+ signature=signature,
318
+ )
319
320
@override
321
def clear_cache(self) -> None: