@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 176 lines 6 KB
Raw
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 """Data sent via torch.multiprocessing.Queue for zero-copy IPC.
33
34 Contains the tensor_id and the actual tensor. The tensor is
35 shared in memory (GPU or CPU) for efficient inter-process communication.
36 """
37
38 sender_id: str
39 message_id: int
40 tensor_id: int
41 tensor: torch.Tensor
42
43
44 class TensorIpcSender(OOBTensorConsumer):
45 """Send-side logic for tensor IPC via torch.multiprocessing.Queue.
46
47 Uses a single queue targeting rank 0 (the only rank that consumes
48 multimodal tensors during TP>1 / PP>1. Note: DP>1 not supported).
49 """
50
51 def __init__(self, queue: TensorIpcQueue):
52 self.queue = queue
53 self._tensor_id_counter = 0
54 self._message_counter = 0
55 self._sender_id = uuid.uuid4().hex[:8]
56
57 def set_target_engine(self, target_engine: int) -> None:
58 if target_engine != 0:
59 raise IndexError(
60 "TensorIpcSender only supports a single queue; "
61 f"got target engine {target_engine}"
62 )
63
64 def new_message(self) -> None:
65 self._message_counter += 1
66 self._tensor_id_counter = 0
67
68 def __call__(self, tensor: torch.Tensor) -> dict[str, Any] | None:
69 """Send tensor via queue, return its handle. Returns None if failed."""
70 try:
71 # Move tensor to shared memory for IPC
72 # This is required for proper inter-process communication
73 if not tensor.is_shared():
74 tensor = tensor.share_memory_()
75
76 metadata = {
77 "sender_id": self._sender_id,
78 "message_id": self._message_counter,
79 "tensor_id": self._tensor_id_counter,
80 }
81
82 self._tensor_id_counter += 1
83
84 ipc_data = TensorIpcData(**metadata, tensor=tensor) # type: ignore[arg-type]
85
86 # Use a timeout to avoid blocking indefinitely
87 self.queue.put(ipc_data, timeout=10.0)
88
89 logger.debug(
90 "Sent tensor %s for (shape=%s, device=%s) "
91 "via IPC queue (shared memory)",
92 metadata,
93 tensor.shape,
94 tensor.device,
95 )
96
97 return metadata
98 except Exception as e:
99 logger.warning(
100 "Failed to send tensor via IPC queue: %s. "
101 "Falling back to standard serialization.",
102 e,
103 )
104 return None
105
106
107 @dataclasses.dataclass
108 class _Sender:
109 current_message_id: int = -1
110 tensors: dict[int, dict[int, torch.Tensor]] = field(default_factory=dict)
111
112
113 class TensorIpcReceiver:
114 """Receive-side logic for tensor IPC via torch.multiprocessing.Queue.
115
116 Wraps the queue receive logic previously embedded in MsgpackDecoder.
117 """
118
119 def __init__(self, queue: TensorIpcQueue):
120 self.queue = queue
121 self._tensor_buffers = defaultdict[str, _Sender](_Sender)
122
123 def __call__(
124 self, dtype: str, shape: tuple[int, ...], meta: dict[str, Any]
125 ) -> torch.Tensor:
126 """Retrieve a tensor from torch.multiprocessing.Queue.
127
128 Uses a drain-and-buffer pattern: drains all available tensors from
129 the queue, buffering them, until the requested tensor is found.
130 Works for CUDA and CPU.
131 """
132 # Create lookup key from handle
133 sender_id: str = meta["sender_id"]
134 message_id: int = meta["message_id"]
135 tensor_id: int = meta["tensor_id"]
136
137 # Drain all available tensors. We save them regardless if this is
138 # the one we're waiting for as they may arrive out of order from
139 # multiple producers.
140 while True:
141 sender = self._tensor_buffers.get(sender_id)
142 if sender is not None:
143 tensors = sender.tensors
144 tensor = tensors.get(message_id, {}).pop(tensor_id, None)
145 if tensor is not None:
146 if sender.current_message_id != message_id:
147 while tensors and (mid := next(iter(tensors))) < message_id:
148 if sender.tensors.pop(mid):
149 logger.warning(
150 "Discarding %d stale tensors from sender %s",
151 sender_id,
152 )
153 sender.current_message_id = message_id
154 logger.debug(
155 "Received tensor %s from sender %s for (shape=%s, device=%s) "
156 "via IPC queue (shared memory)",
157 (message_id, tensor_id),
158 sender_id,
159 tensor.shape,
160 tensor.device,
161 )
162 return tensor
163
164 ipc_data: TensorIpcData = self.queue.get(timeout=10.0)
165
166 # Store tensor
167 sender = self._tensor_buffers[ipc_data.sender_id]
168 if sender.current_message_id > ipc_data.message_id:
169 logger.warning(
170 "Ignoring stale tensor from sender %s", ipc_data.sender_id
171 )
172 continue
173
174 sender.tensors.setdefault(ipc_data.message_id, {})[ipc_data.tensor_id] = (
175 ipc_data.tensor
176 )