SonicMoE: A Hardware-Efficient and Software-Extensible Blueprint for Fine-Grained MoEs

2026-04-21 · Source: Tri Dao · Field: Technology & Digital — Artificial Intelligence & Machine Learning, Software Development & Engineering · Depth: Expert, extended

Summary

SonicMoE is a new hardware-efficient and software-extensible blueprint designed to accelerate fine-grained Mixture-of-Experts (MoE) models, now supporting NVIDIA Blackwell GPUs (B200/B300) in addition to Hopper (H100). It addresses two key challenges in fine-grained MoEs: activation memory scaling with expert granularity and increasing I/O costs. SonicMoE achieves this by an I/O-aware algorithmic redesign that avoids caching O(TKd)-sized tensors, making activation memory independent of expert granularity and eliminating large HBM round-trips. This approach yields 1.87-4.04x speedup over existing kernels like ScatterMoE and MoMoE. The system utilizes a unified software abstraction built on QuACK, which localizes architecture-specific optimizations, and exploits Blackwell hardware features like Cluster Launch Control (CLC) and 2CTA MMA to overlap I/O with computation. Benchmarks show SonicMoE consistently leads across six open-source MoE configurations, achieving 54%/35% higher forward/backward TFLOPS than DeepGEMM and 21% higher forward TFLOPS than Triton's official example.

Key takeaway

For Machine Learning Engineers or AI Scientists developing or deploying fine-grained Mixture-of-Experts models, you should consider integrating SonicMoE to significantly improve training efficiency. Its algorithmic redesign and hardware-aware optimizations on NVIDIA Hopper and Blackwell GPUs reduce activation memory and I/O bottlenecks, delivering up to 4.04x speedup. This allows you to scale MoE models with higher granularity and sparsity without incurring prohibitive hardware costs, enabling more advanced model architectures.

Key insights

SonicMoE optimizes fine-grained MoE training by redesigning algorithms and leveraging hardware to reduce activation memory and I/O costs.

Principles

• Activation memory can be independent of expert granularity.
• I/O costs in memory-bound kernels can be hidden by computation.
• Unified software abstraction simplifies hardware porting.

Method

SonicMoE redesigns the backward pass to avoid caching O(TKd)-sized tensors, fuses gather/activation operations, and reorders contractions. It uses QuACK for modular kernel development and exploits Blackwell's CLC and 2CTA MMA for I/O overlap.

In practice

• Use QuACK for modular GPU kernel development.
• Implement gather fusion to improve L2 cache hit rates.
• Overlap I/O with MMA computation for memory-bound kernels.

Topics

• Mixture-of-Experts
• GPU Acceleration
• NVIDIA Blackwell
• QuACK Library
• Kernel Optimization

Code references

• Dao-AILab/quack
• Dao-AILab/sonic-moe
• tilde-research/MoMoE-impl
• NVIDIA/cutlass
• Dao-AILab/gram-newton-schulz

Best for: AI Engineer, NLP Engineer, Research Scientist, AI Scientist, Machine Learning Engineer, AI Hardware Engineer

Related on AIssential

• Accelerating Mixture-of-Experts Execution with FarSkip-Collective Models — FarSkip-Collective accelerates MoE models by overlapping computation and communication, eliminating idle GPU time during distributed execution.
• Run Step 3.7 Flash on NVIDIA GPUs with Enterprise-Ready Multimodal AI — Step 3.7 Flash is an enterprise-ready multimodal MoE model optimized for NVIDIA infrastructure.
• GPT-OSS 20B: A Sparse MoE Pretraining Benchmark for MLPerf Training v6.0 — GPT-OSS 20B offers a stable, accessible benchmark for sparse Mixture-of-Experts LLM pretraining in MLPerf Training v6.0.
• Path-Constrained Mixture-of-Experts — Path-Constrained Mixture-of-Experts (PathMoE) improves MoE efficiency and performance by explicitly managing and amplifying token routing paths across layers.
• MIT Proved 90% of Every AI Model Is Dead Weight. It Took 8 Years for the Hardware to Catch Up. — Most neural network weights are redundant; efficient subnetworks can be trained directly with modern hardware and tooling.
• Mixture-of-Parallelisms: Towards Memory-Efficient Training Stack for Mixture-of-Experts Models — Mixture-of-Parallelisms (MoP) optimizes MoE model training by combining diverse parallelism techniques for memory and throughput efficiency.

See Counsel's argued verdicts on the open AI decisions leaders are weighing →

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by Tri Dao.