Papers
arxiv:2605.23911

Cross-Platform Fused MoE Dispatch in Triton: Portable Expert Routing Without CUDA

Published on Apr 7
Authors:

Abstract

TritonMoE presents a portable, fused Mixture-of-Experts dispatch kernel written in Triton that achieves high throughput across different hardware platforms while maintaining correctness and reducing memory traffic through optimized GEMM operations.

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.

Community

Sign up or log in to comment

Get this paper in your agent:

hf papers read 2605.23911
Don't have the latest CLI?
curl -LsSf https://hf.co/cli/install.sh | bash

Models citing this paper 0

No model linking this paper

Cite arxiv.org/abs/2605.23911 in a model README.md to link it from this page.

Datasets citing this paper 0

No dataset linking this paper

Cite arxiv.org/abs/2605.23911 in a dataset README.md to link it from this page.

Spaces citing this paper 0

No Space linking this paper

Cite arxiv.org/abs/2605.23911 in a Space README.md to link it from this page.

Collections including this paper 0

No Collection including this paper

Add this paper to a collection to link it from this page.