[Model][Perf] Build Beam's fused all-reduce ahead of time under csrc/beam

Move the Beam all-reduce kernel out of the model folder and drop its
runtime FlashInfer JIT. The MNNVL Lamport device helpers it uses are
vendored from FlashInfer into csrc/mnnvl_lamport.cuh and the kernel lives
in csrc/beam/allreduce_norm.cu (SM90/SM100, --fmad=false). FlashInfer
still maps the multicast workspace (McastGPUBuffer), but the Lamport
protocol, its initialization, and its flags are owned by vLLM, so the
build no longer depends on FlashInfer's headers. The large-batch fallback
uses the vendored two-shot kernel.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

CMakeLists.txt

@@ -1150,6 +1150,27 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       "Building fused GDN decode for archs: ${FUSED_GDN_DECODE_ARCHS}")
   endif()
 
+  # Beam's fused all-reduce needs NVLink multicast and thread-block clusters.
+  # Its norms must not contract into FMAs to keep the FP32 boundary exact.
+  if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
+    cuda_archs_loose_intersection(BEAM_ALLREDUCE_ARCHS
+      "9.0a;10.0f;10.3f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(BEAM_ALLREDUCE_ARCHS
+      "9.0a;10.0a;10.3a" "${CUDA_ARCHS}")
+  endif()
+  if(BEAM_ALLREDUCE_ARCHS)
+    set(BEAM_ALLREDUCE_SRC "csrc/beam/allreduce_norm.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${BEAM_ALLREDUCE_SRC}"
+      CUDA_ARCHS "${BEAM_ALLREDUCE_ARCHS}")
+    set_property(SOURCE ${BEAM_ALLREDUCE_SRC} APPEND PROPERTY
+      COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:--fmad=false>")
+    list(APPEND VLLM_STABLE_EXT_SRC "${BEAM_ALLREDUCE_SRC}")
+    message(STATUS
+      "Building Beam all-reduce for archs: ${BEAM_ALLREDUCE_ARCHS}")
+  endif()
+
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
     cuda_archs_loose_intersection(KIMI_K3_ATTN_RES_ARCHS
       "10.0f" "${CUDA_ARCHS}")
@@ -1251,6 +1272,10 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
       target_compile_definitions(_C_stable_libtorch PRIVATE
         VLLM_ENABLE_KIMI_K3_ATTN_RES=1)
     endif()
+    if(BEAM_ALLREDUCE_ARCHS)
+      target_compile_definitions(_C_stable_libtorch PRIVATE
+        VLLM_ENABLE_BEAM_ALLREDUCE=1)
+    endif()
     # Needed by CUTLASS kernels
     target_compile_definitions(_C_stable_libtorch PRIVATE
       CUTLASS_ENABLE_DIRECT_CUDA_DRIVER_CALL=1)

csrc/beam/allreduce_norm.cu

@@ -1,13 +1,17 @@
-// SPDX-License-Identifier: Apache-2.0
-// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-// Communication helpers and launch geometry from NVIDIA/FlashInfer's MNNVL
-// all-reduce (Apache-2.0); Beam supplies the normalization epilogue.
+// Beam's TP all-reduce fused with its FP32 norm-residual-norm epilogue and,
+// optionally, the TRTLLM-Gen MoE top-k finalize. Built on the MNNVL Lamport
+// primitives in mnnvl_lamport.cuh; launch geometry follows FlashInfer's MNNVL
+// all-reduce.
 
-#include "flashinfer/comm/trtllm_mnnvl_allreduce.cuh"
-#include "tvm_ffi_utils.h"
+#include "../torch_utils.h"
+#include "../mnnvl_lamport.cuh"
 
+#include <torch/headeronly/core/ScalarType.h>
+
+namespace vllm {
 namespace beam {
-namespace ar = flashinfer::trtllm_mnnvl_allreduce;
+
+namespace ar = vllm::mnnvl;
 namespace cg = cooperative_groups;
 using T = __nv_bfloat16;
 using Vec = ar::utils::PackedVec<float4, T>;
@@ -277,78 +281,103 @@ __global__ void __launch_bounds__(1024) finishTwoshot(Params p) {
   epilogue<UseCGA>(p, value, packedIdx, offset);
   cudaGridDependencySynchronize();
   flag.waitAndUpdate(
-      {static_cast<uint32_t>(flashinfer::round_up(p.ar.numTokens, p.ar.nRanks) *
+      {static_cast<uint32_t>(ar::round_up(p.ar.numTokens, p.ar.nRanks) *
                              p.ar.tokenDim * sizeof(T)),
        static_cast<uint32_t>(p.ar.numTokens * p.ar.tokenDim * sizeof(T)), 0,
        0});
 }
 
 void checkCuda(cudaError_t status) {
-  TVM_FFI_ICHECK_EQ(status, cudaSuccess) << cudaGetErrorString(status);
+  STD_TORCH_CHECK(status == cudaSuccess, cudaGetErrorString(status));
 }
 
-void launch(TensorView input, TensorView residual,
-            tvm::ffi::Optional<TensorView> postWeight,
-            tvm::ffi::Optional<TensorView> nextWeight, TensorView hidden,
-            tvm::ffi::Optional<TensorView> normalized, int64_t pointers,
-            TensorView flags, int64_t multicast, int64_t local, int64_t world,
-            int64_t rank, double postEps, double nextEps, double xScale,
-            double rScale, bool postNorm, bool nextNorm, bool scaled,
-            bool oneShot, tvm::ffi::Optional<TensorView> experts,
-            tvm::ffi::Optional<TensorView> expertWeights,
-            tvm::ffi::Optional<TensorView> permutation, double routedScale,
-            bool smallBatch) {
-  ffi::CUDADeviceGuard guard(input.device().device_id);
-  auto stream = get_stream(input.device());
-  TVM_FFI_ICHECK_EQ(encode_dlpack_dtype(input.dtype()), bfloat16_code);
-  TVM_FFI_ICHECK_EQ(input.size(1) % 8, 0);
+ar::AllReduceKernelParams<T> workspaceParams(torch::stable::Tensor const& input,
+                                             int64_t buffer_ptrs,
+                                             torch::stable::Tensor& flags,
+                                             int64_t multicast_ptr,
+                                             int64_t world, int64_t rank) {
+  STD_TORCH_CHECK(
+      input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
+      "Beam all-reduce requires BF16 input");
+  STD_TORCH_CHECK(
+      input.dim() == 2 && input.is_contiguous() && input.size(1) % 8 == 0,
+      "Beam all-reduce requires a contiguous [tokens, 8k] input");
+  STD_TORCH_CHECK(world == 2 || world == 4 || world == 8,
+                  "Beam all-reduce supports TP2/4/8");
+  ar::AllReduceKernelParams<T> p{};
+  p.shardPtr = static_cast<T const*>(input.data_ptr());
+  p.inputPtrs = reinterpret_cast<T**>(buffer_ptrs);
+  p.mcastPtr = reinterpret_cast<T*>(multicast_ptr);
+  p.bufferFlags = static_cast<uint32_t*>(flags.data_ptr());
+  p.nRanks = world;
+  p.rank = rank;
+  p.numTokens = input.size(0);
+  p.tokenDim = input.size(1);
+  return p;
+}
+
+}  // namespace beam
+}  // namespace vllm
+
+void beam_allreduce_norm(
+    torch::stable::Tensor const& input, torch::stable::Tensor const& residual,
+    std::optional<torch::stable::Tensor> post_weight,
+    std::optional<torch::stable::Tensor> next_weight,
+    torch::stable::Tensor& hidden,
+    std::optional<torch::stable::Tensor> normalized, int64_t buffer_ptrs,
+    torch::stable::Tensor& flags, int64_t multicast_ptr, int64_t local_ptr,
+    int64_t world, int64_t rank, double post_eps, double next_eps,
+    double x_scale, double r_scale, bool post_norm, bool next_norm, bool scaled,
+    bool one_shot, bool small_batch,
+    std::optional<torch::stable::Tensor> experts,
+    std::optional<torch::stable::Tensor> expert_weights,
+    std::optional<torch::stable::Tensor> permutation, double routed_scale) {
+  using namespace vllm::beam;
+  int const device = input.get_device_index();
+  torch::stable::accelerator::DeviceGuard const device_guard(device);
+  cudaStream_t stream = get_current_cuda_stream(device);
   Params p{};
+  p.ar = workspaceParams(input, buffer_ptrs, flags, multicast_ptr, world, rank);
+  p.ar.localBufferPtr = reinterpret_cast<T*>(local_ptr);
   p.ar.outputPtr = normalized.has_value()
-                       ? static_cast<T*>(normalized.value().data_ptr())
+                       ? static_cast<T*>(normalized->data_ptr())
                        : nullptr;
   p.ar.prenormedPtr = static_cast<T*>(hidden.data_ptr());
-  p.ar.shardPtr = static_cast<T const*>(input.data_ptr());
   p.ar.residualInPtr = static_cast<T const*>(residual.data_ptr());
-  p.ar.gammaPtr = postWeight.has_value()
-                      ? static_cast<T const*>(postWeight.value().data_ptr())
+  p.ar.gammaPtr = post_weight.has_value()
+                      ? static_cast<T const*>(post_weight->data_ptr())
                       : nullptr;
-  p.nextWeight = nextWeight.has_value()
-                     ? static_cast<T const*>(nextWeight.value().data_ptr())
+  p.ar.epsilon = post_eps;
+  p.nextWeight = next_weight.has_value()
+                     ? static_cast<T const*>(next_weight->data_ptr())
                      : nullptr;
-  p.ar.inputPtrs = reinterpret_cast<T**>(pointers);
-  p.ar.mcastPtr = reinterpret_cast<T*>(multicast);
-  p.ar.localBufferPtr = reinterpret_cast<T*>(local);
-  p.ar.bufferFlags = static_cast<uint32_t*>(flags.data_ptr());
-  p.ar.nRanks = world;
-  p.ar.rank = rank;
-  p.ar.numTokens = input.size(0);
-  p.ar.tokenDim = input.size(1);
-  p.ar.epsilon = postEps;
-  p.nextEps = nextEps;
-  p.xScale = xScale;
-  p.rScale = rScale;
-  p.postNorm = postNorm;
-  p.nextNorm = nextNorm;
+  p.nextEps = next_eps;
+  p.xScale = x_scale;
+  p.rScale = r_scale;
+  p.postNorm = post_norm;
+  p.nextNorm = next_norm;
   p.scaled = scaled;
   if (experts.has_value()) {
-    TVM_FFI_ICHECK(expertWeights.has_value() && permutation.has_value());
-    p.experts = static_cast<T const*>(experts.value().data_ptr());
-    p.expertWeights = expertWeights.value().data_ptr();
-    p.permutation = static_cast<int const*>(permutation.value().data_ptr());
-    p.topK = permutation.value().size(1);
+    STD_TORCH_CHECK(expert_weights.has_value() && permutation.has_value());
+    p.experts = static_cast<T const*>(experts->data_ptr());
+    p.expertWeights = expert_weights->data_ptr();
+    p.permutation = static_cast<int const*>(permutation->data_ptr());
+    p.topK = permutation->size(1);
     p.floatWeights =
-        encode_dlpack_dtype(expertWeights.value().dtype()) == float32_code;
-    p.routedScale = routedScale;
+        expert_weights->scalar_type() == torch::headeronly::ScalarType::Float;
+    p.routedScale = routed_scale;
   }
+  if (p.ar.numTokens == 0) return;
   auto [threads, clusters, loads] =
-      ar::utils::adjustGridConfig(p.ar.numTokens, p.ar.tokenDim, 8, true);
-  TVM_FFI_ICHECK(!smallBatch || (oneShot && !p.experts));
-  if (smallBatch) {
+      ar::utils::adjustGridConfig(p.ar.numTokens, p.ar.tokenDim, 8, true,
+                                  get_device_prop()->multiProcessorCount);
+  STD_TORCH_CHECK(!small_batch || (one_shot && !p.experts));
+  if (small_batch) {
     clusters = 1;
     threads = p.ar.tokenDim / 8;
   }
-  TVM_FFI_ICHECK_EQ(loads, 1);
-  TVM_FFI_ICHECK_EQ(threads * clusters * 8, p.ar.tokenDim);
+  STD_TORCH_CHECK(loads == 1);
+  STD_TORCH_CHECK(threads * clusters * 8 == p.ar.tokenDim);
   cudaLaunchAttribute attrs[2]{};
   attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
   attrs[0].val.programmaticStreamSerializationAllowed = 1;
@@ -359,11 +388,11 @@ void launch(TensorView input, TensorView residual,
   config.blockDim = dim3(threads);
   config.stream = stream;
   config.attrs = attrs;
-  config.numAttrs = smallBatch ? 1 : 2;
-  if (oneShot) {
+  config.numAttrs = small_batch ? 1 : 2;
+  if (one_shot) {
 #define LAUNCH(W)                                                          \
   case W:                                                                  \
-    if (smallBatch)                                                        \
+    if (small_batch)                                                       \
       checkCuda(cudaLaunchKernelEx(&config, &oneshot<W, false, true>, p)); \
     else if (p.experts)                                                    \
       checkCuda(cudaLaunchKernelEx(&config, &oneshot<W, true>, p));        \
@@ -374,14 +403,11 @@ void launch(TensorView input, TensorView residual,
       LAUNCH(2);
       LAUNCH(4);
       LAUNCH(8);
-      default:
-        TVM_FFI_LOG_AND_THROW(ValueError) << "Beam supports TP2/4/8";
     }
 #undef LAUNCH
   } else {
     cudaLaunchConfig_t scatter{};
-    scatter.gridDim =
-        dim3(p.ar.numTokens, flashinfer::ceil_div(p.ar.tokenDim, 1024));
+    scatter.gridDim = dim3(p.ar.numTokens, ar::ceil_div(p.ar.tokenDim, 1024));
     scatter.blockDim = dim3(128);
     scatter.stream = stream;
     scatter.attrs = attrs;
@@ -398,8 +424,6 @@ void launch(TensorView input, TensorView residual,
       LAUNCH(2);
       LAUNCH(4);
       LAUNCH(8);
-      default:
-        TVM_FFI_LOG_AND_THROW(ValueError) << "Beam supports TP2/4/8";
     }
 #undef LAUNCH
     if (clusters > 1) {
@@ -410,5 +434,39 @@ void launch(TensorView input, TensorView residual,
     }
   }
 }
-}  // namespace beam
-TVM_FFI_DLL_EXPORT_TYPED_FUNC(beam_allreduce, beam::launch);
+
+void beam_allreduce(torch::stable::Tensor const& input,
+                    torch::stable::Tensor& output, int64_t buffer_ptrs,
+                    torch::stable::Tensor& flags, int64_t multicast_ptr,
+                    int64_t world, int64_t rank) {
+  using namespace vllm::beam;
+  int const device = input.get_device_index();
+  torch::stable::accelerator::DeviceGuard const device_guard(device);
+  auto p =
+      workspaceParams(input, buffer_ptrs, flags, multicast_ptr, world, rank);
+  p.outputPtr = static_cast<T*>(output.data_ptr());
+  if (p.numTokens == 0) return;
+  cudaLaunchAttribute attrs[1]{};
+  attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization;
+  attrs[0].val.programmaticStreamSerializationAllowed = 1;
+  cudaLaunchConfig_t config{};
+  config.gridDim = dim3(p.numTokens, ar::ceil_div(p.tokenDim / 8, 128));
+  config.blockDim = dim3(128);
+  config.stream = get_current_cuda_stream(device);
+  config.attrs = attrs;
+  config.numAttrs = 1;
+  switch (world) {
+    case 2:
+      checkCuda(cudaLaunchKernelEx(&config,
+                                   &ar::twoshotAllreduceKernel<2, T, true>, p));
+      break;
+    case 4:
+      checkCuda(cudaLaunchKernelEx(&config,
+                                   &ar::twoshotAllreduceKernel<4, T, true>, p));
+      break;
+    case 8:
+      checkCuda(cudaLaunchKernelEx(&config,
+                                   &ar::twoshotAllreduceKernel<8, T, true>, p));
+      break;
+  }
+}

csrc/mnnvl_lamport.cuh

@@ -0,0 +1,778 @@
+/*
+ * Copyright (c) 2022-2024, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// MNNVL Lamport all-reduce primitives adapted from FlashInfer's
+// include/flashinfer/comm/trtllm_mnnvl_allreduce.cuh (v0.7.0), itself from
+// TensorRT-LLM. Workspace layout: three Lamport buffers per rank, initialized
+// to FP32 -0.0, plus a uint32 flag array
+// {current, dirty, bytes per buffer, dirty stages, bytes to clear[4], count}.
+
+#pragma once
+
+#include <cooperative_groups.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <array>
+#include <tuple>
+#include <type_traits>
+#include <utility>
+
+namespace vllm {
+namespace mnnvl {
+
+template <typename T>
+struct AllReduceKernelParams {
+  T* outputPtr;
+  T* prenormedPtr;
+  T const* shardPtr;
+  T const* residualInPtr;
+  T const* gammaPtr;
+  T** inputPtrs;
+  T* mcastPtr;
+  T* localBufferPtr;
+  int nRanks;
+  int rank;
+  int numTokens;
+  int tokenDim;
+  float epsilon;
+  uint32_t* bufferFlags;
+};
+
+template <typename T1, typename T2>
+__forceinline__ __device__ __host__ constexpr T1 ceil_div(const T1 x,
+                                                          const T2 y) noexcept {
+  return (x + y - 1) / y;
+}
+
+template <typename T1, typename T2>
+__forceinline__ __device__ __host__ constexpr T1 round_up(const T1 x,
+                                                          const T2 y) noexcept {
+  return ceil_div(x, y) * y;
+}
+
+namespace utils {
+
+constexpr uint16_t kNEGZERO_FP16 = 0x8000U;
+constexpr uint32_t kNEGZERO_FP32 = 0x80000000U;
+constexpr int kMaxClusterSize = 8;
+
+template <typename T>
+union Fp16BitCast {
+  T mFp;
+  uint16_t mInt;
+
+  constexpr Fp16BitCast() : mInt(0) {}
+
+  constexpr Fp16BitCast(T val) : mFp(val) {}
+
+  constexpr Fp16BitCast(uint16_t val) : mInt(val) {}
+};
+
+template <typename T>
+inline __device__ float toFloat(T val) {
+  return val;
+}
+
+template <>
+inline __device__ float toFloat<__nv_bfloat16>(__nv_bfloat16 val) {
+  return __bfloat162float(val);
+}
+template <>
+inline __device__ float toFloat<__nv_half>(__nv_half val) {
+  return __half2float(val);
+}
+
+template <typename T>
+inline __device__ T fromFloat(float val) {
+  return val;
+}
+
+template <>
+inline __device__ __nv_bfloat16 fromFloat<__nv_bfloat16>(float val) {
+  return __float2bfloat16(val);
+}
+
+template <>
+inline __device__ __nv_half fromFloat<__nv_half>(float val) {
+  return __float2half(val);
+}
+
+template <typename T>
+static constexpr __device__ __host__ T negZero() {
+  if constexpr (std::is_same_v<T, float>) {
+    return -0.0F;
+  } else if constexpr (std::is_same_v<T, __nv_bfloat16> ||
+                       std::is_same_v<T, __nv_half>) {
+    return Fp16BitCast<T>(kNEGZERO_FP16).mFp;
+  } else {
+    static_assert(sizeof(T) == 0, "negativeZero not specialized for this type");
+  }
+  return T{};
+}
+
+// The Lamport sentinel is a bit pattern. Compare bit-exact: an FTZ float
+// compare would also match negative subnormals and deadlock the poll loop.
+template <typename T>
+static inline __device__ bool isNegZero(T val) {
+  if constexpr (std::is_same_v<T, float>) {
+    return __float_as_uint(val) == kNEGZERO_FP32;
+  } else if constexpr (std::is_same_v<T, __nv_bfloat16> ||
+                       std::is_same_v<T, __nv_half>) {
+    return Fp16BitCast<T>(val).mInt == kNEGZERO_FP16;
+  } else {
+    static_assert(sizeof(T) == 0, "isNegZero not specialized for this type");
+  }
+  return false;
+}
+
+template <typename PackedType, typename T>
+constexpr __device__ __host__ PackedType getPackedLamportInit() {
+  static_assert(sizeof(PackedType) % sizeof(T) == 0,
+                "PackedType size must be divisible by T size");
+  constexpr int kNumElements = sizeof(PackedType) / sizeof(T);
+
+  union PackedT {
+    PackedType mPacked;
+    std::array<T, kNumElements> mElements;
+
+    constexpr PackedT() : mElements{} {
+      for (int i = 0; i < kNumElements; i++) {
+        mElements[i] = negZero<T>();
+      }
+    }
+  };
+
+  PackedT initValue{};
+  return initValue.mPacked;
+}
+
+struct LamportBufferLayout {
+  uint32_t numStages = 1;
+  uint32_t bytesPerBuffer = 0;
+  static constexpr uint32_t sNumLamportBuffers = 3;
+
+  [[nodiscard]] __device__ __host__ size_t getTotalBytes() const {
+    return numStages * static_cast<size_t>(bytesPerBuffer / numStages) *
+           sNumLamportBuffers;
+  }
+
+  [[nodiscard]] __device__ __host__ void* getStagePtr(
+      void* bufferBasePtr, uint32_t lamportIndex, uint32_t stageIndex) const {
+    return reinterpret_cast<void*>(
+        reinterpret_cast<char*>(bufferBasePtr) +
+        static_cast<size_t>((lamportIndex * numStages + stageIndex) *
+                            static_cast<size_t>(bytesPerBuffer / numStages)));
+  }
+};
+
+namespace cg = cooperative_groups;
+
+// PackedType is the Lamport buffer access width (LDG.128 or LDG.64).
+template <typename PackedType = float4, bool UseCGA = false>
+__device__ struct __attribute__((aligned(32))) LamportFlags {
+ public:
+  __device__ explicit LamportFlags(uint32_t* bufferFlags,
+                                   uint32_t numStages = 1)
+      : mBufferFlagsPtr(bufferFlags), mFlagAccessPtr(&bufferFlags[8]) {
+    mCurBufferLayout.numStages = numStages;
+    uint4 flag = reinterpret_cast<uint4*>(bufferFlags)[0];
+    mCurrentIndex = flag.x;
+    mDirtyIndex = flag.y;
+    mCurBufferLayout.bytesPerBuffer = flag.z;
+    mDirtyBufferLayout.bytesPerBuffer = flag.z;
+    mDirtyBufferLayout.numStages = flag.w;
+    *reinterpret_cast<uint4*>(&mBytesToClear) =
+        reinterpret_cast<uint4*>(bufferFlags)[1];
+  }
+
+  [[nodiscard]] __device__ void* getCurLamportBuf(void* bufferBasePtr,
+                                                  int stageIdx = 0) const {
+    return mCurBufferLayout.getStagePtr(bufferBasePtr, mCurrentIndex, stageIdx);
+  }
+
+  // Refill a dirty stage (-1 for all) with the sentinel. Every thread of a
+  // 1D-block grid must participate.
+  __device__ void clearDirtyLamportBuf(void* bufferBasePtr, int stageIdx = -1) {
+    uint32_t globalCtaIdx = blockIdx.x * gridDim.y + blockIdx.y;
+    uint32_t globalTid = globalCtaIdx * blockDim.x + threadIdx.x;
+    uint32_t numThreads = gridDim.x * gridDim.y * blockDim.x;
+
+    if (stageIdx == -1) {
+      for (uint32_t i = 0; i < mDirtyBufferLayout.numStages; i++) {
+        clearPackedBuf(bufferBasePtr, globalTid, numThreads, mBytesToClear[i],
+                       mDirtyIndex, i);
+      }
+    } else if (stageIdx < mDirtyBufferLayout.numStages) {
+      clearPackedBuf(bufferBasePtr, globalTid, numThreads,
+                     mBytesToClear[stageIdx], mDirtyIndex, stageIdx);
+    }
+  }
+
+  __device__ void ctaArrive() {
+    if constexpr (UseCGA) {
+      cg::cluster_group cluster = cg::this_cluster();
+      __cluster_barrier_arrive();
+      if (cluster.block_rank() == 0 && threadIdx.x < 32) {
+        __cluster_barrier_wait();
+        arriveCounter(threadIdx.x);
+      }
+      return;
+    }
+    uint32_t const barrierThreads =
+        round_up(static_cast<uint32_t>(blockDim.x), 32u);
+    // A named barrier orders payload stores before the arrival counter
+    // without a full __syncthreads().
+    if (threadIdx.x < 32) {
+      asm volatile("barrier.cta.sync 1, %0;" ::"r"(barrierThreads) : "memory");
+      arriveCounter(threadIdx.x);
+    } else {
+      asm volatile("barrier.cta.arrive 1, %0;" ::"r"(barrierThreads)
+                   : "memory");
+    }
+  }
+
+  __device__ void waitAndUpdate(uint4 bytesToClearPerStage) {
+    bool isLastCtaT0{false};
+    int targetCount{0};
+    cg::grid_group grid = cg::this_grid();
+    // The last thread may exit early, so thread 0 updates the flags.
+    isLastCtaT0 = grid.thread_rank() == 0;
+    if constexpr (UseCGA) {
+      cg::cluster_group cluster = cg::this_cluster();
+      targetCount = gridDim.x * gridDim.y * gridDim.z / cluster.num_blocks();
+    } else {
+      targetCount = gridDim.x * gridDim.y * gridDim.z;
+    }
+    if (isLastCtaT0) {
+      uint4* flagPtr = reinterpret_cast<uint4*>(mBufferFlagsPtr);
+      while (*reinterpret_cast<uint32_t volatile*>(mFlagAccessPtr) <
+             targetCount) {
+      }
+      // 'Current' becomes 'Dirty'.
+      flagPtr[0] = {(mCurrentIndex + 1) % 3, mCurrentIndex,
+                    mCurBufferLayout.bytesPerBuffer,
+                    mCurBufferLayout.numStages};
+      flagPtr[1] = bytesToClearPerStage;
+      *mFlagAccessPtr = 0;
+    }
+  }
+
+ private:
+  uint32_t* mBufferFlagsPtr;
+  uint32_t* mFlagAccessPtr;
+
+  uint32_t mCurrentIndex, mDirtyIndex;
+  alignas(16) std::array<uint32_t, 4> mBytesToClear;
+  LamportBufferLayout mCurBufferLayout, mDirtyBufferLayout;
+
+  __device__ void arriveCounter(int tid) {
+    if (tid == 0) {
+#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000))
+      asm volatile("red.async.release.global.gpu.add.u32 [%0], %1;" ::"l"(
+                       mFlagAccessPtr),
+                   "r"(1)
+                   : "memory");
+#elif (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700))
+      asm volatile(
+          "red.release.global.gpu.add.u32 [%0], %1;" ::"l"(mFlagAccessPtr),
+          "r"(1)
+          : "memory");
+#else
+      atomicAdd(mFlagAccessPtr, 1);
+#endif
+    }
+  }
+
+  inline __device__ void clearPackedBuf(void* bufferBasePtr, uint32_t globalTid,
+                                        uint32_t numThreads,
+                                        uint32_t bytesToClear,
+                                        uint8_t dirtyIndex, uint8_t stageIdx) {
+    uint32_t clearBoundary =
+        ceil_div<uint32_t>(bytesToClear, sizeof(PackedType));
+    for (uint32_t packedIdx = globalTid; packedIdx < clearBoundary;
+         packedIdx += numThreads) {
+      reinterpret_cast<PackedType*>(mDirtyBufferLayout.getStagePtr(
+          bufferBasePtr, dirtyIndex, stageIdx))[packedIdx] =
+          getPackedLamportInit<PackedType, float>();
+    }
+  }
+};
+
+template <typename PackedType, typename T>
+union PackedVec {
+  PackedType packed;
+  T elements[sizeof(PackedType) / sizeof(T)];
+};
+
+template <typename PackedType, typename T>
+inline __device__ void sanitizeLamportPayload(PackedVec<PackedType, T>& value) {
+#pragma unroll
+  for (int i = 0; i < sizeof(PackedType) / sizeof(T); i++) {
+    if (isNegZero(value.elements[i])) {
+      value.elements[i] = fromFloat<T>(0.f);
+    }
+  }
+}
+
+template <typename PackedType, typename T>
+inline __device__ PackedType loadPacked(T* ptr) {
+  return *reinterpret_cast<PackedType*>(ptr);
+}
+
+template <typename PackedType, typename T>
+inline __device__ const PackedType loadPacked(T const* ptr) {
+  return *reinterpret_cast<PackedType const*>(ptr);
+}
+
+template <typename PackedType>
+union VolatilePackedLoad {
+  PackedType packed;
+  uint32_t words[sizeof(PackedType) / sizeof(uint32_t)];
+};
+
+template <typename PackedType>
+inline __device__ VolatilePackedLoad<PackedType> loadPackedVolatile(
+    void const* ptr) {
+  static_assert(sizeof(PackedType) == 0, "Not implemented");
+  return {};
+}
+
+template <>
+inline __device__ VolatilePackedLoad<float4> loadPackedVolatile<float4>(
+    void const* ptr) {
+  VolatilePackedLoad<float4> returnValue;
+  asm volatile("ld.volatile.global.v4.u32 {%0, %1, %2, %3}, [%4];\n"
+               : "=r"(returnValue.words[0]), "=r"(returnValue.words[1]),
+                 "=r"(returnValue.words[2]), "=r"(returnValue.words[3])
+               : "l"(ptr)
+               : "memory");
+  return returnValue;
+}
+
+template <typename PackedType>
+inline __device__ bool isLamportDirty(
+    VolatilePackedLoad<PackedType> const& value) {
+  // PTX only guarantees 64-bit atomicity, so every word must be checked.
+  bool dirty = false;
+#pragma unroll
+  for (int i = 0; i < sizeof(PackedType) / sizeof(uint32_t); i++) {
+    dirty |= value.words[i] == kNEGZERO_FP32;
+  }
+  return dirty;
+}
+
+template <int Rank, uint8_t WorldSize, uint8_t LocalRank, typename T,
+          typename PackedType, int kELTS_PER_THREAD>
+inline __device__ bool pollOneshotRemoteRank(
+    PackedVec<PackedType, T>* remoteValues, T* stagePtrLocal, int token,
+    int tokenDim, int packedIdx) {
+  if constexpr (Rank == LocalRank) {
+    return true;
+  } else {
+    auto loaded = loadPackedVolatile<PackedType>(
+        &stagePtrLocal[token * tokenDim * WorldSize + Rank * tokenDim +
+                       packedIdx * kELTS_PER_THREAD]);
+    remoteValues[Rank].packed = loaded.packed;
+    return !isLamportDirty(loaded);
+  }
+}
+
+template <uint8_t WorldSize, uint8_t LocalRank, typename T, typename PackedType,
+          int kELTS_PER_THREAD, int... Ranks>
+inline __device__ bool pollOneshotRemoteRanks(
+    PackedVec<PackedType, T>* remoteValues, T* stagePtrLocal, int token,
+    int tokenDim, int packedIdx, std::integer_sequence<int, Ranks...>) {
+  bool valid = true;
+  ((valid &= pollOneshotRemoteRank<Ranks, WorldSize, LocalRank, T, PackedType,
+                                   kELTS_PER_THREAD>(
+        remoteValues, stagePtrLocal, token, tokenDim, packedIdx)),
+   ...);
+  return valid;
+}
+
+template <typename T, typename PackedType, int kELTS_PER_THREAD>
+inline __device__ void accumulatePacked(float (&accum)[kELTS_PER_THREAD],
+                                        PackedVec<PackedType, T> const& value) {
+#pragma unroll
+  for (int i = 0; i < kELTS_PER_THREAD; i++) {
+    accum[i] += toFloat<T>(value.elements[i]);
+  }
+}
+
+template <int Rank, uint8_t LocalRank, typename T, typename PackedType,
+          int kELTS_PER_THREAD>
+inline __device__ void accumulateOneshotRank(
+    float (&accum)[kELTS_PER_THREAD],
+    PackedVec<PackedType, T> const* remoteValues,
+    PackedVec<PackedType, T> const& localValue) {
+  if constexpr (Rank == LocalRank) {
+    accumulatePacked<T, PackedType, kELTS_PER_THREAD>(accum, localValue);
+  } else {
+    accumulatePacked<T, PackedType, kELTS_PER_THREAD>(accum,
+                                                      remoteValues[Rank]);
+  }
+}
+
+template <uint8_t LocalRank, typename T, typename PackedType,
+          int kELTS_PER_THREAD, int... Ranks>
+inline __device__ void accumulateOneshotRanks(
+    float (&accum)[kELTS_PER_THREAD],
+    PackedVec<PackedType, T> const* remoteValues,
+    PackedVec<PackedType, T> const& localValue,
+    std::integer_sequence<int, Ranks...>) {
+  (accumulateOneshotRank<Ranks, LocalRank, T, PackedType, kELTS_PER_THREAD>(
+       accum, remoteValues, localValue),
+   ...);
+}
+
+// Polls every peer each iteration rather than waiting on them in turn.
+template <uint8_t WorldSize, uint8_t LocalRank, typename T, typename PackedType,
+          int kELTS_PER_THREAD>
+inline __device__ void waitOneshotRemoteRanks(
+    PackedVec<PackedType, T>* remoteValues, T* stagePtrLocal, int token,
+    int tokenDim, int packedIdx) {
+  static_assert(LocalRank < WorldSize);
+  while (1) {
+    bool const valid = pollOneshotRemoteRanks<WorldSize, LocalRank, T,
+                                              PackedType, kELTS_PER_THREAD>(
+        remoteValues, stagePtrLocal, token, tokenDim, packedIdx,
+        std::make_integer_sequence<int, WorldSize>{});
+    if (valid) {
+      break;
+    }
+  }
+}
+
+// Every rank sums in rank order, so all ranks produce identical results.
+template <uint8_t WorldSize, uint8_t LocalRank, typename T, typename PackedType,
+          int kELTS_PER_THREAD>
+inline __device__ PackedVec<PackedType, T> reduceOneshotDeterministic(
+    PackedVec<PackedType, T> const* remoteValues,
+    PackedVec<PackedType, T> const& localValue) {
+  static_assert(LocalRank < WorldSize);
+  float accum[kELTS_PER_THREAD];
+#pragma unroll
+  for (int i = 0; i < kELTS_PER_THREAD; i++) {
+    accum[i] = 0.f;
+  }
+  accumulateOneshotRanks<LocalRank, T, PackedType, kELTS_PER_THREAD>(
+      accum, remoteValues, localValue,
+      std::make_integer_sequence<int, WorldSize>{});
+
+  PackedVec<PackedType, T> packedAccum;
+#pragma unroll
+  for (int i = 0; i < kELTS_PER_THREAD; i++) {
+    packedAccum.elements[i] = fromFloat<T>(accum[i]);
+  }
+  return packedAccum;
+}
+
+template <uint8_t WorldSize, int kRankChunk, typename T, typename PackedType,
+          int kELTS_PER_THREAD>
+inline __device__ PackedVec<PackedType, T> reduceLamportRanksChunked(
+    T* buffer, int token, int tokenDim, int tokenOffset) {
+  static_assert(kRankChunk > 0, "kRankChunk must be positive");
+  static_assert(WorldSize % kRankChunk == 0,
+                "WorldSize must be divisible by kRankChunk");
+  // Chunking keeps large world sizes from holding every payload live.
+  float accum[kELTS_PER_THREAD];
+#pragma unroll
+  for (int i = 0; i < kELTS_PER_THREAD; i++) {
+    accum[i] = 0.f;
+  }
+
+#pragma unroll 1
+  for (int rankBase = 0; rankBase < WorldSize; rankBase += kRankChunk) {
+    float chunkAccum[kELTS_PER_THREAD];
+    while (1) {
+      bool valid = true;
+#pragma unroll
+      for (int i = 0; i < kELTS_PER_THREAD; i++) {
+        chunkAccum[i] = 0.f;
+      }
+#pragma unroll
+      for (int rr = 0; rr < kRankChunk; rr++) {
+        int const r = rankBase + rr;
+        auto loaded = loadPackedVolatile<PackedType>(
+            &buffer[token * tokenDim * WorldSize + r * tokenDim + tokenOffset]);
+        PackedVec<PackedType, T> value;
+        value.packed = loaded.packed;
+        valid &= !isLamportDirty(loaded);
+        accumulatePacked<T, PackedType, kELTS_PER_THREAD>(chunkAccum, value);
+      }
+      if (valid) {
+        break;
+      }
+    }
+#pragma unroll
+    for (int i = 0; i < kELTS_PER_THREAD; i++) {
+      accum[i] += chunkAccum[i];
+    }
+  }
+
+  PackedVec<PackedType, T> packedAccum;
+#pragma unroll
+  for (int i = 0; i < kELTS_PER_THREAD; i++) {
+    packedAccum.elements[i] = fromFloat<T>(accum[i]);
+  }
+  return packedAccum;
+}
+
+uint32_t constexpr kWARP_SIZE = 32U;
+uint32_t constexpr kLOG2_WARP_SIZE = 5U;
+uint32_t constexpr kLANE_ID_MASK = 0x1f;
+uint32_t constexpr kFINAL_MASK = 0xffffffff;
+
+template <typename T>
+inline __device__ T warpReduceSumFull(T val) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1) {
+    val += __shfl_xor_sync(kFINAL_MASK, val, mask, kWARP_SIZE);
+  }
+  return val;
+}
+
+template <typename T>
+inline __device__ T warpReduceSumPartial(T val) {
+  int laneId = threadIdx.x & kLANE_ID_MASK;
+  // Only the last warp calls this.
+  int warpSize = blockDim.x - (threadIdx.x & ~(kWARP_SIZE - 1));
+  unsigned int active_mask = (1U << warpSize) - 1;
+
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1) {
+    int targetLane = laneId ^ mask;
+    auto tmp = __shfl_xor_sync(active_mask, val, mask, kWARP_SIZE);
+    val += targetLane < warpSize ? tmp : 0;
+  }
+  return val;
+}
+
+// SYNC shares the sum with every thread; otherwise only thread 0 has it.
+template <typename T, bool SYNC = false>
+inline __device__ T blockReduceSumPartial(T val) {
+  __shared__ T smem[kWARP_SIZE];
+  int laneId = threadIdx.x & kLANE_ID_MASK;
+  int warpId = threadIdx.x >> kLOG2_WARP_SIZE;
+  int warpNum = (blockDim.x + kWARP_SIZE - 1) >> kLOG2_WARP_SIZE;
+
+  val = (warpId == warpNum - 1) ? warpReduceSumPartial(val)
+                                : warpReduceSumFull(val);
+  if (laneId == 0) {
+    smem[warpId] = val;
+  }
+  __syncthreads();
+
+  if (warpId == 0) {
+    val = (laneId < warpNum) ? smem[laneId] : (T)0.f;
+    val = (warpNum == 1) ? warpReduceSumPartial(val) : warpReduceSumFull(val);
+
+    if constexpr (SYNC) {
+      if (laneId == 0) {
+        smem[warpId] = val;
+      }
+    }
+  }
+  if constexpr (SYNC) {
+    __syncthreads();
+    val = smem[0];
+  }
+  return val;
+}
+
+template <typename T>
+inline __device__ T blockReduceSumFull(T val) {
+  __shared__ T smem[kWARP_SIZE];
+  int lane_id = threadIdx.x & kLANE_ID_MASK;
+  int warp_id = threadIdx.x >> kLOG2_WARP_SIZE;
+  int warp_num = blockDim.x >> kLOG2_WARP_SIZE;
+
+  val = warpReduceSumFull(val);
+  if (lane_id == 0) {
+    smem[warp_id] = val;
+  }
+  __syncthreads();
+
+  val = (lane_id < warp_num) ? smem[lane_id] : (T)0.f;
+  val = warpReduceSumFull(val);
+
+  return val;
+}
+
+template <typename T, bool SYNC = false>
+inline __device__ T blockReduceSum(T val) {
+  bool hasPartialWarp = (blockDim.x & kLANE_ID_MASK) != 0;
+  if (hasPartialWarp) {
+    return blockReduceSumPartial<T, SYNC>(val);
+  } else {
+    return blockReduceSumFull<T>(val);
+  }
+}
+
+// Grid configuration for the fused one-shot and RMSNorm kernels.
+// Returns (block_size, cluster_size, loads_per_thread).
+inline std::tuple<int, int, int> adjustGridConfig(int numTokens, int dim,
+                                                  int eltsPerThread,
+                                                  bool useCluster,
+                                                  int smCount) {
+  int clusterSize = useCluster ? kMaxClusterSize : 1;
+  int blockSize = 128;
+  int threadsNeeded = ceil_div(dim, eltsPerThread);
+  int loadsPerThread = 1;
+
+  blockSize = ceil_div(threadsNeeded, clusterSize);
+  if (useCluster) {
+    // Shrink the cluster until the row partitions cleanly across CTAs.
+    while (threadsNeeded % clusterSize != 0 && clusterSize > 1) {
+      clusterSize /= 2;
+    }
+    int const maxDivisibleClusterSize = clusterSize;
+    blockSize = ceil_div(threadsNeeded, clusterSize);
+    // Trade cluster width for CTAs of at least 128 threads.
+    while (blockSize < 128 && clusterSize >= 2) {
+      blockSize *= 2;
+      clusterSize /= 2;
+    }
+    // Enough tokens already fill the GPU: prefer wider CTAs.
+    while (numTokens * clusterSize > smCount && clusterSize > 1 &&
+           blockSize <= 512) {
+      blockSize *= 2;
+      clusterSize /= 2;
+    }
+    // The grid still underfills the GPU: restore cluster width.
+    while (clusterSize < maxDivisibleClusterSize) {
+      int const candidateClusterSize = clusterSize * 2;
+      int const candidateBlockSize =
+          ceil_div(threadsNeeded, candidateClusterSize);
+      if (candidateBlockSize < 64 ||
+          numTokens * candidateClusterSize > smCount) {
+        break;
+      }
+      clusterSize = candidateClusterSize;
+      blockSize = candidateBlockSize;
+    }
+  }
+  // Keep very wide rows within 1024 threads per CTA.
+  while (blockSize > 1024) {
+    if (useCluster && clusterSize < kMaxClusterSize) {
+      clusterSize = clusterSize << 1;
+    } else if (loadsPerThread < 8) {
+      loadsPerThread += 1;
+    } else {
+      break;
+    }
+    blockSize = ceil_div(threadsNeeded, clusterSize * loadsPerThread);
+  }
+  return {blockSize, clusterSize, loadsPerThread};
+}
+
+}  // namespace utils
+
+enum MNNVLTwoShotStage : uint8_t {
+  SCATTER = 0,
+  BROADCAST = 1,
+  NUM_STAGES = 2,
+};
+
+// Scatters token rows to their owner rank, reduces owned rows in rank order,
+// and broadcasts them through multicast. WaitForResults also gathers the
+// reduced rows into outputPtr; otherwise a follow-up kernel consumes them.
+template <uint8_t WorldSize, typename T, bool WaitForResults = true,
+          typename PackedType = float4>
+__global__ __launch_bounds__(128) void twoshotAllreduceKernel(
+    AllReduceKernelParams<T> params) {
+  using namespace utils;
+  constexpr int kELTS_PER_THREAD = sizeof(PackedType) / sizeof(T);
+  constexpr uint32_t kELT_SIZE = sizeof(T);
+
+  int const packedIdx = blockIdx.y * blockDim.x + threadIdx.x;
+  int const token = blockIdx.x;
+  int const tokenOffset = packedIdx * kELTS_PER_THREAD;
+  int const threadOffset = token * params.tokenDim + tokenOffset;
+  bool const inBounds = tokenOffset < params.tokenDim;
+  int const destRank = token % WorldSize;
+  int const destTokenOffset = token / WorldSize;
+
+  cudaGridDependencySynchronize();
+
+  LamportFlags<PackedType> flag(params.bufferFlags,
+                                MNNVLTwoShotStage::NUM_STAGES);
+  T* scatterBufLocal = reinterpret_cast<T*>(flag.getCurLamportBuf(
+      params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER));
+  T* scatterBufDest = reinterpret_cast<T*>(flag.getCurLamportBuf(
+      params.inputPtrs[destRank], MNNVLTwoShotStage::SCATTER));
+  T* broadcastBufW = reinterpret_cast<T*>(
+      flag.getCurLamportBuf(params.mcastPtr, MNNVLTwoShotStage::BROADCAST));
+  T* broadcastBufR = reinterpret_cast<T*>(flag.getCurLamportBuf(
+      params.inputPtrs[params.rank], MNNVLTwoShotStage::BROADCAST));
+
+  PackedVec<PackedType, T> val;
+  if (inBounds) {
+    val.packed = loadPacked<PackedType>(&params.shardPtr[threadOffset]);
+    sanitizeLamportPayload<PackedType, T>(val);
+
+    reinterpret_cast<PackedType*>(
+        &scatterBufDest[destTokenOffset * params.tokenDim * WorldSize +
+                        params.rank * params.tokenDim])[packedIdx] = val.packed;
+  }
+
+  cudaTriggerProgrammaticLaunchCompletion();
+  flag.clearDirtyLamportBuf(params.inputPtrs[params.rank],
+                            MNNVLTwoShotStage::SCATTER);
+
+  if (inBounds && destRank == params.rank) {
+    int const localToken = token / WorldSize;
+    constexpr int kRankChunk = WorldSize < 16 ? WorldSize : 16;
+    PackedVec<PackedType, T> packedAccum =
+        reduceLamportRanksChunked<WorldSize, kRankChunk, T, PackedType,
+                                  kELTS_PER_THREAD>(
+            scatterBufLocal, localToken, params.tokenDim, tokenOffset);
+    // A reduced value can round to the sentinel.
+    sanitizeLamportPayload<PackedType, T>(packedAccum);
+    reinterpret_cast<PackedType*>(
+        &broadcastBufW[token * params.tokenDim])[packedIdx] =
+        packedAccum.packed;
+  }
+
+  flag.clearDirtyLamportBuf(params.inputPtrs[params.rank],
+                            MNNVLTwoShotStage::BROADCAST);
+
+  if constexpr (WaitForResults) {
+    // Out-of-bounds threads still arrive so waitAndUpdate sees every CTA.
+    flag.ctaArrive();
+
+    if (inBounds) {
+      auto loaded =
+          loadPackedVolatile<PackedType>(&broadcastBufR[threadOffset]);
+      while (isLamportDirty(loaded)) {
+        loaded = loadPackedVolatile<PackedType>(&broadcastBufR[threadOffset]);
+      }
+      reinterpret_cast<PackedType*>(&params.outputPtr[threadOffset])[0] =
+          loaded.packed;
+    }
+
+    flag.waitAndUpdate(
+        {static_cast<uint32_t>(round_up(params.numTokens, WorldSize) *
+                               params.tokenDim * kELT_SIZE),
+         static_cast<uint32_t>(params.numTokens * params.tokenDim * kELT_SIZE),
+         0, 0});
+  }
+}
+
+}  // namespace mnnvl
+}  // namespace vllm

csrc/ops.h

@@ -475,6 +475,26 @@ void fused_kda_chunk(
     std::optional<torch::stable::Tensor> has_initial_state);
 #endif
 
+#ifdef VLLM_ENABLE_BEAM_ALLREDUCE
+void beam_allreduce_norm(
+    torch::stable::Tensor const& input, torch::stable::Tensor const& residual,
+    std::optional<torch::stable::Tensor> post_weight,
+    std::optional<torch::stable::Tensor> next_weight,
+    torch::stable::Tensor& hidden,
+    std::optional<torch::stable::Tensor> normalized, int64_t buffer_ptrs,
+    torch::stable::Tensor& flags, int64_t multicast_ptr, int64_t local_ptr,
+    int64_t world, int64_t rank, double post_eps, double next_eps,
+    double x_scale, double r_scale, bool post_norm, bool next_norm, bool scaled,
+    bool one_shot, bool small_batch,
+    std::optional<torch::stable::Tensor> experts,
+    std::optional<torch::stable::Tensor> expert_weights,
+    std::optional<torch::stable::Tensor> permutation, double routed_scale);
+void beam_allreduce(torch::stable::Tensor const& input,
+                    torch::stable::Tensor& output, int64_t buffer_ptrs,
+                    torch::stable::Tensor& flags, int64_t multicast_ptr,
+                    int64_t world, int64_t rank);
+#endif
+
 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
 void kimi_k3_attn_res(torch::stable::Tensor& prefix,
                       std::optional<torch::stable::Tensor> delta,

csrc/torch_bindings.cpp

@@ -577,6 +577,22 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
       "Tensor? state_indices=None, Tensor? has_initial_state=None) -> ()");
 #endif
 
+#ifdef VLLM_ENABLE_BEAM_ALLREDUCE
+  // Beam's MNNVL all-reduce fused with its norm-residual-norm boundary.
+  ops.def(
+      "beam_allreduce_norm(Tensor input, Tensor residual, "
+      "Tensor? post_weight, Tensor? next_weight, Tensor! hidden, "
+      "Tensor!? normalized, int buffer_ptrs, Tensor! flags, "
+      "int multicast_ptr, int local_ptr, int world, int rank, "
+      "float post_eps, float next_eps, float x_scale, float r_scale, "
+      "bool post_norm, bool next_norm, bool scaled, bool one_shot, "
+      "bool small_batch, Tensor? experts, Tensor? expert_weights, "
+      "Tensor? permutation, float routed_scale) -> ()");
+  ops.def(
+      "beam_allreduce(Tensor input, Tensor! output, int buffer_ptrs, "
+      "Tensor! flags, int multicast_ptr, int world, int rank) -> ()");
+#endif
+
 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
   ops.def(
       "kimi_k3_attn_res("
@@ -874,6 +890,11 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
   ops.impl("fused_kda_chunk", TORCH_BOX(&fused_kda_chunk));
 #endif
 
+#ifdef VLLM_ENABLE_BEAM_ALLREDUCE
+  ops.impl("beam_allreduce_norm", TORCH_BOX(&beam_allreduce_norm));
+  ops.impl("beam_allreduce", TORCH_BOX(&beam_allreduce));
+#endif
+
 #ifdef VLLM_ENABLE_KIMI_K3_ATTN_RES
   ops.impl("kimi_k3_attn_res", TORCH_BOX(&kimi_k3_attn_res));
 #endif

setup.py

@@ -1623,7 +1623,6 @@ package_data = {
     "vllm": [
         "py.typed",
         "libs/*.so*",
-        "models/beam/nvidia/*.cu",
         "model_executor/layers/fused_moe/configs/*.json",
         "model_executor/layers/quantization/utils/configs/*.json",
         "entrypoints/serve/instrumentator/static/*.js",

tests/models/test_beam_construction.py

@@ -277,7 +277,7 @@ def test_moe_finalize_deferral_requires_replicated_tokens(
         (4, 8, 8192, 65536, -1, 8192, False),
     ],
 )
-def test_decoder_finalizes_large_ep_batches_before_flashinfer_allreduce(
+def test_decoder_finalizes_large_ep_batches_before_fused_allreduce(
     monkeypatch,
     world_size,
     ep_size,

tests/models/test_beam_numerics.py

@@ -458,6 +458,14 @@ def test_fused_kernels_reuse_specializations_across_token_counts(
     torch.accelerator.synchronize()
 
 
+def _skip_without_beam_allreduce():
+    if not hasattr(torch.ops._C, "beam_allreduce_norm"):
+        pytest.skip("Beam all-reduce kernel not built")
+    mnnvl = pytest.importorskip("flashinfer.comm.mnnvl")
+    if not mnnvl.is_multicast_supported(0):
+        pytest.skip("NVLink multicast required")
+
+
 def _beam_allreduce_worker(rank, world_size, port, hidden_size, host="127.0.0.1"):
     from datetime import timedelta
 
@@ -554,9 +562,7 @@ def test_allreduce_boundary_handles_batch_changes_and_graph_replay(
 
     if torch.accelerator.device_count() < world_size:
         pytest.skip(f"Requires {world_size} GPUs")
-    mnnvl = pytest.importorskip("flashinfer.comm.mnnvl")
-    if not mnnvl.is_multicast_supported(0):
-        pytest.skip("NVLink multicast required")
+    _skip_without_beam_allreduce()
     spawn(
         _beam_allreduce_worker,
         args=(world_size, get_open_port(), hidden_size),
@@ -713,9 +719,7 @@ def test_moe_finalize_allreduce_preserves_local_rounding_and_replay(
 
     if torch.accelerator.device_count() < world_size:
         pytest.skip(f"Requires {world_size} GPUs")
-    mnnvl = pytest.importorskip("flashinfer.comm.mnnvl")
-    if not mnnvl.is_multicast_supported(0):
-        pytest.skip("NVLink multicast required")
+    _skip_without_beam_allreduce()
     spawn(
         _beam_finalize_worker,
         args=(world_size, get_open_port(), expert_parallel, hidden_size),

vllm/models/beam/nvidia/allreduce.py

@@ -1,50 +1,26 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""CUDA JIT MNNVL collective with Beam's FP32 sandwich-norm epilogue."""
-
-import functools
-import hashlib
-from pathlib import Path
+"""MNNVL all-reduce with Beam's FP32 sandwich-norm epilogue."""
 
 import torch
+import torch.distributed as dist
 
 from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
 from vllm.utils.flashinfer import has_flashinfer
-from vllm.utils.math_utils import cdiv
+from vllm.utils.math_utils import cdiv, round_up
 
-
-@functools.cache
-def _module():
-    from flashinfer.jit import env as jit_env
-    from flashinfer.jit.core import current_compilation_context, gen_jit_spec
-
-    source = Path(__file__).with_suffix(".cu")
-    communication_header = (
-        jit_env.FLASHINFER_INCLUDE_DIR / "flashinfer/comm/trtllm_mnnvl_allreduce.cuh"
-    )
-    digest = hashlib.sha256(
-        source.read_bytes() + communication_header.read_bytes()
-    ).hexdigest()[:16]
-    flags = current_compilation_context.get_nvcc_flags_list(
-        supported_major_versions=[9, 10]
-    )
-    spec = gen_jit_spec(
-        f"beam_allreduce_{digest}",
-        [source],
-        extra_cuda_cflags=[*flags, "--fmad=false"],
-        extra_include_paths=[jit_env.FLASHINFER_CSRC_DIR],
-    )
-    spec.extra_cuda_cflags.remove("-use_fast_math")
-    return spec.build_and_load()
+_NUM_LAMPORT_BUFFERS = 3
+# FlashInfer's MNNVL one-shot cutoff on the BF16 payload staged across ranks.
+_ONESHOT_MAX_BYTES = 1024 * 1024
 
 
 class BeamAllReduceWorkspace:
-    """Reuse FlashInfer's multicast allocation and Lamport generation protocol."""
+    """Multicast Lamport buffers and generation flags for the Beam kernels."""
 
     @classmethod
     def create(
         cls,
-        group: torch.distributed.ProcessGroup,
+        group: dist.ProcessGroup,
         hidden_size: int,
         max_tokens: int,
         dtype: torch.dtype,
@@ -55,6 +31,7 @@ class BeamAllReduceWorkspace:
             or dtype != torch.bfloat16
             or hidden_size % 1024
             or not 1024 <= hidden_size <= 8192
+            or not hasattr(torch.ops._C, "beam_allreduce_norm")
             or not has_flashinfer()
         ):
             return None
@@ -72,44 +49,65 @@ class BeamAllReduceWorkspace:
         return cls(group, hidden_size, max_tokens)
 
     def __init__(
-        self, group: torch.distributed.ProcessGroup, hidden_size: int, max_tokens: int
+        self, group: dist.ProcessGroup, hidden_size: int, max_tokens: int
     ) -> None:
-        from flashinfer.comm import create_allreduce_fusion_workspace
-        from flashinfer.comm.mnnvl import TorchDistBackend
-
         self.world_size = group.size()
-        self.rank = torch.distributed.get_rank(group)
+        self.rank = dist.get_rank(group)
         self.hidden_size = hidden_size
         self.max_tokens = max_tokens
         # Scale one-shot cutoffs by the BF16 payload staged across all ranks.
         token_bytes = hidden_size * 2
         rank_token_bytes = self.world_size * token_bytes
+        self.oneshot_max_tokens = _ONESHOT_MAX_BYTES // rank_token_bytes
         self.small_batch_max_tokens = 3200 * 1024 // rank_token_bytes
         self.ep_finalize_oneshot_max_tokens = 2560 * 1024 // rank_token_bytes
-        # Two-shot reserves two rows per token; one-shot needs one per rank.
-        workspace_tokens = max(
-            max_tokens,
-            min(max_tokens, self.small_batch_max_tokens) * self.world_size // 2,
-        )
-        self.workspace = create_allreduce_fusion_workspace(
-            backend="mnnvl",
-            world_size=self.world_size,
-            rank=self.rank,
-            max_token_num=workspace_tokens,
-            hidden_dim=hidden_size,
-            dtype=torch.bfloat16,
-            comm_backend=TorchDistBackend(group=group),
-            group=group,
-        )
-        self.module = _module()
-        self.allocation_bytes = (
-            self.workspace.handle.mcast_device_memory.get_allocation_size()
-        )
         self.separate_norm_min_tokens = cdiv(
             10 * 1024**2 * self.world_size, token_bytes
         )
         self.ep_finalize_max_tokens = 640 * 1024**2 // token_bytes
 
+        # One-shot stages one row per rank; two-shot two rows per owned token.
+        oneshot_tokens = min(max_tokens, self.small_batch_max_tokens)
+        buffer_bytes = max(
+            oneshot_tokens * rank_token_bytes,
+            2 * cdiv(max_tokens, self.world_size) * rank_token_bytes,
+        )
+        # Both two-shot stages stay 16-byte aligned.
+        buffer_bytes = round_up(buffer_bytes, 32)
+        if buffer_bytes >= 2**32:
+            raise ValueError(
+                f"Beam all-reduce needs {buffer_bytes} bytes per Lamport "
+                "buffer, which exceeds the 32-bit flag field"
+            )
+        # FlashInfer only maps the multicast memory; the Lamport protocol and
+        # its flags are ours. Small batches reduce ~0.3us faster on its mapping
+        # than on torch symmetric memory.
+        from flashinfer.comm.mnnvl import McastGPUBuffer, TorchDistBackend
+
+        device = torch.device("cuda", torch.accelerator.current_device_index())
+        self.buffer = McastGPUBuffer(
+            _NUM_LAMPORT_BUFFERS * buffer_bytes,
+            self.world_size,
+            self.rank,
+            device,
+            TorchDistBackend(group=group),
+        )
+        self.multicast_ptr = self.buffer.get_multicast_ptr()
+        if not self.multicast_ptr:
+            raise RuntimeError("Beam all-reduce requires NVLink multicast")
+        self.local_ptr = self.buffer.get_unicast_ptr(self.rank)
+        self.buffer_ptrs = self.buffer.get_buffer_ptrs_dev()
+        self.buffer.lamport_initialize(self.rank, torch.float32)
+        # [current, dirty, bytes per buffer, dirty stages, bytes to clear x4,
+        # arrival count]
+        self.flags = torch.tensor(
+            [0, 2, buffer_bytes, 0, 0, 0, 0, 0, 0],
+            dtype=torch.uint32,
+            device=device,
+        )
+        torch.accelerator.synchronize()
+        dist.barrier(group=group)
+
 
 def beam_allreduce_norm_residual_norm(
     x: torch.Tensor,
@@ -131,15 +129,17 @@ def beam_allreduce_norm_residual_norm(
     assert x.shape[1] == workspace.hidden_size
     assert x.shape[0] <= workspace.max_tokens
     if x.shape[0] >= workspace.separate_norm_min_tokens:
-        from flashinfer.comm import AllReduceFusionPattern, allreduce_fusion
-
         from .layernorm import beam_norm_residual_norm
 
-        reduced = allreduce_fusion(
-            input=x,
-            workspace=workspace.workspace,
-            pattern=AllReduceFusionPattern.kAllReduce,
-            launch_with_pdl=True,
+        reduced = torch.empty_like(x)
+        torch.ops._C.beam_allreduce(
+            x,
+            reduced,
+            workspace.buffer_ptrs,
+            workspace.flags,
+            workspace.multicast_ptr,
+            workspace.world_size,
+            workspace.rank,
         )
         return beam_norm_residual_norm(
             reduced,
@@ -169,8 +169,6 @@ def _cuda_allreduce_norm_residual_norm(
     routed_scale: float = 1.0,
     oneshot_max_tokens: int = 0,
 ):
-    from flashinfer.comm.trtllm_mnnvl_ar import MNNVLAllreduceFusionStrategy
-
     assert x.dtype == residual.dtype == torch.bfloat16
     assert x.is_contiguous() and residual.is_contiguous()
     m, n = x.shape
@@ -178,22 +176,18 @@ def _cuda_allreduce_norm_residual_norm(
     hidden = torch.empty_like(x)
     normalized = torch.empty_like(x) if next_eps is not None else None
     if m:
-        strategy = MNNVLAllreduceFusionStrategy.select_strategy(
-            workspace.world_size, m, n, x.dtype
-        )
         small_batch = unfinalized is None and m <= workspace.small_batch_max_tokens
-        w = workspace.workspace
-        workspace.module.beam_allreduce(
+        torch.ops._C.beam_allreduce_norm(
             x,
             residual,
             post_weight,
             next_weight,
             hidden,
             normalized,
-            w.uc_ptrs_dev,
-            w.buffer_flags,
-            w.mc_ptr,
-            w.uc_ptr_local,
+            workspace.buffer_ptrs,
+            workspace.flags,
+            workspace.multicast_ptr,
+            workspace.local_ptr,
             workspace.world_size,
             workspace.rank,
             post_eps or 0.0,
@@ -202,16 +196,14 @@ def _cuda_allreduce_norm_residual_norm(
             post_eps is not None,
             next_eps is not None,
             scales is not None,
-            small_batch
-            or m <= oneshot_max_tokens
-            or strategy == MNNVLAllreduceFusionStrategy.ONESHOT,
+            small_batch or m <= max(oneshot_max_tokens, workspace.oneshot_max_tokens),
+            small_batch,
             unfinalized.gemm2_permuted if unfinalized is not None else None,
             unfinalized.expert_weights if unfinalized is not None else None,
             unfinalized.expanded_idx_to_permuted_idx
             if unfinalized is not None
             else None,
             routed_scale,
-            small_batch,
         )
     return hidden, normalized