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 15, 03:30 UTC
[Bugfix][V1][P/D]Fix the issue of occasional garbled output for P2pNcclConnector (#20263)
Signed-off-by: Abatom <abzhonghua@gmail.com>
Zhonghua Deng committed
Jul 1, 2025 at 07:45 UTC
ded1fb635b7c1504a83fc7c195a5bf47d31c1bef
1 file changed
+5
-4
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py
+5
-4
@@ -310,10 +310,11 @@ class P2pNcclEngine:
310
elif data["cmd"] == "PUT":
311
tensor_id = data["tensor_id"]
312
try:
313
- tensor = torch.empty(data["shape"],
314
- dtype=getattr(
315
- torch, data["dtype"]),
316
- device=self.device)
313
+ with torch.cuda.stream(self.recv_stream):
314
+ tensor = torch.empty(data["shape"],
315
+ dtype=getattr(
316
+ torch, data["dtype"]),
317
+ device=self.device)
318
self.router_socket.send_multipart(
319
[remote_address, b"0"])
320
comm, rank = self.comms[remote_address.decode()]