SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations (2512.14080v1)

Abstract: Mixture of Experts (MoE) models have emerged as the de facto architecture for scaling up LLMs 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.

• The paper presents a memory-optimal backward algorithm that reduces activation memory by up to 45% in large-scale MoE models.
• It introduces IO-compute overlapping GPU kernels and a tile-aware token rounding routing strategy to mitigate padding waste and achieve 88% of the cuBLAS theoretical TFLOPS.
• Methodological enhancements enable efficient scaling of fine-grained, sparse MoEs without additional hardware penalties, opening avenues for trillion-parameter models.

SonicMoE: Accelerating Mixture-of-Experts Architectures through IO and Tile-aware Kernel Optimizations

Introduction

The proliferation of Mixture-of-Experts (MoE) architectures has enabled large-scale LLMs to achieve enhanced parameterization without a commensurate increase in computational cost. Fine-grained and highly sparse MoE configurations have demonstrated improved model quality per FLOP, reflected as trends in the latest open-source LLM families. However, this granularity and sparsity exacerbate hardware inefficiencies: increased activation memory, memory-bound kernel regimes due to intensified IO, and wasteful padding-induced compute in Grouped GEMM operations. This paper introduces SonicMoE, a robust kernel and model co-design methodology that provides: (1) minimal-memory MoE computation compatible with scalable granularity, (2) IO-compute overlapping GPU kernels tailored for Hopper and Blackwell architectures, and (3) a tile-aware token rounding routing strategy that strictly bounds wasted computational padding.

Background and Characterization

MoE layers substitute dense MLPs in transformer channels, dispatching tokens via routers to a constant number $K$ of $E$ experts, and aggregating outputs using group-wise operations. On modern GPUs, efficient MoE inference and training rely on Grouped GEMM kernels with variable token dimensions per expert ($\mathbf{M}$), incurring nontrivial padding overheads when $T_e$ is not a multiple of hardware tile sizes.

Recent MoE scaling laws favor higher expert granularity ($G = d / n$, for embedding $d$ and expert intermediate $n$) and greater sparsity ($K/E$ ratio), both leading to lower kernel arithmetic intensity and heightened hardware memory bottleneck. Figure 1 demonstrates SonicMoE’s constant per-layer activation memory despite increased granularity and a substantial margin over other baselines (0.20–1.59$\times$). In throughput, SonicMoE sustains 88% of the cuBLAS theoretical upper bound across configurations.

Figure 1: SonicMoE's per-layer activation memory usage remains constant as expert granularity increases and exhibits superior memory efficiency and throughput relative to all baselines under a 30B MoE configuration on H100.

Grouped GEMM operations underpin MoE kernel efficiency, with Figure 2 illustrating token gathering, projection, and aggregation stages. The token routing method (TC top-K, EC, etc.) not only defines compute patterns but also directly impacts task accuracy and hardware utilization.

Figure 2: Schematic of MoE forward pass leveraging Grouped GEMM, illustrating expert input gathering and processing.

Memory-Efficient Backward Algorithm

One major innovation in SonicMoE is the elimination of backward activation memory scaling with granularity. Standard backward computation for fine-grained MoE requires caching large $Y$ and $X_e$ tensors ($O(TKd)$ size). SonicMoE instead reconstructs backward paths to avoid such caches, reducing per-layer memory by up to 45% for 7B models and more for larger configurations (Figure 3). This aligns activation memory with dense model baselines and enables scalable expert increases without hardware penalty.

Figure 3: Peak activation memory usage per layer across model scales; SonicMoE’s memory remains flat under increased granularity, compared to sharp increases for others.

Key to this is recomputing intermediates in-register and fusing gather operations—SonicMoE eschews storing $Y$ and $dY$, and caches only minimal $X$, $H$, and routing metadata on HBM. This operational change underlies the substantial reduction in backward memory pressure.

IO and Compute-Overlap Kernel Design

With granularity and sparsity pushing MoE workloads into memory-bound regimes, SonicMoE makes strategic kernel enhancements:

• Gather Fusion: Inputs to Grouped GEMM are gathered and loaded asynchronously from global memory to shared memory in a single fused step, eliminating redundant loads and pre-processing. Figure 4 visualizes SonicMoE’s kernel workflow with distinct gather, computation, and aggregation steps mapped to launched kernels.

  Figure 4: SonicMoE's computational kernel workflow with 8 kernel launches across forward and backward passes, and critical gather/store points.

