@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 80b70884eb
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 6, 05:45 UTC

Add tensor IPC transfer mechanism for multimodal data (#32104)

Signed-off-by: Brandon Pelfrey <bpelfrey@nvidia.com> Signed-off-by: Brandon Pelfrey <brandonpelfrey@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Brandon Pelfrey committed Mar 21, 2026 at 13:10 UTC 80b70884ebd6b5a297d4b87e05651ea2d7c9d116
13 files changed +1430 -25
tests/v1/test_serial_utils.py
+145
@@ -278,3 +278,148 @@ def test_custom_class_serialization_disallowed_without_pickle():
278 with pytest.raises(TypeError):
279 # Attempt to encode the custom class
280 encoder.encode(obj)
281 +
282 +
283 +@dataclass
284 +class RequestWithTensor:
285 + """Mock request with non-multimodal tensor field like EngineCoreRequest."""
286 +
287 + prompt_embeds: torch.Tensor | None
288 + data: str
289 +
290 +
291 +def test_non_multimodal_tensor_with_ipc():
292 + """Test that non-multimodal tensor fields work correctly with IPC enabled.
293 +
294 + This reproduces the bug where fields like prompt_embeds: torch.Tensor | None
295 + would fail to decode when IPC is enabled because _decode_tensor expected a
296 + raw tensor tuple but received a msgpack-decoded TensorIpcHandle list.
297 + """
298 + import torch.multiprocessing as torch_mp
299 +
300 + from vllm.v1.engine.tensor_ipc import TensorIpcReceiver, TensorIpcSender
301 +
302 + # Create tensor queues for IPC
303 + tensor_queues = [torch_mp.Queue()]
304 +
305 + # Create encoder with IPC sender
306 + sender = TensorIpcSender(tensor_queues[0])
307 + encoder = MsgpackEncoder(oob_tensor_consumer=sender)
308 +
309 + # Create decoder with IPC receiver
310 + receiver = TensorIpcReceiver(tensor_queues[0])
311 + decoder = MsgpackDecoder(RequestWithTensor, oob_tensor_provider=receiver)
312 +
313 + # Create a request with a non-multimodal tensor
314 + original_tensor = torch.randn(5, 10, dtype=torch.float32)
315 + request = RequestWithTensor(prompt_embeds=original_tensor, data="test_data")
316 +
317 + # Encode the request - this should send the tensor via IPC
318 + encoded = encoder.encode(request)
319 +
320 + # Verify encoding succeeded
321 + assert len(encoded) > 0
322 +
323 + # Decode the request - this should retrieve the tensor from IPC queue
324 + # Previously this would fail because the decoder tried to unpack the
325 + # handle list as raw tensor bytes metadata.
326 + decoded = decoder.decode(encoded)
327 +
328 + # Verify the decoded request matches the original
329 + assert isinstance(decoded, RequestWithTensor)
330 + assert decoded.data == "test_data"
331 + assert decoded.prompt_embeds is not None
332 + assert torch.allclose(decoded.prompt_embeds, original_tensor), (
333 + "Decoded tensor does not match the original tensor."
334 + )
335 +
336 +
337 +def test_non_multimodal_tensor_with_ipc_none_value():
338 + """Test that None values for tensor fields work correctly with IPC enabled."""
339 + import torch.multiprocessing as torch_mp
340 +
341 + from vllm.v1.engine.tensor_ipc import TensorIpcReceiver, TensorIpcSender
342 +
343 + # Create tensor queues for IPC
344 + tensor_queues = [torch_mp.Queue()]
345 +
346 + # Create encoder with IPC sender
347 + sender = TensorIpcSender(tensor_queues[0])
348 + encoder = MsgpackEncoder(oob_tensor_consumer=sender)
349 +
350 + # Create decoder with IPC receiver
351 + receiver = TensorIpcReceiver(tensor_queues[0])
352 + decoder = MsgpackDecoder(RequestWithTensor, oob_tensor_provider=receiver)
353 +
354 + # Create a request with None for the tensor field
355 + request = RequestWithTensor(prompt_embeds=None, data="test_data_with_none")
356 +
357 + # Encode and decode the request
358 + encoded = encoder.encode(request)
359 + decoded = decoder.decode(encoded)
360 +
361 + # Verify the decoded request matches the original
362 + assert isinstance(decoded, RequestWithTensor)
363 + assert decoded.data == "test_data_with_none"
364 + assert decoded.prompt_embeds is None
365 +
366 +
367 +def test_multiple_senders_single_receiver_ipc():
368 + """Test N senders sharing a queue with a single receiver via msgpack.
369 +
370 + Simulates the real vLLM topology where multiple API server frontends
371 + each have their own MsgpackEncoder + TensorIpcSender, all putting
372 + tensors onto the same torch.mp queue, and a single engine core
373 + decodes them with one MsgpackDecoder + TensorIpcReceiver.
374 + """
375 + import torch.multiprocessing as torch_mp
376 +
377 + from vllm.v1.engine.tensor_ipc import TensorIpcReceiver, TensorIpcSender
378 +
379 + num_senders = 3
380 + num_messages_per_sender = 2
381 + tensor_queue = torch_mp.Queue()
382 +
383 + # Create N independent senders (each gets its own uuid-based sender_id)
384 + senders = []
385 + encoders = []
386 + for _ in range(num_senders):
387 + s = TensorIpcSender(tensor_queue)
388 + senders.append(s)
389 + encoders.append(MsgpackEncoder(oob_tensor_consumer=s))
390 +
391 + # Single receiver
392 + receiver = TensorIpcReceiver(tensor_queue)
393 + decoder = MsgpackDecoder(RequestWithTensor, oob_tensor_provider=receiver)
394 +
395 + # Encode messages from all senders, interleaving the order
396 + # so that tensors from different senders land on the queue interleaved.
397 + encoded_payloads: list[tuple[int, int, torch.Tensor, list]] = []
398 + for msg_idx in range(num_messages_per_sender):
399 + for sender_idx in range(num_senders):
400 + tensor = torch.full(
401 + (sender_idx + 1, msg_idx + 2),
402 + float(sender_idx * 100 + msg_idx),
403 + dtype=torch.float32,
404 + )
405 + req = RequestWithTensor(
406 + prompt_embeds=tensor,
407 + data=f"s{sender_idx}_m{msg_idx}",
408 + )
409 + encoded = encoders[sender_idx].encode(req)
410 + encoded_payloads.append((sender_idx, msg_idx, tensor, encoded))
411 +
412 + # Decode all messages — the receiver must correctly match each
413 + # tensor handle to the right TensorIpcData from the shared queue.
414 + for sender_idx, msg_idx, original_tensor, encoded in encoded_payloads:
415 + decoded = decoder.decode(encoded)
416 + assert isinstance(decoded, RequestWithTensor)
417 + assert decoded.data == f"s{sender_idx}_m{msg_idx}"
418 + assert decoded.prompt_embeds is not None
419 + assert decoded.prompt_embeds.shape == original_tensor.shape, (
420 + f"Shape mismatch for sender {sender_idx} msg {msg_idx}: "
421 + f"{decoded.prompt_embeds.shape} != {original_tensor.shape}"
422 + )
423 + assert torch.allclose(decoded.prompt_embeds, original_tensor), (
424 + f"Value mismatch for sender {sender_idx} msg {msg_idx}"
425 + )
tests/v1/test_tensor_ipc_queue.py new
+943
@@ -0,0 +1,943 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Tests for tensor IPC queue functionality."""
5 +
6 +import contextlib
7 +import multiprocessing as mp
8 +from dataclasses import dataclass
9 +from multiprocessing.synchronize import Barrier as BarrierType
10 +from multiprocessing.synchronize import Event as EventType
11 +from typing import Any
12 +
13 +import pytest
14 +import torch
15 +import torch.multiprocessing as torch_mp
16 +
17 +from vllm.v1.engine.tensor_ipc import (
18 + TensorIpcData,
19 + TensorIpcReceiver,
20 + TensorIpcSender,
21 +)
22 +from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
23 +
24 +
25 +@pytest.fixture(scope="module", autouse=True)
26 +def setup_multiprocessing():
27 + """Set multiprocessing start method to 'spawn' for compatibility."""
28 + with contextlib.suppress(RuntimeError):
29 + # Already set, which is fine
30 + torch_mp.set_start_method("spawn", force=True)
31 + yield
32 +
33 +
34 +@dataclass
35 +# Use a typed container so the test covers the real vLLM path where tensor IPC
36 +# handles are encoded and decoded as fields nested inside larger msgpack payloads.
37 +class TensorEnvelope:
38 + tensor: torch.Tensor
39 + label: str
40 +
41 +
42 +def encoder_process(
43 + tensor_queue: torch_mp.Queue,
44 + payload_queue: mp.Queue,
45 + result_queue: mp.Queue,
46 + tensor_data: dict[str, Any],
47 + ready_event: EventType,
48 + retrieval_done: EventType,
49 +):
50 + """Process that msgpack-encodes and sends tensors via IPC."""
51 + try:
52 + sender = TensorIpcSender(tensor_queue)
53 + encoder = MsgpackEncoder(oob_tensor_consumer=sender)
54 +
55 + if torch.cuda.is_available():
56 + device = "cuda:0"
57 + tensor = torch.randn(
58 + *tensor_data["shape"], dtype=tensor_data["dtype"], device=device
59 + )
60 + else:
61 + # Fall back to CPU for testing
62 + device = "cpu"
63 + tensor = torch.randn(*tensor_data["shape"], dtype=tensor_data["dtype"])
64 +
65 + message = TensorEnvelope(tensor=tensor, label="cuda-msgpack")
66 + encoded = encoder.encode(message)
67 + payload_queue.put(encoded, timeout=10.0)
68 +
69 + ready_event.set()
70 +
71 + result_queue.put(
72 + {
73 + "success": True,
74 + "encoded_length": len(encoded),
75 + "device": str(device),
76 + "tensor_shape": tuple(tensor.shape),
77 + }
78 + )
79 + retrieval_done.wait(timeout=30.0)
80 + except Exception as e:
81 + import traceback
82 +
83 + ready_event.set()
84 + retrieval_done.set()
85 + result_queue.put(
86 + {"success": False, "error": str(e), "traceback": traceback.format_exc()}
87 + )
88 +
89 +
90 +def decoder_process(
91 + tensor_queue: torch_mp.Queue,
92 + payload_queue: mp.Queue,
93 + result_queue: mp.Queue,
94 + expected_shape: tuple,
95 + encoder_ready: EventType,
96 + retrieval_done: EventType,
97 +):
98 + """Process that msgpack-decodes tensors received via IPC."""
99 + try:
100 + if not encoder_ready.wait(timeout=10.0):
101 + raise TimeoutError("Encoder did not signal ready")
102 +
103 + encoded = payload_queue.get(timeout=5.0)
104 + receiver = TensorIpcReceiver(tensor_queue)
105 + decoder = MsgpackDecoder(TensorEnvelope, oob_tensor_provider=receiver)
106 + decoded = decoder.decode(encoded)
107 +
108 + result_queue.put(
109 + {
110 + "success": True,
111 + "tensor_shape": tuple(decoded.tensor.shape),
112 + "device": str(decoded.tensor.device),
113 + "label": decoded.label,
114 + "matches_expected": tuple(decoded.tensor.shape) == expected_shape,
115 + }
116 + )
117 + except Exception as e:
118 + import traceback
119 +
120 + retrieval_done.set()
121 + result_queue.put(
122 + {"success": False, "error": str(e), "traceback": traceback.format_exc()}
123 + )
124 + else:
125 + retrieval_done.set()
126 +
127 +
128 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
129 +def test_cuda_tensor_queue_basic():
130 + """Test CUDA tensor IPC through the msgpack encoder/decoder path."""
131 + tensor_queue = torch_mp.Queue()
132 + payload_queue: mp.Queue = mp.Queue()
133 + result_queue: mp.Queue = mp.Queue()
134 + encoder_ready = mp.Event()
135 + retrieval_done = mp.Event()
136 +
137 + tensor_shape = (4, 8, 16)
138 + tensor_dtype = torch.float32
139 +
140 + encoder_proc = mp.Process(
141 + target=encoder_process,
142 + args=(
143 + tensor_queue,
144 + payload_queue,
145 + result_queue,
146 + {"shape": tensor_shape, "dtype": tensor_dtype},
147 + encoder_ready,
148 + retrieval_done,
149 + ),
150 + )
151 + encoder_proc.start()
152 +
153 + decoder_proc = mp.Process(
154 + target=decoder_process,
155 + args=(
156 + tensor_queue,
157 + payload_queue,
158 + result_queue,
159 + tensor_shape,
160 + encoder_ready,
161 + retrieval_done,
162 + ),
163 + )
164 + decoder_proc.start()
165 +
166 + encoder_result = result_queue.get(timeout=10.0)
167 + decoder_result = result_queue.get(timeout=10.0)
168 +
169 + encoder_proc.join(timeout=5.0)
170 + decoder_proc.join(timeout=5.0)
171 +
172 + # Verify results
173 + assert encoder_result["success"], (
174 + f"Encoder failed: {encoder_result.get('error')}\n"
175 + f"{encoder_result.get('traceback', '')}"
176 + )
177 + assert decoder_result["success"], (
178 + f"Decoder failed: {decoder_result.get('error')}\n"
179 + f"{decoder_result.get('traceback', '')}"
180 + )
181 + assert decoder_result["matches_expected"], "Tensor shape mismatch"
182 + assert "cuda" in decoder_result["device"], "Tensor not on CUDA device"
183 + assert decoder_result["label"] == "cuda-msgpack"
184 +
185 +
186 +def test_cpu_tensor_fallback():
187 + """Test that CPU tensors use standard serialization path."""
188 + encoder = MsgpackEncoder()
189 +
190 + # Create a CPU tensor
191 + tensor = torch.randn(3, 4, dtype=torch.float32)
192 +
193 + # Encode the tensor (should use standard path, not queue)
194 + encoded = encoder.encode({"test_tensor": tensor})
195 +
196 + # Verify encoding succeeded
197 + assert len(encoded) > 0
198 + assert isinstance(encoded, (list, tuple))
199 +
200 + # Basic check: no queue should be used, so tensor goes through standard path
201 + # This is mainly to ensure no exceptions are raised
202 +
203 +
204 +def test_msgpack_encoder_decoder_with_ipc():
205 + """Test the full msgpack + tensor IPC path in one process."""
206 + tensor_queue = torch_mp.Queue()
207 + sender = TensorIpcSender(tensor_queue)
208 + encoder = MsgpackEncoder(oob_tensor_consumer=sender)
209 + receiver = TensorIpcReceiver(tensor_queue)
210 + decoder = MsgpackDecoder(TensorEnvelope, oob_tensor_provider=receiver)
211 +
212 + # Use CPU here to exercise the msgpack + sender/receiver integration
213 + # without relying on same-process CUDA IPC behavior.
214 + tensor = torch.randn(2, 3)
215 +
216 + message = TensorEnvelope(tensor=tensor, label="test")
217 + encoded = encoder.encode(message)
218 + assert len(encoded) > 0
219 +
220 + decoded = decoder.decode(encoded)
221 + assert isinstance(decoded, TensorEnvelope)
222 + assert decoded.label == "test"
223 + assert torch.allclose(decoded.tensor, tensor)
224 +
225 +
226 +def test_decoder_buffer_management():
227 + """Test receiver's tensor buffer management when draining queue."""
228 + tensor_queue = torch_mp.Queue()
229 +
230 + sender_id = "test_sender"
231 + message_id = 1
232 +
233 + # Put multiple tensors in queue using TensorIpcData
234 + tensors_data = [
235 + (0, torch.randn(2, 3)),
236 + (1, torch.randn(4, 5)),
237 + (2, torch.randn(6, 7)),
238 + ]
239 +
240 + for tensor_id, tensor in tensors_data:
241 + ipc_data = TensorIpcData(
242 + sender_id=sender_id,
243 + message_id=message_id,
244 + tensor_id=tensor_id,
245 + tensor=tensor,
246 + )
247 + tensor_queue.put(ipc_data)
248 +
249 + # Create receiver directly
250 + receiver = TensorIpcReceiver(tensor_queue)
251 +
252 + # Request tensor_id=2 (should buffer tensor_id=0 and tensor_id=1)
253 + handle = {"sender_id": sender_id, "message_id": message_id, "tensor_id": 2}
254 +
255 + result = receiver("float32", (6, 7), handle)
256 + assert result.shape == (6, 7)
257 +
258 + # Verify buffer has tensor_id 0 and 1
259 + sender = receiver._tensor_buffers[sender_id]
260 + tensors = sender.tensors.get(message_id, {})
261 + assert 0 in tensors
262 + assert 1 in tensors
263 +
264 + # Request buffered tensor
265 + handle2 = {"sender_id": sender_id, "message_id": message_id, "tensor_id": 0}
266 +
267 + result2 = receiver("float32", (2, 3), handle2)
268 + assert result2.shape == (2, 3)
269 + # tensor_id 0 should be removed from buffer
270 + sender = receiver._tensor_buffers[sender_id]
271 + tensors = sender.tensors.get(message_id, {})
272 + assert 0 not in tensors
273 +
274 +
275 +def api_server_worker(
276 + server_id: int,
277 + tensor_queue: torch_mp.Queue,
278 + result_queue: mp.Queue,
279 + barrier: BarrierType,
280 + retrieval_done: EventType,
281 +):
282 + """Worker simulating an API server sending tensors."""
283 + try:
284 + # Each server sends a unique tensor
285 + tensor = torch.ones(server_id + 1, server_id + 2) * server_id
286 + sender_id = f"server_{server_id}"
287 +
288 + # Wait for all servers to be ready
289 + barrier.wait()
290 +
291 + # Send tensor using TensorIpcData
292 + ipc_data = TensorIpcData(
293 + sender_id=sender_id,
294 + message_id=0,
295 + tensor_id=0,
296 + tensor=tensor,
297 + )
298 + tensor_queue.put(ipc_data)
299 +
300 + result_queue.put({"server_id": server_id, "success": True})
301 +
302 + # Keep process alive until main process has retrieved all tensors
303 + # This prevents shared memory handles from being invalidated
304 + retrieval_done.wait(timeout=30.0)
305 + except Exception as e:
306 + import traceback
307 +
308 + result_queue.put(
309 + {
310 + "server_id": server_id,
311 + "success": False,
312 + "error": str(e),
313 + "traceback": traceback.format_exc(),
314 + }
315 + )
316 +
317 +
318 +def test_multiple_api_servers_to_engine():
319 + """Test multiple API servers sending to one engine core via multiprocessing."""
320 + num_api_servers = 3
321 + tensor_queue = torch_mp.Queue()
322 + result_queue: mp.Queue = mp.Queue()
323 + barrier = mp.Barrier(num_api_servers)
324 + retrieval_done = mp.Event()
325 +
326 + # Start multiple API server processes
327 + processes = []
328 + for server_id in range(num_api_servers):
329 + proc = mp.Process(
330 + target=api_server_worker,
331 + args=(server_id, tensor_queue, result_queue, barrier, retrieval_done),
332 + )
333 + proc.start()
334 + processes.append(proc)
335 +
336 + # Collect results from all servers
337 + results = []
338 + for _ in range(num_api_servers):
339 + result = result_queue.get(timeout=10.0)
340 + results.append(result)
341 +
342 + # Verify all servers succeeded
343 + for result in results:
344 + assert result["success"], (
345 + f"Server {result['server_id']} failed: {result.get('error')}"
346 + )
347 +
348 + # Verify all tensors are in queue
349 + received_tensors = []
350 + for _ in range(num_api_servers):
351 + ipc_data = tensor_queue.get(timeout=1.0)
352 + received_tensors.append((ipc_data.sender_id, ipc_data.tensor))
353 +
354 + assert len(received_tensors) == num_api_servers
355 +
356 + # Verify tensor content (order may vary with multiprocessing)
357 + tensor_by_sender = {sid: t for sid, t in received_tensors}
358 + for server_id in range(num_api_servers):
359 + expected_id = f"server_{server_id}"
360 + assert expected_id in tensor_by_sender, (
361 + f"Missing tensor from server {server_id}"
362 + )
363 + expected_tensor = torch.ones(server_id + 1, server_id + 2) * server_id
364 + assert torch.allclose(tensor_by_sender[expected_id], expected_tensor)
365 +
366 + # Signal workers that retrieval is complete
367 + retrieval_done.set()
368 +
369 + # Wait for all processes to complete
370 + for proc in processes:
371 + proc.join(timeout=5.0)
372 +
373 +
374 +def mixed_tensor_encoder_process(
375 + tensor_queue: torch_mp.Queue,
376 + result_queue: mp.Queue,
377 + ready_event: EventType,
378 + retrieval_done: EventType,
379 +):
380 + """Process that encodes mixed CPU/CUDA tensors."""
381 + try:
382 + sender = TensorIpcSender(tensor_queue)
383 + _encoder = MsgpackEncoder(oob_tensor_consumer=sender)
384 +
385 + # Create only CUDA tensor for IPC (CPU will be serialized)
386 + # But actually, let's just send CUDA tensor directly
387 + cuda_tensor = torch.randn(4, 5, device="cuda:0")
388 +
389 + # Manually send via IPC to test the mechanism
390 + cuda_tensor_shared = cuda_tensor.share_memory_()
391 +
392 + ipc_data = TensorIpcData(
393 + sender_id="mixed_encoder",
394 + message_id=0,
395 + tensor_id=0,
396 + tensor=cuda_tensor_shared,
397 + )
398 + tensor_queue.put(ipc_data, timeout=10.0)
399 +
400 + ready_event.set()
401 +
402 + result_queue.put({"success": True, "sent_cuda": True})
403 +
404 + # Keep process alive until decoder has retrieved the tensor
405 + retrieval_done.wait(timeout=30.0)
406 + except Exception as e:
407 + import traceback
408 +
409 + ready_event.set()
410 + result_queue.put(
411 + {"success": False, "error": str(e), "traceback": traceback.format_exc()}
412 + )
413 +
414 +
415 +def mixed_tensor_decoder_process(
416 + tensor_queue: torch_mp.Queue,
417 + result_queue: mp.Queue,
418 + encoder_ready: EventType,
419 + retrieval_done: EventType,
420 +):
421 + """Process that retrieves mixed tensors from queue."""
422 + try:
423 + # Wait for encoder to finish
424 + if not encoder_ready.wait(timeout=10.0):
425 + raise TimeoutError("Encoder did not signal ready")
426 +
427 + # Try to get CUDA tensor from queue
428 + ipc_data = tensor_queue.get(timeout=5.0)
429 +
430 + result_queue.put(
431 + {
432 + "success": True,
433 + "is_cuda": ipc_data.tensor.is_cuda,
434 + "shape": tuple(ipc_data.tensor.shape),
435 + }
436 + )
437 +
438 + # Signal that retrieval is complete
439 + retrieval_done.set()
440 + except Exception as e:
441 + import traceback
442 +
443 + retrieval_done.set() # Signal even on failure
444 + result_queue.put(
445 + {"success": False, "error": str(e), "traceback": traceback.format_exc()}
446 + )
447 +
448 +
449 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
450 +def test_mixed_cpu_cuda_tensors():
451 + """Test encoding with mixed CPU and CUDA tensors using multiprocessing."""
452 + tensor_queue = torch_mp.Queue()
453 + result_queue: mp.Queue = mp.Queue()
454 + encoder_ready = mp.Event()
455 + retrieval_done = mp.Event()
456 +
457 + # Start encoder process
458 + encoder_proc = mp.Process(
459 + target=mixed_tensor_encoder_process,
460 + args=(tensor_queue, result_queue, encoder_ready, retrieval_done),
461 + )
462 + encoder_proc.start()
463 +
464 + # Start decoder process
465 + decoder_proc = mp.Process(
466 + target=mixed_tensor_decoder_process,
467 + args=(tensor_queue, result_queue, encoder_ready, retrieval_done),
468 + )
469 + decoder_proc.start()
470 +
471 + # Get results
472 + encoder_result = result_queue.get(timeout=10.0)
473 + decoder_result = result_queue.get(timeout=10.0)
474 +
475 + encoder_proc.join(timeout=5.0)
476 + decoder_proc.join(timeout=5.0)
477 +
478 + # Verify encoder succeeded
479 + assert encoder_result["success"], (
480 + f"Encoder failed: {encoder_result.get('error')}\n"
481 + f"{encoder_result.get('traceback', '')}"
482 + )
483 +
484 + # Verify decoder succeeded and got CUDA tensor
485 + assert decoder_result["success"], (
486 + f"Decoder failed: {decoder_result.get('error')}\n"
487 + f"{decoder_result.get('traceback', '')}"
488 + )
489 + assert decoder_result["is_cuda"], "Retrieved tensor is not on CUDA"
490 + assert decoder_result["shape"] == (4, 5), (
491 + f"Unexpected shape: {decoder_result['shape']}"
492 + )
493 +
494 +
495 +def cpu_tensor_ipc_encoder_process(
496 + tensor_queue: torch_mp.Queue,
497 + result_queue: mp.Queue,
498 + tensor_shape: tuple,
499 + ready_event: EventType,
500 + retrieval_done: EventType,
501 +):
502 + """Process that encodes and sends CPU tensors via IPC queue."""
503 + try:
504 + # Create encoder with IPC enabled for all tensors
505 + sender = TensorIpcSender(tensor_queue)
506 + encoder = MsgpackEncoder(oob_tensor_consumer=sender)
507 +
508 + # Create a CPU tensor
509 + tensor = torch.randn(*tensor_shape, dtype=torch.float32)
510 +
511 + # Encode the tensor (should use IPC queue, not standard serialization)
512 + encoded = encoder.encode({"test_tensor": tensor})
513 +
514 + # Signal that encoding is complete
515 + ready_event.set()
516 +
517 + result_queue.put(
518 + {
519 + "success": True,
520 + "encoded_length": len(encoded),
521 + "device": str(tensor.device),
522 + "tensor_shape": tuple(tensor.shape),
523 + }
524 + )
525 +
526 + # Keep process alive until decoder has retrieved the tensor
527 + # This is necessary for CPU tensor shared memory to remain valid
528 + retrieval_done.wait(timeout=30.0)
529 + except Exception as e:
530 + import traceback
531 +
532 + ready_event.set()
533 + result_queue.put(
534 + {"success": False, "error": str(e), "traceback": traceback.format_exc()}
535 + )
536 +
537 +
538 +def cpu_tensor_ipc_decoder_process(
539 + tensor_queue: torch_mp.Queue,
540 + result_queue: mp.Queue,
541 + expected_shape: tuple,
542 + encoder_ready: EventType,
543 + retrieval_done: EventType,
544 +):
545 + """Process that decodes and receives CPU tensors from IPC queue."""
546 + try:
547 + # Wait for encoder to finish sending
548 + if not encoder_ready.wait(timeout=10.0):
549 + raise TimeoutError("Encoder did not signal ready")
550 +
551 + # Get tensor from queue
552 + ipc_data = tensor_queue.get(timeout=5.0)
553 +
554 + result_queue.put(
555 + {
556 + "success": True,
557 + "tensor_id": ipc_data.tensor_id,
558 + "tensor_shape": tuple(ipc_data.tensor.shape),
559 + "device": str(ipc_data.tensor.device),
560 + "matches_expected": tuple(ipc_data.tensor.shape) == expected_shape,
561 + "is_cpu": ipc_data.tensor.device.type == "cpu",
562 + }
563 + )
564 +
565 + # Signal that retrieval is complete
566 + retrieval_done.set()
567 + except Exception as e:
568 + import traceback
569 +
570 + retrieval_done.set() # Signal even on failure
571 + result_queue.put(
572 + {"success": False, "error": str(e), "traceback": traceback.format_exc()}
573 + )
574 +
575 +
576 +def test_cpu_tensor_ipc():
577 + """Test CPU tensor sharing via IPC queue when mm_tensor_ipc is enabled."""
578 + # Set up single queue and synchronization
579 + tensor_queue = torch_mp.Queue()
580 + result_queue: mp.Queue = mp.Queue()
581 + encoder_ready = mp.Event()
582 + retrieval_done = mp.Event()
583 +
584 + tensor_shape = (3, 5, 7)
585 +
586 + # Start encoder process
587 + encoder_proc = mp.Process(
588 + target=cpu_tensor_ipc_encoder_process,
589 + args=(
590 + tensor_queue,
591 + result_queue,
592 + tensor_shape,
593 + encoder_ready,
594 + retrieval_done,
595 + ),
596 + )
597 + encoder_proc.start()
598 +
599 + # Start decoder process
600 + decoder_proc = mp.Process(
601 + target=cpu_tensor_ipc_decoder_process,
602 + args=(
603 + tensor_queue,
604 + result_queue,
605 + tensor_shape,
606 + encoder_ready,
607 + retrieval_done,
608 + ),
609 + )
610 + decoder_proc.start()
611 +
612 + # Wait for processes and collect results
613 + encoder_result = result_queue.get(timeout=10.0)
614 + decoder_result = result_queue.get(timeout=10.0)
615 +
616 + encoder_proc.join(timeout=5.0)
617 + decoder_proc.join(timeout=5.0)
618 +
619 + # Verify results
620 + assert encoder_result["success"], (
621 + f"Encoder failed: {encoder_result.get('error')}\n"
622 + f"{encoder_result.get('traceback', '')}"
623 + )
624 + assert decoder_result["success"], (
625 + f"Decoder failed: {decoder_result.get('error')}\n"
626 + f"{decoder_result.get('traceback', '')}"
627 + )
628 + assert decoder_result["matches_expected"], "Tensor shape mismatch"
629 + assert decoder_result["is_cpu"], "Tensor not on CPU device"
630 +
631 +
632 +def test_ipc_disabled_mode():
633 + """Test that IPC is disabled when no sender is provided."""
634 + tensor_queues = [torch_mp.Queue()]
635 +
636 + # Create encoder without IPC sender (IPC disabled)
637 + encoder = MsgpackEncoder()
638 +
639 + # Create a CPU tensor
640 + cpu_tensor = torch.randn(2, 3, dtype=torch.float32)
641 +
642 + # Encode the tensor (should use standard serialization, not IPC)
643 + encoded = encoder.encode({"test_tensor": cpu_tensor})
644 +
645 + # Verify encoding succeeded
646 + assert len(encoded) > 0
647 + assert isinstance(encoded, (list, tuple))
648 +
649 + # Verify queue is empty (no IPC was used)
650 + assert tensor_queues[0].empty(), "Tensor queue should be empty when IPC is disabled"
651 +
652 + # If CUDA is available, test with CUDA tensor too
653 + if torch.cuda.is_available():
654 + cuda_tensor = torch.randn(4, 5, device="cuda:0")
655 + encoded_cuda = encoder.encode({"cuda_tensor": cuda_tensor})
656 + assert len(encoded_cuda) > 0
657 + assert tensor_queues[0].empty(), (
658 + "Tensor queue should be empty for CUDA tensor when IPC is disabled"
659 + )
660 +
661 +
662 +@dataclass
663 +class MultiTensorMessage:
664 + """Message with multiple tensors to test multi-tensor IPC."""
665 +
666 + t1: torch.Tensor
667 + t2: torch.Tensor
668 + sender_label: str
669 +
670 +
671 +def concurrent_sender_process(
672 + tensor_queue: torch_mp.Queue,
673 + payload_queue: mp.Queue,
674 + result_queue: mp.Queue,
675 + sender_index: int,
676 + num_messages: int,
677 + barrier: BarrierType,
678 + retrieval_done: EventType,
679 +):
680 + """Process that acts as one of N concurrent senders."""
681 + try:
682 + sender = TensorIpcSender(tensor_queue)
683 + encoder = MsgpackEncoder(oob_tensor_consumer=sender)
684 +
685 + # Wait for all senders to be ready before sending
686 + barrier.wait(timeout=10.0)
687 +
688 + encoded_payloads = []
689 + for msg_idx in range(num_messages):
690 + # Each sender creates uniquely-shaped tensors so we can
691 + # verify correct routing on the receiver side.
692 + t1 = torch.full((sender_index + 1, 3), float(msg_idx), dtype=torch.float32)
693 + t2 = torch.full(
694 + (2, sender_index + 2), float(msg_idx + 100), dtype=torch.float64
695 + )
696 + msg = MultiTensorMessage(
697 + t1=t1,
698 + t2=t2,
699 + sender_label=f"sender_{sender_index}_msg_{msg_idx}",
700 + )
701 + encoded = encoder.encode(msg)
702 + encoded_payloads.append(encoded)
703 +
704 + # Send all encoded payloads via the regular (non-tensor) queue
705 + for encoded in encoded_payloads:
706 + payload_queue.put(encoded, timeout=10.0)
707 +
708 + result_queue.put(
709 + {
710 + "success": True,
711 + "sender_index": sender_index,
712 + "num_sent": num_messages,
713 + }
714 + )
715 +
716 + # Keep alive so shared-memory handles remain valid
717 + retrieval_done.wait(timeout=30.0)
718 + except Exception as e:
719 + import traceback
720 +
721 + result_queue.put(
722 + {
723 + "success": False,
724 + "sender_index": sender_index,
725 + "error": str(e),
726 + "traceback": traceback.format_exc(),
727 + }
728 + )
729 +
730 +
731 +def test_concurrent_senders_single_receiver():
732 + """Test N concurrent senders sharing one queue with a single receiver.
733 +
734 + Each sender encodes multiple messages (each containing two tensors) via
735 + its own MsgpackEncoder + TensorIpcSender. A single TensorIpcReceiver
736 + on the receiving side must correctly drain-and-buffer interleaved
737 + TensorIpcData items from the shared queue and match them back to the
738 + right message handles during decode.
739 + """
740 + num_senders = 4
741 + num_messages_per_sender = 3
742 + tensor_queue = torch_mp.Queue()
743 + payload_queue: mp.Queue = mp.Queue()
744 + result_queue: mp.Queue = mp.Queue()
745 + barrier = mp.Barrier(num_senders)
746 + retrieval_done = mp.Event()
747 +
748 + # Launch sender processes
749 + processes = []
750 + for i in range(num_senders):
751 + proc = mp.Process(
752 + target=concurrent_sender_process,
753 + args=(
754 + tensor_queue,
755 + payload_queue,
756 + result_queue,
757 + i,
758 + num_messages_per_sender,
759 + barrier,
760 + retrieval_done,
761 + ),
762 + )
763 + proc.start()
764 + processes.append(proc)
765 +
766 + # Collect send confirmations
767 + send_results = []
768 + for _ in range(num_senders):
769 + send_results.append(result_queue.get(timeout=15.0))
770 + for r in send_results:
771 + assert r["success"], (
772 + f"Sender {r['sender_index']} failed: {r.get('error')}\n"
773 + f"{r.get('traceback', '')}"
774 + )
775 +
776 + # Now decode all messages from the main process using a single receiver
777 + receiver = TensorIpcReceiver(tensor_queue)
778 + decoder = MsgpackDecoder(MultiTensorMessage, oob_tensor_provider=receiver)
779 +
780 + decoded_messages: list[MultiTensorMessage] = []
781 + total = num_senders * num_messages_per_sender
782 + for _ in range(total):
783 + encoded = payload_queue.get(timeout=10.0)
784 + decoded = decoder.decode(encoded)
785 + assert isinstance(decoded, MultiTensorMessage)
786 + decoded_messages.append(decoded)
787 +
788 + # Signal senders they can exit
789 + retrieval_done.set()
790 +
791 + # Group by sender_label prefix to verify all messages arrived
792 + by_sender: dict[int, list[MultiTensorMessage]] = {}
793 + for msg in decoded_messages:
794 + # label format: "sender_{i}_msg_{j}"
795 + parts = msg.sender_label.split("_")
796 + sender_idx = int(parts[1])
797 + by_sender.setdefault(sender_idx, []).append(msg)
798 +
799 + assert len(by_sender) == num_senders, (
800 + f"Expected {num_senders} senders, got {len(by_sender)}"
801 + )
802 +
803 + for sender_idx in range(num_senders):
804 + msgs = sorted(by_sender[sender_idx], key=lambda m: m.sender_label)
805 + assert len(msgs) == num_messages_per_sender, (
806 + f"Sender {sender_idx}: expected {num_messages_per_sender} "
807 + f"messages, got {len(msgs)}"
808 + )
809 + for msg_idx, msg in enumerate(msgs):
810 + assert msg.sender_label == f"sender_{sender_idx}_msg_{msg_idx}"
811 + # Verify tensor shapes match what the sender created
812 + assert msg.t1.shape == (sender_idx + 1, 3)
813 + assert msg.t2.shape == (2, sender_idx + 2)
814 + # Verify tensor values
815 + assert torch.allclose(msg.t1, torch.full_like(msg.t1, float(msg_idx)))
816 + assert torch.allclose(msg.t2, torch.full_like(msg.t2, float(msg_idx + 100)))
817 +
818 + for proc in processes:
819 + proc.join(timeout=5.0)
820 +
821 +
822 +def test_concurrent_senders_interleaved_buffer():
823 + """Test receiver buffering when tensors from multiple senders interleave.
824 +
825 + Manually enqueue TensorIpcData from two senders in an interleaved order
826 + and verify the receiver correctly buffers and retrieves each tensor by
827 + its (sender_id, message_id, tensor_id) handle.
828 + """
829 + tensor_queue = torch_mp.Queue()
830 +
831 + # Sender A: 2 tensors for message 1
832 + a_t0 = torch.randn(2, 3)
833 + a_t1 = torch.randn(4, 5)
834 + # Sender B: 2 tensors for message 1
835 + b_t0 = torch.randn(6, 7)
836 + b_t1 = torch.randn(8, 9)
837 +
838 + # Interleave: B_t0, A_t0, B_t1, A_t1
839 + for sid, mid, tid, t in [
840 + ("B", 1, 0, b_t0),
841 + ("A", 1, 0, a_t0),
842 + ("B", 1, 1, b_t1),
843 + ("A", 1, 1, a_t1),
844 + ]:
845 + tensor_queue.put(
846 + TensorIpcData(sender_id=sid, message_id=mid, tensor_id=tid, tensor=t)
847 + )
848 +
849 + receiver = TensorIpcReceiver(tensor_queue)
850 +
851 + # Request A_t1 first — receiver must drain and buffer B_t0, A_t0, B_t1
852 + result = receiver(
853 + "float32", a_t1.shape, {"sender_id": "A", "message_id": 1, "tensor_id": 1}
854 + )
855 + assert torch.equal(result, a_t1)
856 +
857 + # Now request B_t0 from buffer
858 + result = receiver(
859 + "float32", b_t0.shape, {"sender_id": "B", "message_id": 1, "tensor_id": 0}
860 + )
861 + assert torch.equal(result, b_t0)
862 +
863 + # Request A_t0 from buffer
864 + result = receiver(
865 + "float32", a_t0.shape, {"sender_id": "A", "message_id": 1, "tensor_id": 0}
866 + )
867 + assert torch.equal(result, a_t0)
868 +
869 + # Request B_t1 from buffer
870 + result = receiver(
871 + "float64", b_t1.shape, {"sender_id": "B", "message_id": 1, "tensor_id": 1}
872 + )
873 + assert torch.equal(result, b_t1)
874 +
875 + # All buffers should be drained
876 + for sid in ("A", "B"):
877 + tensors = receiver._tensor_buffers[sid].tensors.get(1, {})
878 + assert len(tensors) == 0, f"Sender {sid} buffer not empty: {tensors}"
879 +
880 +
881 +def test_mixed_cpu_cuda_with_ipc_enabled():
882 + """Test that encoder is configured correctly for IPC with all tensor types."""
883 + if not torch.cuda.is_available():
884 + pytest.skip("CUDA not available")
885 +
886 + tensor_queue = torch_mp.Queue()
887 +
888 + # Create sender and encoder with IPC enabled
889 + sender = TensorIpcSender(tensor_queue)
890 + encoder = MsgpackEncoder(oob_tensor_consumer=sender)
891 +
892 + # Verify sender configuration
893 + assert encoder.oob_tensor_consumer is not None, "Consumer should be set"
894 +
895 + # Note: Actual IPC transfer only works across processes
896 + # (tested in test_cpu_tensor_ipc)
897 + # This test just verifies the configuration is correct
898 +
899 +
900 +def test_tensor_cleanup_after_decode():
901 + """Test that tensors are removed from tracking after successful decode."""
902 + # Create a tensor queue
903 + tensor_queue = torch_mp.Queue()
904 +
905 + # Create and encode a tensor
906 + tensor = torch.randn(5, 5)
907 + # Move to shared memory for IPC
908 + if not tensor.is_shared():
909 + tensor.share_memory_()
910 +
911 + # Manually create a TensorIpcData and put it in the queue
912 + sender_id = "test_sender"
913 + message_id = 0
914 + tensor_id = 0
915 + ipc_data = TensorIpcData(
916 + sender_id=sender_id,
917 + message_id=message_id,
918 + tensor_id=tensor_id,
919 + tensor=tensor,
920 + )
921 + tensor_queue.put(ipc_data)
922 +
923 + # Create receiver directly
924 + receiver = TensorIpcReceiver(tensor_queue)
925 +
926 + handle = {
927 + "sender_id": sender_id,
928 + "message_id": message_id,
929 + "tensor_id": tensor_id,
930 + }
931 +
932 + # Receive the tensor - this should retrieve it from the queue
933 + decoded_tensor = receiver(
934 + str(tensor.dtype).removeprefix("torch."), tensor.shape, handle
935 + )
936 +
937 + # Verify the tensor was decoded
938 + assert decoded_tensor.shape == tensor.shape, "Decoded tensor should match shape"
939 +
940 + # Verify the tensor was removed from buffer after decode
941 + sender = receiver._tensor_buffers[sender_id]
942 + tensors = sender.tensors.get(message_id, {})
943 + assert tensor_id not in tensors, "Tensor should be removed from buffer"
vllm/config/model.py
+25 -1
@@ -14,7 +14,12 @@ import vllm.envs as envs
14 from vllm.config.model_arch import (
15 ModelArchitectureConfig,
16 )
17 -from vllm.config.multimodal import MMCacheType, MMEncoderTPMode, MultiModalConfig
17 +from vllm.config.multimodal import (
18 + MMCacheType,
19 + MMEncoderTPMode,
20 + MMTensorIPC,
21 + MultiModalConfig,
22 +)
23 from vllm.config.pooler import PoolerConfig
24 from vllm.config.scheduler import RunnerType
25 from vllm.config.utils import config, getattr_iter
@@ -310,6 +315,7 @@ class ModelConfig:
315 interleave_mm_strings: InitVar[bool | None] = None
316 skip_mm_profiling: InitVar[bool | None] = None
317 video_pruning_rate: InitVar[float | None] = None
318 + mm_tensor_ipc: InitVar[MMTensorIPC] = None
319
320 def compute_hash(self) -> str:
321 """
@@ -430,6 +436,7 @@ class ModelConfig:
436 interleave_mm_strings: bool | None,
437 skip_mm_profiling: bool | None,
438 video_pruning_rate: float | None,
439 + mm_tensor_ipc: MMTensorIPC,
440 ) -> None:
441 # Keep set served_model_name before maybe_model_redirect(self.model)
442 self.served_model_name = get_served_model_name(
@@ -612,6 +619,7 @@ class ModelConfig:
619 interleave_mm_strings=interleave_mm_strings,
620 skip_mm_profiling=skip_mm_profiling,
621 video_pruning_rate=video_pruning_rate,
622 + mm_tensor_ipc=mm_tensor_ipc,
623 )
624
625 mm_config_kwargs = {
@@ -1112,6 +1120,22 @@ class ModelConfig:
1120 f"({parallel_config.decode_context_parallel_size})."
1121 )
1122
1123 + # torch_shm uses a single IPC queue to rank 0; DP>1 is
1124 + # incompatible because API servers can't know which
1125 + # CoreEngine the scheduler will assign work to. TP>1 is
1126 + # also not supported because this requires broadcasting
1127 + # MM tensors between all TP ranks.
1128 + if (
1129 + self.multimodal_config is not None
1130 + and self.multimodal_config.mm_tensor_ipc == "torch_shm"
1131 + and parallel_config.world_size_across_dp > 1
1132 + ):
1133 + raise ValueError(
1134 + "mm_tensor_ipc='torch_shm' is not supported with "
1135 + "data_parallel_size > 1 or tensor_parallel_size > 1 "
1136 + "or pipeline_parallel_size > 1."
1137 + )
1138 +
1139 def get_sliding_window(self) -> int | None:
1140 """Get the sliding window size from the HF text config if present."""
1141 return getattr(self.hf_text_config, "sliding_window", None)
vllm/config/multimodal.py
+6
@@ -59,6 +59,7 @@ class MultiModalDummyOptionsBuiltins(TypedDict, total=False):
59
60 MMEncoderTPMode = Literal["weights", "data"]
61 MMCacheType = Literal["shm", "lru"]
62 +MMTensorIPC = Literal["direct_rpc", "torch_shm"]
63 MMDummyOptions: TypeAlias = dict[str, BaseDummyOptions]
64 """
65 A dictionary containing an entry for each modality type of dummy data.
@@ -172,6 +173,11 @@ class MultiModalConfig:
173 Value sits in range [0;1) and determines fraction of media tokens
174 from each video to be pruned.
175 """
176 + mm_tensor_ipc: MMTensorIPC = "direct_rpc"
177 + """IPC (inter-process communication) method for multimodal tensors.
178 + - "direct_rpc": Use msgspec serialization via RPC
179 + - "torch_shm": Use torch.multiprocessing shared memory for zero-copy IPC
180 + Defaults to "direct_rpc". """
181
182 @field_validator("limit_per_prompt", mode="before")
183 @classmethod
vllm/config/vllm.py
+11
@@ -766,6 +766,17 @@ class VllmConfig:
766 else:
767 self.parallel_config.disable_nccl_for_dp_synchronization = False
768
769 + if (
770 + self.model_config is not None
771 + and self.model_config.multimodal_config is not None
772 + and self.model_config.multimodal_config.mm_tensor_ipc == "torch_shm"
773 + and os.environ.get("VLLM_WORKER_MULTIPROC_METHOD") != "spawn"
774 + ):
775 + raise ValueError(
776 + "torch_shm is known to fail without "
777 + "VLLM_WORKER_MULTIPROC_METHOD set to spawn"
778 + )
779 +
780 from vllm.platforms import current_platform
781
782 if (
vllm/engine/arg_utils.py
+6 -1
@@ -79,7 +79,7 @@ from vllm.config.model import (
79 RunnerOption,
80 TokenizerMode,
81 )
82 -from vllm.config.multimodal import MMCacheType, MMEncoderTPMode
82 +from vllm.config.multimodal import MMCacheType, MMEncoderTPMode, MMTensorIPC
83 from vllm.config.observability import DetailedTraceModules
84 from vllm.config.parallel import (
85 All2AllBackend,
@@ -509,6 +509,7 @@ class EngineArgs:
509 io_processor_plugin: str | None = None
510 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling
511 video_pruning_rate: float | None = MultiModalConfig.video_pruning_rate
512 + mm_tensor_ipc: MMTensorIPC = MultiModalConfig.mm_tensor_ipc
513 # LoRA fields
514 enable_lora: bool = False
515 max_loras: int = LoRAConfig.max_loras
@@ -1097,6 +1098,9 @@ class EngineArgs:
1098 multimodal_group.add_argument(
1099 "--video-pruning-rate", **multimodal_kwargs["video_pruning_rate"]
1100 )
1101 + multimodal_group.add_argument(
1102 + "--mm-tensor-ipc", **multimodal_kwargs["mm_tensor_ipc"]
1103 + )
1104
1105 # LoRA related configs
1106 lora_kwargs = get_kwargs(LoRAConfig)
@@ -1423,6 +1427,7 @@ class EngineArgs:
1427 override_attention_dtype=self.override_attention_dtype,
1428 logits_processors=self.logits_processors,
1429 video_pruning_rate=self.video_pruning_rate,
1430 + mm_tensor_ipc=self.mm_tensor_ipc,
1431 io_processor_plugin=self.io_processor_plugin,
1432 )
1433
vllm/entrypoints/cli/serve.py
+2 -1
@@ -290,7 +290,7 @@ def run_multi_api_server(args: argparse.Namespace):
290
291 with launch_core_engines(
292 vllm_config, executor_class, log_stats, addresses, num_api_servers
293 - ) as (local_engine_manager, coordinator, addresses):
293 + ) as (local_engine_manager, coordinator, addresses, tensor_queue):
294 # Construct common args for the APIServerProcessManager up-front.
295 api_server_manager_kwargs = dict(
296 target_server_fn=run_api_server_worker_proc,
@@ -303,6 +303,7 @@ def run_multi_api_server(args: argparse.Namespace):
303 stats_update_address=coordinator.get_stats_publish_address()
304 if coordinator
305 else None,
306 + tensor_queue=tensor_queue,
307 )
308
309 # For dp ranks > 0 in external/hybrid DP LB modes, we must delay the
vllm/v1/engine/core.py
+17 -7
@@ -13,6 +13,7 @@ from enum import IntEnum
13 from functools import partial
14 from inspect import isclass, signature
15 from logging import DEBUG
16 +from multiprocessing.queues import Queue
17 from typing import Any, TypeVar, cast
18
19 import msgspec
@@ -59,6 +60,7 @@ from vllm.v1.engine import (
60 UtilityOutput,
61 UtilityResult,
62 )
63 +from vllm.v1.engine.tensor_ipc import TensorIpcReceiver
64 from vllm.v1.engine.utils import (
65 EngineHandshakeMetadata,
66 EngineZmqAddresses,
@@ -788,6 +790,7 @@ class EngineCoreProc(EngineCore):
790 executor_class: type[Executor],
791 log_stats: bool,
792 client_handshake_address: str | None = None,
793 + tensor_queue: Queue | None = None,
794 *,
795 engine_index: int = 0,
796 ):
@@ -802,6 +805,12 @@ class EngineCoreProc(EngineCore):
805 self.engines_running = False
806 self.shutdown_state = EngineShutdownState.RUNNING
807
808 + # Receiver for tensor IPC
809 + self.tensor_ipc_receiver: TensorIpcReceiver | None = None
810 + if tensor_queue is not None:
811 + self.tensor_ipc_receiver = TensorIpcReceiver(tensor_queue)
812 + logger.info("Using tensor IPC queue for multimodal tensor sharing")
813 +
814 with self._perform_handshakes(
815 handshake_address,
816 identity,
@@ -1340,9 +1349,11 @@ class EngineCoreProc(EngineCore):
1349 ):
1350 """Input socket IO thread."""
1351
1343 - # Msgpack serialization decoding.
1344 - add_request_decoder = MsgpackDecoder(EngineCoreRequest)
1345 - generic_decoder = MsgpackDecoder()
1352 + # Msgpack serialization decoding with optional tensor IPC receiver.
1353 + add_request_decoder = MsgpackDecoder(
1354 + EngineCoreRequest, oob_tensor_provider=self.tensor_ipc_receiver
1355 + )
1356 + generic_decoder = MsgpackDecoder(oob_tensor_provider=self.tensor_ipc_receiver)
1357
1358 with ExitStack() as stack, zmq.Context() as ctx:
1359 input_sockets = [
@@ -1418,10 +1429,7 @@ class EngineCoreProc(EngineCore):
1429 self.input_queue.put_nowait((request_type, request))
1430
1431 def process_output_sockets(
1421 - self,
1422 - output_paths: list[str],
1423 - coord_output_path: str | None,
1424 - engine_index: int,
1432 + self, output_paths: list[str], coord_output_path: str | None, engine_index: int
1433 ):
1434 """Output socket IO thread."""
1435
@@ -1580,6 +1588,7 @@ class DPEngineCoreProc(EngineCoreProc):
1588 executor_class: type[Executor],
1589 log_stats: bool,
1590 client_handshake_address: str | None = None,
1591 + tensor_queue: Queue | None = None,
1592 ):
1593 assert vllm_config.model_config.is_moe, (
1594 "DPEngineCoreProc should only be used for MoE models"
@@ -1605,6 +1614,7 @@ class DPEngineCoreProc(EngineCoreProc):
1614 log_stats,
1615 client_handshake_address,
1616 engine_index=dp_rank,
1617 + tensor_queue=tensor_queue,
1618 )
1619
1620 def _init_data_parallel(self, vllm_config: VllmConfig):
vllm/v1/engine/core_client.py
+17 -4
@@ -12,6 +12,7 @@ from collections import defaultdict, deque
12 from collections.abc import Awaitable, Callable, Sequence
13 from concurrent.futures import Future
14 from dataclasses import dataclass
15 +from multiprocessing.queues import Queue
16 from threading import Thread
17 from typing import Any, TypeAlias, TypeVar
18
@@ -45,6 +46,7 @@ from vllm.v1.engine import (
46 from vllm.v1.engine.coordinator import DPCoordinator
47 from vllm.v1.engine.core import EngineCore, EngineCoreProc
48 from vllm.v1.engine.exceptions import EngineDeadError
49 +from vllm.v1.engine.tensor_ipc import TensorIpcSender
50 from vllm.v1.engine.utils import (
51 CoreEngineActorManager,
52 CoreEngineProcManager,
@@ -477,9 +479,6 @@ class MPClient(EngineCoreClient):
479 client_addresses: dict[str, str] | None = None,
480 ):
481 self.vllm_config = vllm_config
480 - # Serialization setup.
481 - self.encoder = MsgpackEncoder()
482 - self.decoder = MsgpackDecoder(EngineCoreOutputs)
482
483 # ZMQ setup.
484 sync_ctx = zmq.Context(io_threads=2)
@@ -501,11 +500,14 @@ class MPClient(EngineCoreClient):
500 enable_input_socket_handover = parallel_config.enable_elastic_ep
501
502 self.stats_update_address: str | None = None
503 + tensor_queue: Queue | None = None
504 if client_addresses:
505 # Engines are managed externally to this client.
506 input_address = client_addresses["input_address"]
507 output_address = client_addresses["output_address"]
508 self.stats_update_address = client_addresses.get("stats_update_address")
509 + # Tensor queues passed via client_addresses for multi-API-server case
510 + tensor_queue = client_addresses.get("tensor_queue") # type: ignore[assignment]
511 self.input_socket = self.resources.input_socket = make_zmq_socket(
512 self.ctx,
513 input_address,
@@ -532,7 +534,7 @@ class MPClient(EngineCoreClient):
534
535 with launch_core_engines(
536 vllm_config, executor_class, log_stats, addresses
535 - ) as (engine_manager, coordinator, addresses):
537 + ) as (engine_manager, coordinator, addresses, tensor_queue):
538 self.resources.coordinator = coordinator
539 self.resources.engine_manager = engine_manager
540
@@ -542,6 +544,17 @@ class MPClient(EngineCoreClient):
544 coordinator.get_stats_publish_address()
545 )
546
547 + # Serialization setup with tensor queues for multimodal tensor IPC.
548 + tensor_ipc_sender: TensorIpcSender | None = None
549 + model_config = getattr(vllm_config, "model_config", None)
550 + if model_config is not None and model_config.multimodal_config is not None:
551 + mm_tensor_ipc = model_config.multimodal_config.mm_tensor_ipc
552 + if mm_tensor_ipc == "torch_shm" and tensor_queue is not None:
553 + tensor_ipc_sender = TensorIpcSender(tensor_queue)
554 +
555 + self.encoder = MsgpackEncoder(oob_tensor_consumer=tensor_ipc_sender)
556 + self.decoder = MsgpackDecoder(EngineCoreOutputs)
557 +
558 dp_size = parallel_config.data_parallel_size
559 dp_rank = parallel_config.data_parallel_index
560 dp_local_size = parallel_config.data_parallel_size_local
vllm/v1/engine/tensor_ipc.py new
+178
@@ -0,0 +1,178 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""Tensor IPC transport via torch.multiprocessing.Queue.
5 +
6 +This module contains the queue-based transport logic for sharing tensors
7 +between processes (e.g., API server -> engine core). The msgpack layer
8 +emits/consumes lightweight :class:`TensorIpcData` values, while transport
9 +state such as request association, handle generation, queue routing, buffering,
10 +and cleanup lives here.
11 +"""
12 +
13 +import dataclasses
14 +import uuid
15 +from collections import defaultdict
16 +from dataclasses import field
17 +from multiprocessing.queues import Queue as MPQueue
18 +from typing import Any
19 +
20 +import torch
21 +
22 +from vllm.logger import init_logger
23 +from vllm.v1.serial_utils import OOBTensorConsumer
24 +
25 +logger = init_logger(__name__)
26 +
27 +TensorIpcQueue = MPQueue
28 +
29 +
30 +@dataclasses.dataclass
31 +class TensorIpcData:
32 + """
33 + Data sent via torch.multiprocessing.Queue for zero-copy IPC.
34 +
35 + Contains the tensor_id and the actual tensor. The tensor is
36 + shared in memory (GPU or CPU) for efficient inter-process communication.
37 + """
38 +
39 + sender_id: str
40 + message_id: int
41 + tensor_id: int
42 + tensor: torch.Tensor
43 +
44 +
45 +class TensorIpcSender(OOBTensorConsumer):
46 + """Send-side logic for tensor IPC via torch.multiprocessing.Queue.
47 +
48 + Uses a single queue targeting rank 0 (the only rank that consumes
49 + multimodal tensors during TP>1 / PP>1. Note: DP>1 not supported).
50 + """
51 +
52 + def __init__(self, queue: TensorIpcQueue):
53 + self.queue = queue
54 + self._tensor_id_counter = 0
55 + self._message_counter = 0
56 + self._sender_id = uuid.uuid4().hex[:8]
57 +
58 + def set_target_engine(self, target_engine: int) -> None:
59 + if target_engine != 0:
60 + raise IndexError(
61 + "TensorIpcSender only supports a single queue; "
62 + f"got target engine {target_engine}"
63 + )
64 +
65 + def new_message(self) -> None:
66 + self._message_counter += 1
67 + self._tensor_id_counter = 0
68 +
69 + def __call__(self, tensor: torch.Tensor) -> dict[str, Any] | None:
70 + """Send tensor via queue, return its handle. Returns None if failed."""
71 + try:
72 + # Move tensor to shared memory for IPC
73 + # This is required for proper inter-process communication
74 + if not tensor.is_shared():
75 + tensor = tensor.share_memory_()
76 +
77 + metadata = {
78 + "sender_id": self._sender_id,
79 + "message_id": self._message_counter,
80 + "tensor_id": self._tensor_id_counter,
81 + }
82 +
83 + self._tensor_id_counter += 1
84 +
85 + ipc_data = TensorIpcData(**metadata, tensor=tensor) # type: ignore[arg-type]
86 +
87 + # Use a timeout to avoid blocking indefinitely
88 + self.queue.put(ipc_data, timeout=10.0)
89 +
90 + logger.debug(
91 + "Sent tensor %s for (shape=%s, device=%s) "
92 + "via IPC queue (shared memory)",
93 + metadata,
94 + tensor.shape,
95 + tensor.device,
96 + )
97 +
98 + return metadata
99 + except Exception as e:
100 + logger.warning(
101 + "Failed to send tensor via IPC queue: %s. "
102 + "Falling back to standard serialization.",
103 + e,
104 + )
105 + return None
106 +
107 +
108 +@dataclasses.dataclass
109 +class _Sender:
110 + current_message_id: int = -1
111 + tensors: dict[int, dict[int, torch.Tensor]] = field(default_factory=dict)
112 +
113 +
114 +class TensorIpcReceiver:
115 + """Receive-side logic for tensor IPC via torch.multiprocessing.Queue.
116 +
117 + Wraps the queue receive logic previously embedded in MsgpackDecoder.
118 + """
119 +
120 + def __init__(self, queue: TensorIpcQueue):
121 + self.queue = queue
122 + self._tensor_buffers = defaultdict[str, _Sender](_Sender)
123 +
124 + def __call__(
125 + self, dtype: str, shape: tuple[int, ...], meta: dict[str, Any]
126 + ) -> torch.Tensor:
127 + """Retrieve a tensor from torch.multiprocessing.Queue.
128 +
129 + Uses a drain-and-buffer pattern: drains all available tensors from
130 + the queue, buffering them, until the requested tensor is found.
131 + Works for CUDA and CPU.
132 + """
133 +
134 + # Create lookup key from handle
135 + sender_id: str = meta["sender_id"]
136 + message_id: int = meta["message_id"]
137 + tensor_id: int = meta["tensor_id"]
138 +
139 + # Drain all available tensors. We save them regardless if this is
140 + # the one we're waiting for as they may arrive out of order from
141 + # multiple producers.
142 + while True:
143 + sender = self._tensor_buffers.get(sender_id)
144 + if sender is not None:
145 + tensors = sender.tensors
146 + tensor = tensors.get(message_id, {}).pop(tensor_id, None)
147 + if tensor is not None:
148 + if sender.current_message_id != message_id:
149 + while tensors and (mid := next(iter(tensors))) < message_id:
150 + if sender.tensors.pop(mid):
151 + logger.warning(
152 + "Discarding %d stale tensors from sender %s",
153 + sender_id,
154 + )
155 + sender.current_message_id = message_id
156 + logger.debug(
157 + "Received tensor %s from sender %s for (shape=%s, device=%s) "
158 + "via IPC queue (shared memory)",
159 + (message_id, tensor_id),
160 + sender_id,
161 + tensor.shape,
162 + tensor.device,
163 + )
164 + return tensor
165 +
166 + ipc_data: TensorIpcData = self.queue.get(timeout=10.0)
167 +
168 + # Store tensor
169 + sender = self._tensor_buffers[ipc_data.sender_id]
170 + if sender.current_message_id > ipc_data.message_id:
171 + logger.warning(
172 + "Ignoring stale tensor from sender %s", ipc_data.sender_id
173 + )
174 + continue
175 +
176 + sender.tensors.setdefault(ipc_data.message_id, {})[ipc_data.tensor_id] = (
177 + ipc_data.tensor
178 + )
vllm/v1/engine/utils.py
+15 -2
@@ -10,6 +10,7 @@ from dataclasses import dataclass
10 from enum import Enum, auto
11 from multiprocessing import Process, connection
12 from multiprocessing.process import BaseProcess
13 +from multiprocessing.queues import Queue
14 from typing import TYPE_CHECKING
15 from unittest.mock import patch
16
@@ -95,6 +96,7 @@ class CoreEngineProcManager:
96 executor_class: type[Executor],
97 log_stats: bool,
98 client_handshake_address: str | None = None,
99 + tensor_queue: Queue | None = None,
100 ):
101 context = get_mp_context()
102 common_kwargs = {
@@ -103,6 +105,7 @@ class CoreEngineProcManager:
105 "handshake_address": handshake_address,
106 "executor_class": executor_class,
107 "log_stats": log_stats,
108 + "tensor_queue": tensor_queue,
109 }
110
111 if client_handshake_address:
@@ -864,6 +867,7 @@ def launch_core_engines(
867 CoreEngineProcManager | CoreEngineActorManager | None,
868 DPCoordinator | None,
869 EngineZmqAddresses,
870 + Queue | None,
871 ]
872 ]:
873 """Launch engine and DP coordinator processes as needed."""
@@ -878,6 +882,14 @@ def launch_core_engines(
882
883 offline_mode = local_start_index is not None
884
885 + # Create a single tensor IPC queue for sharing multimodal tensors between
886 + # API servers and engine core. Returns a single queue since we only support
887 + # DP=1 for this data flow.
888 + tensor_queue: Queue | None = None
889 + multimodal_config = vllm_config.model_config.multimodal_config
890 + if multimodal_config is not None and multimodal_config.mm_tensor_ipc == "torch_shm":
891 + tensor_queue = get_mp_context().Queue()
892 +
893 # Run the DP Coordinator process with rank 0 when in online DP mode.
894 # The coordinator is needed for:
895 # 1. Internal/hybrid LB: collecting and publishing queue stats for load balancing
@@ -913,7 +925,7 @@ def launch_core_engines(
925 log_stats=log_stats,
926 )
927
916 - yield engine_actor_manager, coordinator, addresses
928 + yield engine_actor_manager, coordinator, addresses, tensor_queue
929 return
930
931 if offline_mode:
@@ -975,11 +987,12 @@ def launch_core_engines(
987 local_engine_count=local_engine_count,
988 start_index=dp_rank,
989 local_start_index=local_start_index or 0,
990 + tensor_queue=tensor_queue,
991 )
992 else:
993 local_engine_manager = None
994
982 - yield local_engine_manager, coordinator, addresses
995 + yield local_engine_manager, coordinator, addresses, tensor_queue
996
997 # Now wait for engines to start.
998 wait_for_engine_startup(
vllm/v1/serial_utils.py
+59 -8
@@ -4,6 +4,7 @@
4 import dataclasses
5 import importlib
6 import pickle
7 +from abc import ABC, abstractmethod
8 from collections.abc import Callable, Sequence
9 from functools import partial
10 from inspect import isclass
@@ -53,6 +54,27 @@ MMF_CLASS_TO_FACTORY: dict[type[BaseMultiModalField], str] = {
54 bytestr: TypeAlias = bytes | bytearray | memoryview | zmq.Frame
55
56
57 +class OOBTensorConsumer(ABC):
58 + @abstractmethod
59 + def __call__(self, tensor: torch.Tensor) -> dict | None:
60 + """
61 + Called with tensors for the current message.
62 + Returns None to reject the tensor (falls back to regular serialization),
63 + otherwise a dict with arbitrary placeholder data to be included
64 + in the serialized message.
65 + """
66 + return None
67 +
68 + @abstractmethod
69 + def new_message(self) -> None:
70 + """Called at the start of each new encoded message."""
71 + pass
72 +
73 +
74 +# dtype, shape, metadata -> tensor
75 +OOBTensorProvider = Callable[[str, tuple[int, ...], dict], torch.Tensor]
76 +
77 +
78 def _log_insecure_serialization_warning():
79 logger.warning_once(
80 "Allowing insecure serialization using pickle due to "
@@ -119,9 +141,16 @@ class MsgpackEncoder:
141
142 By default, arrays below 256B are serialized inline Larger will get sent
143 via dedicated messages. Note that this is a per-tensor limit.
144 +
145 + When a ``oob_tensor_consumer`` is provided, tensors (CUDA and CPU) will be
146 + offered to it for out-of-band handling.
147 """
148
124 - def __init__(self, size_threshold: int | None = None):
149 + def __init__(
150 + self,
151 + size_threshold: int | None = None,
152 + oob_tensor_consumer: OOBTensorConsumer | None = None,
153 + ):
154 if size_threshold is None:
155 size_threshold = envs.VLLM_MSGPACK_ZERO_COPY_THRESHOLD
156 self.encoder = msgpack.Encoder(enc_hook=self.enc_hook)
@@ -130,11 +159,14 @@ class MsgpackEncoder:
159 # pass custom data to the hook otherwise.
160 self.aux_buffers: list[bytestr] | None = None
161 self.size_threshold = size_threshold
162 + self.oob_tensor_consumer = oob_tensor_consumer
163 if envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
164 _log_insecure_serialization_warning()
165
166 def encode(self, obj: Any) -> Sequence[bytestr]:
167 try:
168 + if self.oob_tensor_consumer is not None:
169 + self.oob_tensor_consumer.new_message()
170 self.aux_buffers = bufs = [b""]
171 bufs[0] = self.encoder.encode(obj)
172 # This `bufs` list allows us to collect direct pointers to backing
@@ -147,6 +179,8 @@ class MsgpackEncoder:
179
180 def encode_into(self, obj: Any, buf: bytearray) -> Sequence[bytestr]:
181 try:
182 + if self.oob_tensor_consumer is not None:
183 + self.oob_tensor_consumer.new_message()
184 self.aux_buffers = [buf]
185 bufs = self.aux_buffers
186 self.encoder.encode_into(obj, buf)
@@ -222,17 +256,19 @@ class MsgpackEncoder:
256
257 def _encode_tensor(
258 self, obj: torch.Tensor
225 - ) -> tuple[str, tuple[int, ...], int | memoryview]:
226 - assert self.aux_buffers is not None
259 + ) -> tuple[str, tuple[int, ...], int | dict | memoryview]:
260 + oob_consumer = self.oob_tensor_consumer
261 # view the tensor as a contiguous 1D array of bytes
228 - arr_data = tensor_data(obj)
229 - if obj.nbytes < self.size_threshold:
262 + if obj.nbytes < self.size_threshold and obj.is_cpu:
263 # Smaller tensors are encoded inline, just like ndarrays.
231 - data = msgpack.Ext(CUSTOM_TYPE_RAW_VIEW, arr_data)
264 + data = msgpack.Ext(CUSTOM_TYPE_RAW_VIEW, tensor_data(obj))
265 + elif oob_consumer is not None and (data := oob_consumer(obj)) is not None:
266 + assert isinstance(data, dict)
267 else:
268 # Otherwise encode index of backing buffer to avoid copy.
269 + assert self.aux_buffers is not None
270 data = len(self.aux_buffers)
235 - self.aux_buffers.append(arr_data)
271 + self.aux_buffers.append(tensor_data(obj))
272 dtype = str(obj.dtype).removeprefix("torch.")
273 return dtype, obj.shape, data
274
@@ -279,9 +315,17 @@ class MsgpackDecoder:
315
316 Note that unlike vanilla `msgspec` Decoders, this interface is generally
317 not thread-safe when encoding tensors / numpy arrays.
318 +
319 + ``oob_tensor_provider`` must be used when an OOBTensorConsumer is used on the
320 + encoder side.
321 """
322
284 - def __init__(self, t: Any | None = None, share_mem: bool = True):
323 + def __init__(
324 + self,
325 + t: Any | None = None,
326 + share_mem: bool = True,
327 + oob_tensor_provider: OOBTensorProvider | None = None,
328 + ):
329 self.share_mem = share_mem
330 self.pin_tensors = is_pin_memory_available()
331 args = () if t is None else (t,)
@@ -289,6 +333,7 @@ class MsgpackDecoder:
333 *args, ext_hook=self.ext_hook, dec_hook=self.dec_hook
334 )
335 self.aux_buffers: Sequence[bytestr] = ()
336 + self.oob_tensor_provider = oob_tensor_provider
337 if envs.VLLM_ALLOW_INSECURE_SERIALIZATION:
338 _log_insecure_serialization_warning()
339
@@ -353,6 +398,12 @@ class MsgpackDecoder:
398
399 def _decode_tensor(self, arr: Any) -> torch.Tensor:
400 dtype, shape, data = arr
401 + if isinstance(data, dict):
402 + assert self.oob_tensor_provider, (
403 + "Received OOB tensor but tensor provider is not set"
404 + )
405 + return self.oob_tensor_provider(dtype, shape, data)
406 +
407 is_aux = isinstance(data, int)
408 buffer = self.aux_buffers[data] if is_aux else data
409 buffer = buffer if isinstance(buffer, memoryview) else memoryview(buffer)
vllm/v1/utils.py
+6 -1
@@ -10,6 +10,7 @@ from contextlib import AbstractContextManager
10 from dataclasses import dataclass
11 from multiprocessing import connection
12 from multiprocessing.process import BaseProcess
13 +from multiprocessing.queues import Queue
14 from typing import (
15 TYPE_CHECKING,
16 Any,
@@ -173,6 +174,7 @@ class APIServerProcessManager:
174 input_addresses: list[str],
175 output_addresses: list[str],
176 stats_update_address: str | None = None,
177 + tensor_queue: Queue | None = None,
178 ):
179 """Initialize and start API server worker processes.
180
@@ -185,6 +187,7 @@ class APIServerProcessManager:
187 input_addresses: Input addresses for each API server
188 output_addresses: Output addresses for each API server
189 stats_update_address: Optional stats update address
190 + tensor_queue: Optional tensor IPC queue for sharing MM tensors
191 """
192 self.listen_address = listen_address
193 self.sock = sock
@@ -205,6 +208,8 @@ class APIServerProcessManager:
208 }
209 if stats_update_address is not None:
210 client_config["stats_update_address"] = stats_update_address
211 + if tensor_queue is not None:
212 + client_config["tensor_queue"] = tensor_queue
213
214 proc = spawn_context.Process(
215 target=target_server_fn,
@@ -419,7 +424,7 @@ def tensor_data(tensor: torch.Tensor) -> memoryview:
424 Returns:
425 A memoryview of the tensor data as uint8.
426 """
422 - return tensor.flatten().contiguous().view(torch.uint8).numpy().data
427 + return tensor.flatten().cpu().contiguous().view(torch.uint8).numpy().data
428
429
430 @dataclass