revert custom reduce scattre (no lamport)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

csrc/custom_all_reduce.cu

@@ -104,53 +104,6 @@ void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
   }
 }
 
-void reduce_scatter(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
-                    fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) {
-  auto fa = reinterpret_cast<vllm::CustomAllreduce*>(_fa);
-  const at::cuda::OptionalCUDAGuard device_guard(device_of(inp));
-  auto stream = c10::cuda::getCurrentCUDAStream().stream();
-
-  TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
-  TORCH_CHECK(inp.numel() == out.numel() * fa->world_size_);
-  TORCH_CHECK(_is_weak_contiguous(inp));
-  TORCH_CHECK(_is_weak_contiguous(out));
-  auto input_size = inp.numel() * inp.element_size();
-  auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
-  if (reg_buffer) {
-    TORCH_CHECK_LE(input_size, reg_buffer_sz_bytes);
-    AT_CUDA_CHECK(cudaMemcpyAsync(reg_buffer, inp.data_ptr(), input_size,
-                                  cudaMemcpyDeviceToDevice, stream));
-  } else {
-    reg_buffer = inp.data_ptr();
-  }
-  switch (out.scalar_type()) {
-    case at::ScalarType::Float: {
-      fa->reduce_scatter<float>(stream, reinterpret_cast<float*>(reg_buffer),
-                                reinterpret_cast<float*>(out.data_ptr()),
-                                inp.numel());
-      break;
-    }
-    case at::ScalarType::Half: {
-      fa->reduce_scatter<half>(stream, reinterpret_cast<half*>(reg_buffer),
-                               reinterpret_cast<half*>(out.data_ptr()),
-                               inp.numel());
-      break;
-    }
-#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__))
-    case at::ScalarType::BFloat16: {
-      fa->reduce_scatter<nv_bfloat16>(
-          stream, reinterpret_cast<nv_bfloat16*>(reg_buffer),
-          reinterpret_cast<nv_bfloat16*>(out.data_ptr()), inp.numel());
-      break;
-    }
-#endif
-    default:
-      throw std::runtime_error(
-          "custom reduce_scatter only supports float32, float16 and "
-          "bfloat16");
-  }
-}
-
 void dispose(fptr_t _fa) {
   delete reinterpret_cast<vllm::CustomAllreduce*>(_fa);
 }

csrc/custom_all_reduce.cuh

@@ -313,24 +313,6 @@ __global__ void __launch_bounds__(512, 1)
   barrier_at_end<ngpus, true>(sg, self_sg, rank);
 }
 
-template <typename T, int ngpus>
-__global__ void __launch_bounds__(512, 1)
-    cross_device_reduce_scatter(RankData* _dp, RankSignals sg, Signal* self_sg,
-                                T* __restrict__ result, int rank,
-                                int chunk_size) {
-  using P = typename packed_t<T>::P;
-  using A = typename packed_t<T>::A;
-  auto dp = *_dp;
-  int offset = rank * chunk_size;
-  barrier_at_start<ngpus>(sg, self_sg, rank);
-  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < chunk_size;
-       idx += gridDim.x * blockDim.x) {
-    ((P*)result)[idx] =
-        packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], offset + idx);
-  }
-  barrier_at_end<ngpus, true>(sg, self_sg, rank);
-}
-
 template <typename P>
 DINLINE P* get_tmp_buf(Signal* sg) {
   return (P*)(((Signal*)sg) + 1);
@@ -634,69 +616,6 @@ class CustomAllreduce {
 #undef KL
   }
 
-  template <typename T>
-  void reduce_scatter(cudaStream_t stream, T* input, T* output, int size,
-                      int threads = 512, int block_limit = defaultBlockLimit) {
-    auto d = packed_t<T>::P::size;
-    if (size % d != 0)
-      throw std::runtime_error(
-          "custom reduce_scatter currently requires input length to be "
-          "multiple of " +
-          std::to_string(d));
-    if (size % world_size_ != 0)
-      throw std::runtime_error(
-          "custom reduce_scatter requires input length to be divisible by "
-          "world_size");
-    if (block_limit > kMaxBlocks)
-      throw std::runtime_error("max supported block limit is " +
-                               std::to_string(kMaxBlocks) + ". Got " +
-                               std::to_string(block_limit));
-
-    RankData* ptrs;
-    cudaStreamCaptureStatus status;
-    CUDACHECK(cudaStreamIsCapturing(stream, &status));
-    if (status == cudaStreamCaptureStatusActive) {
-      ptrs = d_rank_data_base_ + graph_unreg_buffers_.size();
-      graph_unreg_buffers_.push_back(input);
-    } else {
-      auto it = buffers_.find(input);
-      if (it == buffers_.end())
-        throw std::runtime_error(
-            "buffer address " +
-            std::to_string(reinterpret_cast<uint64_t>(input)) +
-            " is not registered!");
-      ptrs = it->second;
-    }
-
-    int chunk_size = size / world_size_ / d;
-    int blocks = std::min(block_limit, (chunk_size + threads - 1) / threads);
-
-#define KL(ngpus)                                                        \
-  cross_device_reduce_scatter<T, ngpus><<<blocks, threads, 0, stream>>>( \
-      ptrs, sg_, self_sg_, output, rank_, chunk_size);
-
-    switch (world_size_) {
-      case 2:
-        KL(2);
-        break;
-      case 4:
-        KL(4);
-        break;
-      case 6:
-        KL(6);
-        break;
-      case 8:
-        KL(8);
-        break;
-      default:
-        throw std::runtime_error(
-            "custom reduce_scatter only supports num gpus in (2,4,6,8). "
-            "Actual num gpus = " +
-            std::to_string(world_size_));
-    }
-#undef KL
-  }
-
   ~CustomAllreduce() {
     for (auto [_, ptr] : ipc_handles_) {
       CUDACHECK(cudaIpcCloseMemHandle(ptr));
@@ -710,4 +629,4 @@ class CustomAllreduce {
  * template void vllm::CustomAllreduce::allreduce<half>(cudaStream_t, half *,
  half *, int, int, int);
 */
-}  // namespace vllm
+}  // namespace vllm

csrc/ops.h

@@ -280,8 +280,6 @@ fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
                       bool fully_connected);
 void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
                 fptr_t reg_buffer, int64_t reg_buffer_sz_bytes);
