SonicMoE: Ускорение MoE с помощью оптимизаций с учётом ввода-вывода и структуры тайлов

SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations

December 16, 2025

Авторы: Wentao Guo, Mayank Mishra, Xinle Cheng, Ion Stoica, Tri Dao

cs.AI

Аннотация

Модели типа «Смесь экспертов» (MoE) стали фактически стандартной архитектурой для масштабирования языковых моделей без значительного увеличения вычислительных затрат. Современные MoE-модели демонстрируют явную тенденцию к высокой гранулярности экспертов (меньшая промежуточная размерность эксперта) и большей разреженности (постоянное число активируемых экспертов при увеличении общего их количества), что повышает качество модели на один FLOP. Однако мелкозернистые MoE страдают от увеличения объема памяти для активаций и снижения аппаратной эффективности из-за более высоких затрат на ввод-вывод, в то время как более разреженные MoE сталкиваются с бесполезными вычислениями из-за дополнения (padding) в групповых ядрах GEMM. В ответ на это мы предлагаем эффективный по памяти алгоритм для вычисления прямого и обратного проходов MoE с минимальным кэшированием активаций для обратного прохода. Мы также разрабатываем GPU-ядро, которое совмещает операции ввода-вывода с вычислениями, принося пользу всем MoE-архитектурам. Наконец, мы предлагаем новый метод «округления токенов», который минимизирует бесполезные вычисления, вызванные дополнением в групповых ядрах GEMM. В результате наш метод SonicMoE сокращает объем памяти для активаций на 45% и обеспечивает увеличение пропускной способности вычислений в 1.86 раза на GPU Hopper по сравнению с ядром MoE в формате BF16 от ScatterMoE для мелкозернистой 7B MoE. Конкретно, SonicMoE на 64 H100 обеспечивает пропускную способность обучения в 213 миллиардов токенов в день, что сопоставимо с 225 миллиардами токенов в день у ScatterMoE на 96 H100 для обучения 7B MoE-модели с использованием FSDP-2 в кодовой базе lm-engine. В условиях высокой разреженности MoE наш алгоритм округления токенов с учетом блоков (tile-aware) дает дополнительное ускорение времени выполнения ядра в 1.16 раза по сравнению с классической маршрутизацией top-K, сохраняя при этом схожую производительность на downstream-задачах. Мы открываем исходный код всех наших ядер, чтобы обеспечить более быструю обучение 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.