Cross-Platform Fused MoE Dispatch in Triton: Portable Expert Routing Without CUDA
On the map: Mixture-of-experts → Efficient MoE training and inference systems
Abstract (arXiv)
Mixture-of-Experts (MoE) architectures power the majority of frontier large language models, but their inference is bottlenecked by irregular memory access patterns and expert routing overhead. Existing optimized MoE kernels (Megablocks, Tutel, FasterMoE) are implemented in CUDA and locked to NVIDIA hardware. We present TritonMoE, a fused MoE dispatch kernel written entirely in OpenAI Triton that performs the complete forward pass -- router scoring, token permutation, expert GEMMs, and weighted output combination -- using only portable Triton primitives. Our key optimization is a fused gate+up GEMM kernel that computes both SwiGLU projections from shared L2-cached input tiles with in-register SiLU activation, eliminating 35% of global memory traffic. On an NVIDIA A100, TritonMoE achieves 89-131% of the throughput of the CUDA-optimized Megablocks at inference batch sizes (<= 512 tokens) across Mixtral-8x7B, DeepSeek-V3, and Qwen2-MoE configurations. All 162 correctness tests pass on both NVIDIA A100 and AMD MI300X with zero code changes, validating cross-platform portability. We additionally characterize sensitivity to routing imbalance under Zipfian-skewed expert assignments and identify the regime -- 64+ experts under extreme skew -- where our fixed-tile scheduling underperforms Megablocks' block-sparse layout, motivating dynamic block-to-expert assignment as future work. Code is available at https://github.com/bassrehab/triton-kernels.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Placed in Efficient MoE training and inference systems, one of 12 directions in Mixture-of-experts. Builds on 6 of its 13 references that are in the map. First posted on arXiv 7 Apr 2026.
Builds on
6 of 13 references.
Closest ancestors · 2
Same direction, newest first
- 2022MegaBlocks: Efficient Sparse Training with Mixture-of-Experts29 Nov
- 2022Tutel: Adaptive Mixture-of-Experts at Scale7 Jun
Other references · 4
- 2023Efficient Memory Management for Large Language Model Serving with PagedAttentioncited by 5,767
- 2024Qwen2 Technical Reportcited by 2,076
- 2024Mixtral of Expertscited by 1,785
- 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,601
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2026Decoding the Skew: Distribution-Aware MoE Inference with Adaptive Kernel Dispatch
- 2024ExpertFlow: Efficient Mixture-of-Experts Inference via Predictive Expert Caching and Token Scheduling
- 2025Accelerating MoE Model Inference with Expert Sharding
- 2025OD-MoE: On-Demand Expert Loading for Cacheless Edge-Distributed MoE Inference
- 2025ExpertFlow: Adaptive Expert Scheduling and Memory Coordination for Efficient MoE Inference
- 2026Every Expert Counts: ExactMoE for Memory-Efficient W4A16 Inference
- 2025Efficient MoE Serving in the Memory-Bound Regime: Balance Activated Experts, Not Tokens
- 2026Expert Coupling in MoE Pretraining: Reducing All-to-All Overhead with Correlated Placement and Token Shuffling