SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations

TL;DR

SonicMoE is a Hopper/Blackwell-optimized MoE training kernel suite from the Dao-AILab group that attacks fine-grained-and-sparse MoEs on three fronts: an activation-memory-minimal backward pass that decouples per-layer activation footprint from expert granularity, an IO-aware Grouped GEMM with gather-fused HBM loads and heavy epilogue fusion that overlaps memory traffic with matmul, and a “token rounding” routing tweak that rounds per-expert token counts to GEMM tile multiples so sparse MoEs stop wasting compute on padding. On a fine-grained 7B MoE, SonicMoE cuts activation memory by 45% and runs 1.86× faster than ScatterMoE’s BF16 kernel on H100, with an additional 1.16× kernel speedup from token rounding at high sparsity. Implemented in CuTe-DSL with a PyTorch interface; open-source at Dao-AILab/sonic-moe.

Key claims

Fine-grained MoEs (small expert intermediate, many experts) are increasingly memory-bandwidth-bound: arithmetic intensity scales as roughly $1/(\text{granularity}/\text{activation-ratio})$ in the standard formulation, so increasing granularity at fixed FLOPs linearly increases IO cost [§2.2, Fig. 3, Table 1].
A standard ScatterMoE-style backward materializes activations whose size scales linearly with expert granularity; SonicMoE rewrites the backward computation graph so $\mathbf{d}\mathbf{z}_1$ and gathered $\mathbf{x}$ are never written to HBM, caching only $\mathbf{x}$ and $\mathbf{z}_2$ plus routing metadata — the same activation footprint as a dense model with the same active parameter count [§3.2, Fig. 13].
Activation memory per MoE layer drops by up to 45% for a fine-grained 7B MoE, and the savings are independent of expert granularity — empirically verified from 1.4B to 120B configurations [§3.2, Fig. 13].
Gather fusion with the GMEM-to-SMEM load (using cp.async rather than a separate gather kernel) is preserved across both forward and backward in SonicMoE; ScatterMoE and MoMoE only fuse it in the forward, and the two un-fused backward gathers consume 19.6% (ScatterMoE) and 20.6% (MoMoE) of total backward time on the 7B configuration [§4.1.1, Fig. 6, Table 2].
Epilogue fusion combines (1) SwiGLU/dSwiGLU into up-proj/down-proj epilogues and (2) the down-proj activation-gradient kernel additionally computes $\mathbf{d}\mathbf{x}_1$ and $\mathbf{d}\mathbf{x}_2$ together; this replaces ScatterMoE’s separate down-proj-act (0.43 ms) + $\mathbf{d}\mathbf{x}$ (0.24 ms) + dSwiGLU (0.33 ms) — total 0.99 ms — with a single 0.47 ms kernel [§4.1.2, Fig. 6].
2-CTA cluster GEMM on Blackwell requires a custom relay-warp mechanism so the leader CTA’s MMA can wait on gather completion in both CTAs (cp.async only signals within its own CTA) [§4.1.1 Blackwell paragraph].
Token Rounding (TR) routing rounds per-expert token counts to multiples of the Grouped GEMM tile size (typically 128) while bounding the max deviation from top- $k$ token-choice routing to one tile per expert; eliminates padding-induced wasted FLOPs in sparse MoEs without changing total tokens in expectation [§2.3, §“Summary of contributions”].
At high sparsity (4× expert count on a 30B MoE), TR delivers a 1.16× kernel-time speedup over vanilla top- $k$ token-choice routing while preserving inference quality at 1.4B-parameter scale [Abstract, §“Summary of contributions”].
End-to-end 7B MoE training on 64 H100s with SonicMoE+FSDP-2 reaches 213B tokens/day — comparable to ScatterMoE’s 225B tokens/day on 96 H100s [Abstract].
DeepGEMM++ (DeepGEMM Grouped GEMM + separate optimized gather + separate SwiGLU) is still slower than SonicMoE on forward up-proj despite individually tuned kernels, because the unfused gather contributes a separate kernel launch and extra HBM round-trip [§4.1.1 Hopper paragraph, Fig. 6].
Hopper backward speedups over baselines on a fine-grained 7B MoE: +43% over DeepGEMM++ forward, +83% over ScatterMoE backward, +115% over MoMoE backward [§“Summary of contributions”].
The kernel suite is written in CuTe-DSL exposed via PyTorch — same Dao-AILab DSL stack as FlashAttention-4 and Gram Newton-Schulz; ships as a drop-in for ScatterMoE [§“Summary of contributions”, §1].

Method

SonicMoE launches 8 modular Grouped-GEMM kernels per MoE block (forward: up-proj, down-proj, expert-aggregation; backward: down-proj-act-grad, up-proj-act-grad, aggregation-grad, two weight-grad kernels). Three architecture–software co-design moves carry the speedup.

