PaperlayerView exact revision on arXiv

Source-linked AI summary

SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations

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

arXiv:2512.14080v2cs.LGcs.AI

TL;DR

Granular and sparse MoEs improve quality per FLOP but incur activation-memory, IO, and grouped-GEMM padding inefficiencies. SonicMoE addresses these issues with memory-efficient computation, IO-overlapped GPU kernels, and tile-aware token rounding, achieving lower memory use and higher throughput across Hopper and Blackwell evaluations.

  • Problem

    Granular and sparse MoEs face increased activation memory, higher IO costs, and wasted grouped-GEMM computation from padding despite improving quality per FLOP.

  • Method

    SonicMoE co-designs an activation-efficient MoE computation algorithm, IO-overlapped GPU kernels, and token-rounding routing aligned with GEMM tile sizes.

  • Results

    SonicMoE reduces fine-grained 7B MoE activation memory by up to 45% on H100 GPUs and improves throughput across Hopper and Blackwell comparisons.

  • Takeaways & Limitations

    Tile-aware token rounding adds speedup without accuracy loss, while the combined design improves MoE training efficiency across granular and sparse settings.

  • Takeaways & Limitations

    The analysis assumes fine-grained MoEs under iso-FLOPs and iso-parameter settings, and removing temporary materialization would introduce determinism, numerical-accuracy, and communication-compatibility issues.

Abstract

from arXiv · show

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 a 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. On Blackwell GPUs, SonicMoE also achieves a 25% and 15% relative speedup on the forward and backward pass respectively compared to a highly optimized DeepGEMM baseline on OLMoE-sized 7B MoE models. 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 on Hopper GPUs. We open-source all our kernels.

1 Introduction

SonicMoE co-designs memory-efficient MoE computation, IO-overlapped GPU kernels, and tile-aware routing to address the hardware inefficiencies of increasingly granular and sparse MoEs.

  • 1 Introduction: Increasing expert granularity and sparsity raises activation memory, IO cost, and grouped-GEMM padding waste, pushing MoE training toward a memory-bound regime.These inefficiencies motivate hardware- and routing-aware kernel design.
  • 1 Introduction: SonicMoE combines memory-efficient activation handling, IO-overlapped kernels, and token rounding to improve training efficiency for granular and sparse MoEs.The design targets activation memory, memory-bandwidth bottlenecks, and grouped-GEMM padding waste.
  • 1 Introduction: 45% lower per-layer activation memory is achieved for a fine-grained 7B MoE on H100 GPUs without increasing FLOPs.The method avoids caching activations needed for router-gradient computation while preserving mathematical equivalence.
  • 1 Introduction: 43% higher forward TFLOPS than DeepGEMM and 83% and 115% higher backward TFLOPS than ScatterMoE and MoMoE, respectively, are reported for a fine-grained 7B MoE on H100 GPUs.On B300 GPUs, the approach also reports 25% higher forward and 15% higher backward TFLOPS than DeepGEMM for OLMoE-sized 7B MoEs.
  • 1 Introduction: 16% higher end-to-end MoE compute throughput is obtained with token rounding than vanilla top-K routing in highly sparse 1.4B-parameter training on H100 GPUs.Rounding per-expert token counts to GEMM-tile multiples bounds assignment deviation by one tile and reduces padding waste while preserving token-choice accuracy.
  • 1 Introduction: SonicMoE is released with a PyTorch interface and permissive license to support researchers and practitioners.The implementation is mainly written in CuTe-DSL.

2 Background

