[Kernel] Fuse DeepEP-v2 decode globalize+align into one launch

CMakeLists.txt

@@ -1287,7 +1287,8 @@ set(VLLM_MOE_EXT_SRC
 if(VLLM_GPU_LANG STREQUAL "CUDA")
   list(APPEND VLLM_MOE_EXT_SRC
     "csrc/moe/moe_wna16.cu"
-    "csrc/moe/grouped_topk_kernels.cu")
+    "csrc/moe/grouped_topk_kernels.cu"
+    "csrc/moe/fused_globalize_align_block_size.cu")
 endif()
 
 if(VLLM_GPU_LANG STREQUAL "CUDA")

csrc/moe/fused_globalize_align_block_size.cu

@@ -0,0 +1,266 @@
+// Single-launch fusion of the DeepEP-v2 decode metadata chain
+// (_globalize_recv_topk_idx + moe_align_block_size + count_and_sort) for the
+// humming INDEXED / CUDA-graph decode path. recv values are LOCAL expert ids;
+// the kernel bins by local id while storing back the global id, sized by local
+// experts. Hybrid launch: one non-cooperative block for small decode, else a
+// cooperative multi-SM grid (launched via cudaLaunchKernelEx so grid.sync() is
+// CUDA-graph capturable). Barriers: zero+sentinel -> shared histogram
+// (distributed reduce) -> block-0 prefix + expert_ids fill ->
+// scatter+globalize.
+
+#include <cooperative_groups.h>
+#include <cuda_runtime.h>
+
+#include <torch/csrc/stable/macros.h>
+#include <torch/csrc/stable/accelerator.h>
+#include <torch/csrc/stable/ops.h>
+#include <torch/csrc/stable/tensor.h>
+#include <torch/headeronly/core/ScalarType.h>
+
+#include "torch_utils.h"
+
+namespace cg = cooperative_groups;
+
+namespace vllm {
+namespace moe {
+
+static constexpr int FGAS_THREADS = 256;
+static constexpr long FGAS_SB_MAX_EFFECTIVE_WORK = 4096;
+
+// COOP selects the phase barrier: grid.sync() (multi-block) vs __syncthreads().
+// topk / block_size are runtime args: any (topk, block_size) is supported, and
+// the kernel is barrier/memory-bound so the divisions are not on the hot path.
+template <bool COOP>
+__device__ __forceinline__ void fused_gas_body(
+    long* __restrict__ topk_idx, int* __restrict__ sorted_ids,
+    int* __restrict__ expert_ids, int* __restrict__ num_tokens_post_pad,
+    int* __restrict__ counts, int* __restrict__ scatter_idx, int num_recv,
+    int rank_expert_offset, int global_num_experts, int numel,
+    int max_num_tokens_padded, int max_num_m_blocks, int local_num_experts,
+    int topk, int block_size) {
+  extern __shared__ int sh[];  // [local_num_experts] per-block histogram
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+  auto BAR = [&]() {
+    if constexpr (COOP)
+      cg::this_grid().sync();
+    else
+      __syncthreads();
+  };
+
+  // Phase 0: zero the histogram + sentinel-fill sorted_ids.
+  for (int e = tid; e < local_num_experts; e += stride) counts[e] = 0;
+  for (int i = tid; i < max_num_tokens_padded; i += stride)
+    sorted_ids[i] = numel;
+  BAR();
+
+  // Phase 1: per-block shared histogram, one global atomic per (block, expert).
+  // Reduce stays distributed across blocks (centralizing it in block 0 is
+  // slower).
+  for (int e = threadIdx.x; e < local_num_experts; e += blockDim.x) sh[e] = 0;
+  __syncthreads();
+  for (int i = tid; i < numel; i += stride) {
+    int val = (int)topk_idx[i];
+    int row = i / topk;
+    bool valid = (val >= 0) && (val < local_num_experts) &&
+                 (val + rank_expert_offset < global_num_experts) &&
+                 (row < num_recv);
+    if (valid) atomicAdd(&sh[val], 1);
+  }
+  __syncthreads();
+  for (int e = threadIdx.x; e < local_num_experts; e += blockDim.x)
+    if (sh[e]) atomicAdd(&counts[e], sh[e]);
+  BAR();
+
+  // Phase 2 (block 0). Thread 0 runs the tiny serial prefix scan; then all of
+  // block 0's threads split the expert_ids fill grid-stride.
+  if (blockIdx.x == 0) {
+    int* sh_start = sh;  // [local_num_experts] padded start offsets
+    int& sh_run = sh[local_num_experts];
+    if (threadIdx.x == 0) {
+      int run = 0;
+      for (int e = 0; e < local_num_experts; e++) {
+        int c = counts[e];
+        int nb = (c + block_size - 1) / block_size;
+        sh_start[e] = run;
+        counts[e] = run;  // exclusive prefix; empty experts skip
+        run += nb * block_size;
+      }
+      sh_run = run;
+      *num_tokens_post_pad = run;
+    }
+    __syncthreads();
+    int run = sh_run;
+    // owner(j) = highest expert whose padded start <= j*block_size (skips empty
+    // experts); blocks past the last token get -1.
+    for (int j = threadIdx.x; j < max_num_m_blocks; j += blockDim.x) {
+      int off = j * block_size;
+      if (off >= run) {
+        expert_ids[j] = -1;
+        continue;
+      }
+      int owner = 0;
+      for (int e = 1; e < local_num_experts; e++)
+        if (sh_start[e] <= off) owner = e;
+      expert_ids[j] = owner;
+    }
+  }
+  BAR();
+
+  // Phase 3: scatter token indices at cursor[val]++ and globalize in place.
+  for (int i = tid; i < numel; i += stride) {
+    int val = (int)topk_idx[i];
+    int g = val + rank_expert_offset;
+    int row = i / topk;
+    bool valid = (val >= 0) && (val < local_num_experts) &&
+                 (g < global_num_experts) && (row < num_recv);
+    if (scatter_idx != nullptr) scatter_idx[i] = valid ? i : -1;
+    if (valid) {
+      int pos = atomicAdd(&counts[val], 1);
+      sorted_ids[pos] = i;
+    }
+    topk_idx[i] = valid ? g : -1;
+  }
+}
+
+template <bool COOP>
+__global__ void fused_gas_kernel(
+    long* __restrict__ topk_idx, const int* __restrict__ psum,
+    int* __restrict__ sorted_ids, int* __restrict__ expert_ids,
+    int* __restrict__ num_tokens_post_pad, int* __restrict__ counts,
+    int* __restrict__ scatter_idx, int P, int rank_expert_offset,
+    int global_num_experts, int numel, int max_num_tokens_padded,
+    int max_num_m_blocks, int local_num_experts, int topk, int block_size) {
+  // num_recv read on-device (psum[P-1]) -> baked into replay, cudagraph-safe.
+  fused_gas_body<COOP>(topk_idx, sorted_ids, expert_ids, num_tokens_post_pad,
+                       counts, scatter_idx, psum[P - 1], rank_expert_offset,
+                       global_num_experts, numel, max_num_tokens_padded,
+                       max_num_m_blocks, local_num_experts, topk, block_size);
+}
+
+static int fgas_sm_count() {
+  // Device-constant: query the driver once, never on the capture/replay path.
+  static int sm_count = -1;
+  if (sm_count < 0) {
+    int dev = 0;
+    cudaGetDevice(&dev);
+    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev);
+    if (sm_count < 1) sm_count = 1;
+  }
+  return sm_count;
+}
+
+static int fgas_cooperative_blocks(int bps, long work) {
+  int max_blocks = bps * fgas_sm_count();
+  if (max_blocks < 1) max_blocks = 1;
+  long need = (work + FGAS_THREADS - 1) / FGAS_THREADS;  // no over-launch
+  if (need < 1) need = 1;
+  return need < (long)max_blocks ? (int)need : max_blocks;
+}
+
+static void fgas_launch(long* p_topk, const int* p_psum, int* p_sorted,
+                        int* p_expert, int* p_num, int* p_counts,
+                        int* p_scatter, int P, int reo, int gne, int numel,
+                        int mntp, int mnmb, int local_e, int topk,
+                        int block_size, cudaStream_t stream) {
+  long work = numel > mntp ? numel : mntp;
+  long effective_work = (long)mntp + 2L * numel + (long)mnmb * local_e;
+  // Histogram [local_e] + one scalar (phase-2 padded-run total).
+  size_t smem = (size_t)(local_e + 1) * sizeof(int);
+
+  // Small decode: one non-cooperative block, __syncthreads barriers. Account
+  // for both numel passes and the expert_ids owner scan in the crossover.
+  if (effective_work <= FGAS_SB_MAX_EFFECTIVE_WORK) {
+    fused_gas_kernel<false><<<1, FGAS_THREADS, smem, stream>>>(
+        p_topk, p_psum, p_sorted, p_expert, p_num, p_counts, p_scatter, P, reo,
+        gne, numel, mntp, mnmb, local_e, topk, block_size);
+    STD_CUDA_CHECK(cudaGetLastError());
+    return;
+  }
+
+  // Large decode: cooperative multi-SM. Occupancy (fixed per rank) queried
+  // once.
+  static int bps = -1;
+  if (bps < 0)
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+        &bps, (const void*)fused_gas_kernel<true>, FGAS_THREADS, smem);
+  int blocks = fgas_cooperative_blocks(bps, work);
+
+  // Cooperative launch via cudaLaunchKernelEx + the cooperative launch
+  // attribute: unlike the legacy cudaLaunchCooperativeKernel, this form IS
+  // stream-capturable, so grid.sync() works inside a captured CUDA graph.
+  cudaLaunchConfig_t config = {};
+  config.gridDim = dim3(blocks);
+  config.blockDim = dim3(FGAS_THREADS);
+  config.dynamicSmemBytes = smem;
+  config.stream = stream;
+  cudaLaunchAttribute attr = {};
+  attr.id = cudaLaunchAttributeCooperative;
+  attr.val.cooperative = 1;
+  config.attrs = &attr;
+  config.numAttrs = 1;
+  STD_CUDA_CHECK(cudaLaunchKernelEx(&config, fused_gas_kernel<true>, p_topk,
+                                    p_psum, p_sorted, p_expert, p_num, p_counts,
+                                    p_scatter, P, reo, gne, numel, mntp, mnmb,
+                                    local_e, topk, block_size));
+}
+
+}  // namespace moe
+}  // namespace vllm
+
+// recv_topk_idx (int64, [N, topk]) is globalized in place and the align/sort
+// outputs (int32) are filled. counts is an internal write-cursor allocated
+// here.
+void fused_globalize_align_block_size(
+    torch::stable::Tensor topk_idx, torch::stable::Tensor psum,
+    int64_t rank_expert_offset, int64_t global_num_experts,
+    int64_t local_num_experts, int64_t block_size,
+    torch::stable::Tensor sorted_ids, torch::stable::Tensor expert_ids,
+    torch::stable::Tensor num_tokens_post_pad,
+    std::optional<torch::stable::Tensor> scatter_idx) {
+  STD_TORCH_CHECK(topk_idx.scalar_type() == torch::headeronly::ScalarType::Long,
+                  "fused_globalize_align_block_size: topk_idx must be int64");
+  STD_TORCH_CHECK(
+      local_num_experts <= 1024,
+      "fused_globalize_align_block_size: local_num_experts <= 1024");
+
+  const torch::stable::accelerator::DeviceGuard device_guard(
+      topk_idx.get_device_index());
+  const cudaStream_t stream =
+      get_current_cuda_stream(topk_idx.get_device_index());
+
+  int numel = (int)topk_idx.numel();
+  int topk = (int)topk_idx.size(1);
+  int mntp = (int)sorted_ids.size(0);
+  int mnmb = (int)expert_ids.size(0);
+  int local_e = (int)local_num_experts;
+  int reo = (int)rank_expert_offset;
+  int gne = (int)global_num_experts;
+  int P = (int)psum.size(0);
+
+  // Internal per-local-expert histogram / write-cursor.
+  torch::stable::Tensor counts = torch::stable::new_empty(
+      topk_idx, {local_num_experts}, torch::headeronly::ScalarType::Int);
+
+  long* p_topk = reinterpret_cast<long*>(topk_idx.mutable_data_ptr());
+  const int* p_psum = reinterpret_cast<const int*>(psum.const_data_ptr());
+  int* p_sorted = reinterpret_cast<int*>(sorted_ids.mutable_data_ptr());
+  int* p_expert = reinterpret_cast<int*>(expert_ids.mutable_data_ptr());
+  int* p_num = reinterpret_cast<int*>(num_tokens_post_pad.mutable_data_ptr());
+  int* p_counts = reinterpret_cast<int*>(counts.mutable_data_ptr());
+  int* p_scatter = nullptr;
+  if (scatter_idx.has_value()) {
+    STD_TORCH_CHECK(
+        scatter_idx->scalar_type() == torch::headeronly::ScalarType::Int &&
+            scatter_idx->device() == topk_idx.device() &&
+            scatter_idx->is_contiguous() &&
+            scatter_idx->numel() == topk_idx.numel(),
+        "scatter_idx must be contiguous int32 on the same device and match "
+        "topk_idx.numel()");
+    p_scatter = reinterpret_cast<int*>(scatter_idx->mutable_data_ptr());
+  }
+
+  vllm::moe::fgas_launch(p_topk, p_psum, p_sorted, p_expert, p_num, p_counts,
+                         p_scatter, P, reo, gne, numel, mntp, mnmb, local_e,
+                         topk, (int)block_size, stream);
+}

