[Tianlei Wu (tianleiwu)]

Description

• Extend isFusedMoeRoutingSupported in onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu to query onnxruntime::cuda::CanUseFusedRouteMaps(num_tokens, num_experts, experts_per_token) when (num_tokens == 1 && num_experts <= 32) does not apply.
• Update launchFusedRoutingBuildExpertMaps to dispatch to onnxruntime::cuda::LaunchFusedRouteMaps<T> (FusedSoftmaxTopKMapsKernel) for half and __nv_bfloat16 router logits when num_tokens > 1 or num_experts > 32, while preserving fusedRoutingBuildExpertMapsSingleTokenKernel for num_tokens == 1 && num_experts <= 32.

Motivation and Context

• FusedSoftmaxTopKMapsKernel (LaunchFusedRouteMaps in qmoe_fused_route.cuh) was already used by the INT2/INT3 fast path in moe_quantization.cc for 1 <= num_tokens <= 8, 64 <= num_experts <= 512, experts_per_token <= 16, but isFusedMoeRoutingSupported in moe_kernels.cu only allowed num_tokens == 1 && num_experts <= 32.
• Consequently, all INT4 QMoE layers with num_experts = 512 and speculative verification batches (S = 8) fell back to two separate kernel launches per layer (SoftmaxTopKMergeKernel + fusedBuildExpertMapsSortFirstToken, ~9.0 µs/layer).
• Wiring LaunchFusedRouteMaps into CutlassMoeFCRunner replaces both launches with a single-CTA FusedSoftmaxTopKMapsKernel (4.31 µs/layer on H200 at S = 8, E = 512, K = 10), eliminating 1 kernel launch and ~4.7 µs per INT4 MoE layer (48 launches / ~0.22 ms per 48-layer decode cycle).

[Copilot]

🟡 Changes recommended

The eligibility helper is undefined, the launcher call is incompatible with the existing API, and the new path lacks regression coverage.

3 open issues

• Declare and namespace CanUseFusedRouteMaps eligibility predicate · New
• Use declared routing API and handle fused-launcher fallback · New
• Add regression tests for fused routing and fallback paths · New

What changed in this PR

This PR aims to reduce CUDA INT4 QMoE routing overhead by integrating fused routing and expert-map generation into CutlassMoeFCRunner for small token batches and larger expert counts.

Changes:

• Extends fused-routing eligibility beyond the existing single-token, ≤32-expert path.
• Adds FP16/BF16 dispatch to LaunchFusedRouteMaps, while preserving the existing small-expert kernel.
• Passes the token count into the routing launcher.

File	Description
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_kernels.cu	Extends routing eligibility and connects the runner to fused route-map generation.

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Architecture-dependent CUDA dispatch and fallback behavior require confirmed build and GPU validation before final approval.

No open issues 🎉

3 resolved since last review

• Use declared routing API and handle fused-launcher fallback
• Declare and namespace CanUseFusedRouteMaps eligibility predicate
• Add regression tests for fused routing and fallback paths

_{Review effort: Balanced}

[kunal-vaishnavi]

From Copilot:

Completed an implementation pass and a second edge-case/regression pass at 3290281. No actionable finding.

Checked architecture/alignment rejection, float-logit fallback, expert-map ordering and permutation conventions, routing-weight normalization, expert-parallel eligibility, and routing-snapshot timing. A declined fused launch explicitly computes softmax/top-k before the standard expert-map prologue, rather than leaving the caller's routing buffers uninitialized. Existing map/output-scale tests were inspected.

Validation: source and existing-test review only for this PR. No PR-specific ORT build, provider/Python suite execution, cross-architecture qualification, or performance benchmark was performed. This is a source-review outcome, not a formal approval or a claim that all tests and CI passed.