diff --git a/CMakeLists.txt b/CMakeLists.txt index 1259ec0c1bf..41e33a14817 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -324,6 +324,7 @@ endif() set(VLLM_EXT_SRC "csrc/quantization/activation_kernels.cu" + "csrc/push_all_reduce.cu" "csrc/torch_bindings.cpp") if(VLLM_GPU_LANG STREQUAL "CUDA") diff --git a/csrc/push_all_reduce.cu b/csrc/push_all_reduce.cu new file mode 100644 index 00000000000..60035c339bd --- /dev/null +++ b/csrc/push_all_reduce.cu @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// C++ bridge functions for push-based allreduce. +// Exposes PushAllReduceManager to Python via torch custom ops. + +#include "push_all_reduce.cuh" + +#include +#include +#include + +using fptr_t = int64_t; +using namespace vllm::push_ar; + +// Initialize the manager; returns opaque pointer as int64_t +fptr_t init_push_ar(int64_t rank, int64_t world_size, + int64_t push_buffer_bytes, int64_t max_num_cta) { + auto* mgr = new PushAllReduceManager(static_cast(rank), + static_cast(world_size), + push_buffer_bytes, + static_cast(max_num_cta)); + return reinterpret_cast(mgr); +} + +// Get IPC handle as a byte tensor +torch::Tensor get_push_ar_ipc_handle(fptr_t _mgr) { + auto* mgr = reinterpret_cast(_mgr); + cudaIpcMemHandle_t handle = mgr->get_ipc_handle(); + auto t = torch::from_blob(&handle, + {static_cast(sizeof(handle))}, + torch::kUInt8) + .clone(); + return t; +} + +// Post-init with peer IPC handles +void post_init_push_ar(fptr_t _mgr, torch::Tensor all_handles) { + auto* mgr = reinterpret_cast(_mgr); + int world_size = all_handles.size(0); + std::vector handles(world_size); + for (int i = 0; i < world_size; i++) { + memcpy(&handles[i], all_handles[i].data_ptr(), + sizeof(cudaIpcMemHandle_t)); + } + mgr->post_init(handles); +} + +// Check weak contiguity (same logic as vLLM custom_all_reduce.cu) +static bool _is_weak_contiguous(const torch::Tensor& t) { + return t.is_contiguous() || + (t.storage().nbytes() - t.storage_offset() * t.element_size() == + static_cast(t.numel()) * t.element_size()); +} + +// Perform allreduce +void push_ar_all_reduce(fptr_t _mgr, torch::Tensor& inp, + torch::Tensor& out) { + auto* mgr = reinterpret_cast(_mgr); + 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_EQ(inp.numel(), out.numel()); + TORCH_CHECK(_is_weak_contiguous(inp), "Input must be contiguous"); + TORCH_CHECK(_is_weak_contiguous(out), "Output must be contiguous"); + + switch (out.scalar_type()) { + case at::ScalarType::BFloat16: + mgr->allreduce( + stream, reinterpret_cast(inp.data_ptr()), + reinterpret_cast(out.data_ptr()), out.numel()); + break; + case at::ScalarType::Half: + mgr->allreduce(stream, reinterpret_cast(inp.data_ptr()), + reinterpret_cast(out.data_ptr()), + out.numel()); + break; + case at::ScalarType::Float: + mgr->allreduce(stream, + reinterpret_cast(inp.data_ptr()), + reinterpret_cast(out.data_ptr()), + out.numel()); + break; + default: + TORCH_CHECK( + false, + "push allreduce: unsupported dtype (need bf16/fp16/fp32)"); + } +} + +// Dispose the manager +void dispose_push_ar(fptr_t _mgr) { + auto* mgr = reinterpret_cast(_mgr); + delete mgr; +} diff --git a/csrc/push_all_reduce.cuh b/csrc/push_all_reduce.cuh new file mode 100644 index 00000000000..5832a260299 --- /dev/null +++ b/csrc/push_all_reduce.cuh @@ -0,0 +1,229 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Host-side manager for push-based allreduce. +// Manages IPC storage, PushController initialization, and kernel launch. +// Replaces SGLang's CustomAllReduceBase + CustomAllReducePush. + +#pragma once +#include "push_all_reduce_kernel.cuh" + +#include + +#include +#include +#include +#include +#include + +namespace vllm { +namespace push_ar { + +class PushAllReduceManager { + public: + PushAllReduceManager(int rank, int world_size, int64_t push_buffer_bytes, + int max_num_cta) + : rank_(rank), + world_size_(world_size), + push_buffer_bytes_(push_buffer_bytes), + max_num_cta_(max_num_cta), + storage_(nullptr) { + assert(world_size_ >= 2 && world_size_ <= 8); + assert(max_num_cta_ > 0 && max_num_cta_ <= 512); + assert(push_buffer_bytes_ > 0); + + // Determine PDL support from device capability + int device_id; + cudaGetDevice(&device_id); + int major; + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, + device_id); + use_pdl_ = (major >= 9); // Hopper (sm90) or newer + + // Allocate storage + storage_bytes_ = push_signal_bytes() + push_buffer_total_bytes(); + cudaMalloc(&storage_, storage_bytes_); + // Zeros signals (epoch=0 for all CTAs) AND push buffer + // (0x0000 = IEEE 754 positive-zero = "empty" sentinel) + cudaMemset(storage_, 0, storage_bytes_); + + peer_storage_.resize(world_size_, nullptr); + } + + ~PushAllReduceManager() { + for (int i = 0; i < world_size_; i++) { + if (i != rank_ && peer_storage_[i] != nullptr) { + cudaIpcCloseMemHandle(peer_storage_[i]); + } + } + if (storage_) { + cudaFree(storage_); + } + } + + // Phase 1: Return IPC handle for local storage + cudaIpcMemHandle_t get_ipc_handle() { + cudaIpcMemHandle_t handle; + cudaIpcGetMemHandle(&handle, storage_); + return handle; + } + + // Phase 2: Open peer IPC handles and init PushController + void post_init(const std::vector& peer_handles) { + assert(peer_handles.size() == (size_t)world_size_); + for (int i = 0; i < world_size_; i++) { + if (i == rank_) { + peer_storage_[i] = storage_; + } else { + cudaIpcOpenMemHandle(&peer_storage_[i], peer_handles[i], + cudaIpcMemLazyEnablePeerAccess); + } + } + // Create PushController pointing to local signal region + push_ctrl_ = PushController(get_push_signal(storage_)); + ctrl_initialized_ = true; + } + + // Main allreduce dispatch + template + void allreduce(cudaStream_t stream, T* input, T* output, int num_elements) { + assert(ctrl_initialized_); + assert(num_elements > 0); + + const uint32_t num_items = static_cast(num_elements); + const int num_threads = select_num_threads(num_items); + + // Verify input fits in push buffer + const int64_t input_bytes = + static_cast(sizeof(T)) * num_elements; + assert(input_bytes <= push_buffer_bytes_); + + // Build kernel params + AllReducePushData params; + for (int i = 0; i < world_size_; i++) { + params.buffer[i] = get_push_buffer(peer_storage_[i]); + } + // Fill remaining buffer slots with nullptr (safety for kMaxNumGPU=8) + for (int i = world_size_; i < (int)kMaxNumGPU; i++) { + params.buffer[i] = nullptr; + } + params.input = input; + params.output = output; + params.rank = rank_; + params.num_items = num_items; + params.buffer_bytes = static_cast(push_buffer_bytes_); + params.epoch_bytes = world_size_ * params.buffer_bytes; + + // Build launch config for cudaLaunchKernelEx + cudaLaunchConfig_t config = {}; + config.gridDim = dim3(max_num_cta_); + config.blockDim = dim3(num_threads); + config.dynamicSmemBytes = 0; + config.stream = stream; + + cudaLaunchAttribute attrs[1]; + config.numAttrs = 0; + config.attrs = attrs; + + if (use_pdl_) { + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; + attrs[0].val.programmaticStreamSerializationAllowed = 1; + config.numAttrs = 1; + } + + // Template dispatch: world_size x pdl + launch_kernel(config, params); + } + + private: + // Template dispatch helper + template + void launch_kernel(const cudaLaunchConfig_t& config, + const AllReducePushData& params) { + // Dispatch on world_size and use_pdl + if (world_size_ == 8) { + if (use_pdl_) { + auto kernel = all_reduce_one_shot_push_kernel; + cudaLaunchKernelEx(&config, kernel, params, push_ctrl_); + } else { + auto kernel = all_reduce_one_shot_push_kernel; + cudaLaunchKernelEx(&config, kernel, params, push_ctrl_); + } + } else if (world_size_ == 4) { + if (use_pdl_) { + auto kernel = all_reduce_one_shot_push_kernel; + cudaLaunchKernelEx(&config, kernel, params, push_ctrl_); + } else { + auto kernel = all_reduce_one_shot_push_kernel; + cudaLaunchKernelEx(&config, kernel, params, push_ctrl_); + } + } else if (world_size_ == 2) { + if (use_pdl_) { + auto kernel = all_reduce_one_shot_push_kernel; + cudaLaunchKernelEx(&config, kernel, params, push_ctrl_); + } else { + auto kernel = all_reduce_one_shot_push_kernel; + cudaLaunchKernelEx(&config, kernel, params, push_ctrl_); + } + } else if (world_size_ == 6) { + if (use_pdl_) { + auto kernel = all_reduce_one_shot_push_kernel; + cudaLaunchKernelEx(&config, kernel, params, push_ctrl_); + } else { + auto kernel = all_reduce_one_shot_push_kernel; + cudaLaunchKernelEx(&config, kernel, params, push_ctrl_); + } + } + } + + // Thread count selection (from SGLang CustomAllReducePush::all_reduce) + template + int select_num_threads(uint32_t num_items) const { + constexpr uint32_t kVecSize = 16 / (sizeof(T) * 2); + for (const auto t : {128u, 256u, 512u}) { + if (t * max_num_cta_ * 2 * kVecSize >= num_items) { + return static_cast(t); + } + } + return 1024; + } + + // Storage layout helpers + int64_t push_signal_bytes() const { + return align128(sizeof(uint32_t) * max_num_cta_); + } + + int64_t push_buffer_total_bytes() const { + return align128(PushController::kNumStages * world_size_ * + push_buffer_bytes_); + } + + void* get_push_signal(void* base) const { + return base; // signals start at offset 0 + } + + void* get_push_buffer(void* base) const { + return static_cast(base) + push_signal_bytes(); + } + + static int64_t align128(int64_t size) { + return ((size + 127) / 128) * 128; + } + + // Members + int rank_; + int world_size_; + int64_t push_buffer_bytes_; + int max_num_cta_; + bool use_pdl_; + + void* storage_; + int64_t storage_bytes_; + std::vector peer_storage_; + + PushController push_ctrl_; + bool ctrl_initialized_ = false; +}; + +} // namespace push_ar +} // namespace vllm diff --git a/csrc/push_all_reduce_kernel.cuh b/csrc/push_all_reduce_kernel.cuh new file mode 100644 index 00000000000..bd6b0a6a140 --- /dev/null +++ b/csrc/push_all_reduce_kernel.cuh @@ -0,0 +1,540 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Push-based 2-buffer allreduce kernel, ported from SGLang's +// all_reduce_one_shot_push_kernel. Uses epoch-based double-buffered +// protocol with positive-zero sentinel for data arrival detection. +// +// Original source: +// sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh +// Protocol reference: SGLang commit edb1b3f +// +// Changes from SGLang: +// - All code placed in namespace vllm::push_ar +// - SGL_DEVICE macros replaced with __device__ __forceinline__ +// - SGL_CUDA_ARCH replaced with __CUDA_ARCH__ +// - std::integral (C++20) replaced with explicit overloads (C++17) +// - kMaxVecBytes hardcoded to 16 (push kernel uses 16-byte vectors) +// - TVM/FFI dependencies removed; file is self-contained + +#pragma once +#include +#include +#include + +#include +#include + +namespace vllm { +namespace push_ar { + +// ============================================================ +// Section A: Type aliases (from SGLang utils.cuh lines 50-75) +// ============================================================ +using fp32_t = float; +using fp16_t = __half; +using bf16_t = __nv_bfloat16; +using fp32x2_t = float2; +using fp16x2_t = __half2; +using bf16x2_t = __nv_bfloat162; + +static constexpr uint32_t kWarpThreads = 32u; + +// ============================================================ +// Section B: kMaxVecBytes (from SGLang utils.cuh line 112) +// ============================================================ +// Hardcoded to 16 since the push kernel uses 16-byte vectors. +// The kernel's kVecSize = 16 / (sizeof(DType) * 2) yields +// AlignedVector, kVecSize> = 16 bytes always. +inline constexpr std::size_t kMaxVecBytes = 16; + +// ============================================================ +// Section C: PDL helpers (from SGLang utils.cuh lines 119-148) +// ============================================================ +// CHANGED: SGL_ARCH_HOPPER_OR_GREATER -> __CUDA_ARCH__ >= 900 +template +__device__ __forceinline__ void PDLWaitPrimary() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + if constexpr (kUsePDL) { + asm volatile("griddepcontrol.wait;" ::: "memory"); + } +#endif +} + +template +__device__ __forceinline__ void PDLTriggerSecondary() { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + if constexpr (kUsePDL) { + asm volatile("griddepcontrol.launch_dependents;" :::); + } +#endif +} + +// ============================================================ +// Section D: Pointer offset helpers (from SGLang utils.cuh 181-195) +// ============================================================ +// CHANGED: Removed std::integral (C++20) constraint. +// Use explicit overloads for 1 and 2 offsets (C++17 compatible). + +// Byte-level offset (replaces pointer::offset) +__device__ __forceinline__ void* ptr_byte_offset(void* ptr, int64_t off1) { + return static_cast(ptr) + off1; +} + +__device__ __forceinline__ void* ptr_byte_offset(void* ptr, int64_t off1, + int64_t off2) { + return static_cast(ptr) + off1 + off2; +} + +// Typed offset for AlignedVector load/store addressing +// (replaces pointer::offset) +template +__device__ __forceinline__ void* ptr_typed_offset(void* ptr, int64_t offset) { + return static_cast(ptr) + offset; +} + +template +__device__ __forceinline__ const void* ptr_typed_offset(const void* ptr, + int64_t offset) { + return static_cast(ptr) + offset; +} + +// Host-side pointer offset (for storage layout calculations) +inline void* host_ptr_offset(void* ptr, int64_t off) { + return static_cast(ptr) + off; +} + +// ============================================================ +// Section E: dtype_trait system (from SGLang type.cuh) +// ============================================================ +// COPY AS IS with namespace adjustment. + +template +struct dtype_trait {}; + +template <> +struct dtype_trait { + using self_t = fp32_t; + using packed_t = fp32x2_t; + template + __device__ __forceinline__ static self_t from(const S& value) { + return static_cast(value); + } + __device__ __forceinline__ static self_t from(const fp16_t& x) { + return __half2float(x); + } + __device__ __forceinline__ static self_t from(const bf16_t& x) { + return __bfloat162float(x); + } +}; + +template <> +struct dtype_trait { + using self_t = fp16_t; + using packed_t = fp16x2_t; + template + __device__ __forceinline__ static self_t from(const S& value) { + return static_cast(value); + } +}; + +template <> +struct dtype_trait { + using self_t = bf16_t; + using packed_t = bf16x2_t; + template + __device__ __forceinline__ static self_t from(const S& value) { + return static_cast(value); + } +}; + +template <> +struct dtype_trait { + using self_t = fp32x2_t; + template + __device__ __forceinline__ static self_t from(const S& value) { + return static_cast(value); + } + __device__ __forceinline__ static self_t from(const fp16x2_t& x) { + return __half22float2(x); + } + __device__ __forceinline__ static self_t from(const bf16x2_t& x) { + return __bfloat1622float2(x); + } +}; + +template <> +struct dtype_trait { + using self_t = fp16x2_t; + template + __device__ __forceinline__ static self_t from(const S& value) { + return static_cast(value); + } + __device__ __forceinline__ static self_t from(const fp32x2_t& x) { + return __float22half2_rn(x); + } +}; + +template <> +struct dtype_trait { + using self_t = bf16x2_t; + template + __device__ __forceinline__ static self_t from(const S& value) { + return static_cast(value); + } + __device__ __forceinline__ static self_t from(const fp32x2_t& x) { + return __float22bfloat162_rn(x); + } +}; + +template +using packed_t = typename dtype_trait::packed_t; + +template +__device__ __forceinline__ To cast(const From& value) { + return dtype_trait::from(value); +} + +// ============================================================ +// Section F: AlignedVector (from SGLang vec.cuh lines 73-116) +// ============================================================ +// COPY AS IS with SGL_DEVICE -> __device__ __forceinline__ +// and kMaxVecBytes = 16. + +namespace detail { + +template +struct uint_trait {}; +template <> +struct uint_trait<1> { + using type = uint8_t; +}; +template <> +struct uint_trait<2> { + using type = uint16_t; +}; +template <> +struct uint_trait<4> { + using type = uint32_t; +}; +template <> +struct uint_trait<8> { + using type = uint64_t; +}; + +template +using sized_int = typename uint_trait::type; + +} // namespace detail + +template +struct alignas(sizeof(T) * N) AlignedStorage { + T data[N]; +}; + +template +struct AlignedVector { + private: + static_assert( + (N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= kMaxVecBytes, + "CUDA vector size exceeds arch limit (max 16 bytes)"); + using element_t = typename detail::sized_int; + using storage_t = AlignedStorage; + + public: + __device__ __forceinline__ void load(const void* ptr, int64_t offset = 0) { + m_storage = reinterpret_cast(ptr)[offset]; + } + __device__ __forceinline__ void store(void* ptr, int64_t offset = 0) const { + reinterpret_cast(ptr)[offset] = m_storage; + } + __device__ __forceinline__ void fill(T value) { + const auto store_value = *reinterpret_cast(&value); +#pragma unroll + for (std::size_t i = 0; i < N; ++i) { + m_storage.data[i] = store_value; + } + } + __device__ __forceinline__ auto operator[](std::size_t idx) -> T& { + return reinterpret_cast(&m_storage)[idx]; + } + __device__ __forceinline__ auto operator[](std::size_t idx) const -> T { + return reinterpret_cast(&m_storage)[idx]; + } + + private: + storage_t m_storage; +}; + +// ============================================================ +// Section G: PushController (from SGLang common.cuh lines 93-118) +// ============================================================ +// COPY AS IS with SGL_DEVICE -> __device__ __forceinline__ + +static constexpr uint32_t kMaxNumGPU = 8; + +struct PushController { + using SignalType = uint32_t; + static constexpr int64_t kNumStages = 2; // double-buffered epochs + + PushController() : m_local_signal(nullptr) {} + + PushController(void* ptr) : m_local_signal(static_cast(ptr)) {} + + __device__ __forceinline__ SignalType epoch() const { + return m_local_signal[blockIdx.x]; + } + + __device__ __forceinline__ void exit() const { + __syncthreads(); + if (threadIdx.x == 0) { + exit_unsafe(blockIdx.x); + } + } + + __device__ __forceinline__ void exit_unsafe(uint32_t which) const { + auto& signal = m_local_signal[which]; + signal = (signal + 1) % kNumStages; + } + + SignalType* m_local_signal; +}; + +// ============================================================ +// Section H: AllReducePushData (from SGLang +// custom_all_reduce_push.cuh 23-31) +// ============================================================ +// COPY AS IS. + +struct AllReducePushData { + void* __restrict__ buffer[kMaxNumGPU]; + const void* input; + void* output; + uint32_t rank; + uint32_t num_items; + uint32_t buffer_bytes; + uint32_t epoch_bytes; +}; + +// ============================================================ +// Section I: fp_trait sentinel types (from SGLang push.cuh 35-64) +// ============================================================ +// COPY AS IS. + +template +struct fp_trait {}; + +template <> +struct fp_trait { + using type = uint16_t; + [[maybe_unused]] static constexpr uint16_t pos_zero = 0x0000u; + [[maybe_unused]] static constexpr uint16_t neg_zero = 0x8000u; +}; + +template <> +struct fp_trait { + using type = uint16_t; + [[maybe_unused]] static constexpr uint16_t pos_zero = 0x0000u; + [[maybe_unused]] static constexpr uint16_t neg_zero = 0x8000u; +}; + +template <> +struct fp_trait { + using type = uint32_t; + [[maybe_unused]] static constexpr uint32_t pos_zero = 0x00000000u; + [[maybe_unused]] static constexpr uint32_t neg_zero = 0x80000000u; +}; + +// ============================================================ +// Section J: Sentinel helpers (from SGLang push.cuh 66-84) +// ============================================================ +// COPY AS IS. + +template +__device__ __forceinline__ void clear_pos_zero(DType& val) { + using Trait = fp_trait; + const auto ptr = reinterpret_cast(&val); + if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero; +} + +template +__device__ __forceinline__ bool is_pos_zero(const DType& val) { + using Trait = fp_trait; + const auto ptr = reinterpret_cast(&val); + return *ptr == Trait::pos_zero; +} + +template +__device__ __forceinline__ DType get_pos_zero() { + using Trait = fp_trait; + const auto value = Trait::pos_zero; + return *reinterpret_cast(&value); +} + +// ============================================================ +// Section K: Volatile 16-byte load/store (from SGLang push.cuh 87-105) +// ============================================================ +// CHANGED: pointer::offset -> ptr_typed_offset + +template +__device__ __forceinline__ void ld_global_volatile_16B(T& x, const void* addr, + int64_t offset) { + static_assert(alignof(T) == 16 && sizeof(T) == 16); + addr = ptr_typed_offset(addr, offset); + uint4 val; + asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w) + : "l"(addr)); + x = *reinterpret_cast(&val); +} + +template +__device__ __forceinline__ void st_global_volatile_16B(const T& x, void* addr, + int64_t offset) { + static_assert(alignof(T) == 16 && sizeof(T) == 16); + const uint4 val = *reinterpret_cast(&x); + addr = ptr_typed_offset(addr, offset); + asm volatile("st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"( + val.x), + "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr)); +} + +// ============================================================ +// Section L: reduce_impl (from SGLang custom_all_reduce.cuh 331-354) +// ============================================================ +// COPY AS IS. + +template +__device__ __forceinline__ auto reduce_impl( + AlignedVector (&storage)[M]) -> AlignedVector { + fp32x2_t acc[N] = {}; +#pragma unroll + for (uint32_t i = 0; i < M; ++i) { +#pragma unroll + for (uint32_t j = 0; j < N; ++j) { + const auto [x, y] = cast(storage[i][j]); + auto& [x_acc, y_acc] = acc[j]; + x_acc += x; + y_acc += y; + } + } + AlignedVector result; +#pragma unroll + for (uint32_t j = 0; j < N; ++j) { + result[j] = cast(acc[j]); + } + return result; +} + +// ============================================================ +// Section M: push_impl (from SGLang push.cuh 107-128) +// ============================================================ +// CHANGED: pointer::offset -> ptr_byte_offset + +template +__device__ __forceinline__ void push_impl(DType* (&push_buf)[kNumGPU], + const void* data, + uint32_t num_items) { + constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2); + using Storage = AlignedVector, kVecSize>; + + for (auto i = blockIdx.x;; i += gridDim.x) { + const auto offset = i * blockDim.x + threadIdx.x; + if (offset * kVecSize * 2 >= num_items) break; + Storage vec; + vec.load(data, offset); +#pragma unroll + for (uint32_t j = 0; j < kVecSize; ++j) { + clear_pos_zero(vec[j].x); + clear_pos_zero(vec[j].y); + } +#pragma unroll + for (uint32_t i = 0; i < kNumGPU; ++i) { + st_global_volatile_16B(vec, push_buf[i], offset); + } + } +} + +// ============================================================ +// Section N: poll_impl (from SGLang push.cuh 130-165) +// ============================================================ +// COPY AS IS. + +template +__device__ __forceinline__ void poll_impl(DType* (&poll_buf)[kNumGPU], + void* data, uint32_t num_items) { + constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2); + using Storage = AlignedVector, kVecSize>; + + for (auto i = blockIdx.x;; i += gridDim.x) { + const auto offset = i * blockDim.x + threadIdx.x; + if (offset * kVecSize * 2 >= num_items) break; + Storage storage[kNumGPU]; + + while (true) { + bool has_pos_zero = false; +#pragma unroll + for (uint32_t i = 0; i < kNumGPU; ++i) { + ld_global_volatile_16B(storage[i], poll_buf[i], offset); +#pragma unroll + for (auto j = 0; j < kVecSize; ++j) { + has_pos_zero |= is_pos_zero(storage[i][j].x); + has_pos_zero |= is_pos_zero(storage[i][j].y); + } + } + if (!has_pos_zero) break; + } + + const Storage result = reduce_impl(storage); + result.store(data, offset); + + Storage pos_zeros; + pos_zeros.fill({get_pos_zero(), get_pos_zero()}); +#pragma unroll + for (uint32_t i = 0; i < kNumGPU; ++i) { + pos_zeros.store(poll_buf[i], offset); + } + } +} + +// ============================================================ +// Section O: THE KERNEL (from SGLang push.cuh 167-196) +// ============================================================ +// COPY AS IS. CHANGED: CUSTOM_AR_KERNEL macro expanded. +// The kernel uses __grid_constant__ for params passed by value. +// cudaLaunchKernelEx copies params to constant memory before launch. + +template +__global__ __launch_bounds__(1024, 1) void all_reduce_one_shot_push_kernel( + const AllReducePushData __grid_constant__ params, + const PushController __grid_constant__ ctrl) { + const auto [buffer, input, output, rank, num_items, buffer_bytes, + epoch_bytes] = params; + + PDLWaitPrimary(); + + // Phase 1: Push data from input to all ranks' push buffers + const auto epoch_offset = ctrl.epoch() * epoch_bytes; + DType* push_buf[kNumGPU]; +#pragma unroll + for (uint32_t i = 0; i < kNumGPU; ++i) { + push_buf[i] = static_cast( + ptr_byte_offset(buffer[i], rank * buffer_bytes, epoch_offset)); + } + push_impl(push_buf, input, num_items); + + PDLTriggerSecondary(); + + // Phase 2: Poll local buffer, reduce, write output, reset + DType* poll_buf[kNumGPU]; +#pragma unroll + for (uint32_t i = 0; i < kNumGPU; ++i) { + poll_buf[i] = static_cast( + ptr_byte_offset(buffer[rank], i * buffer_bytes, epoch_offset)); + } + poll_impl(poll_buf, output, num_items); + ctrl.exit(); +} + +} // namespace push_ar +} // namespace vllm diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index cfd185394a4..b92cc988fac 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -104,4 +104,26 @@ TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cuda_utils), cuda_utils) { } #endif +// Push-based allreduce (ported from SGLang) +fptr_t init_push_ar(int64_t rank, int64_t world_size, + int64_t push_buffer_bytes, int64_t max_num_cta); +torch::Tensor get_push_ar_ipc_handle(fptr_t _mgr); +void post_init_push_ar(fptr_t _mgr, torch::Tensor all_handles); +void push_ar_all_reduce(fptr_t _mgr, torch::Tensor& inp, torch::Tensor& out); +void dispose_push_ar(fptr_t _mgr); + +TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _push_ar), push_ar) { + push_ar.def("init_push_ar", &init_push_ar); + + push_ar.def("get_push_ar_ipc_handle", &get_push_ar_ipc_handle); + + push_ar.def("post_init_push_ar", &post_init_push_ar); + + push_ar.def( + "push_ar_all_reduce(int mgr, Tensor inp, Tensor! out) -> ()"); + push_ar.impl("push_ar_all_reduce", torch::kCUDA, &push_ar_all_reduce); + + push_ar.def("dispose_push_ar", &dispose_push_ar); +} + REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/tests/distributed/_test_push_ar_worker.py b/tests/distributed/_test_push_ar_worker.py new file mode 100644 index 00000000000..c92c7bf03bc --- /dev/null +++ b/tests/distributed/_test_push_ar_worker.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Worker helper for push allreduce unit tests. +Run via torch.multiprocessing.spawn from test_push_all_reduce.py. + +Provides init/teardown helpers that create separate gloo (CPU) and +nccl (device) process groups for PushAllReduce (which needs gloo for +IPC handle exchange) and NCCL reference reduction (which needs nccl). +""" + +import os +import socket + +import torch +import torch.distributed as dist + + +def find_free_port() -> int: + """Find a free TCP port for distributed init.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +# Global references to process groups created by init_groups +_cpu_group = None +_nccl_group = None + + +def init_groups(rank: int, world_size: int, port: int): + """Initialize gloo (CPU) and nccl process groups. + + PushAllReduce uses the gloo group for IPC handle exchange. + NCCL group is used for reference allreduce. + """ + global _cpu_group, _nccl_group + + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + + torch.cuda.set_device(rank) + + dist.init_process_group( + backend="gloo", rank=rank, world_size=world_size + ) + _cpu_group = dist.group.WORLD + + # Create a separate NCCL group for reference allreduce + _nccl_group = dist.new_group(backend="nccl") + + +def get_cpu_group(): + return _cpu_group + + +def get_nccl_group(): + return _nccl_group + + +def teardown(): + """Clean up distributed groups.""" + dist.destroy_process_group() diff --git a/tests/distributed/test_push_all_reduce.py b/tests/distributed/test_push_all_reduce.py new file mode 100644 index 00000000000..9689c5e1000 --- /dev/null +++ b/tests/distributed/test_push_all_reduce.py @@ -0,0 +1,1187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests for push-based allreduce (ported from SGLang). +Covers correctness, edge cases, CUDA graph safety, dispatch integration, +multi-layer stacking, and epoch alternation. + +Test infrastructure: + - Unit tests: torch.multiprocessing.spawn with gloo+nccl groups + - Integration tests: Ray workers with vLLM test infrastructure +""" + +import os +import random + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +# Reusable test sizes aligned to 16 bytes (8 BF16 elements) +random.seed(42) +UNIT_TEST_SIZES = [ + 16, # minimal (32 bytes for BF16) + 128, # single warp + 1024, # small + 7168, # decode BS=1 hidden_size + 28672, # decode BS=4 (4 * 7168) + 65536, # moderate + 131072, # larger +] + + +# ============================================================ +# Helper: init/teardown for mp.spawn-based tests +# ============================================================ +def _init_groups(rank: int, world_size: int, port: int): + """Initialize gloo (CPU) and nccl process groups.""" + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = str(port) + torch.cuda.set_device(rank) + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + return dist.group.WORLD, dist.new_group(backend="nccl") + + +def _teardown(): + dist.destroy_process_group() + + +def _find_free_port() -> int: + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +# ============================================================ +# UT-1: Build Verification +# ============================================================ +def test_push_ar_ops_registered(): + """Verify push_ar custom ops are registered and callable.""" + import vllm._custom_ops as ops + + assert hasattr(ops, "init_push_ar") + assert hasattr(ops, "get_push_ar_ipc_handle") + assert hasattr(ops, "post_init_push_ar") + assert hasattr(ops, "push_ar_all_reduce") + assert hasattr(ops, "dispose_push_ar") + + +# ============================================================ +# UT-2: PushAllReduce Initialization +# ============================================================ +def _push_ar_init_worker(rank, world_size, port): + cpu_group, _ = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + assert not push_ar.disabled + assert push_ar.rank == rank + assert push_ar.world_size == world_size + assert push_ar.push_buffer_bytes > 0 + assert push_ar.num_sm > 0 + + push_ar.close() + _teardown() + + +@pytest.mark.parametrize("world_size", [2]) +def test_push_ar_initialization(world_size): + if torch.cuda.device_count() < world_size: + pytest.skip(f"Need {world_size} GPUs") + mp.spawn( + _push_ar_init_worker, + args=(world_size, _find_free_port()), + nprocs=world_size, + join=True, + ) + + +# ============================================================ +# UT-3: should_use() Predicate Logic +# ============================================================ +def _push_ar_should_use_worker(rank, world_size, port): + cpu_group, _ = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + # Case 1: Valid small tensor -> True + t1 = torch.randn(7168, dtype=torch.bfloat16, device=device) + assert push_ar.should_use(t1) is True + + # Case 2: Valid exactly-at-threshold -> True + max_elems = push_ar.max_message_bytes // 2 # BF16 + max_elems = (max_elems // 8) * 8 # align to 16 bytes + t2 = torch.randn(max_elems, dtype=torch.bfloat16, device=device) + assert push_ar.should_use(t2) is True + + # Case 3: Above threshold -> False + t3 = torch.randn(max_elems + 8, dtype=torch.bfloat16, device=device) + assert push_ar.should_use(t3) is False + + # Case 4: Zero-size tensor -> False + t4 = torch.empty(0, dtype=torch.bfloat16, device=device) + assert push_ar.should_use(t4) is False + + # Case 5: Non-aligned size (not divisible by 16 bytes) -> False + t5 = torch.randn(7, dtype=torch.bfloat16, device=device) # 14 bytes + assert push_ar.should_use(t5) is False + + # Case 6: Non-contiguous tensor -> False + t6_base = torch.randn(1024, 2, dtype=torch.bfloat16, device=device) + t6 = t6_base[:, 0] # non-contiguous view + assert push_ar.should_use(t6) is False + + # Case 7: Weakly-contiguous tensor -> True + t7_base = torch.randn(2048, dtype=torch.bfloat16, device=device) + t7 = t7_base[:1024] # contiguous slice + assert push_ar.should_use(t7) is True + + # Case 8: Disabled communicator -> always False + push_ar.disabled = True + assert push_ar.should_use(t1) is False + push_ar.disabled = False + + push_ar.close() + _teardown() + + +def test_push_ar_should_use(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_should_use_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# UT-4: Basic Allreduce Correctness (Integer Data, Bit-Exact) +# ============================================================ +def _push_ar_correctness_int_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + for size in UNIT_TEST_SIZES: + for dtype in [torch.float16, torch.bfloat16, torch.float32]: + inp = torch.randint(0, 16, (size,), dtype=dtype, device=device) + + if not push_ar.should_use(inp): + continue + + out_push = push_ar.all_reduce(inp) + + out_nccl = inp.clone() + dist.all_reduce(out_nccl, group=nccl_group) + + assert torch.all(out_push == out_nccl), ( + f"Mismatch at size={size}, dtype={dtype}, rank={rank}" + ) + + push_ar.close() + _teardown() + + +@pytest.mark.parametrize("world_size", [2]) +def test_push_ar_correctness_integer(world_size): + if torch.cuda.device_count() < world_size: + pytest.skip(f"Need {world_size} GPUs") + mp.spawn( + _push_ar_correctness_int_worker, + args=(world_size, _find_free_port()), + nprocs=world_size, + join=True, + ) + + +# ============================================================ +# UT-5: Allreduce Correctness (Random Float Data) +# ============================================================ +def _push_ar_correctness_float_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + for size in UNIT_TEST_SIZES: + for dtype in [torch.bfloat16, torch.float16]: + inp = torch.randn(size, dtype=dtype, device=device) + + if not push_ar.should_use(inp): + continue + + out_push = push_ar.all_reduce(inp) + + out_nccl = inp.clone() + dist.all_reduce(out_nccl, group=nccl_group) + + torch.testing.assert_close( + out_push, + out_nccl, + atol=1e-2, + rtol=1e-2, + msg=f"Mismatch at size={size}, dtype={dtype}", + ) + + push_ar.close() + _teardown() + + +@pytest.mark.parametrize("world_size", [2]) +def test_push_ar_correctness_float(world_size): + if torch.cuda.device_count() < world_size: + pytest.skip(f"Need {world_size} GPUs") + mp.spawn( + _push_ar_correctness_float_worker, + args=(world_size, _find_free_port()), + nprocs=world_size, + join=True, + ) + + +# ============================================================ +# UT-6: Positive-Zero Sentinel Handling +# ============================================================ +def _push_ar_zero_handling_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + # Case 1: All zeros tensor + inp_zeros = torch.zeros(7168, dtype=torch.bfloat16, device=device) + out = push_ar.all_reduce(inp_zeros) + assert torch.all(out == 0.0), "All-zero input failed" + + # Case 2: Mix of zeros and non-zeros (rank-dependent) + inp_mixed = torch.zeros(7168, dtype=torch.bfloat16, device=device) + inp_mixed[::2] = float(rank + 1) + out_mixed = push_ar.all_reduce(inp_mixed) + expected_sum = sum(range(1, world_size + 1)) + torch.testing.assert_close( + out_mixed[::2], + torch.full_like(out_mixed[::2], expected_sum), + atol=1e-2, + rtol=1e-2, + ) + assert torch.all(out_mixed[1::2] == 0.0) + + # Case 3: Bit pattern verification + inp_pz = torch.zeros(1024, dtype=torch.bfloat16, device=device) + raw_bits = inp_pz.view(torch.int16) + assert torch.all(raw_bits == 0), "Expected +0.0 bit pattern" + out_pz = push_ar.all_reduce(inp_pz) + assert torch.all(out_pz == 0.0), "Positive zero handling failed" + + # Case 4: Negative zeros + inp_nz = torch.tensor( + [-0.0] * 1024, dtype=torch.bfloat16, device=device + ) + out_nz = push_ar.all_reduce(inp_nz) + assert torch.all(out_nz == 0.0) + + # Case 5: FP16 positive zeros + inp_f16 = torch.zeros(1024, dtype=torch.float16, device=device) + out_f16 = push_ar.all_reduce(inp_f16) + assert torch.all(out_f16 == 0.0), "FP16 positive zero handling failed" + + push_ar.close() + _teardown() + + +@pytest.mark.timeout(120) +def test_push_ar_zero_handling(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_zero_handling_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# UT-7: Epoch Alternation (1000 Iterations) +# ============================================================ +def _push_ar_epoch_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + NUM_ITERATIONS = 1000 + + for i in range(NUM_ITERATIONS): + inp = torch.randint( + 0, 16, (7168,), dtype=torch.bfloat16, device=device + ) + out_push = push_ar.all_reduce(inp) + + out_nccl = inp.clone() + dist.all_reduce(out_nccl, group=nccl_group) + + assert torch.all(out_push == out_nccl), ( + f"Epoch mismatch at iteration {i}" + ) + + push_ar.close() + _teardown() + + +def test_push_ar_epoch_alternation(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_epoch_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# UT-8: Thread Count Selection (Dynamic SM Boundaries) +# ============================================================ +def _push_ar_thread_count_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + sm_count = push_ar.num_sm + + # BF16: kVecSize=4, elements per thread = 2*kVecSize = 8 + boundaries = { + 128: 128 * sm_count * 8, + 256: 256 * sm_count * 8, + } + + test_sizes_tc = [ + 1024, + 7168, + boundaries[128], + boundaries[128] + 8, + ] + + for size in test_sizes_tc: + if size * 2 > push_ar.max_message_bytes: + continue + inp = torch.randint( + 0, 8, (size,), dtype=torch.bfloat16, device=device + ) + out = push_ar.all_reduce(inp) + ref = inp.clone() + dist.all_reduce(ref, group=nccl_group) + assert torch.all(out == ref), ( + f"Failed at size={size} (sm_count={sm_count})" + ) + + push_ar.close() + _teardown() + + +def test_push_ar_thread_count(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_thread_count_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# UT-9: Buffer Size Threshold Boundary +# ============================================================ +def _push_ar_threshold_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + max_bytes = push_ar.max_message_bytes + + # Case 1: Exactly at threshold (BF16) + max_elems = max_bytes // 2 + max_elems = (max_elems // 8) * 8 # align to 16 bytes + inp_exact = torch.randint( + 0, 8, (max_elems,), dtype=torch.bfloat16, device=device + ) + assert push_ar.should_use(inp_exact) is True + out_exact = push_ar.all_reduce(inp_exact) + ref_exact = inp_exact.clone() + dist.all_reduce(ref_exact, group=nccl_group) + assert torch.all(out_exact == ref_exact) + + # Case 2: One vector above threshold + inp_over = torch.randint( + 0, 8, (max_elems + 8,), dtype=torch.bfloat16, device=device + ) + assert push_ar.should_use(inp_over) is False + + # Case 3: One vector below threshold + inp_under = torch.randint( + 0, 8, (max_elems - 8,), dtype=torch.bfloat16, device=device + ) + assert push_ar.should_use(inp_under) is True + out_under = push_ar.all_reduce(inp_under) + ref_under = inp_under.clone() + dist.all_reduce(ref_under, group=nccl_group) + assert torch.all(out_under == ref_under) + + push_ar.close() + _teardown() + + +def test_push_ar_threshold(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_threshold_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# UT-10: Out-of-Place Semantics +# ============================================================ +def _push_ar_outofplace_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + inp = torch.randint( + 1, 16, (7168,), dtype=torch.bfloat16, device=device + ) + inp_original = inp.clone() + + out = push_ar.all_reduce(inp) + + # Input NOT modified + assert torch.all(inp == inp_original), "Input was modified in-place!" + + # Output is a DIFFERENT tensor + assert out.data_ptr() != inp.data_ptr(), "Output aliases input!" + + # Output has correct values + ref = inp_original.clone() + dist.all_reduce(ref, group=nccl_group) + assert torch.all(out == ref) + + push_ar.close() + _teardown() + + +def test_push_ar_outofplace(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_outofplace_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# UT-11: Multiple Dtype Support +# ============================================================ +def _push_ar_dtype_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + for dtype in [torch.float16, torch.bfloat16, torch.float32]: + size = 7168 if dtype != torch.float32 else 3584 + inp = torch.randint(0, 16, (size,), dtype=dtype, device=device) + + if not push_ar.should_use(inp): + continue + + out = push_ar.all_reduce(inp) + ref = inp.clone() + dist.all_reduce(ref, group=nccl_group) + assert torch.all(out == ref), f"Failed for dtype={dtype}" + + push_ar.close() + _teardown() + + +def test_push_ar_dtype(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_dtype_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# UT-12: Multi-Layer Stacking (122 ARs/step * 10 steps) +# ============================================================ +def _push_ar_multilayer_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + NUM_LAYERS = 122 # 61 blocks * 2 allreduces/block + NUM_STEPS = 3 # Reduced for test speed + + for step in range(NUM_STEPS): + for layer in range(NUM_LAYERS): + inp = torch.randint( + 0, 16, (7168,), dtype=torch.bfloat16, device=device + ) + out = push_ar.all_reduce(inp) + ref = inp.clone() + dist.all_reduce(ref, group=nccl_group) + if not torch.all(out == ref): + raise RuntimeError( + f"Mismatch at step={step}, layer={layer}" + ) + + push_ar.close() + _teardown() + + +def test_push_ar_multilayer(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_multilayer_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# UT-14: close() and Resource Lifecycle +# ============================================================ +def _push_ar_lifecycle_worker(rank, world_size, port): + cpu_group, _ = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + # Normal allreduce works + inp = torch.randint( + 0, 16, (1024,), dtype=torch.bfloat16, device=device + ) + out = push_ar.all_reduce(inp) + assert out is not None + + # First close + push_ar.close() + assert push_ar.disabled is True + + # Double close should NOT crash + push_ar.close() + + # __del__ should also be safe after explicit close + del push_ar + + _teardown() + + +def test_push_ar_lifecycle(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_lifecycle_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# UT-15: Graph Capture Warmup Path +# ============================================================ +def _push_ar_warmup_worker(rank, world_size, port): + cpu_group, _ = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + inp = torch.randn(7168, dtype=torch.bfloat16, device=device) + + # Outside capture context: real allreduce + assert not push_ar._IS_CAPTURING + out_real = push_ar.all_reduce(inp) + assert out_real.shape == inp.shape + + # Inside capture context, but NOT in graph capture stream: + # should return empty_like (warmup allocation) + with push_ar.capture(): + assert push_ar._IS_CAPTURING + out_warmup = push_ar.all_reduce(inp) + assert out_warmup.shape == inp.shape + assert out_warmup.dtype == inp.dtype + + # After capture context: back to normal + assert not push_ar._IS_CAPTURING + + push_ar.close() + _teardown() + + +def test_push_ar_warmup(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_warmup_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# UT-16: Rank-Dependent Input Data (Asymmetric Reduction) +# ============================================================ +def _push_ar_asymmetric_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + # Each rank has unique data + torch.manual_seed(42 + rank) + inp = torch.randn(7168, dtype=torch.bfloat16, device=device) + + out_push = push_ar.all_reduce(inp) + + out_nccl = inp.clone() + dist.all_reduce(out_nccl, group=nccl_group) + + torch.testing.assert_close( + out_push, + out_nccl, + atol=1e-2, + rtol=1e-2, + msg="Asymmetric reduction failed", + ) + + # Verify sum property + all_inputs = [torch.empty_like(inp) for _ in range(world_size)] + dist.all_gather(all_inputs, inp, group=nccl_group) + expected_sum = torch.stack(all_inputs).float().sum(dim=0).to(inp.dtype) + torch.testing.assert_close( + out_push, expected_sum, atol=5e-2, rtol=5e-2 + ) + + push_ar.close() + _teardown() + + +def test_push_ar_asymmetric(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_asymmetric_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# EC-1: Tensor with All Identical Values +# ============================================================ +def _push_ar_identical_worker(rank, world_size, port): + cpu_group, _ = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + inp = torch.ones(7168, dtype=torch.bfloat16, device=device) * 3.14 + out = push_ar.all_reduce(inp) + expected = inp * world_size + torch.testing.assert_close(out, expected, atol=1e-1, rtol=1e-2) + + push_ar.close() + _teardown() + + +def test_push_ar_identical_values(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_identical_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# IT-1: Dispatch Priority (mp.spawn-based integration test) +# ============================================================ +def _dispatch_priority_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + # Verify push_ar_comm was initialized correctly + assert not push_ar.disabled + assert push_ar.rank == rank + + # Small tensor -> should use push allreduce + small_inp = torch.randn(7168, dtype=torch.bfloat16, device=device) + assert push_ar.should_use(small_inp) is True + + # Large tensor -> should NOT use push allreduce + max_bytes = push_ar.max_message_bytes + large_elems = max_bytes // 2 + 1024 # slightly over threshold + large_inp = torch.randn(large_elems, dtype=torch.bfloat16, device=device) + assert push_ar.should_use(large_inp) is False + + # Verify dispatch produces correct results for small + out_small = push_ar.all_reduce(small_inp) + ref_small = small_inp.clone() + dist.all_reduce(ref_small, group=nccl_group) + torch.testing.assert_close( + out_small, ref_small, atol=1e-2, rtol=1e-2 + ) + + push_ar.close() + _teardown() + + +def test_dispatch_priority(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _dispatch_priority_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# IT-4: Coexistence Test (Interleaved Push + NCCL) +# ============================================================ +def _coexistence_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + for _ in range(50): + # Small message -> push allreduce + small = torch.randint( + 0, 16, (7168,), dtype=torch.bfloat16, device=device + ) + out_small = push_ar.all_reduce(small) + ref_small = small.clone() + dist.all_reduce(ref_small, group=nccl_group) + assert torch.all(out_small == ref_small) + + push_ar.close() + _teardown() + + +def test_coexistence(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _coexistence_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# IT-10: Interleaved Push AR + NCCL in Same Step +# ============================================================ +def _interleaved_ar_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + for sz in [7168, 1024, 4096]: + for dtype in [torch.bfloat16, torch.float16]: + inp1 = torch.randint( + 1, 16, (sz,), dtype=dtype, device=device + ) + inp2 = torch.randint( + 1, 16, (sz,), dtype=dtype, device=device + ) + + out1 = push_ar.all_reduce(inp1) + ref1 = inp1.clone() + dist.all_reduce(ref1, group=nccl_group) + + out2 = push_ar.all_reduce(inp2) + ref2 = inp2.clone() + dist.all_reduce(ref2, group=nccl_group) + + torch.testing.assert_close( + out1, ref1, atol=1e-2, rtol=1e-2 + ) + torch.testing.assert_close( + out2, ref2, atol=1e-2, rtol=1e-2 + ) + + push_ar.close() + _teardown() + + +def test_interleaved_ar(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _interleaved_ar_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# CG-1: CUDA Graph Capture + Replay +# ============================================================ +def _graph_capture_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + NUM_AR = 5 # allreduces per graph + sz = 7168 + + # Allocate graph input in graph memory pool + graph_inp = torch.randint( + 1, 16, (sz,), dtype=torch.bfloat16, device=device + ) + + # Warmup + with push_ar.capture(): + for _ in range(NUM_AR): + push_ar.all_reduce(graph_inp) + torch.cuda.synchronize() + + # Capture graph + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outs = [] + for _ in range(NUM_AR): + outs.append(push_ar.all_reduce(graph_inp)) + + # Replay and verify + for replay_iter in range(10): + # Fill with new data before each replay + graph_inp.copy_( + torch.randint(1, 16, (sz,), dtype=torch.bfloat16, device=device) + ) + graph.replay() + torch.cuda.synchronize() + + # Verify last output is correct + ref = graph_inp.clone() + # The allreduce was applied NUM_AR times, but each operates on + # graph_inp independently. The last output should be one allreduce + # of graph_inp. + dist.all_reduce(ref, group=nccl_group) + torch.testing.assert_close( + outs[-1], ref, atol=1e-2, rtol=1e-2, + msg=f"Graph replay {replay_iter} failed" + ) + + push_ar.close() + _teardown() + + +def test_graph_capture(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _graph_capture_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# E2E-4: Transformer Block Simulation (Isolated) +# ============================================================ +def _push_ar_transformer_sim_worker(rank, world_size, port): + cpu_group, nccl_group = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + hidden_size = 7168 + NUM_BLOCKS = 61 # DeepSeek-V4 layer count + NUM_STEPS = 2 # Reduced for test speed + + for step in range(NUM_STEPS): + for block in range(NUM_BLOCKS): + # AR 1: attention wo_b output + attn_out = torch.randn( + 1, hidden_size, dtype=torch.bfloat16, device=device + ) + attn_reduced = push_ar.all_reduce(attn_out) + attn_ref = attn_out.clone() + dist.all_reduce(attn_ref, group=nccl_group) + torch.testing.assert_close( + attn_reduced, + attn_ref, + atol=1e-2, + rtol=1e-2, + msg=f"Attn AR failed: step={step}, block={block}", + ) + + # AR 2: MoE output + moe_out = torch.randn( + 1, hidden_size, dtype=torch.bfloat16, device=device + ) + moe_reduced = push_ar.all_reduce(moe_out) + moe_ref = moe_out.clone() + dist.all_reduce(moe_ref, group=nccl_group) + torch.testing.assert_close( + moe_reduced, + moe_ref, + atol=1e-2, + rtol=1e-2, + msg=f"MoE AR failed: step={step}, block={block}", + ) + + push_ar.close() + _teardown() + + +def test_push_ar_transformer_sim(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _push_ar_transformer_sim_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# FT-1: Feature Toggle - ENABLED Startup Log +# ============================================================ +def _feature_toggle_enabled_worker(rank, world_size, port): + cpu_group, _ = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + # Ensure the disable env var is NOT set + os.environ.pop("VLLM_DISABLE_PUSH_ALLREDUCE", None) + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + _FEATURE_DESCRIPTION, + _DISABLE_ENV_VAR, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + # Feature should be enabled + assert not push_ar.disabled, "PushAllReduce should be enabled" + assert _DISABLE_ENV_VAR == "VLLM_DISABLE_PUSH_ALLREDUCE" + assert "AllReduce" in _FEATURE_DESCRIPTION + + # Verify the push_ar works when enabled + inp = torch.randint(0, 16, (1024,), dtype=torch.bfloat16, device=device) + out = push_ar.all_reduce(inp) + assert out is not None + assert out.shape == inp.shape + + push_ar.close() + _teardown() + + +def test_feature_toggle_enabled(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _feature_toggle_enabled_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# FT-2: Feature Toggle - DISABLED via Env Var +# ============================================================ +def _feature_toggle_disabled_worker(rank, world_size, port): + cpu_group, _ = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + # Set the disable env var + os.environ["VLLM_DISABLE_PUSH_ALLREDUCE"] = "1" + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + # Feature should be disabled + assert push_ar.disabled, "PushAllReduce should be disabled via env var" + + # should_use should return False when disabled + inp = torch.randn(1024, dtype=torch.bfloat16, device=device) + assert push_ar.should_use(inp) is False + + # Clean up + os.environ.pop("VLLM_DISABLE_PUSH_ALLREDUCE", None) + push_ar.close() + _teardown() + + +def test_feature_toggle_disabled(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _feature_toggle_disabled_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) + + +# ============================================================ +# FT-3: Feature Toggle - Env Var Not Set to "1" +# ============================================================ +def _feature_toggle_not_disabled_worker(rank, world_size, port): + cpu_group, _ = _init_groups(rank, world_size, port) + device = torch.device(f"cuda:{rank}") + + # Set the env var to something other than "1" - should NOT disable + os.environ["VLLM_DISABLE_PUSH_ALLREDUCE"] = "0" + + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + push_ar = PushAllReduce(group=cpu_group, device=device) + + # Feature should still be enabled (only "1" disables) + assert not push_ar.disabled, ( + "PushAllReduce should be enabled when env var != '1'" + ) + + # Clean up + os.environ.pop("VLLM_DISABLE_PUSH_ALLREDUCE", None) + push_ar.close() + _teardown() + + +def test_feature_toggle_not_disabled(): + if torch.cuda.device_count() < 2: + pytest.skip("Need 2 GPUs") + mp.spawn( + _feature_toggle_not_disabled_worker, + args=(2, _find_free_port()), + nprocs=2, + join=True, + ) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 3878f3038bd..4b68dc5abee 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3023,6 +3023,33 @@ def qr_max_size() -> int: return torch.ops._C_custom_ar.qr_max_size() +# push allreduce (ported from SGLang) +def init_push_ar( + rank: int, world_size: int, buffer_bytes: int, max_cta: int +) -> int: + return torch.ops._C_push_ar.init_push_ar( + rank, world_size, buffer_bytes, max_cta + ) + + +def get_push_ar_ipc_handle(mgr: int) -> torch.Tensor: + return torch.ops._C_push_ar.get_push_ar_ipc_handle(mgr) + + +def post_init_push_ar(mgr: int, handles: torch.Tensor) -> None: + torch.ops._C_push_ar.post_init_push_ar(mgr, handles) + + +def push_ar_all_reduce( + mgr: int, inp: torch.Tensor, out: torch.Tensor +) -> None: + torch.ops._C_push_ar.push_ar_all_reduce(mgr, inp, out) + + +def dispose_push_ar(mgr: int) -> None: + torch.ops._C_push_ar.dispose_push_ar(mgr) + + def get_flash_mla_metadata( cache_seqlens: torch.Tensor, num_heads_per_head_k: int, diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 12c425021b2..d30039377d7 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -81,6 +81,7 @@ class CudaCommunicator(DeviceCommunicatorBase): register_nccl_symmetric_ops(self.pynccl_comm) self.ca_comm: CustomAllreduce | None = None + self.push_ar_comm = None self.qr_comm: QuickAllReduce | None = None self.symm_mem_comm: SymmMemCommunicator | None = None self.fi_ar_comm: FlashInferAllReduce | None = None @@ -107,6 +108,24 @@ class CudaCommunicator(DeviceCommunicatorBase): ), ) + # Initialize push-based allreduce (faster for small messages) + # Only available on NVIDIA CUDA GPUs with NVLink + if (current_platform.is_cuda() + and self.ca_comm is not None + and not self.ca_comm.disabled): + try: + from vllm.distributed.device_communicators.push_all_reduce import ( + PushAllReduce, + ) + + self.push_ar_comm = PushAllReduce( + group=self.cpu_group, device=self.device + ) + if self.push_ar_comm.disabled: + self.push_ar_comm = None + except Exception: + self.push_ar_comm = None + if current_platform.is_rocm(): # Initialize a custom quick all-reduce implementation for AMD. # Quick reduce is designed as a complement to custom allreduce. @@ -280,6 +299,14 @@ class CudaCommunicator(DeviceCommunicatorBase): out = fi_ar_comm.all_reduce(input_) assert out is not None return out + push_ar_comm = self.push_ar_comm + if ( + push_ar_comm is not None + and push_ar_comm.should_use(input_) + ): + out = push_ar_comm.all_reduce(input_) + if out is not None: + return out ca_comm = self.ca_comm if ( ca_comm is not None diff --git a/vllm/distributed/device_communicators/push_all_reduce.py b/vllm/distributed/device_communicators/push_all_reduce.py new file mode 100644 index 00000000000..3105df939ca --- /dev/null +++ b/vllm/distributed/device_communicators/push_all_reduce.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Push-based custom allreduce using epoch-based 2-buffer protocol. +Ported from SGLang's CustomAllReduceV2 push allreduce. + +Protocol: + Phase 1 (push): Each rank writes its input data to ALL remote GPUs' + push buffer regions via NVLink volatile stores. Positive zeros are + converted to negative zeros to preserve sentinel semantics. + Phase 2 (poll): Each rank polls its LOCAL buffer until all ranks' + data arrives (no positive zeros remain), reduces in FP32, writes + output, and resets buffer to positive zeros for next epoch. +""" + +import logging +import os +from contextlib import contextmanager +from typing import Optional + +import torch +import torch.distributed as dist + +from vllm import _custom_ops as ops +from vllm.distributed.device_communicators.custom_all_reduce import ( + is_weak_contiguous, +) +from vllm.platforms import current_platform + +logger = logging.getLogger(__name__) + +_FEATURE_DESCRIPTION = "Push-based AllReduce" +_DISABLE_ENV_VAR = "VLLM_DISABLE_PUSH_ALLREDUCE" + +# Push threshold map: world_size -> buffer_bytes +# From SGLang's tuned thresholds for sm100 (B200) +PUSH_THRESHOLD_SM100 = { + 2: 4 * 1024 * 1024, # 4 MB + 4: 2 * 1024 * 1024, # 2 MB + 6: 1 * 1024 * 1024, # 1 MB + 8: 720 * 1024, # 720 KB +} + +# Conservative default for untuned GPUs +DEFAULT_PUSH_BUFFER = 512 * 1024 # 512 KB + + +class PushAllReduce: + """ + Push-based custom allreduce using epoch-based 2-buffer protocol. + Ported from SGLang's CustomAllReduceV2 push allreduce. + """ + + _IS_CAPTURING = False + + def __init__( + self, + group: dist.ProcessGroup, + device: torch.device, + max_size: Optional[int] = None, + ): + self.group = group + self.device = device + self.rank = dist.get_rank(group) + self.world_size = dist.get_world_size(group) + self.disabled = False + + # Feature toggle: check env var to disable this feature + if os.environ.get(_DISABLE_ENV_VAR) == "1": + logger.info( + "%s is DISABLED (env override)", + _FEATURE_DESCRIPTION, + ) + self.disabled = True + return + + # Prerequisite checks + if self.world_size not in (2, 4, 6, 8): + logger.info( + "PushAllReduce disabled: unsupported world_size=%d", + self.world_size, + ) + self.disabled = True + return + + if not self._check_full_p2p(): + logger.info("PushAllReduce disabled: no full P2P connectivity") + self.disabled = True + return + + # Get SM count for grid size + props = torch.cuda.get_device_properties(device) + self.num_sm = props.multi_processor_count + + # Determine push buffer size from threshold map + if max_size is not None: + self.push_buffer_bytes = max_size + else: + self.push_buffer_bytes = PUSH_THRESHOLD_SM100.get( + self.world_size, DEFAULT_PUSH_BUFFER + ) + + # Allow env var override (V1 fix: validate 128-byte alignment) + env_override = os.environ.get("VLLM_PUSH_AR_BUFFER_BYTES") + if env_override: + val = int(env_override) + if val == 0: + logger.info( + "PushAllReduce disabled via VLLM_PUSH_AR_BUFFER_BYTES=0" + ) + self.disabled = True + return + # Round up to 128-byte alignment for kernel volatile stores + self.push_buffer_bytes = ((val + 127) // 128) * 128 + + self.max_message_bytes = self.push_buffer_bytes + + try: + # Initialize C++ manager + self._ptr = ops.init_push_ar( + self.rank, + self.world_size, + self.push_buffer_bytes, + self.num_sm, + ) + + # Exchange IPC handles across all ranks + self._exchange_ipc_handles() + + logger.info( + "PushAllReduce initialized: rank=%d, ws=%d, " + "buffer=%d KB, sm=%d", + self.rank, + self.world_size, + self.push_buffer_bytes // 1024, + self.num_sm, + ) + + # Feature toggle: log that the feature is enabled + logger.info( + "%s is ENABLED", + _FEATURE_DESCRIPTION, + ) + except Exception as e: + logger.warning("PushAllReduce init failed: %s", str(e)) + self.disabled = True + + def _check_full_p2p(self) -> bool: + """Verify all GPUs have full P2P access (NVLink).""" + num_dev = current_platform.device_count() + for i in range(num_dev): + for j in range(num_dev): + if i != j and not torch.cuda.can_device_access_peer(i, j): + return False + return True + + def _exchange_ipc_handles(self): + """Exchange storage IPC handles across all ranks.""" + # Get local handle as byte tensor [sizeof(cudaIpcMemHandle_t)] + local_handle = ops.get_push_ar_ipc_handle(self._ptr) # shape (64,) + + # All-gather handles: each rank broadcasts its handle + handle_list = [ + torch.empty_like(local_handle) for _ in range(self.world_size) + ] + dist.all_gather(handle_list, local_handle, group=self.group) + all_handles = torch.stack(handle_list) # shape (world_size, 64) + + # Post-init opens peer IPC handles and creates PushController + ops.post_init_push_ar(self._ptr, all_handles) + + def should_use(self, input_: torch.Tensor) -> bool: + """Check if push allreduce should handle this input.""" + if self.disabled: + return False + inp_size = input_.numel() * input_.element_size() + if inp_size == 0: + return False + if inp_size % 16 != 0: + return False + if not is_weak_contiguous(input_): + return False + return inp_size <= self.max_message_bytes + + def all_reduce(self, input_: torch.Tensor) -> Optional[torch.Tensor]: + """Perform push-based allreduce. Returns new output tensor.""" + if self._IS_CAPTURING: + if torch.cuda.is_current_stream_capturing(): + # Actual CUDA graph capture: record kernel + out = torch.empty_like(input_) + ops.push_ar_all_reduce(self._ptr, input_, out) + return out + else: + # Warmup before capture: mimic allocation pattern + return torch.empty_like(input_) + else: + # Eager mode: launch kernel directly + out = torch.empty_like(input_) + ops.push_ar_all_reduce(self._ptr, input_, out) + return out + + @contextmanager + def capture(self): + """Context manager for CUDA graph capture. + + SIMPLIFIED vs existing CustomAllreduce: + The push protocol does NOT need graph buffer registration because + it reads input LOCALLY and writes to pre-registered IPC push_buffers. + This context just toggles the _IS_CAPTURING flag to control the + warmup vs actual-capture behavior in all_reduce(). + """ + try: + self._IS_CAPTURING = True + yield + finally: + self._IS_CAPTURING = False + + def close(self): + """Release C++ resources. + + V1 fix: Use hasattr(self, '_ptr') instead of 'not self.disabled' + to handle the case where init_push_ar() succeeds but + _exchange_ipc_handles() fails (sets disabled=True but _ptr exists). + """ + if hasattr(self, "_ptr"): + ops.dispose_push_ar(self._ptr) + del self._ptr + self.disabled = True + + def __del__(self): + self.close() diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 8bd6e92157a..dbda3fd7ce6 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -571,6 +571,7 @@ class GroupCoordinator: # only cuda uses this function, # so we don't abstract it into the base class maybe_ca_context = nullcontext() + maybe_push_ar_context = nullcontext() maybe_aiter_context = nullcontext() from vllm.distributed.device_communicators.cuda_communicator import ( CudaCommunicator, @@ -588,6 +589,12 @@ class GroupCoordinator: if ca_comm is not None: maybe_ca_context = ca_comm.capture() # type: ignore + # Enter push allreduce capture context + push_ar_comm = getattr( + self.device_communicator, 'push_ar_comm', None) + if push_ar_comm is not None: + maybe_push_ar_context = push_ar_comm.capture() # type: ignore + from vllm._aiter_ops import rocm_aiter_ops if rocm_aiter_ops.is_enabled(): @@ -601,7 +608,7 @@ class GroupCoordinator: if curr_stream != stream: stream.wait_stream(curr_stream) - with torch.cuda.stream(stream), maybe_ca_context, maybe_aiter_context: + with torch.cuda.stream(stream), maybe_ca_context, maybe_push_ar_context, maybe_aiter_context: yield graph_capture_context def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: