@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a5b84f1cbf
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Core] Shared memory based object store for Multimodal data caching and IPC (#20452)

Signed-off-by: donglu <donglu@cohere.com>

dongluw committed Sep 12, 2025 at 10:54 UTC a5b84f1cbfeb4030fd7fab4a91d4264b356e20b0
17 files changed +1487 -27
.buildkite/test-pipeline.yaml
+2
@@ -789,6 +789,8 @@ steps:
789 commands:
790 - pytest -v -s distributed/test_comm_ops.py
791 - pytest -v -s distributed/test_shm_broadcast.py
792 + - pytest -v -s distributed/test_shm_buffer.py
793 + - pytest -v -s distributed/test_shm_storage.py
794
795 - label: 2 Node Tests (4 GPUs in total) # 16min
796 timeout_in_minutes: 30
docs/configuration/optimization.md
+26 -5
@@ -230,6 +230,20 @@ Multi-modal IPC caching is automatically enabled when
230 there is a one-to-one correspondence between API (`P0`) and engine core (`P1`) processes,
231 to avoid repeatedly transferring the same multi-modal inputs between them.
232
233 +#### Key-Replicated Cache
234 +
235 +By default, IPC caching uses a **key-replicated cache**, where cache keys exist
236 +in both the API (`P0`) and engine core (`P1`) processes, but the actual cache
237 +data resides only in `P1`.
238 +
239 +#### Shared Memory Cache
240 +
241 +When multiple worker processes are involved (e.g., when TP > 1), a
242 +**shared-memory cache** is more efficient. This can be enabled by setting
243 +`mm_processor_cache_type="shm"`. In this mode, cache keys are stored
244 +on `P0`, while the cache data itself lives in shared memory accessible by all
245 +processes.
246 +
247 ### Configuration
248
249 You can adjust the size of the cache by setting the value of `mm_processor_cache_gb` (default 4 GiB).
@@ -244,6 +258,12 @@ Examples:
258 llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct",
259 mm_processor_cache_gb=8)
260
261 +# Use a shared-memory based IPC cache
262 +llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct",
263 + tensor_parallel_size=2,
264 + mm_processor_cache_type="shm",
265 + mm_processor_cache_gb=8)
266 +
267 # Disable the cache
268 llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct",
269 mm_processor_cache_gb=0)
@@ -253,11 +273,12 @@ llm = LLM(model="Qwen/Qwen2.5-VL-3B-Instruct",
273
274 Based on the configuration, the content of the multi-modal caches on `P0` and `P1` are as follows:
275
256 -| Processor Caching | IPC Caching | `P0` Cache | `P1` Cache | Max. Memory |
257 -|-------------------|-------------|------------|------------|-------------|
258 -| ✅ | ✅ | K | K + V | `mm_processor_cache_gb * data_parallel_size` |
259 -| ✅ | ❌ | K + V | N/A | `mm_processor_cache_gb * api_server_count` |
260 -| ❌ | ❌ | N/A | N/A | `0` |
276 +| mm_processor_cache_type | Cache Type | `P0` Cache | `P1` Engine Cache | `P1` Worker Cache | Max. Memory |
277 +|-------------------|-------------|------------|------------|-------------|-------------|
278 +| lru | Processor Caching | K + V | N/A | N/A | `mm_processor_cache_gb * data_parallel_size` |
279 +| lru | Key-Replicated Caching | K | K + V | N/A | `mm_processor_cache_gb * api_server_count` |
280 +| shm | Shared Memory Caching | K | N/A | V | `mm_processor_cache_gb * api_server_count` |
281 +| N/A | Disabled | N/A | N/A | N/A | `0` |
282
283 K: Stores the hashes of multi-modal items
284 V: Stores the processed tensor data of multi-modal items
tests/distributed/test_shm_buffer.py new
+172
@@ -0,0 +1,172 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import traceback
5 +import unittest
6 +
7 +from vllm.distributed.device_communicators.shm_object_storage import (
8 + SingleWriterShmRingBuffer)
9 +
10 +
11 +class TestSingleWriterShmRingBuffer(unittest.TestCase):
12 + """Test suite for the ring buffer implementation"""
13 +
14 + def setUp(self):
15 + """Set up test fixtures"""
16 + self.buffer_size = 4096
17 + self.ring_buffer = None
18 +
19 + def tearDown(self):
20 + """Clean up after tests"""
21 + if self.ring_buffer:
22 + del self.ring_buffer
23 +
24 + def test_buffer_opening(self):
25 + """Test opening an existing buffer"""
26 + # First create a buffer
27 + self.ring_buffer = SingleWriterShmRingBuffer(
28 + data_buffer_size=self.buffer_size, create=True)
29 +
30 + # Then open it with another instance
31 + reader_buffer = SingleWriterShmRingBuffer(*self.ring_buffer.handle())
32 + self.assertFalse(reader_buffer.is_writer)
33 + self.assertEqual(reader_buffer.shared_memory.name,
34 + self.ring_buffer.shared_memory.name)
35 +
36 + def test_buffer_access(self):
37 + """Test accessing allocated buffers"""
38 + self.ring_buffer = SingleWriterShmRingBuffer(
39 + data_buffer_size=self.buffer_size, create=True)
40 +
41 + size = 100
42 + address, monotonic_id = self.ring_buffer.allocate_buf(size)
43 +
44 + # Write some test data
45 + test_data = b"Hello, World!" * 7 # 91 bytes
46 + with self.ring_buffer.access_buf(address) as (data_buf, metadata):
47 + data_buf[0:len(test_data)] = test_data
48 +
49 + # Read it back
50 + with self.ring_buffer.access_buf(address) as (data_buf2, metadata2):
51 + read_data = bytes(data_buf2[0:len(test_data)])
52 + read_id = metadata2[0]
53 +
54 + self.assertEqual(read_data, test_data)
55 + self.assertEqual(read_id, monotonic_id)
56 +
57 + def test_memory_error_on_full_buffer(self):
58 + """Test that MemoryError is raised when buffer is full"""
59 + small_buffer_size = 200
60 + self.ring_buffer = SingleWriterShmRingBuffer(
61 + data_buffer_size=small_buffer_size, create=True)
62 +
63 + # Fill up the buffer
64 + self.ring_buffer.allocate_buf(100)
65 + self.ring_buffer.allocate_buf(80) # Total: 196 bytes used
66 +
67 + # This should fail
68 + with self.assertRaises(MemoryError):
69 + self.ring_buffer.allocate_buf(1) # Would exceed buffer capacity
70 +
71 + def test_allocation_and_free(self):
72 + """Test allocation and freeing of buffers"""
73 + small_buffer_size = 200
74 + self.ring_buffer = SingleWriterShmRingBuffer(
75 + data_buffer_size=small_buffer_size, create=True)
76 +
77 + size = 80
78 + # Write some data
79 + test_data = b"Repeated test data"
80 + for i in range(5):
81 + address, monotonic_id = self.ring_buffer.allocate_buf(size)
82 + with self.ring_buffer.access_buf(address) as (data_buf, metadata):
83 + data_buf[0:4] = (0).to_bytes(4, "little") # 0 for not in-use
84 + data_buf[4:len(test_data) + 4] = test_data
85 + print(self.ring_buffer.metadata)
86 + freed_ids = self.ring_buffer.free_buf(lambda *args: True)
87 + print(f" Freed IDs: {freed_ids}")
88 + self.assertEqual(freed_ids[0], i)
89 +
90 + def test_clear_buffer(self):
91 + """Test clearing the buffer"""
92 + self.ring_buffer = SingleWriterShmRingBuffer(
93 + data_buffer_size=self.buffer_size, create=True)
94 +
95 + # Allocate some buffers
96 + for _ in range(3):
97 + self.ring_buffer.allocate_buf(100)
98 +
99 + # Clear the buffer
100 + self.ring_buffer.clear()
101 +
102 + # Check that metadata is empty and IDs reset
103 + self.assertEqual(len(self.ring_buffer.metadata), 0)
104 + self.assertEqual(self.ring_buffer.monotonic_id_start, 0)
105 + self.assertEqual(self.ring_buffer.monotonic_id_end, 0)
106 + self.assertEqual(self.ring_buffer.data_buffer_start, 0)
107 + self.assertEqual(self.ring_buffer.data_buffer_end, 0)
108 +
109 +
110 +def main():
111 + """Main function demonstrating usage and running tests"""
112 + print("=== SingleWriterShmRingBuffer Test Suite ===\n")
113 +
114 + # Run unit tests
115 + print("Running unit tests...")
116 + unittest.main(argv=[""], exit=False, verbosity=2)
117 +
118 + print("\n" + "=" * 50)
119 + print("=== Manual Demo ===\n")
120 +
121 + # Manual demonstration
122 + try:
123 + print("Creating ring buffer...")
124 + writer_buffer = SingleWriterShmRingBuffer(data_buffer_size=2048,
125 + create=True)
126 + reader_buffer = SingleWriterShmRingBuffer(*writer_buffer.handle())
127 +
128 + print(f"Buffer created with name: {writer_buffer.shared_memory.name}")
129 +
130 + # Allocate some buffers
131 + print("\nAllocating buffers...")
132 + address_array = []
133 + for i in range(3):
134 + size = 100 + i * 50
135 + try:
136 + writer_buffer.free_buf(lambda *args: True)
137 + address, monotonic_id = writer_buffer.allocate_buf(size)
138 + address_array.append((address, size, monotonic_id))
139 +
140 + # Write some test data
141 + with writer_buffer.access_buf(address) as (data_buf, metadata):
142 + test_message = f"Test message {i}".encode()
143 + data_buf[0:len(test_message)] = test_message
144 +
145 + except MemoryError as e:
146 + print(f" Failed to allocate {size} bytes: {e}")
147 +
148 + print("\nBuffer state:")
149 + print(f" Data buffer start: {writer_buffer.data_buffer_start}")
150 + print(f" Data buffer end: {writer_buffer.data_buffer_end}")
151 + print(f" Monotonic ID start: {writer_buffer.monotonic_id_start}")
152 + print(f" Monotonic ID end: {writer_buffer.monotonic_id_end}")
153 + print(f" Metadata entries: {len(writer_buffer.metadata)}")
154 +
155 + # Try to read back the data
156 + print("\nReading back data...")
157 + for address, size, monotonic_id in address_array:
158 + with reader_buffer.access_buf(address) as (data_buf, metadata):
159 + # Find null terminator or read first 50 chars
160 + data_bytes = bytes(data_buf[0:size])
161 + message = data_bytes.decode()
162 + print(f" ID {monotonic_id}: '{message}'")
163 +
164 + except Exception as e:
165 + print(f"Demo error: {e}")
166 + traceback.print_exc()
167 +
168 + print("\n=== Demo Complete ===")
169 +
170 +
171 +if __name__ == "__main__":
172 + main()
tests/distributed/test_shm_storage.py new
+327
@@ -0,0 +1,327 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import multiprocessing
5 +import random
6 +import time
7 +import traceback
8 +import unittest
9 +from multiprocessing import Lock
10 +
11 +import torch
12 +
13 +# Assuming these are imported from your module
14 +from vllm.distributed.device_communicators.shm_object_storage import (
15 + MsgpackSerde, SingleWriterShmObjectStorage, SingleWriterShmRingBuffer)
16 +from vllm.multimodal.inputs import (MultiModalFieldElem, MultiModalKwargsItem,
17 + MultiModalSharedField)
18 +
19 +
20 +def _dummy_elem(modality: str, key: str, size: int):
21 + return MultiModalFieldElem(
22 + modality=modality,
23 + key=key,
24 + data=torch.empty((size, ), dtype=torch.int8),
25 + field=MultiModalSharedField(1),
26 + )
27 +
28 +
29 +def _dummy_item(modality: str, size_by_key: dict[str, int]):
30 + return MultiModalKwargsItem.from_elems([
31 + _dummy_elem(modality, key, size) for key, size in size_by_key.items()
32 + ])
33 +
34 +
35 +class TestSingleWriterShmObjectStorage(unittest.TestCase):
36 +
37 + def setUp(self):
38 + """Set up test fixtures before each test method."""
39 + ring_buffer = SingleWriterShmRingBuffer(
40 + data_buffer_size=1024 * 100,
41 + create=True, # 10 MB buffer
42 + )
43 + self.storage = SingleWriterShmObjectStorage(
44 + max_object_size=1024 * 10, # 10KB max object
45 + n_readers=2,
46 + ring_buffer=ring_buffer,
47 + serde_class=MsgpackSerde,
48 + reader_lock=Lock(),
49 + )
50 +
51 + def tearDown(self):
52 + """Clean up after each test."""
53 + if self.storage:
54 + del self.storage
55 +
56 + def test_minimal_put_get_cycle(self):
57 + """Test basic put and get operations."""
58 + key = "test_key"
59 + value = _dummy_item("text", {"field1": 10, "field2": 20})
60 +
61 + # Put operation
62 + address, monotonic_id = self.storage.put(key, value)
63 +
64 + # Verify key is in index
65 + self.assertIn(key, self.storage.key_index)
66 + self.assertEqual(self.storage.key_index[key], (address, monotonic_id))
67 + self.assertEqual(self.storage.id_index[monotonic_id], key)
68 +
69 + # Get operation
70 + result = self.storage.get(address, monotonic_id)
71 +
72 + # Verify result
73 + self.assertEqual(result, value)
74 +
75 + def test_put_same_key_twice(self):
76 + """Test behavior when putting the same key multiple times."""
77 + key = "duplicate_key"
78 + value1 = "first value"
79 + value2 = "second value"
80 +
81 + # First put
82 + address1, id1 = self.storage.put(key, value1)
83 + retrieved1 = self.storage.get(address1, id1)
84 + self.assertEqual(retrieved1, value1)
85 +
86 + # should raise an error on second put
87 + with self.assertRaises(ValueError) as context:
88 + self.storage.put(key, value2)
89 +
90 + self.assertIn("already exists in the storage", str(context.exception))
91 +
92 + def test_large_object_rejection(self):
93 + """Test that objects exceeding max_object_size are rejected."""
94 + # Create an object larger than max_object_size
95 + large_data = "x" * (self.storage.max_object_size + 100)
96 +
97 + with self.assertRaises(ValueError) as context:
98 + self.storage.put("large_key", large_data)
99 +
100 + self.assertIn("exceeds max object size", str(context.exception))
101 +
102 + def test_buffer_overflow_and_cleanup(self):
103 + """Test behavior when buffer fills up and needs cleanup."""
104 + # Fill up the buffer with many small objects
105 + stored_items = []
106 +
107 + try:
108 + for i in range(1000): # Try to store many items
109 + key = f"item_{i}"
110 + value = f"data_{i}" * 100 # Make it reasonably sized
111 + address, monotonic_id = self.storage.put(key, value)
112 + stored_items.append((key, value, address, monotonic_id))
113 + except MemoryError:
114 + print(f"Buffer filled after {len(stored_items)} items")
115 +
116 + # Verify that some items are still accessible
117 + accessible_count = 0
118 + for key, original_value, address, monotonic_id in stored_items:
119 + for i in range(self.storage.n_readers):
120 + retrieved = self.storage.get(address, monotonic_id)
121 + if retrieved == original_value:
122 + accessible_count += 1
123 +
124 + self.assertEqual(accessible_count, len(stored_items))
125 +
126 + try:
127 + for i in range(len(stored_items), 1000): # Try to store many items
128 + key = f"item_{i}"
129 + value = f"data_{i}" * 100 # Make it reasonably sized
130 + address, monotonic_id = self.storage.put(key, value)
131 + stored_items.append((key, value, address, monotonic_id))
132 + except MemoryError:
133 + print(f"Buffer filled after {len(stored_items)} items")
134 +
135 + # Verify that some items are still accessibles
136 + for key, original_value, address, monotonic_id in stored_items:
137 + try:
138 + for i in range(self.storage.n_readers):
139 + retrieved = self.storage.get(address, monotonic_id)
140 + if retrieved == original_value:
141 + accessible_count += 1
142 + except ValueError as e:
143 + print(f"Error retrieving {key}: {e}")
144 +
145 + # some items from the first batch may still be accessible
146 + self.assertGreaterEqual(accessible_count, len(stored_items))
147 +
148 + def test_blocking_unread_object(self):
149 + """Test behavior when buffer fills up and needs cleanup."""
150 + # Fill up the buffer with many small objects
151 + stored_items = []
152 +
153 + try:
154 + for i in range(1000): # Try to store many items
155 + key = f"item_{i}"
156 + value = f"data_{i}" * 100 # Make it reasonably sized
157 + address, monotonic_id = self.storage.put(key, value)
158 + stored_items.append((key, value, address, monotonic_id))
159 + except MemoryError:
160 + print(f"Buffer filled after {len(stored_items)} items")
161 +
162 + # read all items except the first one
163 + # to simulate a blocking situation
164 + accessible_count = 0
165 + for key, original_value, address, monotonic_id in stored_items[1:]:
166 + for i in range(self.storage.n_readers):
167 + retrieved = self.storage.get(address, monotonic_id)
168 + if retrieved == original_value:
169 + accessible_count += 1
170 +
171 + self.assertEqual(accessible_count, len(stored_items) - 1)
172 +
173 + try:
174 + key = f"item_{len(stored_items)}"
175 + value = f"data_{len(stored_items)}" * 100
176 + address, monotonic_id = self.storage.put(key, value)
177 + except MemoryError:
178 + print(f"Buffer filled after {len(stored_items)} items")
179 +
180 + # read the first item
181 + for i in range(self.storage.n_readers):
182 + key, original_value, address, monotonic_id = stored_items[0]
183 + retrieved = self.storage.get(address, monotonic_id)
184 + self.assertEqual(retrieved, original_value)
185 +
186 + try:
187 + for i in range(len(stored_items), 1000): # Try to store many items
188 + key = f"item_{i}"
189 + value = f"data_{i}" * 100 # Make it reasonably sized
190 + address, monotonic_id = self.storage.put(key, value)
191 + stored_items.append((key, value, address, monotonic_id))
192 + except MemoryError:
193 + print(f"Buffer filled after {len(stored_items)} items")
194 +
195 + # some items from the first batch may still be accessible
196 + self.assertGreaterEqual(len(stored_items), accessible_count + 10)
197 +
198 + def test_invalid_get_operations(self):
199 + """Test various invalid get operations."""
200 + # Test with non-existent address
201 + with self.assertRaises(ValueError): # Could be various exceptions
202 + self.storage.get(99999, 1)
203 +
204 + # Store something first
205 + address, monotonic_id = self.storage.put("test", "value")
206 +
207 + # Test with wrong monotonic_id
208 + with self.assertRaises(ValueError) as context:
209 + self.storage.get(address, monotonic_id + 100)
210 +
211 + self.assertIn("has been modified or is invalid", \
212 + str(context.exception))
213 +
214 + def test_clear_storage(self):
215 + """Test clearing the storage."""
216 + # Store some items
217 + for i in range(5):
218 + self.storage.put(f"item_{i}", f"value_{i}")
219 +
220 + # Clear the storage
221 + self.storage.clear()
222 +
223 + # Verify that all indices are empty
224 + self.assertEqual(len(self.storage.key_index), 0)
225 + self.assertEqual(len(self.storage.id_index), 0)
226 + self.assertEqual(len(self.storage.ring_buffer.metadata), 0)
227 +
228 + # Verify that new items can be added after clearing
229 + address, monotonic_id = self.storage.put("new_item", "new_value")
230 + self.assertIn("new_item", self.storage.key_index)
231 + self.assertEqual((address, monotonic_id), (0, 0))
232 +
233 +
234 +# Reader process function
235 +def reader_process(process_id, storage_handle, items_to_read):
236 + """Reader process that connects to existing shared memory and reads data."""
237 + reader_storage = SingleWriterShmObjectStorage.create_from_handle(
238 + storage_handle)
239 +
240 + print(f"Reader {process_id} started")
241 +
242 + errors = []
243 +
244 + for key, original_value, address, monotonic_id in items_to_read:
245 + time.sleep(random.random() / 100)
246 + try:
247 + # Read data from shared memory
248 + retrieved_value = reader_storage.get(address, monotonic_id)
249 +
250 + # Verify data integrity
251 + assert retrieved_value == original_value
252 + print(f"Reader {process_id} retrieved {key}: {retrieved_value}")
253 + except Exception as e:
254 + errors.append((key, str(e), type(e).__name__))
255 +
256 +
257 +def run_multiprocess_example():
258 + """Run a minimal working example with real shared memory."""
259 + print("=== Minimal Object Storage Example ===")
260 +
261 + try:
262 + # Create storage instance
263 + ring_buffer = SingleWriterShmRingBuffer(
264 + data_buffer_size=1024 * 100,
265 + create=True, # 10 MB buffer
266 + )
267 + storage = SingleWriterShmObjectStorage(
268 + max_object_size=1024,
269 + n_readers=3,
270 + ring_buffer=ring_buffer,
271 + serde_class=MsgpackSerde,
272 + reader_lock=Lock(),
273 + )
274 +
275 + print(f"Created storage (writer: {storage.is_writer})")
276 +
277 + # Test basic data types
278 + test_data = [
279 + ("user_data", {
280 + "name": "Alice",
281 + "age": 30,
282 + "scores": [95, 87, 92]
283 + }),
284 + ("simple_string", "Hello, World!"),
285 + ("number", 42),
286 + ("list_data", [1, 2, 3, "four", 5.0]),
287 + ]
288 +
289 + stored_items = []
290 +
291 + # Store all data
292 + for key, value in test_data:
293 + print(f"Storing {key}: {value}")
294 + address, monotonic_id = storage.put(key, value)
295 + stored_items.append((key, value, address, monotonic_id))
296 + print(f" -> Stored at address {address}, ID {monotonic_id}")
297 +
298 + print("\n--- Retrieving Data ---")
299 + processes = []
300 + handle = storage.handle()
301 + # initialize lock for reader processes
302 + handle.reader_lock = Lock()
303 + for i in range(storage.n_readers):
304 + p = multiprocessing.Process(target=reader_process,
305 + args=(i, handle, stored_items))
306 + processes.append(p)
307 + p.start()
308 +
309 + for p in processes:
310 + p.join(timeout=10)
311 + if p.is_alive():
312 + p.terminate()
313 + p.join()
314 +
315 + except Exception as e:
316 + print(f"Error in minimal example: {e}")
317 + traceback.print_exc()
318 +
319 +
320 +if __name__ == "__main__":
321 + # Run the minimal example first
322 + run_multiprocess_example()
323 + print("\n" + "=" * 50 + "\n")
324 +
325 + # Run the test suite
326 + print("Running comprehensive test suite...")
327 + unittest.main(verbosity=2, exit=False)
tests/multimodal/test_cache.py
+4 -4
@@ -10,8 +10,8 @@ from vllm.config import ModelConfig, ParallelConfig, VllmConfig
10 from vllm.multimodal.cache import (MultiModalCache,
11 MultiModalProcessorCacheItem,
12 MultiModalProcessorCacheItemMetadata,
13 - processor_cache_from_config,
14 - receiver_cache_from_config)
13 + engine_receiver_cache_from_config,
14 + processor_cache_from_config)
15 from vllm.multimodal.hasher import MultiModalHasher
16 from vllm.multimodal.inputs import (MultiModalFieldElem, MultiModalKwargsItem,
17 MultiModalKwargsItems,
@@ -115,9 +115,9 @@ def _compare_caches(
115 ):
116 mm_registry = MultiModalRegistry()
117 cache_0_p0 = processor_cache_from_config(config_0, mm_registry)
118 - cache_0_p1 = receiver_cache_from_config(config_0, mm_registry)
118 + cache_0_p1 = engine_receiver_cache_from_config(config_0, mm_registry)
119 cache_1_p0 = processor_cache_from_config(config_1, mm_registry)
120 - cache_1_p1 = receiver_cache_from_config(config_1, mm_registry)
120 + cache_1_p1 = engine_receiver_cache_from_config(config_1, mm_registry)
121
122 cache_size_gb = max(
123 config_0.model_config.mm_processor_cache_gb,
tests/tensorizer_loader/conftest.py
+1
@@ -90,6 +90,7 @@ class DummyExecutor(UniProcExecutor):
90 distributed_init_method=distributed_init_method,
91 is_driver_worker=is_driver_worker,
92 )
93 + self.mm_receiver_cache = None
94 self.collective_rpc("init_worker", args=([kwargs], ))
95 self.collective_rpc("init_device")
96
tools/check_pickle_imports.py
+1
@@ -39,6 +39,7 @@ ALLOWED_FILES = set([
39 'vllm/engine/multiprocessing/client.py',
40 'vllm/distributed/device_communicators/all_reduce_utils.py',
41 'vllm/distributed/device_communicators/shm_broadcast.py',
42 + 'vllm/distributed/device_communicators/shm_object_storage.py',
43 'vllm/engine/multiprocessing/engine.py',
44 'benchmarks/kernels/graph_machete_bench.py',
45 'benchmarks/kernels/benchmark_lora.py',
vllm/config/__init__.py
+20
@@ -262,6 +262,7 @@ def is_init_field(cls: ConfigType, name: str) -> bool:
262 TokenizerMode = Literal["auto", "slow", "mistral", "custom"]
263 ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]
264 MMEncoderTPMode = Literal["weights", "data"]
265 +MMCacheType = Literal["shm", "lru"]
266
267
268 class LogprobsMode(enum.Enum):
@@ -450,6 +451,13 @@ class ModelConfig:
451 `mm_processor_cache_gb * (api_server_count + data_parallel_size)`.
452
453 Set to `0` to disable this cache completely (not recommended)."""
454 + mm_processor_cache_type: MMCacheType = "lru"
455 + """Type of cache to use for the multi-modal preprocessor/mapper. If `shm`,
456 + use shared memory FIFO cache. If `lru`, use mirrored LRU cache."""
457 + mm_shm_cache_max_object_size_mb: int = 128
458 + """Size limit (in MiB) for each object stored in the multi-modal processor
459 + shared memory cache. Only effective when `mm_processor_cache_type` is
460 + `"shm"`."""
461 mm_encoder_tp_mode: MMEncoderTPMode = "weights"
462 """Indicates how to optimize multi-modal encoder inference using
463 tensor parallelism (TP).
@@ -881,6 +889,9 @@ class ModelConfig:
889 media_io_kwargs=self.media_io_kwargs,
890 mm_processor_kwargs=self.mm_processor_kwargs,
891 mm_processor_cache_gb=self.mm_processor_cache_gb,
892 + mm_processor_cache_type=self.mm_processor_cache_type,
893 + mm_shm_cache_max_object_size_mb=self.
894 + mm_shm_cache_max_object_size_mb,
895 mm_encoder_tp_mode=self.mm_encoder_tp_mode,
896 interleave_mm_strings=self.interleave_mm_strings,
897 skip_mm_profiling=self.skip_mm_profiling,
@@ -2448,6 +2459,15 @@ class MultiModalConfig:
2459 Set to `0` to disable this cache completely (not recommended).
2460 """
2461
2462 + mm_processor_cache_type: MMCacheType = "lru"
2463 + """Type of cache to use for the multi-modal preprocessor/mapper. If `shm`,
2464 + use shared memory FIFO cache. If `lru`, use mirrored LRU cache."""
2465 +
2466 + mm_shm_cache_max_object_size_mb: int = 128
2467 + """Size limit (in MiB) for each object stored in the multi-modal processor
2468 + shared memory cache. Only effective when `mm_processor_cache_type` is
2469 + `"shm"`."""
2470 +
2471 mm_encoder_tp_mode: MMEncoderTPMode = "weights"
2472 """
2473 Indicates how to optimize multi-modal encoder inference using
vllm/distributed/device_communicators/shm_object_storage.py new
+635
@@ -0,0 +1,635 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pickle
5 +from abc import ABC, abstractmethod
6 +from collections.abc import Iterable
7 +from contextlib import contextmanager
8 +from dataclasses import dataclass
9 +from itertools import chain
10 +from multiprocessing import shared_memory
11 +from multiprocessing.synchronize import Lock as LockType
12 +from typing import Any, Callable, Optional, Union
13 +from unittest.mock import patch
14 +
15 +import torch
16 +
17 +from vllm.logger import init_logger
18 +
19 +logger = init_logger(__name__)
20 +
21 +
22 +class SingleWriterShmRingBuffer:
23 + """
24 + A single-writer, multiple-reader ring buffer implementation using shared
25 + memory. This class provides a thread-safe ring buffer where one process
26 + can write data while multiple processes/threads can read from it.
27 +
28 + Architecture:
29 + - Uses shared memory for cross-process communication
30 + - Maintains metadata for each allocated buffer chunk in the writer process
31 + - Supports custom "is_free_fn" functions to determine when buffers can be
32 + reused
33 + - Each buffer chunk contains: [4-byte id][4-byte size][actual_data]
34 +
35 + Key Concepts:
36 + - monotonic_id_start/end: Track the range of active buffer IDs
37 + - data_buffer_start/end: Track the physical memory range in use
38 + - Automatic wraparound when reaching buffer end
39 + - Lazy garbage collection based on is_free_fn checks
40 +
41 + Example Usage Scenarios:
42 +
43 + Scenario 1: Simple Linear Allocation
44 + ```
45 + Buffer size: 100 bytes
46 + Initial state: [................................................. ]
47 + ^start=end(0)
48 +
49 + After allocating 20 bytes (id=0):
50 + [id:0|size:20|data........][...................................]
51 + ^start(0) ^end(28)
52 +
53 + After allocating 30 bytes (id=1):
54 + [id:0|size:20|data........][id:1|size:30|data..............][..]
55 + ^start(0) ^end(66)
56 + ```
57 +
58 + Scenario 2: Memory Reclamation
59 + ```
60 + Before freeing (both buffers still in use):
61 + [id:0|size:20|data........][id:1|size:30|data..............][..]
62 + ^start(0) ^end(66)
63 +
64 + After id:0 is marked free by readers:
65 + [FREED.................... ][id:1|size:30|data..............][..]
66 + ^start(28) ^end(66)
67 +
68 + After both are freed:
69 + [FREED..............................................][..]
70 + ^start=end(66)
71 + ```
72 +
73 + Scenario 3: Wraparound Allocation (continuing from Scenario 2)
74 + ```
75 + Starting from after memory reclamation in Scenario 2:
76 + [FREED..............................................][..]
77 + ^start=end(66)
78 +
79 + Allocate 40 bytes (id=2) - only 34 bytes available at end, so wraparound:
80 + [id:2|size:40|data........................][FREED.............][..]
81 + ^end(148) ^start(66)
82 + ```
83 +
84 + Scenario 4: Error Handling - Out of Space
85 + ```
86 + Starting from after wraparound allocation in Scenario 3:
87 + [id:2|size:40|data........................][FREED.............][..]
88 + ^end(148) ^start(66)
89 +
90 + Trying to allocate 20 more bytes:
91 + occupied_size_new = end + size - start = 148 + 28 - 66 > buffer_size(100)
92 + -> Raises MemoryError: "Not enough space in the data buffer"
93 + ```
94 +
95 + Thread Safety:
96 + - Single writer: Only one process/thread should write (allocate_buf)
97 + - Multiple readers: Multiple processes/threads can read (access_buf)
98 + - Reader synchronization handled by is_free_fn callback
99 + - Writer handles garbage collection (free_buf) based on reader feedback
100 +
101 + Memory Layout per Buffer Chunk:
102 + [4-byte monotonic_id][4-byte chunk_size][actual_data...]
103 + ^metadata_start ^data_start
104 +
105 + The monotonic_id ensures data integrity - readers can verify they're
106 + accessing the correct data even after buffer wraparound or reuse.
107 + """
108 +
109 + def __init__(
110 + self,
111 + data_buffer_size: int,
112 + name: Optional[str] = None,
113 + create: bool = False,
114 + ):
115 + self.data_buffer_size = data_buffer_size
116 + self.is_writer = create
117 +
118 + self.ID_NBYTES = 4
119 + self.ID_MAX = 2**31 # exclusive, so 2**31 - 1 is the max value
120 + self.SIZE_NBYTES = 4
121 + # 4 bytes for id, 4 bytes for buffer size
122 + self.MD_SIZE = self.ID_NBYTES + self.SIZE_NBYTES
123 + self.monotonic_id_end = 0
124 + self.monotonic_id_start = 0
125 + self.data_buffer_start = 0
126 + self.data_buffer_end = 0
127 +
128 + if create:
129 + # we are creating a buffer
130 + self.metadata = {
131 + self.monotonic_id_end: self.data_buffer_end
132 + } # monotonic_id -> start address
133 + self.shared_memory = shared_memory.SharedMemory(
134 + create=True, size=self.data_buffer_size, name=name)
135 + else:
136 + # we are opening an existing buffer
137 + # fix to https://stackoverflow.com/q/62748654/9191338
138 + # Python incorrectly tracks shared memory even if it is not
139 + # created by the process. The following patch is a workaround.
140 + with patch(
141 + "multiprocessing.resource_tracker.register",
142 + lambda *args, **kwargs: None,
143 + ):
144 + self.shared_memory = shared_memory.SharedMemory(name=name)
145 + # See https://docs.python.org/3/library/multiprocessing.shared_memory.html # noqa
146 + # Some platforms allocate memory based on page size,
147 + # so the shared memory block size may be larger or equal
148 + # to the requested size. The size parameter is ignored
149 + # when attaching to an existing block.
150 + assert self.shared_memory.size >= self.data_buffer_size
151 +
152 + logger.debug("Shared memory created/opened with name: %s, size: %d",
153 + self.shared_memory.name, self.data_buffer_size)
154 +
155 + def handle(self):
156 + return (
157 + self.data_buffer_size,
158 + self.shared_memory.name,
159 + )
160 +
161 + def clear(self) -> None:
162 + """Clear the ring buffer."""
163 + assert self.is_writer, "Only the writer can clear the buffer."
164 + self.metadata.clear()
165 + self.monotonic_id_end = 0
166 + self.monotonic_id_start = 0
167 + self.data_buffer_start = 0
168 + self.data_buffer_end = 0
169 +
170 + def __del__(self):
171 + if hasattr(self, "shared_memory"):
172 + self.shared_memory.close()
173 + if self.is_writer:
174 + self.shared_memory.unlink()
175 +
176 + def int2byte(self, integer: int) -> bytes:
177 + """Convert an integer to bytes."""
178 + return integer.to_bytes(self.ID_NBYTES, "little", signed=True)
179 +
180 + def byte2int(self, byte_data: bytes) -> int:
181 + """Convert bytes back to an integer."""
182 + return int.from_bytes(byte_data, "little", signed=True)
183 +
184 + def allocate_buf(self, size: int) -> tuple[int, int]:
185 + '''
186 + Allocate a buffer `MD_SIZE` + `size` bytes in the shared memory.
187 + Memory layout:
188 + [4-byte monotonic_id][4-byte size][buffer data...]
189 + '''
190 + assert self.is_writer, "Only the writer can allocate buffers."
191 + assert size > 0, "Size must be greater than 0"
192 + size += self.MD_SIZE # add metadata size to the buffer size
193 + # reset to beginning if the buffer does have enough contiguous space
194 + buffer_end_reset = self.data_buffer_end % self.data_buffer_size
195 + if buffer_end_reset + size > self.data_buffer_size:
196 + buffer_end_reset = (self.data_buffer_end // self.data_buffer_size +
197 + 1) * self.data_buffer_size
198 + else: # no reset needed
199 + buffer_end_reset = self.data_buffer_end
200 +
201 + # check if we have enough space in the data buffer
202 + # i.e. if the new end (self.data_buffer_end + size)
203 + # exceeds the start of the data buffer
204 + occupied_size_new = buffer_end_reset + size - self.data_buffer_start
205 + if occupied_size_new > self.data_buffer_size:
206 + raise MemoryError("Not enough space in the data buffer, "
207 + "try calling free_buf() to free up space")
208 + self.data_buffer_end = buffer_end_reset
209 +
210 + # first 4 bytes as the monotonic id
211 + buf_idx = self.data_buffer_end % self.data_buffer_size
212 + self.shared_memory.buf[buf_idx:buf_idx + self.ID_NBYTES] = \
213 + self.int2byte(self.monotonic_id_end)
214 + # next 4 bytes as the size of the data buffer
215 + self.shared_memory.buf[buf_idx + self.ID_NBYTES: \
216 + buf_idx + self.MD_SIZE] = self.int2byte(size)
217 +
218 + # record metadata
219 + self.metadata[self.monotonic_id_end %
220 + self.ID_MAX] = self.data_buffer_end
221 + # update buffer and monotonic id indices
222 + current_buffer_end = self.data_buffer_end
223 + current_id_end = self.monotonic_id_end
224 + self.data_buffer_end += size
225 + self.monotonic_id_end = (self.monotonic_id_end + 1) % self.ID_MAX
226 + return current_buffer_end, current_id_end
227 +
228 + @contextmanager
229 + def access_buf(self, address: int):
230 + buf_idx = address % self.data_buffer_size
231 +
232 + # read metadata
233 + metadata_buff = self.shared_memory.buf[buf_idx:buf_idx + self.MD_SIZE]
234 + id = self.byte2int(metadata_buff[:self.ID_NBYTES])
235 + size = self.byte2int(metadata_buff[self.ID_NBYTES:self.MD_SIZE])
236 +
237 + # yield the data buffer and metadata
238 + data_buff = self.shared_memory.buf[buf_idx + self.MD_SIZE:buf_idx +
239 + size]
240 + with (memoryview(data_buff) as data_view, ):
241 + yield data_view, (id, size)
242 +
243 + def free_buf(self,
244 + is_free_fn: Callable[[int, memoryview], bool],
245 + nbytes: Optional[int] = None) -> Iterable[int]:
246 + '''
247 + Free a buffer of the given size. This is a no-op in shared memory,
248 + but we need to keep track of the metadata.
249 +
250 + If freed memory spreads across the end and start of the ring buffer,
251 + the actual freed memory will be in two segments. In this case there
252 + still might not be a contiguous space of `nbytes` available.
253 +
254 + Args:
255 + nbytes (int, optional): The size of the buffer to free. If None,
256 + frees the maximum size of the ring buffer.
257 + '''
258 +
259 + assert self.is_writer, "Only the writer can free buffers."
260 + logger.debug(
261 + "Freeing up space in the ring buffer, "
262 + "monotonic_id_start: %d, monotonic_id_end: %d",
263 + self.monotonic_id_start, self.monotonic_id_end)
264 + monotonic_id_before = self.monotonic_id_start
265 + # if nbytes is None, free up the maximum size of the ring buffer
266 + if nbytes is None:
267 + nbytes = self.data_buffer_size
268 + freed_bytes = 0
269 + while self.monotonic_id_start in self.metadata and freed_bytes < nbytes:
270 + address = self.metadata[self.monotonic_id_start]
271 + with self.access_buf(address) as (data_buff, metadata):
272 + if is_free_fn(self.monotonic_id_start, data_buff):
273 + # check passed, we can free the buffer
274 + del self.metadata[self.monotonic_id_start]
275 + self.monotonic_id_start = ((self.monotonic_id_start + 1) %
276 + self.ID_MAX)
277 + self.data_buffer_start = address
278 + freed_bytes += metadata[1]
279 + else:
280 + # there are still readers, we cannot free the buffer
281 + break
282 +
283 + logger.debug(
284 + "Freed %d bytes from the ring buffer, "
285 + "monotonic_id_start: %d, monotonic_id_end: %d", freed_bytes,
286 + self.monotonic_id_start, self.monotonic_id_end)
287 +
288 + # buffer wrap around
289 + if self.data_buffer_start >= self.data_buffer_size:
290 + self.data_buffer_start -= self.data_buffer_size
291 + self.data_buffer_end -= self.data_buffer_size
292 +
293 + monotonic_id_after = self.monotonic_id_start
294 + # id wrap around
295 + if monotonic_id_after >= monotonic_id_before:
296 + return range(monotonic_id_before, monotonic_id_after)
297 + else:
298 + return chain(range(monotonic_id_before, self.ID_MAX),
299 + range(0, monotonic_id_after))
300 +
301 +
302 +class ObjectSerde(ABC):
303 +
304 + @abstractmethod
305 + def serialize(self, value: Any) -> tuple[Any, int, bytes, int]:
306 + """Serialize an object to bytes."""
307 + raise NotImplementedError
308 +
309 + @abstractmethod
310 + def deserialize(self, data: memoryview) -> Any:
311 + """Deserialize bytes back to an object."""
312 + raise NotImplementedError
313 +
314 +
315 +class MsgpackSerde(ObjectSerde):
316 +
317 + def __init__(self):
318 + # Delayed import to avoid circular dependency
319 + from vllm.multimodal.inputs import MultiModalKwargsItem
320 + from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
321 +
322 + self.encoder = MsgpackEncoder()
323 + self.tensor_decoder = MsgpackDecoder(torch.Tensor)
324 + self.mm_decoder = MsgpackDecoder(MultiModalKwargsItem)
325 + self._mm_kwargs_item_cls = MultiModalKwargsItem
326 +
327 + def serialize(
328 + self,
329 + value: Any) -> tuple[Union[bytes, list[bytes]], int, bytes, int]:
330 + len_arr = None
331 + if isinstance(value, (torch.Tensor, self._mm_kwargs_item_cls)):
332 + type_name = type(value).__name__
333 + value = self.encoder.encode(value)
334 + len_arr = [len(s) for s in value]
335 + nbytes = sum(len_arr)
336 + else:
337 + value = pickle.dumps(value, protocol=pickle.HIGHEST_PROTOCOL)
338 + type_name = type(value).__name__
339 + nbytes = len(value)
340 +
341 + object_metadata = (type_name, nbytes, len_arr)
342 + serialized_metadata = pickle.dumps(object_metadata,
343 + protocol=pickle.HIGHEST_PROTOCOL)
344 + return value, nbytes, serialized_metadata, len(serialized_metadata)
345 +
346 + def deserialize(self, data_view: memoryview) -> Any:
347 + # pickle.loads do not read past the end of a pickled object
348 + # within a large buffer, so we can skip storing the metadata size
349 + type_name, nbytes, len_arr = pickle.loads(data_view)
350 + serialized_data = bytearray(data_view[-nbytes:])
351 +
352 + if type_name == torch.Tensor.__name__:
353 + obj = []
354 + start_idx = 0
355 + for length in len_arr:
356 + item_bytes = serialized_data[start_idx:start_idx + length]
357 + obj.append(item_bytes)
358 + start_idx += length
359 + obj = self.tensor_decoder.decode(obj)
360 + elif type_name == self._mm_kwargs_item_cls.__name__:
361 + obj = []
362 + start_idx = 0
363 + for length in len_arr:
364 + item_bytes = serialized_data[start_idx:start_idx + length]
365 + obj.append(item_bytes)
366 + start_idx += length
367 + obj = self.mm_decoder.decode(obj)
368 + elif type_name == bytes.__name__:
369 + obj = pickle.loads(serialized_data)
370 + else:
371 + raise ValueError(
372 + f"Unsupported object type '{type_name}' in metadata")
373 +
374 + return obj
375 +
376 +
377 +@dataclass
378 +class ShmObjectStorageHandle:
379 + max_object_size: int
380 + n_readers: int
381 + ring_buffer_handle: tuple[int, str]
382 + serde_class: type[ObjectSerde]
383 + reader_lock: Optional[LockType]
384 +
385 +
386 +class SingleWriterShmObjectStorage:
387 + """
388 + A single-writer, multiple-reader object storage system built on top of a
389 + shared memory ring buffer. Provides key-value storage with automatic memory
390 + management and cross-process serialization support.
391 +
392 + This storage system follows a FIFO (First-In-First-Out) eviction policy
393 + where the oldest objects are automatically freed when memory runs low.
394 + Memory is reclaimed based on reader reference counting - objects are only
395 + freed when all readers have finished accessing them.
396 +
397 + Architecture:
398 + - Single writer process can put(key, value) objects
399 + - Multiple reader processes can get(address, monotonic_id) objects
400 + - Built on SingleWriterShmRingBuffer for efficient shared memory management
401 + - Thread-safe operations with reader synchronization via locks
402 +
403 + Key Features:
404 + - FIFO Eviction: Oldest objects are evicted first when memory is full
405 + - Reference Counting: Objects are only freed when no readers are
406 + accessing them
407 + - Duplicate Key Handling: Existing keys are not overwritten, just
408 + re-referenced
409 + - Customized Serialization: By default uses Msgpack for efficient
410 + serialization of Python objects, but can be extended for custom types
411 + - Cross-Process Safety: Uses shared memory with proper synchronization
412 + - Automatic Cleanup: Garbage collection happens transparently during
413 + allocation
414 +
415 + Memory Layout per Object:
416 + [4-byte reference_count][metadata_size][serialized_object_data]
417 +
418 + Thread Safety:
419 + - Writer operations (put, clear) are single-threaded by design
420 + - Reader operations (get) are thread-safe with lock-based reference
421 + counting
422 + - Memory reclamation is handled exclusively by the writer process
423 + """
424 +
425 + def __init__(
426 + self,
427 + max_object_size: int,
428 + n_readers: int,
429 + ring_buffer: SingleWriterShmRingBuffer,
430 + serde_class: type[ObjectSerde] = MsgpackSerde,
431 + reader_lock: Optional[LockType] = None,
432 + ):
433 + """
434 + Initialize the object storage.
435 +
436 + Args:
437 + max_object_size: Maximum size for a single object in bytes.
438 + n_readers: Number of reader processes that can access the storage.
439 + ring_buffer: The shared memory ring buffer for storing objects.
440 + serde_class: Serializer/deserializer for objects.
441 + reader_lock: Optional lock for synchronizing reader access.
442 + Raises:
443 + ValueError: If reader_lock is None for readers.
444 + """
445 +
446 + self.max_object_size = max_object_size
447 + self.n_readers = n_readers
448 + self.serde_class = serde_class
449 + self.ser_de = serde_class()
450 + self.ring_buffer = ring_buffer
451 + self.is_writer = self.ring_buffer.is_writer
452 +
453 + self.flag_bytes = 4 # for in-use flag
454 +
455 + if self.is_writer:
456 + # Key-value mapping: key -> (address, monotonic_id)
457 + self.key_index: dict[str, tuple[int, int]] = {}
458 + # Reverse mapping: monotonic_id -> key
459 + self.id_index: dict[int, str] = {}
460 + # Writer flag to track in-use status: monotonic_id -> count
461 + self.writer_flag: dict[int, int] = {}
462 + else:
463 + if reader_lock is None:
464 + raise ValueError("Lock must be provided for readers.")
465 +
466 + self._reader_lock = reader_lock
467 +
468 + def clear(self) -> None:
469 + """Clear the object storage."""
470 + if self.is_writer:
471 + self.ring_buffer.clear()
472 + self.key_index.clear()
473 + self.id_index.clear()
474 + self.writer_flag.clear()
475 + logger.debug("Object storage cleared and reinitialized.")
476 +
477 + def copy_to_buffer(
478 + self,
479 + data: Union[bytes, list[bytes]],
480 + data_bytes: int,
481 + metadata: bytes,
482 + md_bytes: int,
483 + data_view: memoryview,
484 + ) -> None:
485 + data_view[self.flag_bytes:self.flag_bytes + md_bytes] = metadata
486 + if isinstance(data, bytes):
487 + data_view[-data_bytes:] = data
488 + elif isinstance(data, list):
489 + start_idx = self.flag_bytes + md_bytes
490 + for item_bytes in data:
491 + item_size = len(item_bytes)
492 + data_view[start_idx:start_idx + item_size] = item_bytes
493 + start_idx += item_size
494 + else:
495 + raise ValueError(
496 + f"Unsupported data type for serialization: {type(data)}")
497 +
498 + def increment_writer_flag(self, id: int) -> None:
499 + """Set the in-use flag for the writer."""
500 + self.writer_flag[id] = self.writer_flag.get(id, 0) + 1
501 +
502 + def increment_reader_flag(self, data_view: memoryview) -> None:
503 + """Set the in-use flag for the reader."""
504 + # >0 for in-use flag
505 + reader_count = self.ring_buffer.byte2int(data_view)
506 + data_view[:] = self.ring_buffer.int2byte(reader_count + 1)
507 +
508 + def free_unused(self) -> None:
509 + """Free unused buffers in the ring buffer."""
510 + # try to free up 2*max_object_size bytes of space in the ring buffer,
511 + # since the buffer might be fragmented
512 + freed_ids = self.ring_buffer.free_buf(self.default_is_free_check,
513 + 2 * self.max_object_size)
514 + # update the metadata after freeing up space
515 + for freed_id in freed_ids:
516 + key_to_free = self.id_index[freed_id]
517 + del self.key_index[key_to_free]
518 + del self.id_index[freed_id]
519 + del self.writer_flag[freed_id]
520 +
521 + def is_cached(self, key: str) -> bool:
522 + """
523 + Check if the object with the given key is cached.
524 + """
525 + return key in self.key_index
526 +
527 + def get_cached(self, key: str) -> tuple[int, int]:
528 + """
529 + Get the cached object by key if it exists.
530 + """
531 + address, monotonic_id = self.key_index[key]
532 + self.increment_writer_flag(monotonic_id)
533 + return address, monotonic_id
534 +
535 + def put(self, key: str, value: Any) -> tuple[int, int]:
536 + """
537 + Store a key-value pair in the object storage.
538 + Attempts to free max_object_size bytes using FIFO order
539 + when the ring buffer runs out of space during a put() operation.
540 +
541 + Args:
542 + key: String key to identify the object
543 + value: Any serializable Python object
544 +
545 + Raises:
546 + MemoryError: If there's not enough space in the buffer
547 + ValueError: If the serialized object is too large
548 + ValueError: If the key already exists in the storage
549 + """
550 + if key in self.key_index:
551 + raise ValueError(f"Key '{key}' already exists in the storage.")
552 +
553 + object_data, data_bytes, object_metadata, md_bytes = \
554 + self.ser_de.serialize(value)
555 + buffer_size = self.flag_bytes + data_bytes + md_bytes
556 +
557 + # Sanity checks
558 + if buffer_size > self.max_object_size:
559 + raise ValueError(
560 + f"Serialized object size ({buffer_size} bytes) exceeds "
561 + f"max object size ({self.max_object_size} bytes)")
562 +
563 + # Allocate new buffer
564 + try:
565 + address, monotonic_id = self.ring_buffer.allocate_buf(buffer_size)
566 + except MemoryError:
567 + self.free_unused()
568 + # try again after freeing up space
569 + address, monotonic_id = self.ring_buffer.allocate_buf(buffer_size)
570 +
571 + # Write data to buffer
572 + with self.ring_buffer.access_buf(address) as (data_view, metadata):
573 + data_view[:self.flag_bytes] = self.ring_buffer.int2byte(0)
574 + self.copy_to_buffer(object_data, data_bytes, object_metadata,
575 + md_bytes, data_view)
576 + self.increment_writer_flag(monotonic_id)
577 +
578 + # Update key index
579 + self.key_index[key] = (address, monotonic_id)
580 + self.id_index[monotonic_id] = key
581 + return address, monotonic_id
582 +
583 + def get(self, address: int, monotonic_id: int) -> Any:
584 + # Read data from buffer
585 + with self.ring_buffer.access_buf(address) as (data_view, buf_metadata):
586 + # check id from metadata
587 + if buf_metadata[0] != monotonic_id:
588 + raise ValueError(
589 + f"Data for address:id '{address}:{monotonic_id}'"
590 + " has been modified or is invalid.")
591 +
592 + obj = self.ser_de.deserialize(data_view[self.flag_bytes:])
593 +
594 + # decrease the in-use flag for reader reads
595 + if self._reader_lock is not None:
596 + with self._reader_lock:
597 + self.increment_reader_flag(data_view[:self.flag_bytes])
598 + else:
599 + # if self._reader_lock is None, it means we are the writer
600 + # in this case, we do not need to decrease the reader count
601 + assert self.is_writer
602 +
603 + return obj
604 +
605 + def handle(self):
606 + """Get handle for sharing across processes."""
607 + return ShmObjectStorageHandle(
608 + max_object_size=self.max_object_size,
609 + n_readers=self.n_readers,
610 + ring_buffer_handle=self.ring_buffer.handle(),
611 + serde_class=self.serde_class,
612 + reader_lock=self._reader_lock,
613 + )
614 +
615 + @staticmethod
616 + def create_from_handle(
617 + handle: ShmObjectStorageHandle) -> "SingleWriterShmObjectStorage":
618 + logger.debug("Creating storage from handle: %s", handle)
619 + ring_buffer = SingleWriterShmRingBuffer(*handle.ring_buffer_handle)
620 + return SingleWriterShmObjectStorage(
621 + max_object_size=handle.max_object_size,
622 + n_readers=handle.n_readers,
623 + ring_buffer=ring_buffer,
624 + serde_class=handle.serde_class,
625 + reader_lock=handle.reader_lock,
626 + )
627 +
628 + def default_is_free_check(self, id: int, buf: memoryview) -> bool:
629 + """
630 + Default is_free function that checks if the first 4 bytes are zero.
631 + This indicates that the buffer is free.
632 + """
633 + reader_count = int.from_bytes(buf[0:4], "little", signed=True)
634 + writer_count = self.writer_flag[id]
635 + return reader_count >= writer_count * self.n_readers
vllm/engine/arg_utils.py
+15 -2
@@ -27,8 +27,8 @@ from vllm.config import (BlockSize, CacheConfig, CacheDType, CompilationConfig,
27 DistributedExecutorBackend, EPLBConfig,
28 GuidedDecodingBackend, HfOverrides, KVEventsConfig,
29 KVTransferConfig, LoadConfig, LogprobsMode,
30 - LoRAConfig, MambaDType, MMEncoderTPMode, ModelConfig,
31 - ModelDType, ModelImpl, MultiModalConfig,
30 + LoRAConfig, MambaDType, MMCacheType, MMEncoderTPMode,
31 + ModelConfig, ModelDType, ModelImpl, MultiModalConfig,
32 ObservabilityConfig, ParallelConfig, PoolerConfig,
33 PrefixCachingHashAlgo, RunnerOption, SchedulerConfig,
34 SchedulerPolicy, SpeculativeConfig, TaskOption,
@@ -373,6 +373,10 @@ class EngineArgs:
373 MultiModalConfig.mm_processor_kwargs
374 disable_mm_preprocessor_cache: bool = False # DEPRECATED
375 mm_processor_cache_gb: float = MultiModalConfig.mm_processor_cache_gb
376 + mm_processor_cache_type: Optional[MMCacheType] = \
377 + MultiModalConfig.mm_processor_cache_type
378 + mm_shm_cache_max_object_size_mb: int = \
379 + MultiModalConfig.mm_shm_cache_max_object_size_mb
380 mm_encoder_tp_mode: MMEncoderTPMode = MultiModalConfig.mm_encoder_tp_mode
381 io_processor_plugin: Optional[str] = None
382 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling
@@ -782,6 +786,12 @@ class EngineArgs:
786 multimodal_group.add_argument("--disable-mm-preprocessor-cache",
787 action="store_true",
788 deprecated=True)
789 + multimodal_group.add_argument(
790 + "--mm-processor-cache-type",
791 + **multimodal_kwargs["mm_processor_cache_type"])
792 + multimodal_group.add_argument(
793 + "--mm-shm-cache-max-object-size-mb",
794 + **multimodal_kwargs["mm_shm_cache_max_object_size_mb"])
795 multimodal_group.add_argument(
796 "--mm-encoder-tp-mode", **multimodal_kwargs["mm_encoder_tp_mode"])
797 multimodal_group.add_argument(
@@ -998,6 +1008,9 @@ class EngineArgs:
1008 config_format=self.config_format,
1009 mm_processor_kwargs=self.mm_processor_kwargs,
1010 mm_processor_cache_gb=self.mm_processor_cache_gb,
1011 + mm_processor_cache_type=self.mm_processor_cache_type,
1012 + mm_shm_cache_max_object_size_mb=self.
1013 + mm_shm_cache_max_object_size_mb,
1014 mm_encoder_tp_mode=self.mm_encoder_tp_mode,
1015 override_pooler_config=self.override_pooler_config,
1016 logits_processor_pattern=self.logits_processor_pattern,
vllm/envs.py
+7
@@ -175,6 +175,7 @@ if TYPE_CHECKING:
175 VLLM_GPT_OSS_HARMONY_SYSTEM_INSTRUCTIONS: bool = False
176 VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False
177 VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES: bool = True
178 + VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME: str = "VLLM_OBJECT_STORAGE_SHM_BUFFER"
179
180
181 def get_default_cache_root():
@@ -1241,6 +1242,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
1242 # raw bytes. Defaults to True for backward compatibility.
1243 "VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES":
1244 lambda: bool(int(os.getenv("VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES", "1"))),
1245 +
1246 + # Name of the shared memory buffer used for object storage.
1247 + # Only effective when mm_config.mm_processor_cache_type == "shm".
1248 + "VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME":
1249 + lambda: os.getenv("VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME",
1250 + "VLLM_OBJECT_STORAGE_SHM_BUFFER"),
1251 }
1252
1253 # --8<-- [end:env-vars-definition]
vllm/executor/uniproc_executor.py
+10
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import os
5 +from multiprocessing import Lock
6 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
7
8 import torch
@@ -10,9 +11,12 @@ import torch.distributed as dist
11 import vllm.envs as envs
12 from vllm.executor.executor_base import ExecutorBase
13 from vllm.logger import init_logger
14 +from vllm.multimodal import MULTIMODAL_REGISTRY
15 +from vllm.multimodal.cache import worker_receiver_cache_from_config
16 from vllm.utils import (get_distributed_init_method, get_ip, get_open_port,
17 run_method)
18 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
19 +from vllm.v1.executor.utils import get_and_update_mm_cache
20 from vllm.worker.worker_base import WorkerWrapperBase
21
22 logger = init_logger(__name__)
@@ -44,6 +48,8 @@ class UniProcExecutor(ExecutorBase):
48 distributed_init_method=distributed_init_method,
49 is_driver_worker=is_driver_worker,
50 )
51 + self.mm_receiver_cache = worker_receiver_cache_from_config(
52 + self.vllm_config, MULTIMODAL_REGISTRY, Lock())
53 self.collective_rpc("init_worker", args=([kwargs], ))
54 self.collective_rpc("init_device")
55 self.collective_rpc("load_model")
@@ -55,6 +61,8 @@ class UniProcExecutor(ExecutorBase):
61 kwargs: Optional[Dict] = None) -> List[Any]:
62 if kwargs is None:
63 kwargs = {}
64 + if self.mm_receiver_cache is not None and method == "execute_model":
65 + get_and_update_mm_cache(self.mm_receiver_cache, args)
66 answer = run_method(self.driver_worker, method, args, kwargs)
67 return [answer]
68
@@ -128,6 +136,8 @@ class ExecutorWithExternalLauncher(UniProcExecutor):
136 distributed_init_method=distributed_init_method,
137 is_driver_worker=is_driver_worker,
138 )
139 + self.mm_receiver_cache = worker_receiver_cache_from_config(
140 + self.vllm_config, MULTIMODAL_REGISTRY, Lock())
141 self.collective_rpc("init_worker", args=([kwargs], ))
142 self.collective_rpc("init_device")
143 self.collective_rpc("load_model")
vllm/multimodal/cache.py
+211 -9
@@ -3,19 +3,24 @@
3 import sys
4 from abc import ABC, abstractmethod
5 from collections.abc import Mapping, Sequence
6 -from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union
6 +from multiprocessing.synchronize import Lock as LockType
7 +from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union, cast
8
9 import torch
10 from typing_extensions import TypeAlias, override
11
12 +from vllm.distributed.device_communicators.shm_object_storage import (
13 + MsgpackSerde, SingleWriterShmObjectStorage, SingleWriterShmRingBuffer)
14 +from vllm.envs import VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME
15 from vllm.logger import init_logger
12 -from vllm.utils import GiB_bytes, LRUCache
16 +from vllm.utils import GiB_bytes, LRUCache, MiB_bytes
17 from vllm.utils.jsontree import (json_count_leaves, json_map_leaves,
18 json_reduce_leaves)
19
16 -from .inputs import (MultiModalFeatureSpec, MultiModalFieldElem,
17 - MultiModalKwargs, MultiModalKwargsItem,
18 - MultiModalKwargsItems, NestedTensors)
20 +from .inputs import (MultiModalBatchedField, MultiModalFeatureSpec,
21 + MultiModalFieldElem, MultiModalKwargs,
22 + MultiModalKwargsItem, MultiModalKwargsItems,
23 + NestedTensors)
24
25 if TYPE_CHECKING:
26 from vllm.config import ModelConfig, VllmConfig
@@ -389,6 +394,106 @@ class MultiModalProcessorSenderCache(BaseMultiModalProcessorCache):
394 self._cache.clear()
395
396
397 +class ShmObjectStoreSenderCache(BaseMultiModalProcessorCache):
398 + """
399 + The cache which is used on P0 when IPC caching is enabled.
400 +
401 + How to update each item:
402 +
403 + - If the item is already in the cache, clear the input to avoid
404 + unnecessary IPC.
405 +
406 + - If the item is not in the cache, store the data in shared memory.
407 + """
408 +
409 + def __init__(self, vllm_config: "VllmConfig") -> None:
410 + super().__init__()
411 +
412 + self.world_size = vllm_config.parallel_config.world_size
413 + mm_config = vllm_config.model_config.get_multimodal_config()
414 +
415 + ring_buffer = SingleWriterShmRingBuffer(
416 + data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes),
417 + name=VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME,
418 + create=True, # sender is the writer
419 + )
420 + self._shm_cache = SingleWriterShmObjectStorage(
421 + max_object_size=mm_config.mm_shm_cache_max_object_size_mb *
422 + MiB_bytes,
423 + n_readers=self.world_size,
424 + ring_buffer=ring_buffer,
425 + serde_class=MsgpackSerde,
426 + )
427 + # cache (prompt_updates, modality) for P0 only
428 + self._p0_cache: dict[str, tuple[Sequence[ResolvedPromptUpdate],
429 + str]] = {}
430 +
431 + @override
432 + def is_cached_item(self, mm_hash: str) -> bool:
433 + return self._shm_cache.is_cached(mm_hash)
434 +
435 + @override
436 + def get_and_update_item(
437 + self,
438 + mm_item: MultiModalProcessorCacheInItem,
439 + mm_hash: str,
440 + ) -> MultiModalProcessorCacheOutItem:
441 +
442 + if self._shm_cache.is_cached(mm_hash):
443 + address, monotonic_id = self._shm_cache.get_cached(mm_hash)
444 + prompt_updates, modality = self._p0_cache[mm_hash]
445 + return self.address_as_item(address, monotonic_id,
446 + modality), prompt_updates
447 +
448 + assert mm_item is not None, f"Expected a cached item for {mm_hash=}"
449 +
450 + try:
451 + address, monotonic_id = self._shm_cache.put(mm_hash, mm_item[0])
452 + # Try to remove dangling items if p0 cache is too large.
453 + if len(self._p0_cache) >= 2 * len(self._shm_cache.key_index):
454 + self.remove_dangling_items()
455 + self._p0_cache[mm_hash] = mm_item[1], mm_item[0].modality
456 + address_item = self.address_as_item(address, monotonic_id,
457 + mm_item[0].modality)
458 + return address_item, mm_item[1]
459 + except (ValueError, MemoryError) as e:
460 + # put may fail if the object is too large or
461 + # the cache is full.
462 + # In this case we log the error and keep the original mm_input.
463 + logger.debug("Failed to cache mm_input with hash %s: %s", mm_hash,
464 + e)
465 + return mm_item
466 +
467 + @override
468 + def clear_cache(self) -> None:
469 + self._shm_cache.clear()
470 + self._p0_cache.clear()
471 +
472 + def remove_dangling_items(self) -> None:
473 + """Remove items that are no longer in the shared memory cache."""
474 + cached_hashes = self._shm_cache.key_index.keys()
475 + dangling_hashes = set(self._p0_cache.keys()) - cached_hashes
476 + for mm_hash in dangling_hashes:
477 + del self._p0_cache[mm_hash]
478 +
479 + def address_as_item(self, address: int, monotonic_id: int,
480 + modality: str) -> MultiModalKwargsItem:
481 + addr_elem = MultiModalFieldElem(
482 + modality=modality,
483 + key="address",
484 + data=address,
485 + field=MultiModalBatchedField(),
486 + )
487 + id_elem = MultiModalFieldElem(
488 + modality=modality,
489 + key="monotonic_id",
490 + data=monotonic_id,
491 + field=MultiModalBatchedField(),
492 + )
493 + mm_item = MultiModalKwargsItem.from_elems([addr_elem, id_elem])
494 + return mm_item
495 +
496 +
497 def _enable_processor_cache(
498 model_config: "ModelConfig",
499 mm_registry: "MultiModalRegistry",
@@ -408,6 +513,17 @@ def _enable_ipc_cache(vllm_config: "VllmConfig") -> bool:
513 return supports_ipc_cache
514
515
516 +def _enable_mm_input_shm_cache(vllm_config: "VllmConfig") -> bool:
517 + """Whether the shared memory based cache should be enabled."""
518 +
519 + if not _enable_ipc_cache(vllm_config):
520 + return False
521 +
522 + mm_config = vllm_config.model_config.get_multimodal_config()
523 +
524 + return mm_config.mm_processor_cache_type == "shm"
525 +
526 +
527 def processor_cache_from_config(
528 vllm_config: "VllmConfig",
529 mm_registry: "MultiModalRegistry",
@@ -421,7 +537,9 @@ def processor_cache_from_config(
537 if not _enable_ipc_cache(vllm_config):
538 return MultiModalProcessorOnlyCache(model_config)
539
424 - return MultiModalProcessorSenderCache(model_config)
540 + if not _enable_mm_input_shm_cache(vllm_config):
541 + return MultiModalProcessorSenderCache(model_config)
542 + return ShmObjectStoreSenderCache(vllm_config)
543
544
545 def processor_only_cache_from_config(
@@ -491,11 +609,68 @@ class MultiModalReceiverCache(BaseMultiModalReceiverCache):
609 self._cache.clear()
610
611
494 -def receiver_cache_from_config(
612 +class ShmObjectStoreReceiverCache(BaseMultiModalReceiverCache):
613 + """
614 + The cache which is used on P1 Worker Process when IPC caching is enabled.
615 +
616 + How to update each item:
617 +
618 + - If the item has an address, replace the input with the cached item.
619 + - If not, return the input.
620 + """
621 +
622 + def __init__(
623 + self,
624 + vllm_config: "VllmConfig",
625 + shared_worker_lock: LockType,
626 + ) -> None:
627 + super().__init__()
628 +
629 + self.world_size = vllm_config.parallel_config.world_size
630 + mm_config = vllm_config.model_config.get_multimodal_config()
631 +
632 + ring_buffer = SingleWriterShmRingBuffer(
633 + data_buffer_size=int(mm_config.mm_processor_cache_gb * GiB_bytes),
634 + name=VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME,
635 + create=False, # Server is a reader
636 + )
637 + self._shm_cache = SingleWriterShmObjectStorage(
638 + max_object_size=mm_config.mm_shm_cache_max_object_size_mb *
639 + MiB_bytes,
640 + n_readers=self.world_size,
641 + ring_buffer=ring_buffer,
642 + serde_class=MsgpackSerde,
643 + reader_lock=shared_worker_lock,
644 + )
645 +
646 + @override
647 + def get_and_update_item(
648 + self,
649 + mm_item: Optional[MultiModalKwargsItem],
650 + mm_hash: str,
651 + ) -> MultiModalKwargsItem:
652 + assert mm_item is not None, f"Expected an address item for {mm_hash=}"
653 + if "address" in mm_item:
654 + address = cast(int, mm_item["address"].data)
655 + monotonic_id = cast(int, mm_item["monotonic_id"].data)
656 + return self._shm_cache.get(address, monotonic_id)
657 +
658 + return mm_item
659 +
660 + @override
661 + def clear_cache(self) -> None:
662 + self._shm_cache.clear()
663 +
664 +
665 +def engine_receiver_cache_from_config(
666 vllm_config: "VllmConfig",
667 mm_registry: "MultiModalRegistry",
668 ) -> Optional[BaseMultiModalReceiverCache]:
498 - """Return a `BaseMultiModalReceiverCache`, if enabled."""
669 + """
670 + This is used in the engine process.
671 + Return a `BaseMultiModalReceiverCache` only when IPC caching is enabled and
672 + mm_processor_cache_type=="lru".
673 + """
674 model_config = vllm_config.model_config
675
676 if not _enable_processor_cache(model_config, mm_registry):
@@ -504,4 +679,31 @@ def receiver_cache_from_config(
679 if not _enable_ipc_cache(vllm_config):
680 return None
681
507 - return MultiModalReceiverCache(model_config)
682 + if not _enable_mm_input_shm_cache(vllm_config):
683 + return MultiModalReceiverCache(model_config)
684 +
685 + return None
686 +
687 +
688 +def worker_receiver_cache_from_config(
689 + vllm_config: "VllmConfig",
690 + mm_registry: "MultiModalRegistry",
691 + shared_worker_lock: LockType,
692 +) -> Optional[BaseMultiModalReceiverCache]:
693 + """
694 + This is used in the worker process.
695 + Return a `BaseMultiModalReceiverCache` only when IPC caching is enabled and
696 + mm_processor_cache_type=="shm".
697 + """
698 + model_config = vllm_config.model_config
699 +
700 + if not _enable_processor_cache(model_config, mm_registry):
701 + return None
702 +
703 + if not _enable_ipc_cache(vllm_config):
704 + return None
705 +
706 + if not _enable_mm_input_shm_cache(vllm_config):
707 + return None
708 +
709 + return ShmObjectStoreReceiverCache(vllm_config, shared_worker_lock)
vllm/utils/__init__.py
+6
@@ -163,6 +163,12 @@ STR_FLASH_ATTN_VAL: str = "FLASH_ATTN"
163 STR_DUAL_CHUNK_FLASH_ATTN_VAL: str = "DUAL_CHUNK_FLASH_ATTN"
164 STR_INVALID_VAL: str = "INVALID"
165
166 +MB_bytes = 1_000_000
167 +"""The number of bytes in one megabyte (MB)."""
168 +
169 +MiB_bytes = 1 << 20
170 +"""The number of bytes in one mebibyte (MiB)."""
171 +
172 GB_bytes = 1_000_000_000
173 """The number of bytes in one gigabyte (GB)."""
174
vllm/v1/engine/core.py
+2 -2
@@ -23,7 +23,7 @@ from vllm.logger import init_logger
23 from vllm.logging_utils.dump_input import dump_engine_exception
24 from vllm.lora.request import LoRARequest
25 from vllm.multimodal import MULTIMODAL_REGISTRY
26 -from vllm.multimodal.cache import receiver_cache_from_config
26 +from vllm.multimodal.cache import engine_receiver_cache_from_config
27 from vllm.tasks import POOLING_TASKS, SupportedTask
28 from vllm.transformers_utils.config import (
29 maybe_register_config_serialize_by_value)
@@ -131,7 +131,7 @@ class EngineCore:
131 self.use_spec_decode = vllm_config.speculative_config is not None
132
133 self.mm_registry = mm_registry = MULTIMODAL_REGISTRY
134 - self.mm_receiver_cache = receiver_cache_from_config(
134 + self.mm_receiver_cache = engine_receiver_cache_from_config(
135 vllm_config, mm_registry)
136
137 # Setup batch queue for pipeline parallelism.
vllm/v1/executor/multiproc_executor.py
+23 -5
@@ -14,6 +14,7 @@ from enum import Enum, auto
14 from functools import partial
15 from multiprocessing.connection import Connection
16 from multiprocessing.process import BaseProcess
17 +from multiprocessing.synchronize import Lock as LockType
18 from threading import Thread
19 from typing import Any, Callable, Optional, Union, cast
20
@@ -31,10 +32,13 @@ from vllm.distributed.parallel_state import (get_dp_group, get_ep_group,
32 from vllm.executor.multiproc_worker_utils import (
33 set_multiprocessing_worker_envs)
34 from vllm.logger import init_logger
35 +from vllm.multimodal import MULTIMODAL_REGISTRY
36 +from vllm.multimodal.cache import worker_receiver_cache_from_config
37 from vllm.utils import (decorate_logs, get_distributed_init_method,
38 get_loopback_ip, get_mp_context, get_open_port,
39 set_process_title)
40 from vllm.v1.executor.abstract import Executor, FailureCallback
41 +from vllm.v1.executor.utils import get_and_update_mm_cache
42 from vllm.v1.outputs import (AsyncModelRunnerOutput, DraftTokenIds,
43 ModelRunnerOutput)
44 from vllm.worker.worker_base import WorkerWrapperBase
@@ -81,6 +85,8 @@ class MultiprocExecutor(Executor):
85 scheduler_output_handle = self.rpc_broadcast_mq.export_handle()
86
87 # Create workers
88 + context = get_mp_context()
89 + shared_worker_lock = context.Lock()
90 unready_workers: list[UnreadyWorkerProcHandle] = []
91 success = False
92 try:
@@ -92,6 +98,7 @@ class MultiprocExecutor(Executor):
98 rank=rank,
99 distributed_init_method=distributed_init_method,
100 input_shm_handle=scheduler_output_handle,
101 + shared_worker_lock=shared_worker_lock,
102 ))
103
104 # Workers must be created before wait_for_ready to avoid
@@ -380,6 +387,7 @@ class WorkerProc:
387 rank: int,
388 distributed_init_method: str,
389 input_shm_handle: Handle,
390 + shared_worker_lock: LockType,
391 ):
392 self.rank = rank
393 wrapper = WorkerWrapperBase(vllm_config=vllm_config, rpc_rank=rank)
@@ -416,6 +424,10 @@ class WorkerProc:
424 name="WorkerAsyncOutputCopy")
425 self.async_output_copy_thread.start()
426
427 + # Initialize multimodal receiver cache if needed
428 + self.mm_receiver_cache = worker_receiver_cache_from_config(
429 + vllm_config, MULTIMODAL_REGISTRY, shared_worker_lock)
430 +
431 # Initialize device
432 self.worker.init_device()
433
@@ -428,11 +440,12 @@ class WorkerProc:
440
441 @staticmethod
442 def make_worker_process(
431 - vllm_config: VllmConfig,
432 - local_rank: int,
433 - rank: int,
434 - distributed_init_method: str,
435 - input_shm_handle, # Receive SchedulerOutput
443 + vllm_config: VllmConfig,
444 + local_rank: int,
445 + rank: int,
446 + distributed_init_method: str,
447 + input_shm_handle, # Receive SchedulerOutput
448 + shared_worker_lock: LockType,
449 ) -> UnreadyWorkerProcHandle:
450 context = get_mp_context()
451 # (reader, writer)
@@ -449,6 +462,7 @@ class WorkerProc:
462 "input_shm_handle": input_shm_handle,
463 "ready_pipe": (reader, writer),
464 "death_pipe": death_reader,
465 + "shared_worker_lock": shared_worker_lock,
466 }
467 # Run EngineCore busy loop in background process.
468 proc = context.Process(target=WorkerProc.worker_main,
@@ -646,6 +660,10 @@ class WorkerProc:
660 func = getattr(self.worker, method)
661 elif isinstance(method, bytes):
662 func = partial(cloudpickle.loads(method), self.worker)
663 + # retrieve from shm cache if available
664 + if self.mm_receiver_cache is not None \
665 + and func.__name__ == "execute_model":
666 + get_and_update_mm_cache(self.mm_receiver_cache, args)
667 output = func(*args, **kwargs)
668 except Exception as e:
669 # Notes have been introduced in python 3.11
vllm/v1/executor/utils.py new
+25
@@ -0,0 +1,25 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from vllm.multimodal.cache import ShmObjectStoreReceiverCache
4 +from vllm.v1.core.sched.output import SchedulerOutput
5 +
6 +
7 +def get_and_update_mm_cache(
8 + receiver_cache: ShmObjectStoreReceiverCache,
9 + args: tuple[SchedulerOutput],
10 +) -> None:
11 + """
12 + For each MultiModalKwargsItem in SchedulerOutput, fetch from shared memory
13 + cache as needed.
14 +
15 + Args:
16 + receiver_cache: The receiver cache to update.
17 + args: According to the collective_rpc call of execute_model method in
18 + executor, args is a tuple of only one SchedulerOutput element.
19 + """
20 + scheduler_output = args[0]
21 + for request_data in scheduler_output.scheduled_new_reqs:
22 + for i in range(len(request_data.mm_kwargs)):
23 + mm_input = request_data.mm_kwargs[i]
24 + request_data.mm_kwargs[i] = \
25 + receiver_cache.get_and_update_item(mm_input, None)