-void reduce_scatter(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
-                    fptr_t reg_buffer, int64_t reg_buffer_sz_bytes);
 void dispose(fptr_t _fa);
 int64_t meta_size();
 void register_buffer(fptr_t _fa, const std::vector<int64_t>& fake_ipc_ptrs);

csrc/torch_bindings.cpp

@@ -655,10 +655,6 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _custom_ar), custom_ar) {
       "all_reduce(int fa, Tensor inp, Tensor! out, int reg_buffer, "
       "int reg_buffer_sz_bytes) -> ()");
   custom_ar.impl("all_reduce", torch::kCUDA, &all_reduce);
-  custom_ar.def(
-      "reduce_scatter(int fa, Tensor inp, Tensor! out, int reg_buffer, "
-      "int reg_buffer_sz_bytes) -> ()");
-  custom_ar.impl("reduce_scatter", torch::kCUDA, &reduce_scatter);
 
   custom_ar.def("dispose", &dispose);
   custom_ar.def("meta_size", &meta_size);

vllm/_custom_ops.py

@@ -2797,16 +2797,6 @@ def all_reduce(
     torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
 
 
-def reduce_scatter(
-    fa: int,
-    inp: torch.Tensor,
-    out: torch.Tensor,
-    reg_buffer: int,
-    reg_buffer_sz_bytes: int,
-) -> None:
-    torch.ops._C_custom_ar.reduce_scatter(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
-
-
 def dispose(fa: int) -> None:
     torch.ops._C_custom_ar.dispose(fa)
 

vllm/distributed/device_communicators/cuda_communicator.py

@@ -250,10 +250,6 @@ class CudaCommunicator(DeviceCommunicatorBase):
         input_tensor = input_.movedim(0, dim).contiguous()
 
         assert input_tensor.shape[0] % world_size == 0
-
-        # TODO: custom reduce_scatter disabled for now — needs testing
-        # with torch.compile + CUDA graph pipeline before enabling.
-
         chunk_size = input_tensor.shape[0] // world_size
         output_shape = (chunk_size,) + input_tensor.shape[1:]
 

vllm/distributed/device_communicators/custom_all_reduce.py

@@ -281,37 +281,6 @@ class CustomAllreduce:
             # latency) compared to the performance gain of using custom kernels
             return self.all_reduce(input, registered=False)
 
-    def reduce_scatter(
-        self,
-        inp: torch.Tensor,
-        *,
-        out: torch.Tensor | None = None,
-        registered: bool = False,
-    ) -> torch.Tensor:
-        if out is None:
-            out_shape = (inp.shape[0] // self.world_size, *inp.shape[1:])
-            out = torch.empty(out_shape, dtype=inp.dtype, device=inp.device)
-        if registered:
-            ops.reduce_scatter(self._ptr, inp, out, 0, 0)
-        else:
-            ops.reduce_scatter(
-                self._ptr, inp, out, self.buffer_ptrs[self.rank], self.max_size
-            )
-        return out
-
-    def custom_reduce_scatter(self, input: torch.Tensor) -> torch.Tensor | None:
-        """The main reduce_scatter API that provides support for cuda graph."""
-        if self.disabled or not self.should_custom_ar(input):
-            return None
-        if self._IS_CAPTURING:
-            if torch.cuda.is_current_stream_capturing():
-                return self.reduce_scatter(input, registered=True)
-            else:
-                out_shape = (input.shape[0] // self.world_size, *input.shape[1:])
-                return torch.empty(out_shape, dtype=input.dtype, device=input.device)
-        else:
-            return self.reduce_scatter(input, registered=False)
-
     def close(self):
         if not self.disabled and self._ptr:
             if ops is not None: