SonicMoE: Accelerazione delle MoE tramite ottimizzazioni consapevoli delle operazioni di I/O e dei tile

SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations

December 16, 2025

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

cs.AI

Abstract

I modelli Mixture of Experts (MoE) sono emersi come l'architettura di fatto per scalare i modelli linguistici senza aumentare significativamente il costo computazionale. I modelli MoE recenti dimostrano una chiara tendenza verso un'alta granularità degli esperti (dimensione intermedia degli esperti più piccola) e una maggiore sparsità (numero costante di esperti attivati con un numero maggiore di esperti totali), che migliorano la qualità del modello per FLOP. Tuttavia, i MoE a granularità fine soffrono di un aumento dell'ingombro di memoria delle attivazioni e di una ridotta efficienza hardware a causa dei maggiori costi di I/O, mentre i MoE più sparsi soffrono di calcoli sprecati a causa del padding nei kernel Grouped GEMM. In risposta, proponiamo un algoritmo efficiente in memoria per calcolare i passi in avanti e indietro dei MoE con una memorizzazione minima delle attivazioni per il passaggio all'indietro. Progettiamo anche kernel GPU che sovrappongono l'I/O di memoria con il calcolo, a vantaggio di tutte le architetture MoE. Infine, proponiamo un nuovo metodo di "arrotondamento dei token" che minimizza il calcolo sprecato a causa del padding nei kernel Grouped GEMM. Di conseguenza, il nostro metodo SonicMoE riduce la memoria di attivazione del 45% e raggiunge un miglioramento del throughput computazionale di 1.86x sulle GPU Hopper rispetto al kernel MoE BF16 di ScatterMoE per un MoE a granularità fine da 7B. Nello specifico, SonicMoE su 64 H100 raggiunge un throughput di addestramento di 213 miliardi di token al giorno, paragonabile ai 225 miliardi di token al giorno di ScatterMoE su 96 H100 per l'addestramento di un modello MoE da 7B con FSDP-2 utilizzando il codebase lm-engine. In contesti di elevata sparsità MoE, il nostro algoritmo di arrotondamento dei token "tile-aware" produce un ulteriore speedup di 1.16x sul tempo di esecuzione del kernel rispetto al routing top-K standard, mantenendo prestazioni downstream simili. Rilasciamo come open source tutti i nostri kernel per consentire un addestramento più veloce dei modelli 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.