From c8f09e9cf2051a3341b5b032bf9c35d882eaf13f Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Tue, 14 Apr 2026 22:32:05 +0000 Subject: [PATCH] revert custom reduce scattre (no lamport) Signed-off-by: Woosuk Kwon --- csrc/custom_all_reduce.cu | 47 ----------- csrc/custom_all_reduce.cuh | 83 +------------------ csrc/ops.h | 2 - csrc/torch_bindings.cpp | 4 - vllm/_custom_ops.py | 10 --- .../device_communicators/cuda_communicator.py | 4 - .../device_communicators/custom_all_reduce.py | 31 ------- 7 files changed, 1 insertion(+), 180 deletions(-) diff --git a/csrc/custom_all_reduce.cu b/csrc/custom_all_reduce.cu index 5046259ebf2..a38d6fa24a2 100644 --- a/csrc/custom_all_reduce.cu +++ b/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(_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(_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(stream, reinterpret_cast(reg_buffer), - reinterpret_cast(out.data_ptr()), - inp.numel()); - break; - } - case at::ScalarType::Half: { - fa->reduce_scatter(stream, reinterpret_cast(reg_buffer), - reinterpret_cast(out.data_ptr()), - inp.numel()); - break; - } -#if (__CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)) - case at::ScalarType::BFloat16: { - fa->reduce_scatter( - stream, reinterpret_cast(reg_buffer), - reinterpret_cast(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(_fa); } diff --git a/csrc/custom_all_reduce.cuh b/csrc/custom_all_reduce.cuh index 655ca1dfa41..58926f6429d 100644 --- a/csrc/custom_all_reduce.cuh +++ b/csrc/custom_all_reduce.cuh @@ -313,24 +313,6 @@ __global__ void __launch_bounds__(512, 1) barrier_at_end(sg, self_sg, rank); } -template -__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::P; - using A = typename packed_t::A; - auto dp = *_dp; - int offset = rank * chunk_size; - barrier_at_start(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((const P**)&dp.ptrs[0], offset + idx); - } - barrier_at_end(sg, self_sg, rank); -} - template DINLINE P* get_tmp_buf(Signal* sg) { return (P*)(((Signal*)sg) + 1); @@ -634,69 +616,6 @@ class CustomAllreduce { #undef KL } - template - void reduce_scatter(cudaStream_t stream, T* input, T* output, int size, - int threads = 512, int block_limit = defaultBlockLimit) { - auto d = packed_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(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<<>>( \ - 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(cudaStream_t, half *, half *, int, int, int); */ -} // namespace vllm +} // namespace vllm \ No newline at end of file diff --git a/csrc/ops.h b/csrc/ops.h index 9ce6f5b0c49..0a0b6c2d7d0 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -280,8 +280,6 @@ fptr_t init_custom_ar(const std::vector& 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& fake_ipc_ptrs); diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index 2fefc2e59e1..48062c3f47b 100644 --- a/csrc/torch_bindings.cpp +++ b/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); diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index edf4db8d538..d6780185be9 100644 --- a/vllm/_custom_ops.py +++ b/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) diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 1f45376b3f6..a6b9d511ba5 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/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:] diff --git a/vllm/distributed/device_communicators/custom_all_reduce.py b/vllm/distributed/device_communicators/custom_all_reduce.py index c38c1855077..65a19626468 100644 --- a/vllm/distributed/device_communicators/custom_all_reduce.py +++ b/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: