SonicMoE: Aceleración de MoE mediante Optimizaciones Conscientes de E/S y Tiles

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

Resumen

Los modelos Mixture of Experts (MoE) han surgido como la arquitectura de facto para escalar modelos de lenguaje sin aumentar significativamente el coste computacional. Los modelos MoE recientes demuestran una clara tendencia hacia una alta granularidad de expertos (dimensión intermedia del experto más pequeña) y una mayor dispersión (número constante de expertos activados con un mayor número total de expertos), lo que mejora la calidad del modelo por FLOP. Sin embargo, los MoE de grano fino sufren un mayor uso de memoria de activaciones y una reducción de la eficiencia del hardware debido a mayores costes de E/S, mientras que los MoE más dispersos sufren de cálculos desperdiciados debido al relleno (padding) en los kernels Grouped GEMM. En respuesta, proponemos un algoritmo eficiente en memoria para calcular las pasadas hacia adelante y hacia atrás de los MoE con un almacenamiento en caché mínimo de activaciones para la pasada hacia atrás. También diseñamos kernels de GPU que solapan la E/S de memoria con el cálculo, beneficiando a todas las arquitecturas MoE. Finalmente, proponemos un novedoso método de "redondeo de tokens" que minimiza el cómputo desperdiciado debido al relleno en los kernels Grouped GEMM. Como resultado, nuestro método SonicMoE reduce la memoria de activaciones en un 45% y logra una mejora del rendimiento de cómputo de 1.86x en GPUs Hopper en comparación con el kernel MoE BF16 de ScatterMoE para un MoE de grano fino de 7B. Concretamente, SonicMoE en 64 H100s alcanza un rendimiento de entrenamiento de 213 mil millones de tokens por día, comparable a los 225 mil millones de tokens por día de ScatterMoE en 96 H100s para el entrenamiento de un modelo MoE de 7B con FSDP-2 utilizando el código base de lm-engine. Bajo configuraciones de alta dispersión MoE, nuestro algoritmo de redondeo de tokens consciente de los 'tiles' produce una aceleración adicional de 1.16x en el tiempo de ejecución del kernel en comparación con el enrutamiento top-K convencional, manteniendo un rendimiento posterior similar. Liberamos todo nuestro código de kernels como open-source para permitir un entrenamiento más 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.