SonicMoE: Acelerando MoE com Otimizações Conscientes de E/S e Blocos

SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations

December 16, 2025

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

cs.AI

Resumo

Os modelos Mixture of Experts (MoE) emergiram como a arquitetura padrão para escalar modelos de linguagem sem aumentar significativamente o custo computacional. Modelos MoE recentes demonstram uma tendência clara para alta granularidade de especialistas (dimensão intermediária do especialista menor) e maior esparsidade (número constante de especialistas ativados com maior número total de especialistas), o que melhora a qualidade do modelo por FLOP. No entanto, MoEs de granularidade fina sofrem com maior consumo de memória de ativação e eficiência de hardware reduzida devido a maiores custos de E/S, enquanto MoEs mais esparsos sofrem com computações desperdiçadas devido ao preenchimento (padding) nos kernels Grouped GEMM. Em resposta, propomos um algoritmo eficiente em memória para calcular as passagens forward e backward de MoEs com cache de ativação mínimo para a passagem backward. Também projetamos kernels de GPU que sobrepõem E/S de memória com computação, beneficiando todas as arquiteturas MoE. Finalmente, propomos um novo método de "arredondamento de tokens" (token rounding) que minimiza o cálculo desperdiçado devido ao preenchimento nos kernels Grouped GEMM. Como resultado, nosso método SonicMoE reduz a memória de ativação em 45% e alcança uma melhoria de throughput de computação de 1.86x em GPUs Hopper em comparação com o kernel MoE BF16 do ScatterMoE para um MoE de granularidade fina de 7B. Concretamente, o SonicMoE em 64 H100s alcança um throughput de treinamento de 213 bilhões de tokens por dia, comparável aos 225 bilhões de tokens por dia do ScatterMoE em 96 H100s para o treinamento de um modelo MoE de 7B com FSDP-2 usando a codebase lm-engine. Sob configurações de alta esparsidade MoE, nosso algoritmo de arredondamento de tokens com consciência de tile (tile-aware) proporciona um speedup adicional de 1.16x no tempo de execução do kernel em comparação com o roteamento top-K padrão, mantendo desempenho downstream similar. Disponibilizamos todos os nossos kernels em código aberto para permitir um treinamento mais rápido de modelos 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.