MoE layers use routed experts and Grouped GEMM, but increasing granularity and sparsity lowers arithmetic intensity and raises memory and padding inefficiencies. SonicMoE addresses these constraints with tile-aware token rounding that reduces Grouped GEMM padding while preserving token assignments and inference quality.

  • MoE architecture: MoE blocks route tokens to multiple experts, process them independently, and aggregate the expert outputs for the next layer.
  • MoE using Grouped GEMM: Grouped GEMM executes expert-specific matrix multiplications with fixed weight dimensions and variable token dimensions, using either gathered or contiguous inputs.Forward and activation-gradient operations use varlen-M Grouped GEMM, while backward weight gradients use varlen-K Grouped GEMM.
  • Granularity and efficiency: Increasing granularity or sparsity decreases arithmetic intensity because expert-related IO costs grow relative to computation, pushing fine-grained MoEs toward memory-bound execution.The paper assumes fixed embedding dimension and defines fine-grained MoEs as having expert intermediate size smaller than the embedding dimension.
  • MoE routing methods: SonicMoE rounds each expert’s received-token count to nearby Grouped GEMM tile multiples, changing at most one tile per expert to reduce padding waste while preserving inference quality.The method targets wasted FLOPs rather than only padding-related load traffic, unlike prior approaches described in the passage.

3 Memory-efficient MoE algorithm

SonicMoE reorganizes MoE computation into an eight-kernel forward and backward workflow that minimizes cached activations and overlaps memory movement with computation. Its activation footprint is independent of expert granularity and matches the minimum described for backward computation without GEMM recomputation.

  • Kernel workflow: SonicMoE’s eight-kernel workflow launches three forward kernels and five backward kernels for projections, aggregation, activation gradients, input gradients, and weight gradients.Figure 3 distinguishes HBM loads and stores and marks cached activations and intermediate variables.
  • Kernel implementation: The implementation combines an optimized modular Grouped GEMM kernel with an expert aggregation kernel, dispatching configurations and load/store strategies across the eight kernels.
  • Memory-efficient computation: SonicMoE fuses gathers with HBM loads for X and dO, avoiding materialization and activation caching while improving fine-grained MoE throughput.
  • Memory-efficient computation: An alternative backward computation derives dS and dH without requiring Y and dY, avoiding extra FLOPs while reducing activation storage.
  • Activation memory efficiency: SonicMoE caches only X, H, and routing metadata, using 2Td+4TKn bytes per layer and making activation memory independent of expert granularity.The reported usage matches a dense model with the same number of activated parameters and is described as minimal without GEMM-based activation recomputation.

4 IO-aware kernel design

SonicMoE reduces fine-grained MoE IO costs by fusing data movement and epilogue work with computation, while overlapping asynchronous IO with GEMM on Hopper and Blackwell GPUs.

  • 4.1.1 Gather fusion with GMEM-to-SMEM load: Gather fusion combines routed-token gathering with activation loads, avoiding a separate gather-and-pad kernel before Grouped GEMM.On Blackwell, a relay warp forwards cp.async completion across the 2-CTA cluster so the leader MMA warp can wait for both CTAs.
  • 4.1.2 Epilogue fusion: Epilogue fusion combines SwiGLU, dSwiGLU, dH, and dS computations with GEMM output processing to reduce extra IO and kernel time.SonicMoE computes dS through an alternative path that avoids the additional HBM load and activation caching required by competing formulations.
  • 4.2 GEMM MMA Overlapping with Asynchronous IO: Ping-Pong scheduling overlaps one warpgroup’s IO with another’s GEMM on Hopper, sustaining Tensor Core utilization despite heavy epilogues.Heavy HBM stores can otherwise block the next tile’s MMA and reduce TFLOPS by about 20%.
  • 4.2 GEMM MMA Overlapping with Asynchronous IO: Blackwell’s TMEM uses two accumulator stages so epilogue warps process one stage while MMA warps accumulate into the other, improving overlap without register-heavy accumulation.UMMA stores accumulators directly in TMEM and enables concurrent epilogue and MMA work.

5 Token rounding routing

