kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-02-04 10:11:47 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

[None][fix] Avoid write-write race for async pp send. (#10488)

Browse Source 
Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
This commit is contained in:
Yuxian Qiu 2026-01-14 09:39:36 +08:00 committed by GitHub
parent 2acd03030a
commit d3f4fbb742
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
tensorrt_llm/_torch/distributed/communicator.py
View File

@ -839,9 +839,13 @@ class PPCommNCCL:
self.nccl_comm.send(tensor, dest)
return
self.tensor_ready_event.record()
# If the tensor is allocated from non-default memory pool
# like userbuffers, its underlying memory may be reused
# before the send operation is completed.
# We clone the tensor to avoid write-write conflicts.
tensor = tensor.clone()
self.send_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(self.send_stream):
self.tensor_ready_event.wait()
self.nccl_comm.send(tensor, dest)
def recv(self, tensor: torch.Tensor, src: Optional[int] = None):