Repository navigation
perf(moe): fused triton router for sigmoid + selection-bias routing #730
Description
Problem
Five MoE families route with sigmoid scores and a selection-only e_score_correction_bias (DeepSeek-V3 style), and compute it as eager torch ops: glm4_moe, glm_moe_dsa, glm5_next (moe.py _route, three near-identical copies), minimax_m2 and minimax_m3. Each layer runs about 13 small kernels per decode step:
- the router weight is cast to fp32 every step
- an fp32 GEMV
- sigmoid, the bias add, optional group top-k,
torch.topk, gather, sum, divide, the routed scale and dtype casts
The softmax families (qwen3_moe, qwen3_5_moe, qwen4_exp, gpt_oss) already use the fused fused_topk_softmax kernel, and gemma4 uses gemma4_fused_routing, so they do not have this cost.
Measured
GLM-4.5-Air Q2_K GGUF, bs=1 decode, RTX PRO 6000 (sm_120), nsys:
| FreeToken | llama.cpp f2cc728 | |
|---|---|---|
| decode tok/s | 117.9 | 146.9 |
| routing, us per token | 1108 | 246 |
| routing kernels per token | 587 | 90 |
The expert GEMVs take the same time in both engines (2030 vs 2014 us per token). llama.cpp runs the router GEMV plus one fused topk_moe kernel (ggml-cuda/topk-moe.cu). MiniMax-M2 UD-TQ1_0 shows the same pattern: 1315 vs 405 us per token in routing, and 133 vs 152 tok/s.
Proposal
- One shared Triton kernel covering: sigmoid scores, selection bias, optional group-limited top-k (
n_group/topk_group), gathering the unbiased scores, renormalization, routed scale, fp32 outputs. - Graph-capturable, with the padded-row mask like
fused_topk_softmax. - Keep the router weight in fp32 once at load instead of casting it every step.
- Replace the five
_routeimplementations with it. - deepseek_v41's
router_select(kernel/triton/dsv4/router.py) already fuses sigmoid / sqrt-softplus + bias + top-k and could be the starting point. It would need group top-k and renormalization.
Expected
About 1.1-1.3 ms per token saved on GLM-4.5-Air at bs=1 (roughly 136-140 tok/s), and about parity with llama.cpp on MiniMax-M2. The NVFP4 HF checkpoints of the same families also benefit.
Out of scope here
The other gaps found in the same profile: launch gaps between kernels without PDL (~1 ms per token), separate q/k/v bias adds and q8_1 quantizes, and flashinfer's prefill kernel + MergeStates used for decode.
Activity
rainhjrui commentedon Oct 22, 2026
Hi, I've opened a prototype implementation as Draft PR #802 and would appreciate early feedback on its scope and validation approach.
The implementation introduces a shared Triton router for sigmoid-based MoE routing across the five model families mentioned in this issue. It also keeps gate weights resident in FP32 for the three GLM implementations; MiniMax M3 already uses FP32, while MiniMax M2 retains its existing gate precision.
After integrating upstream main at
546c5fe, the router-only benchmark on an RTX 4060 Laptop GPU shows a 3.17x speedup for one activation row and a 1.85x speedup for 93 rows. These measurements include the input cast and gate GEMM, and are not end-to-end model speedups. The draft includes a four-way ablation, raw results, and reproduction commands.The targeted routing tests pass (121 tests), and the selected regression suite passes on both the branch and the same upstream baseline (92 passed, 1 skipped each). Additional real-weight tie cases are preserved as regression fixtures; the documented tie-breaking differences remain a review point.
Full-model end-to-end validation is still missing due to local memory and CPU-offload compatibility limitations for the candidates examined. I have therefore kept the PR as a Draft and do not claim full-model performance or quality validation.
Would you prefer a smaller initial PR focused on the shared kernel and one model integration, or a single PR covering all five families? Would the kernel-level measurements and reproducible tests be sufficient for an initial Draft PR review, with full-model benchmarking performed later on suitable hardware?
Implementation, validation report, and raw evidence: #802
Thanks for your time!