SonicMoE addresses sparse-MoE tile quantization by rounding expert token counts to Grouped GEMM tile multiples, reducing padding waste while preserving routing quality.

  • 5 Token rounding routing: Tile quantization wastes computation because GEMM dimensions must be padded to hardware tile multiples, with the waste becoming nontrivial when expert token counts are small.Sparser MoEs receive fewer tokens per expert, increasing the relative impact of padded tiles.
  • 5 Token rounding routing: Token rounding adjusts each expert’s received-token count to a nearby Mtile multiple, avoiding extra padded GEMM tiles while preserving the original assignments as much as possible.The method alters at most one tile per expert and keeps the maximum deviation from token-choice routing within one tile.
  • 5 Token rounding routing: The routing procedure first performs vanilla top-K token choice, then sorts tokens per expert and either discards selected tokens or pads additional tokens to reach tile-aligned counts.A round-and-sparsify decision chooses between rounding up and down, with nearest-frequency rounding as the default.
  • 5 Token rounding routing: Token rounding maintains robust model quality under sparse MoE training, with performance remaining similar to token-choice routing despite bounded per-expert deviations.The reported robustness holds across the examined sparse-training settings and supports token rounding as a substitute for token-choice routing.
  • 5.1 Training efficiency of sparse MoE: 16% higher kernel TFLOPS than vanilla top-K routing is achieved as expert count increases with constant K in highly sparse MoE training.Token rounding removes tile quantization effects while keeping the comparison at fixed T, n, and K.

6 Experiments

SonicMoE reduces activation memory and improves MoE throughput across H100 and Blackwell configurations, while token rounding preserves model quality and accelerates sparse models.

  • 6.1 Activation memory: 45% lower activation memory is achieved on a fine-grained 7B MoE with n = 256 versus ScatterMoE, with larger savings at 120B scale versus MoMoE.SonicMoE has the lowest peak activation memory across tested model scales on H100 GPUs.
  • 6.2 Training throughput: 43% higher forward TFLOPS than DeepGEMM and 83% and 115% higher backward TFLOPS than ScatterMoE and MoMoE are achieved on the fine-grained 7B H100 configuration.SonicMoE consistently achieves the highest TFLOPS across model scales in Figure 11a.
  • 6.2 Training throughput: 213 billion tokens per day on 64 H100s is comparable to ScatterMoE’s 225 billion tokens per day on 96 H100s for the 7B model with FSDP-2.The measurements use the lm-engine codebase.
  • 6.2 Training throughput: On B300 GPUs, SonicMoE is 25% faster forward and 15% faster backward than DeepGEMM++ for OLMoE-sized 7B MoEs.It generally exceeds 1100 TFLOPS in both passes and is 11.8% faster forward than Triton’s official example.
  • 6.3 Token rounding: Token rounding maintains similar or better task quality than token-choice routing across sparse configurations and can replace token choice during training.Under K/E ≤1/32, it achieves slightly lower validation perplexity and higher or equal average accuracy in the reported settings.
  • 6.3 Token rounding: Token rounding improves model TFLOPS by up to 25.7% forward and 11.8% backward for K/E = 1/64, while delivering 19.6% and 7.9% speedups on a Qwen3-Next configuration.The method is generally robust to rounding subroutine, microbatch size, and tile size when ¯Te/Mtile ≥2.

7 Conclusion

The paper concludes that SonicMoE co-designs memory-efficient computation, IO-overlapped kernels, and tile-aware routing to improve granular and sparse MoE training efficiency.

  • 7 Conclusion: SonicMoE targets the memory and hardware inefficiencies caused by increasing MoE granularity and sparsity.These inefficiencies include larger activation memory, higher IO cost, and padding waste in grouped GEMM.
  • 7 Conclusion: Its token-rounding router aligns expert token counts with grouped-GEMM tile sizes while bounding each expert’s assignment deviation by one tile.The method reduces padding waste while preserving token assignments as much as possible.
  • 7 Conclusion: SonicMoE overlaps memory IO with GEMM computation using asynchronous hardware features, improving utilization when fine-grained MoE GEMMs are small.Asynchronous TMA stores can overlap with TensorCore MMA, unlike synchronous scatter stores on Hopper.
  • 7 Conclusion: The backward computation avoids caching activations needed for router gradients while remaining mathematically equivalent to the original MoE formulation.This choice also avoids extra HBM traffic and reduces cached activation memory associated with storing Y.
  • 7 Conclusion: SonicMoE’s score-gradient path reduces reduction rounds from log2(d) to log2(n), saving at least log2(d/n) rounds.The path also avoids extra loads because e,t and Ae,t are already computed in the dH kernel.

D Efficient top-K sorting kernel for MoE

