[Bugfix][ROCm] Keep TP all_gather on base-class collective (#48446)
Signed-off-by: fai <fangzhouai@gmail.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Fangzhou Ai committed
Jul 12, 2026 at 23:53 UTC
775c1589ea9c7df49a866f900341e3645c5bd60c
1 file changed
+8
vllm/distributed/device_communicators/cuda_communicator.py
+8
@@ -352,6 +352,14 @@ class CudaCommunicator(DeviceCommunicatorBase):
352
if pynccl_comm is None or pynccl_comm.disabled:
353
return super().all_gather(input_, dim)
354
355
+ # On ROCm, the base-class all_gather (all_gather_into_tensor) is faster
356
+ # than the manual pynccl + torch.empty + movedim + reshape path below,
357
+ # which adds a per-call output allocation and (for dim != 0) an extra
358
+ # copy on every step. This is on the hot path for TP forward passes, so
359
+ # keep ROCm on the base-class collective to avoid a decode regression.
360
+ if current_platform.is_rocm():
361
+ return super().all_gather(input_, dim)
362
+
363
input_size = input_.size()
364
output_size = (input_size[0] * self.world_size,) + input_size[1:]
365
output_tensor = torch.empty(