SonicMoE: Beschleunigung von MoE durch E/A- und Tile-bewusste Optimierungen

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

Mixture of Experts (MoE)-Modelle haben sich als de-facto-Architektur etabliert, um Sprachmodelle zu skalieren, ohne die Rechenkosten erheblich zu erhöhen. Aktuelle MoE-Modelle zeigen einen klaren Trend zu hoher Expertengranularität (kleinere Zwischendimension der Experten) und höherer Sparsity (konstante Anzahl aktivierter Experten bei höherer Gesamtzahl an Experten), was die Modellqualität pro FLOP verbessert. Allerdings leiden feingranulare MoEs unter einem erhöhten Aktivierungsspeicherbedarf und reduzierter Hardware-Effizienz aufgrund höherer IO-Kosten, während sparsere MoEs unter Rechenverschwendung durch Padding in gruppierten GEMM-Kerneln leiden. Als Antwort darauf schlagen wir einen speichereffizienten Algorithmus vor, um die Vorwärts- und Rückwärtspass-Berechnungen von MoEs mit minimaler Aktivierungspufferung für den Rückwärtspass durchzuführen. Wir entwickeln auch GPU-Kernel, die Speicher-IO mit Berechnungen überlappen und allen MoE-Architekturen zugutekommen. Schließlich schlagen wir eine neuartige "Token-Rounding"-Methode vor, die die Rechenverschwendung durch Padding in gruppierten GEMM-Kerneln minimiert. Infolgedessen reduziert unsere Methode SonicMoE den Aktivierungsspeicher um 45 % und erreicht einen 1,86-fachen Durchsatzgewinn auf Hopper-GPUs im Vergleich zu ScatterMoEs BF16-MoE-Kernel für ein feingranulares 7B-MoE. Konkret erzielt SonicMoE auf 64 H100s einen Trainingsdurchsatz von 213 Milliarden Token pro Tag, vergleichbar mit ScatterMoEs 225 Milliarden Token pro Tag auf 96 H100s für ein 7B-MoE-Modelltraining mit FSDP-2 unter Verwendung der lm-engine-Codebasis. Unter Einstellungen mit hoher MoE-Sparsity erzielt unser tile-aware Token-Rounding-Algorithmus eine zusätzliche Beschleunigung der Kernel-Ausführungszeit um den Faktor 1,16 im Vergleich zum herkömmlichen Top-K-Routing bei ähnlicher Downstream-Leistung. Wir stellen alle unsere Kernel als Open Source zur Verfügung, um ein schnelleres MoE-Modelltraining zu ermöglichen.

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.