SonicMoE implements a stable, register-based top-K sorting kernel that reduces router overhead for large-token MoE workloads.

  • D Efficient top-K sorting kernel for MoE: PyTorch top-K can consume approximately 40% of router computation time, motivating SonicMoE’s optimized top-K kernel.The kernel supports E ≤4096 and K ≤16 and targets large token counts T.
  • D Efficient top-K sorting kernel for MoE: The kernel parallelizes over tokens, sorts each row’s E routing values with bitonic sort, and selects the first K columns.Column indices are packed into FP32 mantissa bits during sorting.
  • D Efficient top-K sorting kernel for MoE: Register-local intra-thread and intra-warp comparisons provide higher memory bandwidth than PyTorch, Triton, Tilelang, and RTop-K alternatives.Warp shuffles support the bitonic compare-and-merge operations without additional shared-memory traffic.
  • D Efficient top-K sorting kernel for MoE: Packing unique column indices into mantissa bits makes the sorting result stable by eliminating ties during bitonic comparisons.The packed format follows a strategy also used by Triton’s official top-K kernel.

E Referenced tables, figures, and SonicMoE algorithms

The paper documents SonicMoE’s kernel pipeline and contrasts its aggregation and hardware execution choices with existing MoE implementations.

  • Referenced tables, figures, and SonicMoE algorithms: The backward up-projection pipeline computes dX and dW1 through varlen-M and gather-plus-varlen-K grouped GEMMs, followed by expert aggregation for dX.The documented kernels take X, routing metadata, W1, and dH as inputs and output dX and dW1.
  • Referenced tables, figures, and SonicMoE algorithms: Table 4 tracks expert sparsity as K/E and expert granularity as d/n, showing that recent open-source MoEs trend toward higher sparsity and granularity.The table excludes shared experts from the sparsity calculation.
  • Referenced tables, figures, and SonicMoE algorithms: Figure 16 contrasts asynchronous TMA stores, which overlap with TensorCore MMA, with synchronous st.global stores that block the next MMA tile.The TMA-based strategy is supported by a reported average 20.1% speedup over the synchronous scatter-fusion strategy.
  • Referenced tables, figures, and SonicMoE algorithms: SonicMoE stores contiguously packed expert outputs through TMA in the GEMM epilogue, then gathers and sums activated expert outputs per token.ScatterMoE and MoMoE instead fuse scatter with the epilogue and launch a later summation kernel.
  • Referenced tables, figures, and SonicMoE algorithms: The appendix includes kernel comparisons, computational paths, top-K sorting details, and kernel-level throughput ablations.These materials cover gather fusion, expert aggregation, and the effects of SonicMoE’s implemented features.

F.1 Grouped GEMM

SonicMoE’s Grouped GEMM and aggregation kernels improve throughput across Hopper and Blackwell GPUs, including gathered-input cases, by overlapping IO with computation and using tile-aware scheduling.

  • Grouped GEMM with contiguously-packed inputs: 10.0% higher down-proj TFLOPS than DeepGEMM on average on H100 GPUs, while up-proj is 2.7% higher.The down-proj advantage increases when the intermediate size is small; one reported configuration reaches 57.4% higher TFLOPS.
  • Grouped GEMM with contiguously-packed inputs: 12.7% higher down-proj TFLOPS than DeepGEMM on average on B300 GPUs, with 8.1% higher up-proj TFLOPS.SonicMoE also exceeds the official Triton GEMM example by 13.3% for up-proj and 15.6% for down-proj on average.
  • Gathered-input Grouped GEMM: 38.3% higher TFLOPS than DeepGEMM with gather fusion on H100 GPUs, with SonicMoE also exceeding ScatterMoE and MoMoE.The comparison uses gathered varlen-K timing for baselines without equivalent fusion.
  • Gathered-input Grouped GEMM: SonicMoE’s gather fusion changes average TFLOPS by only -0.1% on B300 K-dimension gathers while retaining a wider advantage as expert granularity increases.On H100 K-dimension gathers, fusion changes TFLOPS by 8.5%, and the gap versus separate-gather baselines also widens with granularity.
  • Expert aggregation: 2.92x higher aggregation bandwidth than ScatterMoE on H100 GPUs and 6.72x higher on B300 GPUs.SonicMoE remains at 0.98x the Triton contiguous-sum upper bound on both GPUs; its selected aggregation strategy is 20% faster than the alternative on H100.

