SonicMoE : Accélération des MoE par des optimisations tenant compte des E/S et des tuiles

SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations

December 16, 2025

papers.authors: Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao

cs.AI

papers.abstract

Les modèles Mixture of Experts (MoE) sont devenus l'architecture de facto pour augmenter l'échelle des modèles de langage sans accroître significativement le coût computationnel. Les modèles MoE récents montrent une tendance nette vers une granularité d'experts élevée (dimension intermédiaire des experts plus réduite) et une plus grande parcimonie (nombre constant d'experts activés avec un nombre total d'experts plus important), ce qui améliore la qualité du modèle par FLOP. Cependant, les MoE à granularité fine souffrent d'une empreinte mémoire d'activation accrue et d'une efficacité matérielle réduite en raison de coûts d'E/S plus élevés, tandis que les MoE plus parcimonieux souffrent de calculs gaspillés dus au remplissage (padding) dans les noyaux Grouped GEMM. En réponse, nous proposons un algorithme efficace en mémoire pour calculer les passes avant et arrière des MoE avec une mise en cache minimale des activations pour la passe arrière. Nous concevons également des noyaux GPU qui chevauchent les E/S mémoire avec le calcul, profitant à toutes les architectures MoE. Enfin, nous proposons une nouvelle méthode d'« arrondi des tokens » (token rounding) qui minimise le calcul gaspillé dû au remplissage dans les noyaux Grouped GEMM. En conséquence, notre méthode SonicMoE réduit la mémoire d'activation de 45 % et atteint une amélioration du débit de calcul de 1,86x sur les GPU Hopper par rapport au noyau MoE BF16 de ScatterMoE pour un MoE à granularité fine de 7B. Concrètement, SonicMoE sur 64 H100 atteint un débit d'entraînement de 213 milliards de tokens par jour, comparable aux 225 milliards de tokens par jour de ScatterMoE sur 96 H100 pour l'entraînement d'un modèle MoE 7B avec FSDP-2 en utilisant la base de code lm-engine. Dans des configurations de haute parcimonie MoE, notre algorithme d'arrondi des tokens tenant compte des tuiles (tile-aware) procure une accélération supplémentaire de 1,16x sur le temps d'exécution du noyau par rapport au routage top-K classique, tout en maintenant des performances en aval similaires. Nous ouvrons en libre accès tous nos noyaux pour permettre un entraînement plus rapide des modèles MoE.

English

Mixture of Experts (MoE) models have emerged as the de facto architecture for scaling up language models without significantly increasing the computational cost. Recent MoE models demonstrate a clear trend towards high expert granularity (smaller expert intermediate dimension) and higher sparsity (constant number of activated experts with higher number of total experts), which improve model quality per FLOP. However, fine-grained MoEs suffer from increased activation memory footprint and reduced hardware efficiency due to higher IO costs, while sparser MoEs suffer from wasted computations due to padding in Grouped GEMM kernels. In response, we propose a memory-efficient algorithm to compute the forward and backward passes of MoEs with minimal activation caching for the backward pass. We also design GPU kernels that overlap memory IO with computation benefiting all MoE architectures. Finally, we propose a novel "token rounding" method that minimizes the wasted compute due to padding in Grouped GEMM kernels. As a result, our method SonicMoE reduces activation memory by 45% and achieves a 1.86x compute throughput improvement on Hopper GPUs compared to ScatterMoE's BF16 MoE kernel for a fine-grained 7B MoE. Concretely, SonicMoE on 64 H100s achieves a training throughput of 213 billion tokens per day comparable to ScatterMoE's 225 billion tokens per day on 96 H100s for a 7B MoE model training with FSDP-2 using the lm-engine codebase. Under high MoE sparsity settings, our tile-aware token rounding algorithm yields an additional 1.16x speedup on kernel execution time compared to vanilla top-K routing while maintaining similar downstream performance. We open-source all our kernels to enable faster MoE model training.