From 239e6ff95bf1683dde2c8cfd0f27e8fbd0ba84f5 Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Mon, 20 Apr 2026 19:01:57 +0000 Subject: [PATCH] address comments Signed-off-by: yewentao256 --- .../device_communicators/base_device_communicator.py | 5 ++++- vllm/distributed/device_communicators/cuda_communicator.py | 6 ++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/vllm/distributed/device_communicators/base_device_communicator.py b/vllm/distributed/device_communicators/base_device_communicator.py index e6cf30f0f2d..df84144e712 100644 --- a/vllm/distributed/device_communicators/base_device_communicator.py +++ b/vllm/distributed/device_communicators/base_device_communicator.py @@ -188,6 +188,9 @@ class DeviceCommunicatorBase: "Invalid output shape for all_gather_into_tensor: " f"expected {expected_output_size}, got {tuple(output_tensor.shape)}" ) + assert input_.is_contiguous(), ( + "all_gather_into_tensor requires a contiguous input tensor" + ) assert output_tensor.is_contiguous(), ( "all_gather_into_tensor requires a contiguous output tensor" ) @@ -208,7 +211,7 @@ class DeviceCommunicatorBase: output_size, dtype=input_.dtype, device=input_.device ) # All-gather. - self.all_gather_into_tensor(output_tensor, input_.contiguous()) + self.all_gather_into_tensor(output_tensor, input_) # Reshape output_tensor = output_tensor.reshape((self.world_size,) + input_size) output_tensor = output_tensor.movedim(0, dim) diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 557fba66f85..c83df6d0576 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -240,8 +240,7 @@ class CudaCommunicator(DeviceCommunicatorBase): self, output_tensor: torch.Tensor, input_: torch.Tensor ) -> torch.Tensor: pynccl_comm = self.pynccl_comm - if pynccl_comm is None or pynccl_comm.disabled: - return super().all_gather_into_tensor(output_tensor, input_) + assert pynccl_comm is not None and not pynccl_comm.disabled pynccl_comm.all_gather(output_tensor, input_) return output_tensor @@ -249,8 +248,7 @@ class CudaCommunicator(DeviceCommunicatorBase): self, output_tensor: torch.Tensor, input_tensor: torch.Tensor ) -> torch.Tensor: pynccl_comm = self.pynccl_comm - if pynccl_comm is None or pynccl_comm.disabled: - return super().reduce_scatter_tensor(output_tensor, input_tensor) + assert pynccl_comm is not None and not pynccl_comm.disabled pynccl_comm.reduce_scatter(output_tensor, input_tensor) return output_tensor