Skip to content

GitLab

Unverified Commit ded1fb63 authored by Zhonghua Deng's avatar Zhonghua Deng Committed by GitHub
Browse files

[Bugfix][V1][P/D]Fix the issue of occasional garbled output for P2pNcclConnector (#20263) 


Signed-off-by: default avatarAbatom <abzhonghua@gmail.com>
parent 97d9524f
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 4 deletions
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_engine.py
View file @ ded1fb63
...@@ -310,10 +310,11 @@ class P2pNcclEngine: ...@@ -310,10 +310,11 @@ class P2pNcclEngine:
elif data["cmd"] == "PUT": elif data["cmd"] == "PUT":
tensor_id = data["tensor_id"] tensor_id = data["tensor_id"]
try: try:
tensor = torch.empty(data["shape"], with torch.cuda.stream(self.recv_stream):
dtype=getattr( tensor = torch.empty(data["shape"],
torch, data["dtype"]), dtype=getattr(
device=self.device) torch, data["dtype"]),
device=self.device)
self.router_socket.send_multipart( self.router_socket.send_multipart(
[remote_address, b"0"]) [remote_address, b"0"])
comm, rank = self.comms[remote_address.decode()] comm, rank = self.comms[remote_address.decode()]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment