[CUDA] Fuse MoE softmax/top-k into the expert-map prologue at decode

[Tianlei Wu (tianleiwu)]

Description

At decode (one token per sequence) the MoE/QMoE path launches a standalone softmax/top-k kernel and
then a second single-block kernel that builds the expert permutation maps. Both are one-block
kernels whose cost is launch and drain latency rather than work, so the second launch is nearly pure
overhead. This PR merges them.

fusedRoutingBuildExpertMapsSingleTokenKernel handles the single-token case with one warp: lane e
owns expert e's logit, warp reductions give the softmax, and WarpBitonicSortDescending gives the
top-k. For a single token the permutation is just "the k selected experts ordered by expert id",
which every lane derives with shuffles, so no CUB block rank and no shared memory are needed. The
kernel writes token_selected_experts, token_final_scales and all three permutation maps.

Summary of Changes

• Fused routing prologue. launchFusedRoutingBuildExpertMaps replaces the
  LaunchSoftmaxTopK + fusedBuildExpertMapsSortFirstToken pair on the supported decode shape.
• Shared softmax/top-k math. The routing arithmetic must stay bit-identical to
  SoftmaxTopKWarpBitonicKernel, so SoftmaxScale, SafeInvSum, TopKNormalizeDenom,
  WarpReduceMax and WarpReduceSum move from qmoe_kernels.cu into topk_warp_sort.cuh, and
  both translation units now share one definition.
• Single source of truth for the guard. isFusedMoeRoutingSupported() is host-only and
  deterministic. The QMoE op uses it to decide whether to skip its own softmax/top-k launch, and
  runMoe re-checks it, so the two can never disagree about who computed the routing. It requires
  one token, no expert parallelism, and num_experts <= 32. The FP4 family is excluded in the op
  because its decode fast path consumes expert_indices/expert_scales itself instead of calling
  runMoe.
• Kill switch: ORT_DISABLE_MOE_FUSED_ROUTING=1 restores the two-kernel prologue.

Performance

The win scales with how large a launch is relative to the rest of the call, so it is much bigger on
a smaller GPU. Isolated QMoE node at the gpt-oss-20b decode shape (hidden 2880, intermediate 2880,
32 experts, top-4, batch 1, fp16, int4), arms interleaved within every pass via the kill switch:

RTX 3060 (sm_86), per-channel int4, 12 passes x 1000 reps

arm	mean	min	vs fused
fused routing	255.6 us	250.1 us	—
ORT_DISABLE_MOE_FUSED_ROUTING=1	277.6 us	262.7 us	+8.6%

Fused wins 12 of 12 paired passes (median +7.0%).

RTX 3060 (sm_86), block-wise int4 (block_size=64), 8 passes x 1000 reps

arm	mean	vs fused
fused routing	263.7 us	—
ORT_DISABLE_MOE_FUSED_ROUTING=1	277.7 us	+5.3%

Fused wins 8 of 8 paired passes (median +5.0%).

On an RTX 5060 Ti (sm_120) the same change measured +0.24% end-to-end decode throughput on
gpt-oss-20b — inside that harness's noise floor. A tripwire build confirmed the fused path really
does engage at decode there, so the sm_120 number is a genuine null result rather than a dead path;
the faster GPU simply amortizes the extra launch better.

Testing

On the RTX 3060 (sm_86), with and without ORT_DISABLE_MOE_FUSED_ROUTING=1:

• onnxruntime/test/python/transformers/test_qmoe_cuda.py: 109 passed, 13 skipped
• onnxruntime/test/python/transformers/test_moe_cuda.py: 48 passed, 15 skipped

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

This PR reduces decode-time overhead in the CUDA MoE/QMoE path by fusing the standalone softmax/top-k routing step into the expert-map prologue when the configuration is a single token with ≤32 experts (and no expert parallelism), with an environment-variable kill switch to revert to the prior two-kernel behavior.

Changes:

• Adds a fused single-warp “routing + expert-map build” prologue for the single-token case and wires it through runMoe via FusedRoutingParams.
• Moves softmax/top-k helper math into topk_warp_sort.cuh so the standalone and fused routing paths share bit-identical implementations.
• Updates MoE/QMoE call sites to optionally skip LaunchSoftmaxTopK and instead pass raw router logits into the runner for fused routing.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
onnxruntime/core/providers/cuda/cu_inc/topk_warp_sort.cuh	Centralizes shared softmax/top-k helper math in a CUDA header for bit-identical reuse.
onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu	Removes duplicated helper implementations and reuses the shared helpers from topk_warp_sort.cuh.
onnxruntime/contrib_ops/cuda/moe/moe.cc	Passes default (disabled) fused-routing parameters into runMoe for the non-quantized MoE path.
onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc	Adds gating logic to skip standalone softmax/top-k at supported decode shapes and passes FusedRoutingParams into runMoe.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.h	Extends the runner interface with FusedRoutingParams and declares isFusedMoeRoutingSupported().
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu	Implements the fused routing kernel, kill switch, support predicate, and integrates the fused path into runMoe.