diff --git a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu b/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu index dd19f71b713e6..e703ad696a675 100644 --- a/onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu +++ b/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 +using MoeExpertRadixRank = cub::BlockRadixRankMatch; + template __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; + using BlockRadixRank = MoeExpertRadixRank; extern __shared__ unsigned char temp_storage[]; auto& sort_temp = *reinterpret_cast(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; + using BlockRadixRank = MoeExpertRadixRank; size_t shared_size = sizeof(typename BlockRadixRank::TempStorage); cudaLaunchConfig_t config; @@ -1562,9 +1575,18 @@ 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 +template __global__ void finalizeMoeRoutingKernel(const GemmOutputType* expanded_permuted_rows, OutputType* reduced_unpermuted_output, const ScaleBiasType* bias, const float* scales, const int* unpermuted_row_to_permuted_row, const int* token_selected_experts, const int64_t orig_cols, @@ -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(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 + : &finalizeMoeRoutingKernel; + 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 - : &finalizeMoeRoutingKernel; + ? &finalizeMoeRoutingKernel + : &finalizeMoeRoutingKernel; 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); diff --git a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu b/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu index 8a1993c53a05e..61b06cbeb42ea 100644 --- a/onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu +++ b/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 +__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(threadIdx.y); + if (row >= num_rows) return; + + const int lane = threadIdx.x; + const T* row_logits = logits + static_cast(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(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(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(row) * k + j] = normalize_scales ? (probs[j] / denom) : probs[j]; + topk_indices[static_cast(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<<>>( 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(onnxruntime::cuda::topk::kWarpSize), kRowsPerBlock); + const dim3 warp_grid(static_cast((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<<>>( + logits, topk_scales, topk_indices, num_rows, num_experts, k, normalize_scales); + } else { + SoftmaxTopKWarpTopKKernel<<>>( + logits, topk_scales, topk_indices, num_rows, num_experts, k, normalize_scales); + } + return; } else if (num_experts <= 128) { SoftmaxTopKMergeKernel<<>>( logits, topk_scales, topk_indices, num_rows, num_experts, k, normalize_scales); diff --git a/onnxruntime/core/providers/cuda/cu_inc/topk_warp_sort.cuh b/onnxruntime/core/providers/cuda/cu_inc/topk_warp_sort.cuh index 6c5567ea05046..814e539d0b5a8 100644 --- a/onnxruntime/core/providers/cuda/cu_inc/topk_warp_sort.cuh +++ b/onnxruntime/core/providers/cuda/cu_inc/topk_warp_sort.cuh @@ -163,6 +163,95 @@ struct Greater { } }; +/** + * @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 +__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 +__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 +__device__ __forceinline__ void WarpBitonicTopN(uint64_t (&keys)[N]) { + BitonicSortRegistersDescending(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( + __shfl_xor_sync(0xFFFFFFFFu, static_cast(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(keys); + } +} + /** * @brief Single-warp CUB merge sort of up to `BufferSize` (score, index) pairs * held in shared memory, producing a descending order.