Memory-efficient backward. Standard MoEs cache the gathered input $\mathbf{x}$ and several intermediate activations ( $\mathbf{z}_1$ , $\mathbf{z}_2$ ) for the backward pass; with fine granularity these have size $O(K \cdot d \cdot d_{\text{int}})$ that scales with granularity. SonicMoE identifies an alternative algebraic path for the router gradient that expresses $\mathbf{d}\mathbf{x}_1, \mathbf{d}\mathbf{x}_2$ purely in terms of $\mathbf{x}$ and $\mathbf{z}_2$ (and quantities recomputable in registers during the epilogue), allowing $\mathbf{d}\mathbf{z}_1$ and gathered $\mathbf{x}$ to never touch HBM. The cached footprint matches a dense model with the same activated parameter count — the information-theoretic floor without checkpoint-recomputation.

Gather-fused, epilogue-fused Grouped GEMM. The Grouped GEMM accepts either packed or gathered inputs and fuses the per-expert token-index lookup with the producer-warp HBM-to-SMEM load (via prefetched indices and cooperative fetching among producer warps). Epilogues fuse SwiGLU/dSwiGLU directly into the matmul, and the down-proj-activation-gradient epilogue additionally computes the up-proj gradient in the same kernel, eliminating two HBM round-trips per backward pass. On Blackwell, the 2-CTA cluster MMA requires a dedicated “relay warp” in the follower CTA to bridge cp.async’s intra-CTA-only signaling to the cluster-scope mbarrier the leader’s MMA waits on.

Token rounding routing. After standard top- $k$ token-choice routing, SonicMoE adjusts the per-expert token counts to the nearest multiples of the GEMM tile size (128) by moving up to one tile’s worth of tokens per expert between adjacent ranks. This bounds the deviation from the original routing assignment, preserves total token count in expectation, and exactly eliminates the padding-induced wasted FLOPs that grow as MoEs get sparser.

Results

Activation memory (7B MoE, fine-grained): 45% reduction per layer vs ScatterMoE; flat as a function of granularity, vs ScatterMoE’s linear scaling [Fig. 13].
Kernel throughput (Fig. 6, 7B fine-grained MoE on H100): SonicMoE up-proj 0.55 ms / 559 TFLOPS beats DeepGEMM++ (up-proj + SwiGLU) 0.60 ms total. Down-proj-activation backward 0.47 ms beats ScatterMoE’s down-proj-act + $\mathbf{d}\mathbf{x}$ + dSwiGLU combined at 0.99 ms.
End-to-end training (FSDP-2, 7B MoE, lm-engine): 213 B tokens/day on 64 H100s vs ScatterMoE’s 225 B on 96 H100s — ~1.5× per-GPU throughput.
Backward speedups vs baselines on the fine-grained 7B MoE: +83% over ScatterMoE, +115% over MoMoE [§“Summary of contributions”].
Token Rounding (1.4B sparse MoE): 1.16× compute throughput vs vanilla top- $k$ token-choice routing in highly sparse regimes; downstream model quality preserved on 2B-scale validation.
Hardware-aware speedup scaling: Relative speedup over alternative kernels grows with granularity (Fig. 14) — SonicMoE’s IO-aware optimizations win more as MoEs trend more fine-grained.

Why it’s interesting

SonicMoE is the first filed paper that targets the kernel-level IO bottleneck of fine-grained-sparse MoEs as a first-class problem, and it sits cleanly at the intersection of two active wiki threads. On the IO-aware-kernel side, it extends the IO-Aware Kernel Design template — fuse the producer/consumer pair, exploit algebraic identities to elide intermediates — from FlashAttention, Flash-KMeans, FlashSampling, and Gram-Newton-Schulz into MoE Grouped GEMM. The backward-pass identity that removes $\mathbf{d}\mathbf{z}_1$ from HBM is structurally analogous to FlashAttention’s softmax-block identity that elides the full attention matrix, and the CuTe-DSL implementation matches the Dao-AILab pattern that FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling and Gram Newton-Schulz: A Fast, Hardware-Aware Newton-Schulz Algorithm for Muon established. On the MoE-routing side, Token Rounding adds a hardware-tile-aware constraint to the routing-design vocabulary maintained at MoE Routing Design — a different axis from the soft-gate (Routing-Free Mixture-of-Experts) and LBL-rescoping (Demons in the Detail: On Implementing Load Balancing Loss for Training Specialized Mixture-of-Expert Models) lines, and complementary to Quantile Balancing: A Hyperparameter-Free MoE Load Balancing Method‘s knob-free balancer. The result is also a counterpoint to the systems-level MoE-parallelism recipes in MegaScale-MoE: Large-Scale Communication-Efficient Training of Mixture-of-Experts Models in Production and Scalable Training of Mixture-of-Experts Models with Megatron Core: those resolve the attention-vs-experts parallelism mismatch at the cluster level; SonicMoE resolves the intra-kernel IO mismatch at the single-GPU level. Both are needed at frontier scale.