G.2 Ablation study on different rounding subroutines for token rounding

The token-rounding ablation compares frequency- and score-based rounding choices, showing that quality–efficiency balance favors nearest rounding over always padding or discarding tiles.

  • Rounding subroutines: Token rounding is generally robust to the specific rounding subroutine, as shown by comparisons across nearest, stochastic, balanced, up, and down variants.The evaluation reports validation perplexity and mean accuracy across 11 downstream tasks.
  • Rounding subroutines: NR-f rounds each expert’s frequency to the nearest tile multiple and is the default choice for the main experiments.The decision pads or discards according to which neighboring tile multiple is closer.
  • Always discarding tokens (DOWN): Always discarding tokens gives the shortest MoE kernel runtime but much higher validation perplexity than NR-f, SR-f, and NR-s.The authors caution that the resulting quality degradation may be unacceptable in practice.
  • Always padding tokens (UP): Always padding tokens produces the longest MoE kernel runtime and does not necessarily improve average downstream accuracy.The authors do not recommend always rounding up, despite often lower validation perplexity.

G.3 Ablation study on the effects of microbatch size T and tile size Mtile

Token rounding is robust across microbatch and tile-size choices when the average tokens per expert cover at least two tiles, but quality degrades at one tile.

  • Microbatch size T: Changing microbatch size T changes the average tokens per expert while keeping the minibatch size constant.The routing decision is applied at the microbatch level, so different microbatch granularities produce different token-rounding behavior.
  • Microbatch size T: TR preserves trained MoE quality when ¯Te/Mtile ≥2, whereas ¯Te/Mtile = 1 causes noticeable degradation in perplexity and downstream performance.Even at ¯Te/Mtile = 1, quality remains better than training with EC and then finetuning with TC top-K routing.
  • Tile size Mtile: TR is generally robust to Mtile when ¯Te/Mtile ≥2, but shows noticeable degradation at ¯Te/Mtile = 1 while remaining better than the EC baseline.The tile-size ablation evaluates validation perplexity and mean accuracy across 11 downstream tasks.

H Activation memory and training throughput benchmark configurations

The benchmark configurations vary expert count and model dimensions across four Figure 13 subfigures while fixing the token-rounding tile size at 128.

  • Configuration overview: Mtile is consistently 128 when benchmarking token-rounding speed.The configurations vary T, d, n, K, and E across four subfigure groups.
  • Configuration overview: The four subfigure groups use (T,d,n,K) values of (16384,1536,256,8), (16384,1536,1024,2), (16384,4096,512,8), and (16384,4096,1024,4).Expert count E varies from 16 to 512 depending on the group.

I Hyperparameter details for LM training

The LM training experiments use standardized MoE, attention, routing, optimization, and dataset configurations, with separate procedures for router finetuning and auxiliary-router experiments.

  • The benchmark configurations for the kernel-level studies and LM training experiments are listed in Table 9 and Table 10.
  • Pretraining uses deduplicated FineWeb-Edu with context length 4096, while downstream evaluation covers MMLU and the listed commonsense, reasoning, and science datasets.
  • The models use SwiGLU MoE layers, auxiliary load balancing loss with coefficient 0.01, no router Z loss, and tied token-embedding and LM-head weights.The attention block matches OLMoE’s architecture.
  • Router-only TC finetuning adds 4B tokens, freezes non-router parameters, uses learning rate 2e-4, weight decay 0.01, cosine scheduling, and 1M-token minibatches.Auxiliary load balancing loss is disabled during this finetuning.
  • Auxiliary-router experiments use a two-layer E × E MLP on raw router logits, train averaged multilabel binary cross-entropy scaled by 0.01, and apply masking during evaluation.
  • TC token drop is implemented by discarding tokens selected during TC top-K sorting.
Loading 2512.14080v2…