csrc/moe/moe_ops.h

@@ -44,6 +44,16 @@ void moe_align_block_size(torch::stable::Tensor topk_ids, int64_t num_experts,
                           std::optional<torch::stable::Tensor> maybe_expert_map,
                           std::optional<torch::stable::Tensor> scatter_idx);
 
+#ifdef USE_CUDA
+void fused_globalize_align_block_size(
+    torch::stable::Tensor topk_idx, torch::stable::Tensor psum,
+    int64_t rank_expert_offset, int64_t global_num_experts,
+    int64_t local_num_experts, int64_t block_size,
+    torch::stable::Tensor sorted_ids, torch::stable::Tensor expert_ids,
+    torch::stable::Tensor num_tokens_post_pad,
+    std::optional<torch::stable::Tensor> scatter_idx);
+#endif
+
 void batched_moe_align_block_size(
     int64_t max_tokens_per_batch, int64_t block_size,
     const torch::stable::Tensor& expert_num_tokens,

csrc/moe/torch_bindings.cpp

@@ -42,6 +42,18 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
       "                     Tensor? maybe_expert_map,"
       "                     Tensor(a!)? scatter_idx=None) -> ()");
 
+#ifdef USE_CUDA
+  // Fused globalize + align + sort for the DeepEP-v2 humming INDEXED decode
+  // path (one launch; num_recv read on device from psum[-1], cudagraph-safe).
+  m.def(
+      "fused_globalize_align_block_size(Tensor! topk_idx, Tensor psum,"
+      "                     int rank_expert_offset, int global_num_experts,"
+      "                     int local_num_experts, int block_size,"
+      "                     Tensor! sorted_ids, Tensor! expert_ids,"
+      "                     Tensor! num_tokens_post_pad,"
+      "                     Tensor(a!)? scatter_idx=None) -> ()");
+#endif
+
   // Aligning the number of tokens to be processed by each expert such
   // that it is divisible by the block size, but for the batched case.
   m.def(
@@ -143,6 +155,10 @@ STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) {
   m.impl("topk_softplus_sqrt", TORCH_BOX(&topk_softplus_sqrt));
   m.impl("moe_sum", TORCH_BOX(&moe_sum));
   m.impl("moe_align_block_size", TORCH_BOX(&moe_align_block_size));
+#ifdef USE_CUDA
+  m.impl("fused_globalize_align_block_size",
+         TORCH_BOX(&fused_globalize_align_block_size));
+#endif
   m.impl("batched_moe_align_block_size",
          TORCH_BOX(&batched_moe_align_block_size));
   m.impl("moe_lora_align_block_size", TORCH_BOX(&moe_lora_align_block_size));