• Epilogue Fusion and Ping-Pong Scheduling: Heavy epilogue operations (e.g., SwiGLU and dSwiGLU, $dS$ computation) are fused directly onto MMA blocks. Ping-Pong scheduling enables role-switching between warpgroups, overlapping Tensor Core compute with memory IO to maximize throughput, as depicted in Figure 5.

  Figure 5: Ping-Pong warpgroup scheduling enables concurrent MMA and epilogue operations for IO-bound kernels, optimizing utilization on Hopper.

• Asynchronous Memory Operations: SonicMoE leverages asynchronous TMA load/store for expert aggregation, avoiding synchronous epilogue memory instructions that can stall compute kernels (see Figure 6), and prefers expert-wise gather and sum aggregation over scatter strategies for bandwidth and determinism benefits (Figure 7).

  Figure 6: Asynchronous TMA store allows for higher memory bandwidth and natural overlap with compute compared to synchronous stores, reducing kernel stalls.

  

  Figure 7: Aggregation strategy comparison—SonicMoE's gather-and-sum is faster than scatter-and-add methods used in alternatives.

Tile-aware Token Rounding Routing

As MoEs grow sparser, computational waste due to padding in Grouped GEMM becomes severe (Figure 8). SonicMoE introduces token rounding (TR), a routing method that ensures each expert receives a token count divisible by the tile size. TR operates by adjusting token allocations with bounded deviation per expert (one tile max) from original top-K selection, maintaining expected total tokens and preserving routing fidelity.

Figure 8: Token rounding example ensures expert token assignments are tile-multiples, mitigating padding-induced compute waste.

Compared to vanilla TC top-K routing, TR yields up to 16% higher kernel TFLOPS in highly sparse configurations and exhibits negligible impact on downstream evaluation: Table results manifest parity (or improvement) in perplexity and task accuracy between TR-trained models (evaluated with TC top-K) and standard TC or EC baselines.

Numerical Results and Kernel Benchmarks

Across H100 benchmarks, SonicMoE consistently achieves the best per-layer TFLOPS under both forward and backward passes (Figure 9), outperforming MegaBlocks, ScatterMoE, MoMoE, and DeepGEMM++ by a margin that widens with increased granularity and sparsity. Real-world multi-GPU training matches 225B/day token throughput on 96 H100s using ScatterMoE with only 64 H100s on SonicMoE—an effective 1.5$\times$ hardware efficiency gain.

Figure 9: Forward and backward TFLOPS across MoE kernels in diverse configurations; SonicMoE consistently leads as model size and sparsity increase.

For large-scale open MoEs (1.4B–685B), SonicMoE remains robust, supporting configurations where others fail due to memory limitations. Kernel-level ablations (Figures 12–14) exhibit further gains when substituting TC routing with TR in hardware utilization, especially in configurations with extreme expert count and sparsity.

Figure 10: Single MoE layer TFLOPS in different kernel implementations over models ranging from 7B to 685B; SonicMoE delivers sustained high throughput.

Figure 11: Token rounding routing yields superior model TFLOPS in sparse settings compared to TC top-K routing.

Figure 12: Token rounding routing sustains throughput under scale-up in both expert count and model size.

Practical and Theoretical Implications

SonicMoE’s contributions have practical significance for future MoE model deployment:

• Hardware envelopes dictated by memory bandwidth can be smoothly navigated, making architectural choices (granularity, extreme sparsity) primarily an algorithmic consideration, no longer hardware-constrained.
• TR routing unlocks efficient expert parallelism and MoE scaling, potentially enabling trillion-parameter models with modest compute and memory budgets.
• Kernel fusions and scheduling (gather, epilogue) modeled by SonicMoE suggest directions for ML compilers targeting future AI accelerators beyond Blackwell/Hopper.
• By matching activation memory efficiency to dense baselines and maximizing token throughput, SonicMoE shifts the optimization axis for LLM amortization toward quality per hardware compute hour.

On the theoretical front, SonicMoE demonstrates that fine-grained MoEs do not intrinsically imply memory-prohibitive backward passes when kernel co-design is performed. The token rounding method provides a deterministic, provable bound on allocation divergence, ensuring model performance converges robustly with tile quantization constraints. The negative results regarding scatter-fused aggregation further clarify the precision, repeatability, and reusability needs for distributed MoE training.

Conclusion

SonicMoE introduces a suite of kernel and algorithmic optimizations that remove the principal bottlenecks for training and inference in granular, sparse Mixture-of-Experts architectures. The approach comprises a memory-optimal backward algorithm, IO-overlapped GPU kernel launches, and a tile-aware token rounding routing strategy, each validated numerically and theoretically. The results indicate substantial improvements in memory footprint, hardware throughput, and practical training costs for large-scale MoE models. Future work entails extending these principles to low-precision computation (FP8, MXFP4), communication-compute overlap in expert parallelism, and further co-design between hardware and model for maximizing quality per unit hardware time.