[Model][Perf] Fuse Beam's TP all-reduce with its sandwich norms and MoE finalize (#62366)

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

CMakeLists.txt

@@ -1150,6 +1150,26 @@ 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 kernels use NVLink multicast and thread-block clusters. Their 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_KERNEL_ARCHS
+      "9.0a;10.0f;10.3f" "${CUDA_ARCHS}")
+  else()
+    cuda_archs_loose_intersection(BEAM_KERNEL_ARCHS
+      "9.0a;10.0a;10.3a" "${CUDA_ARCHS}")
+  endif()
+  if(BEAM_KERNEL_ARCHS)
+    set(BEAM_KERNEL_SRCS "csrc/beam/allreduce_norm.cu")
+    set_gencode_flags_for_srcs(
+      SRCS "${BEAM_KERNEL_SRCS}"
+      CUDA_ARCHS "${BEAM_KERNEL_ARCHS}")
+    set_property(SOURCE ${BEAM_KERNEL_SRCS} APPEND PROPERTY
+      COMPILE_OPTIONS "$<$<COMPILE_LANGUAGE:CUDA>:--fmad=false>")
+    list(APPEND VLLM_STABLE_EXT_SRC "${BEAM_KERNEL_SRCS}")
+    message(STATUS "Building Beam kernels for archs: ${BEAM_KERNEL_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 +1271,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_KERNEL_ARCHS)
+      target_compile_definitions(_C_stable_libtorch PRIVATE
+        VLLM_ENABLE_BEAM_KERNELS=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

@@ -0,0 +1,472 @@
+// 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 "../torch_utils.h"
+#include "../mnnvl_lamport.cuh"
+
+#include <torch/headeronly/core/ScalarType.h>
+
+namespace vllm {
+namespace beam {
+
+namespace ar = vllm::mnnvl;
+namespace cg = cooperative_groups;
+using T = __nv_bfloat16;
+using Vec = ar::utils::PackedVec<float4, T>;
+
+struct Params {
+  ar::AllReduceKernelParams<T> ar;
+  T const* nextWeight;
+  float nextEps;
+  float xScale;
+  float rScale;
+  bool postNorm;
+  bool nextNorm;
+  bool scaled;
+  T const* experts;
+  void const* expertWeights;
+  int const* permutation;
+  int topK;
+  bool floatWeights;
+  float routedScale;
+};
+
+template <bool Finalize>
+__device__ Vec loadInput(Params const& p, int token, int packedIdx) {
+  Vec value;
+  if constexpr (!Finalize) {
+    value.packed = ar::utils::loadPacked<float4>(
+        p.ar.shardPtr + int64_t(token) * p.ar.tokenDim + packedIdx * 8);
+    return value;
+  }
+  float sum[8]{};
+  for (int base = 0; base < p.topK; base += 8) {
+    int rows[8];
+    float weights[8];
+#pragma unroll
+    for (int k = 0; k < 8; ++k) {
+      int index = token * p.topK + base + k;
+      rows[k] = base + k < p.topK ? p.permutation[index] : -1;
+      weights[k] = rows[k] < 0 ? 0.f
+                   : p.floatWeights
+                       ? static_cast<float const*>(p.expertWeights)[index]
+                       : ar::utils::toFloat<T>(
+                             static_cast<T const*>(p.expertWeights)[index]);
+    }
+#pragma unroll
+    for (int k = 0; k < 8; ++k) {
+      Vec expert{};
+      if (rows[k] >= 0)
+        expert.packed = ar::utils::loadPacked<float4>(
+            p.experts + int64_t(rows[k]) * p.ar.tokenDim + packedIdx * 8);
+#pragma unroll
+      for (int i = 0; i < 8; ++i)
+        sum[i] = __fadd_rn(
+            sum[i],
+            __fmul_rn(ar::utils::toFloat<T>(expert.elements[i]), weights[k]));
+    }
+  }
+#pragma unroll
+  for (int i = 0; i < 8; ++i) {
+    T rounded = ar::utils::fromFloat<T>(sum[i]);
+    value.elements[i] = ar::utils::fromFloat<T>(
+        __fmul_rn(ar::utils::toFloat<T>(rounded), p.routedScale));
+  }
+  return value;
+}
+
+template <bool UseCGA, int Stage>
+__device__ float invRms(Vec const& value, int dim, float eps) {
+  float sum = 0.f;
+#pragma unroll
+  for (int i = 0; i < 8; ++i) {
+    float v = ar::utils::toFloat<T>(value.elements[i]);
+    sum = __fadd_rn(sum, __fmul_rn(v, v));
+  }
+  sum = ar::utils::blockReduceSum<float, true>(sum);
+  if constexpr (UseCGA) {
+    auto cluster = cg::this_cluster();
+    __shared__ float partial[8];
+    if (threadIdx.x < cluster.num_blocks()) {
+      cluster.map_shared_rank(partial, threadIdx.x)[cluster.block_rank()] = sum;
+    }
+    cluster.sync();
+    sum = 0.f;
+    for (int i = 0; i < cluster.num_blocks(); ++i)
+      sum = __fadd_rn(sum, partial[i]);
+  }
+  __syncthreads();
+  return rsqrtf(sum / dim + eps);
+}
+
+template <bool UseCGA, bool Preloaded = false>
+__device__ void epilogue(Params const& p, Vec value, int packedIdx, int offset,
+                         Vec postWeight = {}, Vec residual = {},
+                         Vec nextWeight = {}) {
+  if (p.postNorm) {
+    float rstd = invRms<UseCGA, 0>(value, p.ar.tokenDim, p.ar.epsilon);
+    Vec weight = postWeight;
+    if constexpr (!Preloaded) {
+      if (p.ar.gammaPtr)
+        weight.packed =
+            ar::utils::loadPacked<float4>(p.ar.gammaPtr + packedIdx * 8);
+    }
+#pragma unroll
+    for (int i = 0; i < 8; ++i) {
+      float v = __fmul_rn(ar::utils::toFloat<T>(value.elements[i]), rstd);
+      if (p.ar.gammaPtr)
+        v = __fmul_rn(v, ar::utils::toFloat<T>(weight.elements[i]));
+      value.elements[i] = ar::utils::fromFloat<T>(v);
+    }
+  }
+  if constexpr (!Preloaded)
+    residual.packed =
+        ar::utils::loadPacked<float4>(p.ar.residualInPtr + offset);
+#pragma unroll
+  for (int i = 0; i < 8; ++i) {
+    T x = value.elements[i], r = residual.elements[i];
+    if (p.scaled) {
+      x = ar::utils::fromFloat<T>(
+          __fmul_rn(ar::utils::toFloat<T>(x), p.xScale));
+      r = ar::utils::fromFloat<T>(
+          __fmul_rn(ar::utils::toFloat<T>(r), p.rScale));
+    }
+    value.elements[i] = ar::utils::fromFloat<T>(
+        __fadd_rn(ar::utils::toFloat<T>(x), ar::utils::toFloat<T>(r)));
+  }
+  *reinterpret_cast<float4*>(p.ar.prenormedPtr + offset) = value.packed;
+  if (p.nextNorm) {
+    float rstd = invRms<UseCGA, 1>(value, p.ar.tokenDim, p.nextEps);
+    Vec weight = nextWeight;
+    if constexpr (!Preloaded) {
+      if (p.nextWeight)
+        weight.packed =
+            ar::utils::loadPacked<float4>(p.nextWeight + packedIdx * 8);
+    }
+#pragma unroll
+    for (int i = 0; i < 8; ++i) {
+      float v = __fmul_rn(ar::utils::toFloat<T>(value.elements[i]), rstd);
+      if (p.nextWeight)
+        v = __fmul_rn(v, ar::utils::toFloat<T>(weight.elements[i]));
+      value.elements[i] = ar::utils::fromFloat<T>(v);
+    }
+    *reinterpret_cast<float4*>(p.ar.outputPtr + offset) = value.packed;
+  }
+}
+
+template <int World, int Rank>
+__device__ Vec reduceLocal(Vec value, T* buffer, int token, int dim,
+                           int packedIdx) {
+  Vec remote[World];
+  ar::utils::waitOneshotRemoteRanks<World, Rank, T, float4, 8>(
+      remote, buffer, token, dim, packedIdx);
+  return ar::utils::reduceOneshotDeterministic<World, Rank, T, float4, 8>(
+      remote, value);
+}
+
+template <int World, bool Finalize = false, bool SmallBatch = false>
+__global__ void __launch_bounds__(1024) oneshot(Params p) {
+  int packedIdx = threadIdx.x;
+  if constexpr (!SmallBatch) packedIdx = cg::this_cluster().thread_rank();
+  int token = blockIdx.x;
+  int offset = token * p.ar.tokenDim + packedIdx * 8;
+  cudaGridDependencySynchronize();
+  ar::utils::LamportFlags<float4, !SmallBatch> flag(p.ar.bufferFlags, 1);
+  auto multicast = static_cast<T*>(flag.getCurLamportBuf(p.ar.mcastPtr, 0));
+  auto local =
+      static_cast<T*>(flag.getCurLamportBuf(p.ar.inputPtrs[p.ar.rank], 0));
+  Vec value = loadInput<Finalize>(p, token, packedIdx);
+  ar::utils::sanitizeLamportPayload<float4, T>(value);
+  reinterpret_cast<float4*>(multicast + token * p.ar.tokenDim * World +
+                            p.ar.rank * p.ar.tokenDim)[packedIdx] =
+      value.packed;
+  Vec postWeight{}, residual{}, nextWeight{};
+  if constexpr (SmallBatch) {
+    // Overlap epilogue inputs with the peer payloads in flight.
+    if (p.postNorm && p.ar.gammaPtr)
+      postWeight.packed =
+          ar::utils::loadPackedVolatile<float4>(p.ar.gammaPtr + packedIdx * 8)
+              .packed;
+    residual.packed =
+        ar::utils::loadPackedVolatile<float4>(p.ar.residualInPtr + offset)
+            .packed;
+    if (p.nextNorm && p.nextWeight)
+      nextWeight.packed =
+          ar::utils::loadPackedVolatile<float4>(p.nextWeight + packedIdx * 8)
+              .packed;
+  }
+  flag.ctaArrive();
+  if constexpr (SmallBatch) {
+    // This workspace uses at most two stages. Constant indices avoid a
+    // dynamically indexed array spilling the Lamport state to local memory.
+    flag.clearDirtyLamportBuf(p.ar.inputPtrs[p.ar.rank], 0);
+    flag.clearDirtyLamportBuf(p.ar.inputPtrs[p.ar.rank], 1);
+  } else {
+    flag.clearDirtyLamportBuf(p.ar.inputPtrs[p.ar.rank], -1);
+  }
+#define REDUCE(R)                                                       \
+  case R:                                                               \
+    if constexpr (World > R)                                            \
+      value = reduceLocal<World, R>(value, local, token, p.ar.tokenDim, \
+                                    packedIdx);                         \
+    break
+  switch (p.ar.rank) {
+    REDUCE(0);
+    REDUCE(1);
+    REDUCE(2);
+    REDUCE(3);
+    REDUCE(4);
+    REDUCE(5);
+    REDUCE(6);
+    REDUCE(7);
+  }
+#undef REDUCE
+  cudaTriggerProgrammaticLaunchCompletion();
+  epilogue<!SmallBatch, SmallBatch>(p, value, packedIdx, offset, postWeight,
+                                    residual, nextWeight);
+  flag.waitAndUpdate({static_cast<uint32_t>(p.ar.numTokens * p.ar.tokenDim *
+                                            World * sizeof(T)),
+                      0, 0, 0});
+}
+
+template <int World>
+__global__ void __launch_bounds__(128) finalizeScatter(Params p) {
+  int packedIdx = blockIdx.y * blockDim.x + threadIdx.x;
+  int token = blockIdx.x;
+  int destRank = token % World;
+  int localToken = token / World;
+  cudaGridDependencySynchronize();
+  ar::utils::LamportFlags<float4> flag(p.ar.bufferFlags, 2);
+  auto local =
+      static_cast<T*>(flag.getCurLamportBuf(p.ar.inputPtrs[p.ar.rank], 0));
+  auto dest =
+      static_cast<T*>(flag.getCurLamportBuf(p.ar.inputPtrs[destRank], 0));
+  auto broadcast = static_cast<T*>(flag.getCurLamportBuf(p.ar.mcastPtr, 1));
+  if (packedIdx * 8 < p.ar.tokenDim) {
+    Vec value = loadInput<true>(p, token, packedIdx);
+    ar::utils::sanitizeLamportPayload<float4, T>(value);
+    reinterpret_cast<float4*>(dest + localToken * p.ar.tokenDim * World +
+                              p.ar.rank * p.ar.tokenDim)[packedIdx] =
+        value.packed;
+  }
+  cudaTriggerProgrammaticLaunchCompletion();
+  flag.clearDirtyLamportBuf(p.ar.inputPtrs[p.ar.rank], 0);
+  if (packedIdx * 8 < p.ar.tokenDim && destRank == p.ar.rank) {
+    Vec value =
+        ar::utils::reduceLamportRanksChunked<World, World, T, float4, 8>(
+            local, localToken, p.ar.tokenDim, packedIdx * 8);
+    ar::utils::sanitizeLamportPayload<float4, T>(value);
+    reinterpret_cast<float4*>(broadcast + token * p.ar.tokenDim)[packedIdx] =
+        value.packed;
+  }
+  flag.clearDirtyLamportBuf(p.ar.inputPtrs[p.ar.rank], 1);
+}
+
+template <bool UseCGA>
+__global__ void __launch_bounds__(1024) finishTwoshot(Params p) {
+  int packedIdx = threadIdx.x;
+  if constexpr (UseCGA) packedIdx = cg::this_cluster().thread_rank();
+  int offset = blockIdx.x * p.ar.tokenDim + packedIdx * 8;
+  cudaTriggerProgrammaticLaunchCompletion();
+  ar::utils::LamportFlags<float4, UseCGA> flag(p.ar.bufferFlags, 2);
+  auto input = static_cast<T*>(flag.getCurLamportBuf(p.ar.localBufferPtr, 1));
+  flag.ctaArrive();
+  auto loaded = ar::utils::loadPackedVolatile<float4>(input + offset);
+  while (ar::utils::isLamportDirty(loaded)) {
+    loaded = ar::utils::loadPackedVolatile<float4>(input + offset);
+  }
+  Vec value{.packed = loaded.packed};
+  epilogue<UseCGA>(p, value, packedIdx, offset);
+  cudaGridDependencySynchronize();
+  flag.waitAndUpdate(
+      {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) {
+  STD_TORCH_CHECK(status == cudaSuccess, cudaGetErrorString(status));
+}
+
+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->data_ptr())
+                       : nullptr;
+  p.ar.prenormedPtr = static_cast<T*>(hidden.data_ptr());
+  p.ar.residualInPtr = static_cast<T const*>(residual.data_ptr());
+  p.ar.gammaPtr = post_weight.has_value()
+                      ? static_cast<T const*>(post_weight->data_ptr())
+                      : nullptr;
+  p.ar.epsilon = post_eps;
+  p.nextWeight = next_weight.has_value()
+                     ? static_cast<T const*>(next_weight->data_ptr())
+                     : nullptr;
+  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()) {
+    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 =
+        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,
+                                  get_device_prop()->multiProcessorCount);
+  STD_TORCH_CHECK(!small_batch || (one_shot && !p.experts));
+  if (small_batch) {
+    clusters = 1;
+    threads = p.ar.tokenDim / 8;
+  }
+  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;
+  attrs[1].id = cudaLaunchAttributeClusterDimension;
+  attrs[1].val.clusterDim = {1, static_cast<unsigned>(clusters), 1};
+  cudaLaunchConfig_t config{};
+  config.gridDim = dim3(p.ar.numTokens, clusters);
+  config.blockDim = dim3(threads);
+  config.stream = stream;
+  config.attrs = attrs;
+  config.numAttrs = small_batch ? 1 : 2;
+  if (one_shot) {
+#define LAUNCH(W)                                                          \
+  case W:                                                                  \
+    if (small_batch)                                                       \
+      checkCuda(cudaLaunchKernelEx(&config, &oneshot<W, false, true>, p)); \
+    else if (p.experts)                                                    \
+      checkCuda(cudaLaunchKernelEx(&config, &oneshot<W, true>, p));        \
+    else                                                                   \
+      checkCuda(cudaLaunchKernelEx(&config, &oneshot<W>, p));              \
+    break
+    switch (world) {
+      LAUNCH(2);
+      LAUNCH(4);
+      LAUNCH(8);
+    }
+#undef LAUNCH
+  } else {
+    cudaLaunchConfig_t scatter{};
+    scatter.gridDim = dim3(p.ar.numTokens, ar::ceil_div(p.ar.tokenDim, 1024));
+    scatter.blockDim = dim3(128);
+    scatter.stream = stream;
+    scatter.attrs = attrs;
+    scatter.numAttrs = 1;
+#define LAUNCH(W)                                                      \
+  case W:                                                              \
+    if (p.experts)                                                     \
+      checkCuda(cudaLaunchKernelEx(&scatter, &finalizeScatter<W>, p)); \
+    else                                                               \
+      checkCuda(cudaLaunchKernelEx(                                    \
+          &scatter, &ar::twoshotAllreduceKernel<W, T, false>, p.ar));  \
+    break
+    switch (world) {
+      LAUNCH(2);
+      LAUNCH(4);
+      LAUNCH(8);
+    }
+#undef LAUNCH
+    if (clusters > 1) {
+      checkCuda(cudaLaunchKernelEx(&config, &finishTwoshot<true>, p));
+    } else {
+      config.numAttrs = 1;
+      checkCuda(cudaLaunchKernelEx(&config, &finishTwoshot<false>, p));
+    }
+  }
+}
+
+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.2), 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_KERNELS
+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_KERNELS
+  // 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_KERNELS
+  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

tests/models/test_beam_numerics.py

@@ -458,6 +458,235 @@ 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
+
+    import torch.distributed as dist
+
+    from vllm.models.beam.nvidia.allreduce import (
+        BeamAllReduceWorkspace,
+        beam_allreduce_norm_residual_norm,
+    )
+
+    torch.accelerator.set_device_index(rank % torch.accelerator.device_count())
+    dist.init_process_group(
+        "nccl",
+        init_method=f"tcp://{host}:{port}",
+        rank=rank,
+        world_size=world_size,
+        timeout=timedelta(seconds=120),
+    )
+    group = dist.new_group(backend="gloo")
+    separate_cutoff = (10 * 1024**2 * world_size + 2 * hidden_size - 1) // (
+        2 * hidden_size
+    )
+    max_tokens = max(513, separate_cutoff + 1)
+    workspace = BeamAllReduceWorkspace(group, hidden_size, max_tokens)
+    torch.manual_seed(42)
+    residual = torch.randn(max_tokens, hidden_size, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(hidden_size, device="cuda", dtype=torch.bfloat16)
+    torch.manual_seed(43 + rank)
+    x = torch.randn_like(residual)
+    # Exercise shared Lamport generations across batch and implementation changes.
+    cutoff = workspace.small_batch_max_tokens
+    batches = [0, 1, cutoff, cutoff + 1, 513, 1]
+    if workspace.separate_norm_min_tokens is not None:
+        cutoff = workspace.separate_norm_min_tokens
+        batches += [cutoff - 1, cutoff, cutoff + 1, 1]
+    for iteration, tokens in enumerate(batches):
+        torch.manual_seed(42 + iteration)
+        residual.normal_()
+        weight.normal_()
+        torch.manual_seed(43 + iteration + rank)
+        x.normal_()
+        x[:, ::7] = -0.0
+        next_norm = iteration % 2 == 0
+        scales = (0.125, 0.9682458365518543)
+        args = (
+            x[:tokens],
+            residual[:tokens],
+            weight,
+            1e-6,
+            weight if next_norm else None,
+            1e-6 if next_norm else None,
+            scales,
+        )
+
+        def check(result, args):
+            partial, r, pw, peps, nw, neps, scales = args
+            # The collective rounds the FP32 sum before the post-norm.
+            reduced = partial.float()
+            dist.all_reduce(reduced)
+            reduced = reduced.bfloat16()
+            expected = _torch_boundary(reduced, r, pw, peps, nw, neps, scales)
+            _assert_boundary_close(result, expected, relative_tolerance=1e-3)
+
+        check(beam_allreduce_norm_residual_norm(*args, workspace), args)
+        torch.accelerator.synchronize()
+        dist.barrier()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = beam_allreduce_norm_residual_norm(*args, workspace)
+        for _ in range(2):
+            x.mul_(0.5)
+            graph.replay()
+            check(captured, args)
+        graph.reset()
+    # A small configured capacity still needs enough space for forced one-shot.
+    cutoff = workspace.small_batch_max_tokens
+    small_workspace = BeamAllReduceWorkspace(group, hidden_size, cutoff)
+    args = (x[:cutoff], residual[:cutoff], weight, 1e-6, weight, 1e-6, scales)
+    check(beam_allreduce_norm_residual_norm(*args, small_workspace), args)
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+@pytest.mark.parametrize(
+    "world_size,hidden_size", [(2, 1024), (2, 8192), (4, 5120), (8, 5120)]
+)
+def test_allreduce_boundary_handles_batch_changes_and_graph_replay(
+    world_size, hidden_size
+):
+    """Changing token counts must not consume stale Lamport slots or deadlock."""
+    from torch.multiprocessing import spawn
+
+    from vllm.utils.network_utils import get_open_port
+
+    if torch.accelerator.device_count() < world_size:
+        pytest.skip(f"Requires {world_size} GPUs")
+    _skip_without_beam_allreduce()
+    spawn(
+        _beam_allreduce_worker,
+        args=(world_size, get_open_port(), hidden_size),
+        nprocs=world_size,
+    )
+
+
+def _beam_finalize_worker(rank, world_size, port, hidden_size):
+    import torch.distributed as dist
+
+    from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
+    from vllm.models.beam.nvidia.allreduce import (
+        BeamAllReduceWorkspace,
+        beam_allreduce_norm_residual_norm,
+        beam_finalize_allreduce_norm_residual_norm,
+    )
+
+    torch.accelerator.set_device_index(rank)
+    dist.init_process_group(
+        "nccl",
+        init_method=f"tcp://127.0.0.1:{port}",
+        rank=rank,
+        world_size=world_size,
+    )
+    group = dist.new_group(backend="gloo")
+    n, k = hidden_size, 8
+    workspace = BeamAllReduceWorkspace(group, n, 4095)
+    for iteration, m in enumerate([0, 1, 127, 513, 2]):
+        torch.manual_seed(100 + iteration)
+        residual = torch.randn(m, n, device="cuda", dtype=torch.bfloat16)
+        weight = torch.randn(n, device="cuda", dtype=torch.bfloat16)
+        torch.manual_seed(200 + rank + iteration)
+        experts = torch.randn(m * k, n, device="cuda", dtype=torch.bfloat16)
+        weights = torch.rand(m, k, device="cuda") / k
+        if iteration % 2:
+            weights = weights.bfloat16()
+        permutation = torch.randperm(m * k, device="cuda", dtype=torch.int32).view(m, k)
+        permutation[::3, -1] = -1
+        x = UnfinalizedMoEOutput(experts, weights, permutation)
+        routed_scale = 2.5 if iteration % 2 else 1.0
+        next_norm = iteration % 2 == 0
+        scales = (0.125, 0.9682458365518543)
+
+        args = (
+            x,
+            residual,
+            weight,
+            1e-6,
+            weight if next_norm else None,
+            1e-6 if next_norm else None,
+            scales,
+            workspace,
+            routed_scale,
+        )
+
+        def run(args=args):
+            return beam_finalize_allreduce_norm_residual_norm(*args)
+
+        def check(actual, args=args):
+            x, residual, pw, peps, nw, neps, scales, _, routed_scale = args
+            m, n = residual.shape
+            local = torch.zeros(m, n, device="cuda")
+            for expert in range(x.expert_weights.shape[1]):
+                rows = x.expanded_idx_to_permuted_idx[:, expert].long()
+                values = x.gemm2_permuted[rows.clamp_min(0)].float()
+                values = torch.where(rows[:, None] >= 0, values, 0)
+                local = local + values * x.expert_weights[:, expert, None].float()
+            local = (local.bfloat16().float() * routed_scale).bfloat16().float()
+            dist.all_reduce(local)
+            expected = _torch_boundary(
+                local.bfloat16(),
+                residual,
+                pw,
+                peps,
+                nw,
+                neps,
+                scales,
+            )
+            _assert_boundary_close(actual, expected)
+
+        check(run())
+        torch.accelerator.synchronize()
+        dist.barrier()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = run(args)
+        for _ in range(2):
+            experts.mul_(0.5)
+            graph.replay()
+            check(captured, args)
+            # Alternate two-stage finalization with the one-stage decode path
+            # on the same Lamport buffers, including during graph replay.
+            small = residual[:8]
+            actual = beam_allreduce_norm_residual_norm(
+                small, small, weight, 1e-6, weight, 1e-6, scales, args[-2]
+            )
+            reduced = small.float()
+            dist.all_reduce(reduced)
+            expected = _torch_boundary(
+                reduced.bfloat16(), small, weight, 1e-6, weight, 1e-6, scales
+            )
+            _assert_boundary_close(actual, expected)
+        graph.reset()
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+@pytest.mark.parametrize("world_size", [2, 4])
+def test_moe_finalize_allreduce_preserves_local_rounding_and_replay(world_size):
+    """Empty local routes and live routing metadata retain BF16 casts on replay."""
+    from torch.multiprocessing import spawn
+
+    from vllm.utils.network_utils import get_open_port
+
+    if torch.accelerator.device_count() < world_size:
+        pytest.skip(f"Requires {world_size} GPUs")
+    _skip_without_beam_allreduce()
+    spawn(
+        _beam_finalize_worker,
+        args=(world_size, get_open_port(), 5120),
+        nprocs=world_size,
+    )
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.parametrize("rows", [1, 2, 4, 8, 12, 16, 24, 64])
 @pytest.mark.parametrize("qkv_rows,gate_rows", [(2048, 1536), (1024, 768)])

vllm/models/beam/nvidia/allreduce.py

@@ -0,0 +1,244 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""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, round_up
+
+_NUM_LAMPORT_BUFFERS = 3
+# FlashInfer's MNNVL one-shot cutoff on the BF16 payload staged across ranks.
+_ONESHOT_MAX_BYTES = 1024 * 1024
+
+
+class BeamAllReduceWorkspace:
+    """Multicast Lamport buffers and generation flags for the Beam kernels."""
+
+    @classmethod
+    def create(
+        cls,
+        group: dist.ProcessGroup,
+        hidden_size: int,
+        max_tokens: int,
+        dtype: torch.dtype,
+    ) -> "BeamAllReduceWorkspace | None":
+        # The CUDA epilogue handles one packed load per thread.
+        if (
+            group.size() not in (2, 4, 8)
+            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
+        from flashinfer.comm.mnnvl import (
+            TorchDistBackend,
+            all_ranks_agree,
+            is_multicast_supported,
+        )
+
+        device = torch.accelerator.current_device_index()
+        supported = torch.cuda.get_device_capability(device)[0] in (9, 10)
+        supported = supported and is_multicast_supported(device)
+        if not all_ranks_agree(TorchDistBackend(group=group), supported):
+            return None
+        return cls(group, hidden_size, max_tokens)
+
+    def __init__(
+        self, group: dist.ProcessGroup, hidden_size: int, max_tokens: int
+    ) -> None:
+        self.world_size = group.size()
+        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.separate_norm_min_tokens = cdiv(
+            10 * 1024**2 * self.world_size, 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,
+    residual: torch.Tensor,
+    post_weight: torch.Tensor,
+    post_eps: float,
+    next_weight: torch.Tensor | None,
+    next_eps: float | None,
+    scales: tuple[float, float],
+    workspace: BeamAllReduceWorkspace,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    assert x.shape == residual.shape
+    for weight in (post_weight, next_weight):
+        if weight is not None:
+            assert weight.dtype == torch.bfloat16 and weight.is_contiguous()
+            assert weight.shape == (workspace.hidden_size,)
+    assert x.dtype == residual.dtype == torch.bfloat16
+    assert x.is_contiguous() and residual.is_contiguous()
+    assert x.shape[1] == workspace.hidden_size
+    assert x.shape[0] <= workspace.max_tokens
+    if x.shape[0] >= workspace.separate_norm_min_tokens:
+        from .layernorm import beam_norm_residual_norm
+
+        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,
+            residual,
+            post_weight,
+            post_eps,
+            next_weight,
+            next_eps,
+            scales,
+            launch_pdl=True,
+        )
+    return _cuda_allreduce_norm_residual_norm(
+        x, residual, post_weight, post_eps, next_weight, next_eps, scales, workspace
+    )
+
+
+def _cuda_allreduce_norm_residual_norm(
+    x,
+    residual,
+    post_weight,
+    post_eps,
+    next_weight,
+    next_eps,
+    scales,
+    workspace,
+    unfinalized: UnfinalizedMoEOutput | None = None,
+    routed_scale: float = 1.0,
+):
+    assert x.dtype == residual.dtype == torch.bfloat16
+    assert x.is_contiguous() and residual.is_contiguous()
+    m, n = x.shape
+    assert n == workspace.hidden_size and m <= workspace.max_tokens
+    hidden = torch.empty_like(x)
+    normalized = torch.empty_like(x) if next_eps is not None else None
+    if m:
+        small_batch = unfinalized is None and m <= workspace.small_batch_max_tokens
+        torch.ops._C.beam_allreduce_norm(
+            x,
+            residual,
+            post_weight,
+            next_weight,
+            hidden,
+            normalized,
+            workspace.buffer_ptrs,
+            workspace.flags,
+            workspace.multicast_ptr,
+            workspace.local_ptr,
+            workspace.world_size,
+            workspace.rank,
+            post_eps or 0.0,
+            next_eps or 0.0,
+            *(scales if scales is not None else (1.0, 1.0)),
+            post_eps is not None,
+            next_eps is not None,
+            scales is not None,
+            small_batch or m <= 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,
+        )
+    return hidden, normalized
+
+
+def beam_finalize_allreduce_norm_residual_norm(
+    x: UnfinalizedMoEOutput,
+    residual: torch.Tensor,
+    post_weight: torch.Tensor,
+    post_eps: float,
+    next_weight: torch.Tensor | None,
+    next_eps: float | None,
+    scales: tuple[float, float],
+    workspace: BeamAllReduceWorkspace,
+    routed_scale: float,
+) -> tuple[torch.Tensor, torch.Tensor | None]:
+    """Finalize local experts before AR, retaining the local BF16 rounding."""
+    m, n = residual.shape
+    for weight in (post_weight, next_weight):
+        if weight is not None:
+            assert weight.dtype == torch.bfloat16 and weight.is_contiguous()
+            assert weight.shape == (n,)
+    assert x.gemm2_permuted.dtype == torch.bfloat16
+    assert x.gemm2_permuted.is_contiguous() and x.gemm2_permuted.shape[1] == n
+    assert x.expert_weights.dtype in (torch.float32, torch.bfloat16)
+    assert x.expanded_idx_to_permuted_idx.dtype == torch.int32
+    assert x.expert_weights.is_contiguous()
+    assert x.expanded_idx_to_permuted_idx.is_contiguous()
+    assert x.expert_weights.shape == x.expanded_idx_to_permuted_idx.shape
+    assert x.expert_weights.shape[0] == m
+    return _cuda_allreduce_norm_residual_norm(
+        residual,
+        residual,
+        post_weight,
+        post_eps,
+        next_weight,
+        next_eps,
+        scales,
+        workspace,
+        x,
+        routed_scale,
+    )

vllm/models/beam/nvidia/attention.py

@@ -48,6 +48,7 @@ class BeamAttention(nn.Module):
         cache_config: CacheConfig | None,
         quant_config: QuantizationConfig | None,
         prefix: str = "",
+        reduce_results: bool = True,
     ) -> None:
         super().__init__()
         self.layer_id = extract_layer_index(prefix)
@@ -104,6 +105,7 @@ class BeamAttention(nn.Module):
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.o_proj",
+            reduce_results=reduce_results,
         )
 
         self.rms_norm_eps = config.rms_norm_eps

vllm/models/beam/nvidia/layernorm.py

@@ -117,10 +117,13 @@ def _norm_residual_norm_kernel(
     R_SCALE,
     NEXT_NORM: tl.constexpr,
     BLOCK: tl.constexpr,
+    PDL: tl.constexpr,
 ):
     row = tl.program_id(0).to(tl.int64)
     col = tl.arange(0, BLOCK)
     dtype = X.dtype.element_ty
+    if PDL:
+        tl.extra.cuda.gdc_wait()
     x = tl.load(X + row * X_STRIDE + col, col < N, other=0).to(tl.float32)
     r = tl.load(R + row * R_STRIDE + col, col < N, other=0).to(dtype).to(tl.float32)
     h, y = _boundary_math(
@@ -150,6 +153,8 @@ def beam_norm_residual_norm(
     next_weight: torch.Tensor | None,
     next_eps: float | None,
     scales: tuple[float, float],
+    *,
+    launch_pdl: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
     """Fuse post-norm, scaled residual add and next pre-norm, keeping each cast.
 
@@ -178,7 +183,9 @@ def beam_norm_residual_norm(
             *scales,
             next_weight is not None,
             triton.next_power_of_2(n),
-            num_warps=4 if n <= 1024 else 8,
+            launch_pdl,
+            num_warps=4 if launch_pdl or n <= 1024 else 8,
             enable_fp_fusion=False,
+            launch_pdl=launch_pdl,
         )
     return hidden, normalized

vllm/models/beam/nvidia/model.py

@@ -11,14 +11,20 @@ selection.
 
 import math
 from collections.abc import Iterable
+from functools import partial
 from typing import cast
 
 import torch
 from torch import nn
 
 from vllm.config import VllmConfig
-from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
+from vllm.distributed import (
+    get_pp_group,
+    get_tensor_model_parallel_world_size,
+    get_tp_group,
+)
 from vllm.model_executor.layers.fused_moe import MoERunner
+from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
 from vllm.model_executor.layers.logits_processor import LogitsProcessor
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
@@ -44,6 +50,11 @@ from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.configs.beam import BeamConfig
 
 from .activation import beam_softcap_logits
+from .allreduce import (
+    BeamAllReduceWorkspace,
+    beam_allreduce_norm_residual_norm,
+    beam_finalize_allreduce_norm_residual_norm,
+)
 from .attention import BeamAttention, use_merged_qkv_gate
 from .layernorm import RMSNorm, beam_norm_residual_norm, beam_rms_norm
 from .moe import BeamMLP, BeamSparseMoeBlock
@@ -69,6 +80,7 @@ class BeamDecoderLayer(nn.Module):
         self,
         vllm_config: VllmConfig,
         prefix: str = "",
+        allreduce_workspace: BeamAllReduceWorkspace | None = None,
     ) -> None:
         super().__init__()
         config = cast(BeamConfig, vllm_config.model_config.hf_config)
@@ -76,9 +88,14 @@ class BeamDecoderLayer(nn.Module):
         quant_config = vllm_config.quant_config
         self.layer_id = extract_layer_index(prefix)
         self.rms_norm_eps = config.rms_norm_eps
+        fuse_attention_reduce = allreduce_workspace is not None
 
         self.self_attn = BeamAttention(
-            config, cache_config, quant_config, prefix=f"{prefix}.self_attn"
+            config,
+            cache_config,
+            quant_config,
+            prefix=f"{prefix}.self_attn",
+            reduce_results=not fuse_attention_reduce,
         )
 
         if config.mlp_layer_types[self.layer_id] == "dense":
@@ -88,10 +105,38 @@ class BeamDecoderLayer(nn.Module):
                 hidden_act=config.hidden_act,
                 quant_config=quant_config,
                 prefix=f"{prefix}.mlp",
+                reduce_results=allreduce_workspace is None,
             )
         else:
             self.mlp = BeamSparseMoeBlock(
-                vllm_config=vllm_config, prefix=f"{prefix}.mlp"
+                vllm_config=vllm_config,
+                prefix=f"{prefix}.mlp",
+                reduce_results=allreduce_workspace is None,
+            )
+
+        fused_boundary = (
+            partial(beam_allreduce_norm_residual_norm, workspace=allreduce_workspace)
+            if allreduce_workspace is not None
+            else beam_norm_residual_norm
+        )
+        self.attention_boundary = (
+            fused_boundary if fuse_attention_reduce else beam_norm_residual_norm
+        )
+        fuse_mlp_reduce = allreduce_workspace is not None and (
+            isinstance(self.mlp, BeamMLP)
+            or self.mlp.experts.moe_config.skip_final_all_reduce
+        )
+        self.mlp_boundary = (
+            fused_boundary if fuse_mlp_reduce else beam_norm_residual_norm
+        )
+        self.finalize_boundary = None
+        if fuse_mlp_reduce and isinstance(self.mlp, BeamSparseMoeBlock):
+            assert allreduce_workspace is not None
+            self.mlp.defer_finalize(allreduce_workspace.max_tokens)
+            self.finalize_boundary = partial(
+                beam_finalize_allreduce_norm_residual_norm,
+                workspace=allreduce_workspace,
+                routed_scale=self.mlp.experts.routed_scaling_factor,
             )
 
         self.input_layernorm = RMSNorm(config.hidden_size)
@@ -114,7 +159,7 @@ class BeamDecoderLayer(nn.Module):
             else normalized_states
         )
         x = self.self_attn(positions, x)
-        hidden_states, x = beam_norm_residual_norm(
+        hidden_states, x = self.attention_boundary(
             x,
             hidden_states,
             self.post_attention_layernorm.weight,
@@ -125,7 +170,11 @@ class BeamDecoderLayer(nn.Module):
         )
         assert x is not None
         x = self.mlp(x)
-        return beam_norm_residual_norm(
+        boundary = self.mlp_boundary
+        if isinstance(x, UnfinalizedMoEOutput):
+            assert self.finalize_boundary is not None
+            boundary = self.finalize_boundary
+        return boundary(
             x,
             hidden_states,
             self.post_feedforward_layernorm.weight,
@@ -203,9 +252,23 @@ class BeamModel(nn.Module, EagleModelMixin):
             )
         else:
             self.embed_tokens = PPMissingLayer()
+        tp_group = get_tp_group()
+        self.allreduce_workspace = BeamAllReduceWorkspace.create(
+            tp_group.cpu_group,
+            config.hidden_size,
+            max(
+                vllm_config.scheduler_config.max_num_batched_tokens,
+                vllm_config.scheduler_config.max_num_seqs,
+            ),
+            vllm_config.model_config.dtype,
+        )
         self.start_layer, self.end_layer, self.layers = make_layers(
             config.num_hidden_layers,
-            lambda prefix: BeamDecoderLayer(vllm_config=vllm_config, prefix=prefix),
+            lambda prefix: BeamDecoderLayer(
+                vllm_config=vllm_config,
+                prefix=prefix,
+                allreduce_workspace=self.allreduce_workspace,
+            ),
             prefix=f"{prefix}.layers",
         )
         self.norm = (

vllm/models/beam/nvidia/moe.py

@@ -6,6 +6,7 @@ from torch import nn
 
 from vllm.config import VllmConfig
 from vllm.model_executor.layers.fused_moe import FusedMoEFactory, GateLinear
+from vllm.model_executor.layers.fused_moe.moe_output import UnfinalizedMoEOutput
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     RowParallelLinear,
@@ -25,6 +26,7 @@ class BeamMLP(nn.Module):
         hidden_act: str,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
+        reduce_results: bool = True,
     ) -> None:
         super().__init__()
         self.gate_up_proj = MergedColumnParallelLinear(
@@ -40,6 +42,7 @@ class BeamMLP(nn.Module):
             bias=False,
             quant_config=quant_config,
             prefix=f"{prefix}.down_proj",
+            reduce_results=reduce_results,
         )
         if hidden_act != "silu":
             raise ValueError(
@@ -57,7 +60,9 @@ class BeamMLP(nn.Module):
 class BeamSparseMoeBlock(nn.Module):
     """Grouped sigmoid ``noaux_tc`` MoE block."""
 
-    def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:
+    def __init__(
+        self, vllm_config: VllmConfig, prefix: str = "", reduce_results: bool = True
+    ) -> None:
         super().__init__()
         config = vllm_config.model_config.hf_config
         parallel_config = vllm_config.parallel_config
@@ -77,6 +82,7 @@ class BeamSparseMoeBlock(nn.Module):
         self.gate.register_parameter("e_score_correction_bias", e_score_correction_bias)
 
         self.experts = FusedMoEFactory(
+            reduce_results=reduce_results,
             gate=self.gate,
             num_experts=config.num_experts,
             top_k=config.num_experts_per_tok,
@@ -99,5 +105,24 @@ class BeamSparseMoeBlock(nn.Module):
             prefix=f"{prefix}.experts",
         )
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def defer_finalize(self, max_tokens: int) -> None:
+        from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
+            TrtLlmNvFp4ExpertsModular,
+            TrtLlmNvFp4ExpertsMonolithic,
+        )
+
+        experts_cls = getattr(
+            self.experts.routed_experts.quant_method, "experts_cls", None
+        )
+        if experts_cls in (TrtLlmNvFp4ExpertsModular, TrtLlmNvFp4ExpertsMonolithic):
+            self.experts.moe_config.defer_moe_finalize(max_tokens)
+
+    def forward(
+        self, hidden_states: torch.Tensor
+    ) -> torch.Tensor | UnfinalizedMoEOutput:
+        if self.experts.moe_config.should_defer_moe_finalize(hidden_states.shape[0]):
+            output = self.experts._forward_impl(hidden_states, hidden_states, None)
+            assert isinstance(output, UnfinalizedMoEOutput)
+            return output
+
         return self.experts(hidden_states=hidden_states, router_logits=hidden_states)