Address review: use registered envs.VLLM_DISABLE_PUSH_ALLREDUCE

Replace direct os.environ.get(_DISABLE_ENV_VAR) == "1" check with
envs.VLLM_DISABLE_PUSH_ALLREDUCE to use the centrally registered
env var from envs.py, which provides validation and caching.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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")

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 <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+#include <torch/all.h>
+
+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<int>(rank),
+                                       static_cast<int>(world_size),
+                                       push_buffer_bytes,
+                                       static_cast<int>(max_num_cta));
+  return reinterpret_cast<fptr_t>(mgr);
+}
+
+// Get IPC handle as a byte tensor
+torch::Tensor get_push_ar_ipc_handle(fptr_t _mgr) {
+  auto* mgr = reinterpret_cast<PushAllReduceManager*>(_mgr);
+  cudaIpcMemHandle_t handle = mgr->get_ipc_handle();
+  auto t = torch::from_blob(&handle,
+                            {static_cast<int64_t>(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<PushAllReduceManager*>(_mgr);
+  int world_size = all_handles.size(0);
+  std::vector<cudaIpcMemHandle_t> 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<size_t>(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<PushAllReduceManager*>(_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<nv_bfloat16>(
+          stream, reinterpret_cast<nv_bfloat16*>(inp.data_ptr()),
+          reinterpret_cast<nv_bfloat16*>(out.data_ptr()), out.numel());
+      break;
+    case at::ScalarType::Half:
+      mgr->allreduce<half>(stream, reinterpret_cast<half*>(inp.data_ptr()),
+                           reinterpret_cast<half*>(out.data_ptr()),
+                           out.numel());
+      break;
+    case at::ScalarType::Float:
+      mgr->allreduce<float>(stream,
+                            reinterpret_cast<float*>(inp.data_ptr()),
+                            reinterpret_cast<float*>(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<PushAllReduceManager*>(_mgr);
+  delete mgr;
+}

csrc/push_all_reduce.cuh

@@ -0,0 +1,248 @@
+// 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 <cuda_runtime.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cstring>
+#include <numeric>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace vllm {
+namespace push_ar {
+
+#define PUSH_AR_CUDACHECK(cmd)                                          \
+  do {                                                                  \
+    cudaError_t e = cmd;                                                \
+    if (e != cudaSuccess) {                                             \
+      throw std::runtime_error(                                         \
+          std::string("push_all_reduce CUDA error at ") + __FILE__ +    \
+          ":" + std::to_string(__LINE__) + " '" +                       \
+          cudaGetErrorString(e) + "'");                                 \
+    }                                                                   \
+  } while (0)
+
+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;
+    PUSH_AR_CUDACHECK(cudaGetDevice(&device_id));
+    int major;
+    PUSH_AR_CUDACHECK(cudaDeviceGetAttribute(
+        &major, cudaDevAttrComputeCapabilityMajor, device_id));
+    use_pdl_ = (major >= 9);  // Hopper (sm90) or newer
+
+    // Allocate storage
+    storage_bytes_ = push_signal_bytes() + push_buffer_total_bytes();
+    PUSH_AR_CUDACHECK(cudaMalloc(&storage_, storage_bytes_));
+    // Zeros signals (epoch=0 for all CTAs) AND push buffer
+    // (0x0000 = IEEE 754 positive-zero = "empty" sentinel)
+    PUSH_AR_CUDACHECK(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;
+    PUSH_AR_CUDACHECK(cudaIpcGetMemHandle(&handle, storage_));
+    return handle;
+  }
+
+  // Phase 2: Open peer IPC handles and init PushController
+  void post_init(const std::vector<cudaIpcMemHandle_t>& 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 {
+        PUSH_AR_CUDACHECK(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 <typename T>
+  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<uint32_t>(num_elements);
+    const int num_threads = select_num_threads<T>(num_items);
+
+    // Verify input fits in push buffer (runtime check, not compiled out)
+    const int64_t input_bytes =
+        static_cast<int64_t>(sizeof(T)) * num_elements;
+    if (input_bytes > push_buffer_bytes_) {
+      throw std::runtime_error(
+          "push_all_reduce: input (" + std::to_string(input_bytes) +
+          " bytes) exceeds push buffer capacity (" +
+          std::to_string(push_buffer_bytes_) + " 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<uint32_t>(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<T>(config, params);
+  }
+
+ private:
+  // Template dispatch helper
+  template <typename T>
+  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<T, 8, true>;
+        cudaLaunchKernelEx(&config, kernel, params, push_ctrl_);
+      } else {
+        auto kernel = all_reduce_one_shot_push_kernel<T, 8, false>;
+        cudaLaunchKernelEx(&config, kernel, params, push_ctrl_);
+      }
+    } else if (world_size_ == 4) {
+      if (use_pdl_) {
+        auto kernel = all_reduce_one_shot_push_kernel<T, 4, true>;
+        cudaLaunchKernelEx(&config, kernel, params, push_ctrl_);
+      } else {
+        auto kernel = all_reduce_one_shot_push_kernel<T, 4, false>;
+        cudaLaunchKernelEx(&config, kernel, params, push_ctrl_);
+      }
+    } else if (world_size_ == 2) {
+      if (use_pdl_) {
+        auto kernel = all_reduce_one_shot_push_kernel<T, 2, true>;
+        cudaLaunchKernelEx(&config, kernel, params, push_ctrl_);
+      } else {
+        auto kernel = all_reduce_one_shot_push_kernel<T, 2, false>;
+        cudaLaunchKernelEx(&config, kernel, params, push_ctrl_);
+      }
+    } else if (world_size_ == 6) {
+      if (use_pdl_) {
+        auto kernel = all_reduce_one_shot_push_kernel<T, 6, true>;
+        cudaLaunchKernelEx(&config, kernel, params, push_ctrl_);
+      } else {
+        auto kernel = all_reduce_one_shot_push_kernel<T, 6, false>;
+        cudaLaunchKernelEx(&config, kernel, params, push_ctrl_);
+      }
+    }
+  }
+
+  // Thread count selection (from SGLang CustomAllReducePush::all_reduce)
+  template <typename T>
+  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<int>(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<char*>(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<void*> peer_storage_;
+
+  PushController push_ctrl_;
+  bool ctrl_initialized_ = false;
+};
+
+}  // namespace push_ar
+}  // namespace vllm

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 <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <cstddef>
+#include <cstdint>
+
+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<packed_t<DType>, 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 <bool kUsePDL>
+__device__ __forceinline__ void PDLWaitPrimary() {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900
+  if constexpr (kUsePDL) {
+    asm volatile("griddepcontrol.wait;" ::: "memory");
+  }
+#endif
+}
+
+template <bool kUsePDL>
+__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<char>)
+__device__ __forceinline__ void* ptr_byte_offset(void* ptr, int64_t off1) {
+  return static_cast<char*>(ptr) + off1;
+}
+
+__device__ __forceinline__ void* ptr_byte_offset(void* ptr, int64_t off1,
+                                                 int64_t off2) {
+  return static_cast<char*>(ptr) + off1 + off2;
+}
+
+// Typed offset for AlignedVector load/store addressing
+// (replaces pointer::offset<T>)
+template <typename T>
+__device__ __forceinline__ void* ptr_typed_offset(void* ptr, int64_t offset) {
+  return static_cast<T*>(ptr) + offset;
+}
+
+template <typename T>
+__device__ __forceinline__ const void* ptr_typed_offset(const void* ptr,
+                                                        int64_t offset) {
+  return static_cast<const T*>(ptr) + offset;
+}
+
+// Host-side pointer offset (for storage layout calculations)
+inline void* host_ptr_offset(void* ptr, int64_t off) {
+  return static_cast<char*>(ptr) + off;
+}
+
+// ============================================================
+// Section E: dtype_trait system (from SGLang type.cuh)
+// ============================================================
+// COPY AS IS with namespace adjustment.
+
+template <typename T>
+struct dtype_trait {};
+
+template <>
+struct dtype_trait<fp32_t> {
+  using self_t = fp32_t;
+  using packed_t = fp32x2_t;
+  template <typename S>
+  __device__ __forceinline__ static self_t from(const S& value) {
+    return static_cast<fp32_t>(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<fp16_t> {
+  using self_t = fp16_t;
+  using packed_t = fp16x2_t;
+  template <typename S>
+  __device__ __forceinline__ static self_t from(const S& value) {
+    return static_cast<fp16_t>(value);
+  }
+};
+
+template <>
+struct dtype_trait<bf16_t> {
+  using self_t = bf16_t;
+  using packed_t = bf16x2_t;
+  template <typename S>
+  __device__ __forceinline__ static self_t from(const S& value) {
+    return static_cast<bf16_t>(value);
+  }
+};
+
+template <>
+struct dtype_trait<fp32x2_t> {
+  using self_t = fp32x2_t;
+  template <typename S>
+  __device__ __forceinline__ static self_t from(const S& value) {
+    return static_cast<fp32x2_t>(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<fp16x2_t> {
+  using self_t = fp16x2_t;
+  template <typename S>
+  __device__ __forceinline__ static self_t from(const S& value) {
+    return static_cast<fp16x2_t>(value);
+  }
+  __device__ __forceinline__ static self_t from(const fp32x2_t& x) {
+    return __float22half2_rn(x);
+  }
+};
+
+template <>
+struct dtype_trait<bf16x2_t> {
+  using self_t = bf16x2_t;
+  template <typename S>
+  __device__ __forceinline__ static self_t from(const S& value) {
+    return static_cast<bf16x2_t>(value);
+  }
+  __device__ __forceinline__ static self_t from(const fp32x2_t& x) {
+    return __float22bfloat162_rn(x);
+  }
+};
+
+template <typename T>
+using packed_t = typename dtype_trait<T>::packed_t;
+
+template <typename To, typename From>
+__device__ __forceinline__ To cast(const From& value) {
+  return dtype_trait<To>::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 <std::size_t N>
+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 <typename T>
+using sized_int = typename uint_trait<sizeof(T)>::type;
+
+}  // namespace detail
+
+template <typename T, std::size_t N>
+struct alignas(sizeof(T) * N) AlignedStorage {
+  T data[N];
+};
+
+template <typename T, std::size_t N>
+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<T>;
+  using storage_t = AlignedStorage<element_t, N>;
+
+ public:
+  __device__ __forceinline__ void load(const void* ptr, int64_t offset = 0) {
+    m_storage = reinterpret_cast<const storage_t*>(ptr)[offset];
+  }
+  __device__ __forceinline__ void store(void* ptr, int64_t offset = 0) const {
+    reinterpret_cast<storage_t*>(ptr)[offset] = m_storage;
+  }
+  __device__ __forceinline__ void fill(T value) {
+    const auto store_value = *reinterpret_cast<element_t*>(&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<T*>(&m_storage)[idx];
+  }
+  __device__ __forceinline__ auto operator[](std::size_t idx) const -> T {
+    return reinterpret_cast<const T*>(&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<SignalType*>(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 <typename T>
+struct fp_trait {};
+
+template <>
+struct fp_trait<bf16_t> {
+  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<fp16_t> {
+  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<float> {
+  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 <typename DType>
+__device__ __forceinline__ void clear_pos_zero(DType& val) {
+  using Trait = fp_trait<DType>;
+  const auto ptr = reinterpret_cast<typename Trait::type*>(&val);
+  if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero;
+}
+
+template <typename DType>
+__device__ __forceinline__ bool is_pos_zero(const DType& val) {
+  using Trait = fp_trait<DType>;
+  const auto ptr = reinterpret_cast<const typename Trait::type*>(&val);
+  return *ptr == Trait::pos_zero;
+}
+
+template <typename DType>
+__device__ __forceinline__ DType get_pos_zero() {
+  using Trait = fp_trait<DType>;
+  const auto value = Trait::pos_zero;
+  return *reinterpret_cast<const DType*>(&value);
+}
+
+// ============================================================
+// Section K: Volatile 16-byte load/store (from SGLang push.cuh 87-105)
+// ============================================================
+// CHANGED: pointer::offset<T> -> ptr_typed_offset<T>
+
+template <typename T>
+__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<T>(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<const T*>(&val);
+}
+
+template <typename T>
+__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<const uint4*>(&x);
+  addr = ptr_typed_offset<T>(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 <typename DType2, size_t N, uint32_t M>
+__device__ __forceinline__ auto reduce_impl(
+    AlignedVector<DType2, N> (&storage)[M]) -> AlignedVector<DType2, N> {
+  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<fp32x2_t>(storage[i][j]);
+      auto& [x_acc, y_acc] = acc[j];
+      x_acc += x;
+      y_acc += y;
+    }
+  }
+  AlignedVector<DType2, N> result;
+#pragma unroll
+  for (uint32_t j = 0; j < N; ++j) {
+    result[j] = cast<DType2>(acc[j]);
+  }
+  return result;
+}
+
+// ============================================================
+// Section M: push_impl (from SGLang push.cuh 107-128)
+// ============================================================
+// CHANGED: pointer::offset -> ptr_byte_offset
+
+template <typename DType, uint32_t kNumGPU>
+__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<packed_t<DType>, 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 <typename DType, uint32_t kNumGPU>
+__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<packed_t<DType>, 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<DType>(), get_pos_zero<DType>()});
+#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 <typename DType, uint32_t kNumGPU, bool kUsePDL>
+__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<kUsePDL>();
+
+  // 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<DType*>(
+        ptr_byte_offset(buffer[i], rank * buffer_bytes, epoch_offset));
+  }
+  push_impl(push_buf, input, num_items);
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  // 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<DType*>(
+        ptr_byte_offset(buffer[rank], i * buffer_bytes, epoch_offset));
+  }
+  poll_impl(poll_buf, output, num_items);
+  ctrl.exit();
+}
+
+}  // namespace push_ar
+}  // namespace vllm

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)

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(("localhost", 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()

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,

vllm/distributed/device_communicators/cuda_communicator.py

@@ -65,6 +65,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
         from vllm.distributed.device_communicators.flashinfer_all_reduce import (
             FlashInferAllReduce,
         )
+        from vllm.distributed.device_communicators.push_all_reduce import (
+            PushAllReduce,
+        )
         from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
         from vllm.distributed.device_communicators.quick_all_reduce import (
             QuickAllReduce,
@@ -81,6 +84,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
                 register_nccl_symmetric_ops(self.pynccl_comm)
 
         self.ca_comm: CustomAllreduce | None = None
+        self.push_ar_comm: PushAllReduce | None = None
         self.qr_comm: QuickAllReduce | None = None
         self.symm_mem_comm: SymmMemCommunicator | None = None
         self.fi_ar_comm: FlashInferAllReduce | None = None
@@ -107,6 +111,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 +302,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
@@ -414,6 +444,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
         if self.pynccl_comm is not None:
             self.pynccl_comm.destroy()
             self.pynccl_comm = None
+        if self.push_ar_comm is not None:
+            self.push_ar_comm.close()
+            self.push_ar_comm = None
         if self.ca_comm is not None:
             self.ca_comm = None
         if self.fi_ar_comm is not None:

vllm/distributed/device_communicators/push_all_reduce.py

@@ -0,0 +1,257 @@
+# 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
+
+import vllm.envs as envs
+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 maps: 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 thresholds for architectures without tuned values
+PUSH_THRESHOLD_DEFAULT = {
+    2: 512 * 1024,  # 512 KB
+    4: 512 * 1024,  # 512 KB
+    6: 512 * 1024,  # 512 KB
+    8: 512 * 1024,  # 512 KB
+}
+
+# Map GPU major compute capability -> threshold table
+_THRESHOLD_BY_ARCH: dict[int, dict[int, int]] = {
+    10: PUSH_THRESHOLD_SM100,  # Blackwell (sm_100)
+}
+
+# Conservative default for untuned GPUs / unrecognized world_size
+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 envs.VLLM_DISABLE_PUSH_ALLREDUCE:
+            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 and architecture for grid size and threshold selection
+        props = torch.cuda.get_device_properties(device)
+        self.num_sm = props.multi_processor_count
+
+        # Determine push buffer size from architecture-specific threshold map
+        if max_size is not None:
+            self.push_buffer_bytes = max_size
+        else:
+            arch_major = props.major
+            threshold_map = _THRESHOLD_BY_ARCH.get(
+                arch_major, PUSH_THRESHOLD_DEFAULT
+            )
+            self.push_buffer_bytes = threshold_map.get(
+                self.world_size, DEFAULT_PUSH_BUFFER
+            )
+            if arch_major not in _THRESHOLD_BY_ARCH:
+                logger.info(
+                    "PushAllReduce: no tuned thresholds for sm_%d%d, "
+                    "using conservative default (%d KB)",
+                    arch_major,
+                    props.minor,
+                    self.push_buffer_bytes // 1024,
+                )
+
+        # 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()

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:

vllm/envs.py

@@ -225,6 +225,7 @@ if TYPE_CHECKING:
     VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False
     VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
     VLLM_ALLREDUCE_USE_FLASHINFER: bool = False
+    VLLM_DISABLE_PUSH_ALLREDUCE: bool = False
     VLLM_TUNED_CONFIG_FOLDER: str | None = None
     VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set()
     VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False
@@ -1653,6 +1654,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_ALLREDUCE_USE_FLASHINFER": lambda: bool(
         int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER", "0"))
     ),
+    # If set to 1, disable push-based allreduce for small tensors.
+    # When disabled, small tensor reductions fall back to the barrier-based
+    # CustomAllreduce path.
+    "VLLM_DISABLE_PUSH_ALLREDUCE": lambda: bool(
+        int(os.getenv("VLLM_DISABLE_PUSH_ALLREDUCE", "0"))
+    ),
     # Experimental: use this to enable MCP tool calling for non harmony models
     "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool(
         int(os.getenv("VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", "0"))