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
[Bugfix] Fix several issues with p2p xPyD in GET type (#23993)
Signed-off-by: Csrayz <jover@cmbchina.com> Signed-off-by: ivyilike <pww123@cmbchina.com> Co-authored-by: ivyilike <pww123@cmbchina.com>
Csrayz committed
Sep 22, 2025 at 22:53 UTC
c10101a3eba053d22842edd95941d49e4324c727
2 files changed
+21
-10
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py
+4
-1
@@ -178,6 +178,9 @@ class P2pNcclConnector(KVConnectorBase_V1):
178
179
# Load the KV for each request each layer
180
for request in metadata.requests:
181
+ request_id = request.request_id
182
+ ip, port = self.parse_request_id(request_id, False)
183
+ remote_address = ip + ":" + str(port + self._rank)
184
for layer_name in forward_context.no_compile_layers:
185
layer = forward_context.no_compile_layers[layer_name]
186
@@ -191,7 +194,7 @@ class P2pNcclConnector(KVConnectorBase_V1):
194
layer = kv_cache[forward_context.virtual_engine]
195
196
kv_cache = self.p2p_nccl_engine.recv_tensor(
194
- request.request_id + "#" + layer_name)
197
+ request.request_id + "#" + layer_name, remote_address)
198
199
if kv_cache is None:
200
logger.warning("🚧kv_cache is None, %s", request.request_id)
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py
+17
-9
@@ -134,7 +134,6 @@ class P2pNcclEngine:
134
# PUT or PUT_ASYNC
135
# tensor_id: torch.Tensor
136
self.send_queue: deque[SendQueueItem] = deque()
137
- self.send_request_id_to_tensor_ids: dict[str, set[str]] = {}
137
if self.send_type == "PUT_ASYNC":
138
self._send_thread = threading.Thread(target=self.send_async,
139
daemon=True)
@@ -143,6 +142,7 @@ class P2pNcclEngine:
142
# tensor_id: torch.Tensor/(addr, dtype, shape)
143
self.recv_store: dict[str, Any] = {}
144
self.recv_request_id_to_tensor_ids: dict[str, set[str]] = {}
145
+ self.send_request_id_to_tensor_ids: dict[str, set[str]] = {}
146
self.socks: dict[str, Any] = {} # remote_address: client socket
147
self.comms: dict[str, Any] = {} # remote_address: (ncclComm_t, rank)
148
@@ -223,18 +223,26 @@ class P2pNcclEngine:
223
# GET
224
with self.send_store_cv:
225
tensor_size = tensor.element_size() * tensor.numel()
226
+ if tensor_size > self.buffer_size_threshold:
227
+ logger.warning(
228
+ "❗[GET]tensor_id:%s, tensor_size:%d, is greater than"
229
+ "buffer size threshold :%d, skip send to %s, rank:%d",
230
+ tensor_id, tensor_size, self.buffer_size_threshold,
231
+ remote_address, self.rank)
232
+ return False
233
while (self.buffer_size + tensor_size
234
> self.buffer_size_threshold):
228
- oldest_tenser_id = next(iter(self.send_store))
229
- oldest_tenser = self.send_store.pop(oldest_tenser_id)
230
- oldest_tenser_size = oldest_tenser.element_size(
231
- ) * oldest_tenser.numel()
232
- self.buffer_size -= oldest_tenser_size
233
- logger.info(
235
+ assert len(self.send_store) > 0
236
+ oldest_tensor_id = next(iter(self.send_store))
237
+ oldest_tensor = self.send_store.pop(oldest_tensor_id)
238
+ oldest_tensor_size = oldest_tensor.element_size(
239
+ ) * oldest_tensor.numel()
240
+ self.buffer_size -= oldest_tensor_size
241
+ logger.debug(
242
"⛔[GET]Send to %s, tensor_id:%s, tensor_size:%d,"
235
- " buffer_size:%d, oldest_tenser_size:%d, rank:%d",
243
+ " buffer_size:%d, oldest_tensor_size:%d, rank:%d",
244
remote_address, tensor_id, tensor_size, self.buffer_size,
237
- oldest_tenser_size, self.rank)
245
+ oldest_tensor_size, self.rank)
246
247
self.send_store[tensor_id] = tensor
248
self.buffer_size += tensor_size