[CUDA] Cut the MoE routing prologue/epilogue cost for small-batch decode

onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu

@@ -338,6 +338,19 @@ bool tryLaunchMoeGemvIntSymmetricInterleavedSwiGLU(
   }
 }
 
+// Ranker used to group the (token, top-k slot) pairs by selected expert.
+//
+// cub::BlockRadixRank keeps a private set of packed digit counters per thread, so its
+// shared memory and its per-thread scan both grow with 2^LOG2_NUM_EXPERTS. At 256
+// experts (LOG2_NUM_EXPERTS == 9) and a decode-sized block of 32 threads that is 512
+// bins spread over a single warp, which costs ~16KB of shared memory, 255 registers per
+// thread and enough spilling that the kernel is dominated by counter bookkeeping rather
+// than by the handful of assignments it actually ranks. The match-based ranker ranks
+// with warp ballots instead of per-thread counters, so its cost tracks the number of
+// keys rather than the number of bins.
+template <int BLOCK_SIZE, int LOG2_NUM_EXPERTS>
+using MoeExpertRadixRank = cub::BlockRadixRankMatch<BLOCK_SIZE, LOG2_NUM_EXPERTS, false>;
+
 template <int BLOCK_SIZE, int EXPERTS_PER_TOKEN, int LOG2_NUM_EXPERTS>
 __global__ void fusedBuildExpertMapsSortFirstTokenKernel(const int* const token_selected_experts,
                                                          int* const permuted_row_to_unpermuted_row, int* const unpermuted_row_to_permuted_row,
@@ -385,7 +398,7 @@ __global__ void fusedBuildExpertMapsSortFirstTokenKernel(const int* const token_
   // that to elide the binary search
 
   // sort the expert map
-  using BlockRadixRank = cub::BlockRadixRank<BLOCK_SIZE, LOG2_NUM_EXPERTS, false>;
+  using BlockRadixRank = MoeExpertRadixRank<BLOCK_SIZE, LOG2_NUM_EXPERTS>;
   extern __shared__ unsigned char temp_storage[];
   auto& sort_temp = *reinterpret_cast<typename BlockRadixRank::TempStorage*>(temp_storage);
 
@@ -437,7 +450,7 @@ bool fusedBuildExpertMapsSortFirstTokenDispatch(const int* token_selected_expert
   const int blocks = (num_tokens + threads - 1) / threads;
   ORT_ENFORCE(blocks == 1, "Current implementation requires single block");
 
-  using BlockRadixRank = cub::BlockRadixRank<BLOCK_SIZE, LOG2_NUM_EXPERTS, false>;
+  using BlockRadixRank = MoeExpertRadixRank<BLOCK_SIZE, LOG2_NUM_EXPERTS>;
   size_t shared_size = sizeof(typename BlockRadixRank::TempStorage);
 
   cudaLaunchConfig_t config;
@@ -1562,6 +1575,15 @@ enum class ScaleMode : int {
 
 constexpr static int FINALIZE_THREADS_PER_BLOCK = 256;
 
+// The routing metadata (selected expert, permuted row, routing scale) only varies with the
+// top-k slot, not with the column, yet the naive reduction loop re-reads it from global memory
+// for every column tile through a chain of two dependent loads (expert map -> permuted row ->
+// expanded row data). With a small grid (one block per token, i.e. a decode step) there is not
+// enough occupancy to hide that chain and the kernel becomes long-scoreboard bound. Staging the
+// metadata in shared memory once per block turns the k expanded-row loads into independent
+// accesses that can all be in flight at the same time.
+constexpr static int FINALIZE_MAX_STAGED_TOPK = FINALIZE_THREADS_PER_BLOCK;
+
 // Final kernel to unpermute and scale
 // This kernel unpermutes the original data, does the k-way reduction and performs the final skip connection.
 template <typename OutputType, class GemmOutputType, class ScaleBiasType, ScaleMode SCALE_MODE>
@@ -1594,22 +1616,48 @@ __global__ void finalizeMoeRoutingKernel(const GemmOutputType* expanded_permuted
   asm volatile("griddepcontrol.wait;");
 #endif
 
+  __shared__ int s_permuted_row[STAGE_ROUTING ? FINALIZE_MAX_STAGED_TOPK : 1];
+  __shared__ int s_expert_id[STAGE_ROUTING ? FINALIZE_MAX_STAGED_TOPK : 1];
+  __shared__ float s_row_scale[STAGE_ROUTING ? FINALIZE_MAX_STAGED_TOPK : 1];
+  if constexpr (STAGE_ROUTING) {
+    for (int64_t k_idx = threadIdx.x; k_idx < experts_per_token; k_idx += FINALIZE_THREADS_PER_BLOCK) {
+      int64_t const k_offset = original_row * experts_per_token + k_idx;
+      int const expert_id = token_selected_experts[k_offset] - start_expert_id;
+      bool const is_local_expert = expert_id >= 0 && expert_id < num_experts_per_node;
+      s_expert_id[k_idx] = expert_id;
+      // -1 marks a slot routed to an expert owned by another rank, which is skipped below.
+      s_permuted_row[k_idx] = is_local_expert ? unpermuted_row_to_permuted_row[original_row + k_idx * num_rows] : -1;
+      s_row_scale[k_idx] = (SCALE_MODE == ScaleMode::NO_SCALE) ? 1.f : scales[k_offset];
+    }
+    __syncthreads();
+  }
+
 #pragma unroll
   for (int elem_index = start_offset; elem_index < num_elems_in_col; elem_index += stride) {
     ComputeElem thread_output;
     thread_output.fill(0);
+#pragma unroll 4
     for (int k_idx = 0; k_idx < experts_per_token; ++k_idx) {
-      const int64_t k_offset = original_row * experts_per_token + k_idx;
-      const int64_t expert_id = token_selected_experts[k_offset] - start_expert_id;
-      if (expert_id < 0 || expert_id >= num_experts_per_node) {
-        continue;
+      int64_t expanded_permuted_row;
+      int64_t expert_id;
+      float row_scale;
+      if constexpr (STAGE_ROUTING) {
+        expanded_permuted_row = s_permuted_row[k_idx];
+        expert_id = s_expert_id[k_idx];
+        row_scale = s_row_scale[k_idx];
+        if (expanded_permuted_row < 0) {
+          continue;
+        }
+      } else {
+        const int64_t k_offset = original_row * experts_per_token + k_idx;
+        expert_id = token_selected_experts[k_offset] - start_expert_id;
+        if (expert_id < 0 || expert_id >= num_experts_per_node) {
+          continue;
+        }
+        expanded_permuted_row = unpermuted_row_to_permuted_row[original_row + k_idx * num_rows];
+        row_scale = (SCALE_MODE == ScaleMode::NO_SCALE) ? 1.f : scales[k_offset];
       }
 
-      const int64_t expanded_original_row = original_row + k_idx * num_rows;
-      const int64_t expanded_permuted_row = unpermuted_row_to_permuted_row[expanded_original_row];
-
-      const float row_scale = (SCALE_MODE == ScaleMode::NO_SCALE) ? 1.f : scales[k_offset];
-
       const auto* expanded_permuted_rows_row_ptr = expanded_permuted_rows_v + expanded_permuted_row * num_elems_in_col;
 
       ComputeElem expert_result = arrayConvert<InputElem, ComputeElem>(expanded_permuted_rows_row_ptr[elem_index]);
@@ -1858,10 +1906,17 @@ void finalizeMoeRoutingKernelLauncher(const GemmOutputType* expanded_permuted_ro
           break;
       }
 #undef LAUNCH_FINALIZE_ONE_ROW
+    } else if (experts_per_token <= FINALIZE_MAX_STAGED_TOPK) {
+      auto func = final_scales
+                      ? &finalizeMoeRoutingKernel<OutputType, GemmOutputType, ScaleBiasType, ScaleMode::DEFAULT, true>
+                      : &finalizeMoeRoutingKernel<OutputType, GemmOutputType, ScaleBiasType, ScaleMode::NO_SCALE, true>;
+      cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales,
+                         unpermuted_row_to_permuted_row, token_selected_experts, cols, experts_per_token, num_experts_per_node_int,
+                         start_expert_id);
     } else {
       auto func = final_scales
-                      ? &finalizeMoeRoutingKernel<OutputType, GemmOutputType, ScaleBiasType, ScaleMode::DEFAULT>
-                      : &finalizeMoeRoutingKernel<OutputType, GemmOutputType, ScaleBiasType, ScaleMode::NO_SCALE>;
+                      ? &finalizeMoeRoutingKernel<OutputType, GemmOutputType, ScaleBiasType, ScaleMode::DEFAULT, false>
+                      : &finalizeMoeRoutingKernel<OutputType, GemmOutputType, ScaleBiasType, ScaleMode::NO_SCALE, false>;
       cudaLaunchKernelEx(&config, func, expanded_permuted_rows, reduced_unpermuted_output, bias_ptr, final_scales,
                          unpermuted_row_to_permuted_row, token_selected_experts, cols, experts_per_token, num_experts_per_node_int,
                          start_expert_id);

onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu

@@ -256,6 +256,75 @@ __global__ void SoftmaxTopKWarpBitonicKernel(const T* logits, float* topk_scales
   }
 }
 
+// Warp-per-row softmax + top-k for expert counts in (64, kWarpSize * kItemsPerLane].
+// Lane `l` owns experts l, l + 32, l + 64, ... (strided so the logit loads stay
+// coalesced), and the whole row - softmax reduction and top-k selection alike - is
+// resolved with warp shuffles in registers: no shared memory, no __syncthreads and
+// no block-wide merge sort.
+//
+// This exists because the MoE decode shapes have very few rows, so the grid is a
+// handful of blocks and the kernel is bound by its dependency chain rather than by
+// throughput. SoftmaxTopKMergeKernel fully sorts all `num_experts` logits through
+// cub::BlockMergeSort, whose per-round MergePath searches and barriers dominate that
+// chain, while the top-k butterfly here costs a fixed ~5 shuffle rounds. Requires
+// k <= kItemsPerLane; tie-breaking matches SoftmaxTopKMergeKernel via the same
+// packed stable sort key.
+template <typename T, int kItemsPerLane, int kRowsPerBlock>
+__global__ void SoftmaxTopKWarpTopKKernel(const T* logits, float* topk_scales, int* topk_indices,
+                                          int num_rows, int num_experts, int k, bool normalize_scales) {
+  constexpr int kWarpSize = onnxruntime::cuda::topk::kWarpSize;
+  const int row = blockIdx.x * kRowsPerBlock + static_cast<int>(threadIdx.y);
+  if (row >= num_rows) return;
+
+  const int lane = threadIdx.x;
+  const T* row_logits = logits + static_cast<size_t>(row) * num_experts;
+
+  // Padding lanes carry (-inf, INT_MAX) so that valid -inf expert scores sort ahead.
+  uint64_t keys[kItemsPerLane];
+  float local_max = onnxruntime::cuda::topk::kNegativeInfinity;
+#pragma unroll
+  for (int j = 0; j < kItemsPerLane; ++j) {
+    const int idx = j * kWarpSize + lane;
+    const float logit = (idx < num_experts) ? static_cast<float>(row_logits[idx])
+                                            : onnxruntime::cuda::topk::kNegativeInfinity;
+    keys[j] = onnxruntime::cuda::topk::PackStableSortKey(logit, (idx < num_experts) ? idx : INT_MAX);
+    local_max = fmaxf(local_max, logit);
+  }
+
+  // Softmax denominator over all experts (needed when normalize_scales is false;
+  // when true it cancels in the top-k normalization but is still correct).
+  const float max_val = WarpReduceMax(local_max);
+  float local_sum = 0.0f;
+#pragma unroll
+  for (int j = 0; j < kItemsPerLane; ++j) {
+    if (j * kWarpSize + lane < num_experts) {
+      local_sum += expf(onnxruntime::cuda::topk::UnpackStableSortScore(keys[j]) - max_val);
+    }
+  }
+  const float inv_sum = SafeInvSum(WarpReduceSum(local_sum));
+
+  // Every lane ends up holding the row's top kItemsPerLane keys, sorted descending.
+  onnxruntime::cuda::topk::WarpBitonicTopN<kItemsPerLane>(keys);
+
+  float probs[kItemsPerLane];
+  float scale_sum = 0.0f;
+#pragma unroll
+  for (int j = 0; j < kItemsPerLane; ++j) {
+    probs[j] = SoftmaxScale(onnxruntime::cuda::topk::UnpackStableSortScore(keys[j]), max_val, inv_sum);
+    if (j < k) scale_sum += probs[j];
+  }
+  const float denom = TopKNormalizeDenom(normalize_scales, scale_sum);
+
+  // The results are replicated across the warp, so lane j simply writes rank j.
+#pragma unroll
+  for (int j = 0; j < kItemsPerLane; ++j) {
+    if (lane == j && j < k) {
+      topk_scales[static_cast<size_t>(row) * k + j] = normalize_scales ? (probs[j] / denom) : probs[j];
+      topk_indices[static_cast<size_t>(row) * k + j] = onnxruntime::cuda::topk::UnpackStableSortIndex(keys[j]);
+    }
+  }
+}
+
 // Warp CUB merge sort softmax + top-k for num_experts <= kBufferSize (64). One
 // warp (32 threads) per block sorts a row's logits held in shared memory. This
 // is the genai-recommended path for sort sizes in (32, 64]. Tie-breaking
@@ -344,6 +413,22 @@ void DispatchSoftmaxTopK(const T* logits, float* topk_scales, int* topk_indices,
       SoftmaxTopKWarpMergeKernel<T, 64><<<grid, onnxruntime::cuda::topk::kWarpSize, 0, stream>>>(
           logits, topk_scales, topk_indices, num_rows, num_experts, k, normalize_scales);
       return;
+    } else if (num_experts <= 256 && k <= 8) {
+      // Warp-per-row register top-k. Pick the smallest per-lane capacity that both
+      // covers num_experts across the warp and can hold the requested k.
+      constexpr int kRowsPerBlock = 4;
+      const dim3 warp_block(static_cast<unsigned>(onnxruntime::cuda::topk::kWarpSize), kRowsPerBlock);
+      const dim3 warp_grid(static_cast<unsigned>((num_rows + kRowsPerBlock - 1) / kRowsPerBlock));
+      const int items_per_lane =
+          std::max(k, (num_experts + onnxruntime::cuda::topk::kWarpSize - 1) / onnxruntime::cuda::topk::kWarpSize);
+      if (items_per_lane <= 4) {
+        SoftmaxTopKWarpTopKKernel<T, 4, kRowsPerBlock><<<warp_grid, warp_block, 0, stream>>>(
+            logits, topk_scales, topk_indices, num_rows, num_experts, k, normalize_scales);
+      } else {
+        SoftmaxTopKWarpTopKKernel<T, 8, kRowsPerBlock><<<warp_grid, warp_block, 0, stream>>>(
+            logits, topk_scales, topk_indices, num_rows, num_experts, k, normalize_scales);
+      }
+      return;
     } else if (num_experts <= 128) {
       SoftmaxTopKMergeKernel<T, 128, 1><<<grid, 128, 0, stream>>>(
           logits, topk_scales, topk_indices, num_rows, num_experts, k, normalize_scales);

onnxruntime/core/providers/cuda/cu_inc/topk_warp_sort.cuh

@@ -163,6 +163,95 @@
   }
 };
 
+/**
+ * @brief Sort `keys[0..N)` descending with an in-register bitonic network.
+ *
+ * N must be a power of two. The network is fully unrolled and touches no memory,
+ * so the whole sort stays in registers.
+ */
+template <int N>
+__device__ __forceinline__ void BitonicSortRegistersDescending(uint64_t (&keys)[N]) {
+  static_assert(N > 0 && (N & (N - 1)) == 0, "N must be a power of two.");
+#pragma unroll
+  for (int k = 2; k <= N; k <<= 1) {
+#pragma unroll
+    for (int j = k >> 1; j > 0; j >>= 1) {
+#pragma unroll
+      for (int i = 0; i < N; ++i) {
+        const int partner = i ^ j;
+        if (partner > i) {
+          const uint64_t a = keys[i];
+          const uint64_t b = keys[partner];
+          const bool a_first = ((i & k) == 0) == (a > b);
+          keys[i] = a_first ? a : b;
+          keys[partner] = a_first ? b : a;
+        }
+      }
+    }
+  }
+}
+
+/**
+ * @brief Sort a bitonic sequence `keys[0..N)` descending (half-cleaner cascade).
+ *
+ * N must be a power of two and the input must already be bitonic.
+ */
+template <int N>
+__device__ __forceinline__ void BitonicCleanDescending(uint64_t (&keys)[N]) {
+  static_assert(N > 0 && (N & (N - 1)) == 0, "N must be a power of two.");
+#pragma unroll
+  for (int j = N >> 1; j > 0; j >>= 1) {
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      const int partner = i ^ j;
+      if (partner > i) {
+        const uint64_t a = keys[i];
+        const uint64_t b = keys[partner];
+        keys[i] = a > b ? a : b;
+        keys[partner] = a > b ? b : a;
+      }
+    }
+  }
+}
+
+/**
+ * @brief Single-warp Top-N over the warp's 32 * N packed keys, entirely in registers.
+ *
+ * Every lane contributes N packed keys (see PackStableSortKey; the packing makes a
+ * plain unsigned comparison equivalent to "larger score first, smaller index on a
+ * tie"). On return every lane holds the N largest keys of the whole warp in
+ * `keys[0..N)`, sorted descending; a caller that wants the top k <= N simply reads
+ * the first k entries. Which N of the 32 * N inputs a lane starts with is
+ * irrelevant, so callers should pick the layout that loads best (e.g. strided by
+ * warp size for coalescing).
+ *
+ * The algorithm is a bitonic sort of each lane's own keys followed by a butterfly
+ * of bitonic Top-N merges over the warp. Merging two descending sequences A and B
+ * as C[i] = max(A[i], B[N-1-i]) yields a bitonic sequence holding exactly the top N
+ * of A union B, which the half-cleaner cascade then sorts. Because that construction
+ * is symmetric in A and B, a single __shfl_xor_sync per key leaves both partners
+ * with the same result, so no shared memory and no barriers are needed.
+ */
+template <int N>
+__device__ __forceinline__ void WarpBitonicTopN(uint64_t (&keys)[N]) {
+  BitonicSortRegistersDescending<N>(keys);
+#pragma unroll
+  for (int step = 1; step < kWarpSize; step <<= 1) {
+    uint64_t partner[N];
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      partner[i] = static_cast<uint64_t>(
+          __shfl_xor_sync(0xFFFFFFFFu, static_cast<unsigned long long>(keys[i]), step));
+    }
+#pragma unroll
+    for (int i = 0; i < N; ++i) {
+      const uint64_t other = partner[N - 1 - i];
+      keys[i] = keys[i] > other ? keys[i] : other;
+    }
+    BitonicCleanDescending<N>(keys);
+  }
+}
+
 /**
  * @brief Single-warp CUB merge sort of up to `BufferSize` (score, index) pairs
  *        held in shared memory, producing a descending order.