[CUDA] Add MoERouter and SwiGLU contrib operators

[Tianlei Wu (tianleiwu)]

Stacked PR. Based on tlwu/20260817/sparse_kv_select (#32124) because it shares quant_sim_common.cuh. Please review that PR first; only the last commit here is new. I will retarget this to main once the base merges.

Description

Adds two CUDA contrib operators for the MoE block. MoERouter turns a gate GEMM's output into the two tensors an expert-parallel QMoE call needs, and SwiGLU exposes the clamped SwiGLU activation that MoE/QMoE already apply internally, for the dense and shared-expert paths that do not run through a grouped expert GEMM.

MoERouter is the piece that makes expert parallelism expressible in ONNX today: the existing QMoE operator takes router_probs and applies its own softmax, so a rank holding only a slice of the experts needs a router row constructed such that that softmax still produces the globally correct weights. Doing that with primitives requires a top-k, a scatter and a masked log per layer.

Summary of Changes

MoERouter

File	Change
onnxruntime/contrib_ops/cuda/math/moe_router.{h,cc}	OpKernel — attributes, validation, output shapes
onnxruntime/contrib_ops/cuda/math/moe_router_impl.{h,cu}	Scoring, selection and log-domain router-row kernels

Scoring and selection are separate, orthogonal attributes so the operator covers the common router designs rather than one model's:

scoring:   sqrt_softplus | softmax | sigmoid
selection: topk | noaux_tc   (noaux_tc adds `bias` before the top-k)

Ties go to the lower expert index. Supplying the optional expert_ids input overrides selection and fixes the choice per token (hash routing); bias is then ignored. Either way the weights are the affinities of the chosen experts, normalized to sum to one.

Expert parallelism is the reason for the two outputs:

• router_probs carries log(w_e) for a chosen local expert and a large negative value elsewhere (-1e30, or -1e4 for float16 which -1e30 does not survive), so QMoE's own softmax over the local block returns w_e / W_local.
• weight_scale is route_scale * W_local, which multiplies that factor back out of the expert output before the all-reduce. A token with no local expert gets a zero scale, which annihilates the degenerate uniform softmax an all-negative row would otherwise produce.

local_expert_start / local_expert_count describe the rank's slice; a single-rank model just sets them to the full expert range.

SwiGLU

File	Change
onnxruntime/contrib_ops/cuda/math/swiglu.{h,cc}	OpKernel
onnxruntime/contrib_ops/cuda/math/swiglu_impl.{h,cu}	Fused clamp + gate kernel

G      = clamp(gate, max=limit)
L      = clamp(up, min=-limit, max=limit)
output = G * Sigmoid(activation_alpha * G) * (L + activation_beta)

• Arithmetic is done in float regardless of T. A limit of zero or less disables both clamps.
• This is deliberately the same activation with the same attribute contract that MoE/QMoE apply internally via their swiglu_limit / activation_alpha / activation_beta attributes, so a model can use the same three values for its routed experts and its dense path.
• up is optional. When omitted, gate carries both halves of one [.., 2 * inter] projection (gate first, then up) and the split is done internally, so a fused sibling GEMM needs no Split node.

Why a standalone operator rather than folding the path into MoE/QMoE: the motivating case is a shared expert that is tensor-parallel sharded while the routed experts are expert-parallel sharded. It therefore has a different (here 8x larger) local intermediate size than a routed expert, is ungated, and runs on every rank — so it cannot be represented as one more always-on expert slot inside the grouped GEMM without putting dense per-token work on a single rank. Accepting TP-sharded shared-expert weights as separate QMoE inputs is a reasonable follow-up; this PR is the piece needed to express the path at all.

Shared

File	Change
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	Schemas, shape inference, formulas
onnxruntime/core/graph/contrib_ops/ms_opset.h	Schema registration
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	CUDA kernel registration for float, float16 and bfloat16
docs/ContribOperators.md, docs/OperatorKernels.md	Regenerated entries

Testing

• onnxruntime/test/python/transformers/test_moe_router_swiglu.py — PyTorch references for both operators over float32/float16/bfloat16: scoring/selection combinations, hash routing via expert_ids, the no-local-expert token, the float16 -1e4 sentinel, SwiGLU with and without up, with and without the clamp, and non-default activation_alpha / activation_beta.

python -m pytest onnxruntime/test/python/transformers/test_moe_router_swiglu.py
# 12 passed, 36 subtests passed

Additive only — no existing operator, schema or kernel is modified; MoE/QMoE are untouched.

Checklist

• Tests added
• No breaking changes
• Documentation updated

[Justin Chu (justinchuby)]

SwiGLU is also declared at https://onnx.ai/onnx/operators/onnx__SwiGLU.html. We could register for the standard domain too?

[kunal-vaishnavi]

This will be a very useful op for expert parallelism. Thanks for adding it! Should we consider expanding its usage for other types of routing besides top-k and hash routing?

[kunal-vaishnavi]

We already have a partial C++ graph fusion for SwiGLU that runs on existing models. Can we update it to work with the SwiGLU op?

[kunal-vaishnavi]

Can we move this into a shared utility for common CUDA kernel calculations?