SonicMoE: Versnelling van MoE met IO- en tegelbewuste optimalisaties

SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations

December 16, 2025

Auteurs: Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao

cs.AI

Samenvatting

Mixture of Experts (MoE)-modellen zijn naar voren gekomen als de de facto architectuur voor het opschalen van taalmodelmodellen zonder de rekenkosten significant te verhogen. Recente MoE-modellen tonen een duidelijke trend richting hoge expertgranulariteit (kleinere expert-intermediaire dimensie) en hogere sparseheid (constant aantal geactiveerde experts met een hoger totaal aantal experts), wat de modelkwaliteit per FLOP verbetert. Fijngranulaire MoE's leiden echter tot een groter geheugenverbruik voor activeringen en een verminderde hardware-efficiëntie door hogere IO-kosten, terwijl sparsere MoE's te maken hebben met verspilde berekeningen door padding in gegroepeerde GEMM-kernels. Als antwoord hierop stellen wij een geheugenefficiënt algoritme voor voor de forward- en backwardpassen van MoE's met minimale caching van activeringen voor de backwardpass. We ontwerpen ook GPU-kernels die geheugen-IO overlappen met berekeningen, wat alle MoE-architecturen ten goede komt. Ten slotte introduceren we een nieuwe "token rounding"-methode die de verspilde rekenkracht door padding in gegroepeerde GEMM-kernels minimaliseert. Hierdoor reduceert onze methode, SonicMoE, het activeringsgeheugen met 45% en behaalt een 1,86x verbetering in rekendoorvoer op Hopper-GPU's in vergelijking met ScatterMoE's BF16 MoE-kernel voor een fijngranulair 7B MoE. Concreet behaalt SonicMoE op 64 H100's een trainingsdoorvoer van 213 miljard tokens per dag, vergelijkbaar met ScatterMoE's 225 miljard tokens per dag op 96 H100's voor de training van een 7B MoE-model met FSDP-2 met behulp van de lm-engine codebase. Onder instellingen met hoge MoE-sparseheid levert ons tile-aware token rounding-algoritme een extra 1,16x versnelling op in kerneluitvoeringstijd in vergelijking met standaard top-K-routering, met behoud van vergelijkbare downstreamprestaties. We maken al onze kernels open source om snellere training van MoE-modellen mogelijk te maken.

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.