SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations

2025-12-15

AI Summary

To address high activation memory, severe I/O bottlenecks, and computational inefficiency from Grouped GEMM padding in fine-grained, highly sparse Mixture-of-Experts (MoE) training, this work introduces three key innovations: (1) a novel token rounding scheduling strategy that mitigates GEMM padding overhead induced by sparsity; (2) a minimal-activation-cache backward algorithm that substantially reduces memory pressure during gradient computation; and (3) a full-architecture-compatible GPU kernel enabling I/O–computation overlap, coupled with a tile-aware scheduling mechanism. Experiments demonstrate a 45% reduction in activation memory and a 1.86× improvement in compute throughput on Hopper GPUs. On a 64× H100 cluster, the system achieves a training throughput of 213B tokens/day—matching the performance of a 96-GPU ScatterMoE baseline. In highly sparse regimes, the optimized kernels deliver a 1.16× speedup over prior approaches.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Learning to Search
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding

Application Category

• Economics, Online Markets and Human Computation: Architectures and workflows that use LLMs for crowd work
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search

Abstract

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.

Problem

• Reduces activation memory in fine-grained MoE models
• Overlaps memory IO with computation for efficiency
• Minimizes wasted compute from padding in sparse MoEs

Innovation

• Memory-efficient algorithm minimizes activation caching
• GPU kernels overlap memory IO with computation
• Token rounding method reduces padding waste in GEMM