[CUDA] Cut the MoE routing prologue/epilogue cost for small-batch decode - #31156
Open
tianleiwu wants to merge 2 commits into
Open
[CUDA] Cut the MoE routing prologue/epilogue cost for small-batch decode#31156tianleiwu wants to merge 2 commits into
tianleiwu wants to merge 2 commits into
Conversation
tianleiwu
force-pushed
the
tlwu/20260730/moe_routing_decode_cost
branch
from
July 30, 2026 09:53
cdd88fa to
72cff94
Compare
The MoE routing kernels run once per layer with a grid of one to four blocks, so they are bound by their own dependency chains rather than by throughput. On the Qwen3.6-35B-A3B-NVFP4 MTP decode shape (4 rows, 256 experts, top-8, hidden 2048) the four of them cost 0.656 ms/step of the 6.033 ms/step of kernel time, with Nsight Compute reporting 0.07-0.4% of both compute and memory throughput: the SMs were essentially idle waiting on dependencies. Three changes, each aimed at the chain that a profile identified: fusedBuildExpertMapsSortFirstTokenKernel used cub::BlockRadixRank, which keeps a private set of packed digit counters per thread. Its shared memory and its per-thread scan both scale with 2^LOG2_NUM_EXPERTS, and the block size is picked from the token count, so 256 experts on a 4-token decode meant 512 bins spread over a single warp: 16KB of shared memory, 255 registers per thread, local memory spilling, and 1784 instructions to rank 32 assignments. Switching to cub::BlockRadixRankMatch, which ranks with warp ballots instead of per-thread counters, drops that to 2.2KB, 64 registers, no spilling and 787 instructions. 4.16 -> 2.22 us. SoftmaxTopK for 64 < num_experts <= 256 fully sorted every logit through cub::BlockMergeSort to read the first k, and the per-round MergePath searches and barriers dominated the chain. SoftmaxTopKWarpTopKKernel instead keeps one row per warp entirely in registers: each lane bitonic-sorts its own keys and the warp runs a butterfly of bitonic top-N merges. Because merging two descending sequences as C[i] = max(A[i], B[N-1-i]) is symmetric in A and B, a single __shfl_xor_sync per key leaves both partners with the same result, so there is no shared memory and no barrier anywhere in the kernel. The new BitonicSortRegistersDescending / BitonicCleanDescending / WarpBitonicTopN primitives live next to the existing warp sort helpers. Tie-breaking is unchanged (same packed stable sort key). 4.52 -> 3.68 us. finalizeMoeRoutingKernel re-read the routing metadata from global memory for every column tile through a chain of two dependent loads, which left it long-scoreboard bound at 7.5 stall cycles per warp. It now stages the metadata in shared memory once per block, as the num_rows == 1 specialisation already did, so the k expanded-row loads are independent and can all be in flight. The staging is a template parameter so the k > 256 fallback keeps the original path without a branch in the hot loop. 4.64 -> 4.18 us. Routing family 0.656 -> 0.529 ms/step, all kernels 6.033 -> 5.893 ms/step. Interleaved end-to-end A/B over five reps: 8.511/8.535/8.501/8.541/8.566 -> 8.467/8.382/8.370/8.374/8.445, i.e. 8.531 -> 8.408 ms/step (-1.4%). Tests: onnxruntime_provider_test '*QMoE*:*MoE*:*TopK*' 72 passed; test_qmoe_nvfp4_cuda.py 22 passed; test_moe_cuda.py 68 passed with the two TestSparseMixer failures reproducing identically on the parent commit.
tianleiwu
force-pushed
the
tlwu/20260730/moe_routing_decode_cost
branch
from
July 30, 2026 21:06
72cff94 to
0b22d02
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
This PR optimizes CUDA Mixture-of-Experts (MoE) routing overhead in small-batch decode scenarios by replacing heavier sort/rank primitives with warp-register top-k and lighter ranking, and by staging routing metadata to reduce dependent global-memory reads in the routing epilogue.
Changes:
- Added a register-only warp bitonic Top-N primitive (
WarpBitonicTopN) totopk_warp_sort.cuh. - Introduced a new warp-per-row softmax + top-k kernel for
num_experts <= 256 && k <= 8inqmoe_kernels.cu, avoiding shared memory and block merge-sort overhead. - Swapped MoE prologue ranking to a match-based radix ranker and added shared-memory staging of routing metadata in
finalizeMoeRoutingKernelto reduce latency at low occupancy.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| onnxruntime/core/providers/cuda/cu_inc/topk_warp_sort.cuh | Adds register-level bitonic sort/clean and a warp Top-N merge primitive used for fast top-k selection. |
| onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu | Adds a warp-per-row softmax + top-k path for (64, 256] experts and small k, dispatching to the new warp Top-N routine. |
| onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu | Replaces block radix rank with a match-based variant and stages routing metadata in shared memory for faster finalize routing at decode. |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
hariharans29
approved these changes
Jul 31, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Cuts the cost of the MoE routing prologue and epilogue at small batch, where those kernels stop
being negligible relative to the expert GEMMs. Touches
moe_gemm/moe_kernels.cu,moe/qmoe_kernels.cuandcore/providers/cuda/cu_inc/topk_warp_sort.cuh.Motivation and Context
At decode batch sizes the routing kernels run on a handful of blocks (grid 1 and grid 4 in the
table below), so they are latency-bound, not throughput-bound, and their fixed cost is a
measurable share of the step. Once the expert GEMVs themselves were optimized, the routing
family became the next item on the profile.
BlockRadixRankkeeps a private set of packed digit counters per thread, which is what makesthe sort-based path expensive at these sizes; the replacement is a warp-level top-k in
core/providers/cuda/cu_inc/topk_warp_sort.cuh. Tie-breaking is unchanged — the same packedstable key is used, so routing decisions are identical.
Measured impact
1x H200 SXM (SM90, 132 SM, ~4.8 TB/s HBM), CUDA 13.0, Qwen3.6-35B-A3B-NVFP4 + MTP
N=3(verify batch
M=4).Per kernel (graph OFF):
finalizeMoeRoutingKernel(grid 4)SoftmaxTopKMergeKernel->WarpTopKKernelfusedBuildExpertMapsSortFirstTokenKernel(grid 1)expandInputRowsKernel(grid 1056)All kernels (both families): 6.033 -> 5.893 ms/step (−0.140 ms/step).
End-to-end (graph ON, 5 interleaved reps per arm):
8.531 -> 8.408 ms/step (−1.4%).
Left on the table (~0.05 ms/step, judged low ROI): 32-bit packed keys in
WarpBitonicTopN, andsplitting
finalizeMoeRouting's columns across more blocks to use more than 4 SMs' L1 bandwidth.Tests
Existing MoE/QMoE unit tests; routing output is unchanged by construction (same stable key), so
these are parity tests rather than new coverage.
Methodology note
End-to-end deltas are quoted as ms/step from a fixed-step measurement, never tok/s: any
numerics change alters the generated sequence and therefore the MTP acceptance rate, which swamps
the speed delta. Per-kernel durations are taken with CUDA graphs off —
nsys --cuda-graph-trace=nodeinflates durations ~35% globally and up to 3.8x for large-gridkernels.
Note
Rebased onto
mainafter #29919 ("Fuse MoE softmax/top-k into the expert-map prologue atdecode") merged. The two are complementary, not competing — no textual or semantic
conflict:
fusedRoutingBuildExpertMapsSingleTokenKernelis gated byisFusedMoeRoutingSupported(), which requiresnum_tokens == 1,ep_size == 1andnum_experts <= kWarpBitonicMaxSize(32). Qwen3.6-35B-A3B has 256 experts, and an MTPverify step has 4 tokens, so that path never engages here.
fusedBuildExpertMapsSortFirstTokenKernel(the general prologue) still usescub::BlockRadixRankonmain— this PR'sBlockRadixRankMatchswap applies to it.num_experts <= 256dispatch still lands onSoftmaxTopKMergeKernel<T, 128, 2>— thisPR's
SoftmaxTopKWarpTopKKernelreplaces it fornum_experts <= 256 && k <= 8.finalizeMoeRoutingKernelwas untouched by [CUDA] Fuse MoE softmax/top-k into the expert-map prologue at decode #29919.#29919 also moved
WarpReduceMax/WarpReduceSum/SafeInvSum/TopKNormalizeDenomintotopk_warp_sort.cuhand re-exported them intoqmoe_kernels.cuwithusingdeclarations, sothe unqualified calls added here still resolve.