onnxruntime
a3c3cf98 - [CUDA][QMoE] Wire LaunchFusedRouteMaps into CutlassMoeFCRunner for S<=8 and E<=512 (#33596)

Commit
3 days ago
[CUDA][QMoE] Wire LaunchFusedRouteMaps into CutlassMoeFCRunner for S<=8 and E<=512 (#33596) ### 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). --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading