address comments

Signed-off-by: yewentao256 <zhyanwentao@126.com>

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)

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