# sonic-moe is a MoE layer that only runs on Hopper and Blackwell, and treats IO as the bottleneck

> SonicMoE is a Mixture-of-Experts implementation built on CuTeDSL, Triton and the QuACK grouped GEMM kernels, aimed at NVIDIA Hopper and Blackwell GPUs, where the design priority is IO rather than arithmetic. It ships with several routing strategies and two weight layouts, and it is a young 0.1.x package.

**Dao-AILab/sonic-moe** — Accelerating MoE with IO and Tile-aware Optimizations

- Repository: https://github.com/Dao-AILab/sonic-moe
- Stars: 781 · Forks: 115
- Language: Python
- License: Apache-2.0
- Published: 2026-10-20 · Updated: 2026-10-20 · Language: en
- Canonical page: https://hysenlabs.com/projects/dao-ailab-sonic-moe

## A MoE implementation with a hard hardware floor

SonicMoE is a Mixture-of-Experts implementation optimized for NVIDIA Hopper, the SM90 generation, Blackwell datacenter parts at SM100 such as the B200 and B300, and Blackwell consumer cards at SM120 such as the RTX 5090. The README calls it a simple but blazing-fast implementation and attributes its approach to IO-aware optimizations.

The prerequisites are as narrow as the hardware list. CUDA 12.9 or newer, and 13.0 or newer for B300 GPUs, Python 3.12 or newer recommended, and PyTorch 2.11 or newer with 2.12 recommended. The package metadata agrees, with requires-python set to >=3.12 and a dependency on torch>=2.11.0.

This is a deliberately specialised project. There is no CPU fallback and no path for older architectures, because the whole point of the design is to use the specific capabilities of recent datacenter GPUs. Anyone evaluating it needs to know that before reading anything else.

The version line is young. The releases are 0.1.2, tagged Hopper and Blackwell GPUs, on 2026-04-20, and 0.1.0, tagged Hopper GPUs, on 2026-02-25, and the last push to the repository was on 2026-08-29. The setup.py in the repository carries the version 0.1.2.post1, so the source tree sits slightly ahead of the last tag.

## IO is the stated bottleneck, not arithmetic

The name and the pitch both point the same way. The repository description is Accelerating MoE with IO and Tile-aware Optimizations, and the README says the implementation uses CuTeDSL and Triton to deliver state-of-the-art performance through IO-aware optimizations.

The README supports that with two figures, one for activation memory usage and one for training throughput, measured on Hopper and on Blackwell. The images carry the numbers rather than the text, so the specific values live in the plots and the blog post linked from the news section rather than in the README prose. The blog post from 2026-04-22 covers the activation memory efficient and IO-aware design and how the project extends it to Blackwell through QuACK's software abstraction.

The claim to keep in mind when reading any of this is that these are the authors' own figures for their own implementation. The project is from the Dao-AILab group, whose work includes the grouped GEMM kernel library this project builds on, and the paper is on arXiv as 2512.14080 with a bibtex entry in the README, so the method is described in a form other people can check.

For an adopter, the interesting question is not whether the numbers are good but whether IO is the bottleneck in your workload.

## CuTeDSL, Triton and QuACK grouped GEMM kernels

The implementation stack is three layers, and the README names each.

CuTeDSL is NVIDIA's Python DSL for writing kernels in CuTe, and Triton is the open source compiler for GPU kernels. The README says the implementation mainly uses these two to deliver the IO-aware optimizations. Triton also shows up in the build tooling, where the Makefile has a debug target that sets TRITON_PRINT_AUTOTUNING=1 so you can watch the autotuner work.

The third layer is the more specific one. The current version of SonicMoE builds on the grouped GEMM kernels from the QuACK library, which is itself built on CUTLASS. The news entries describe this as the route to Blackwell support: on 2026-04-19 the project released Blackwell SM100 support built on QuACK's Grouped GEMM kernels, and the blog post explains how the IO-aware design was extended to Blackwell through QuACK's software abstraction.

That is why the package depends on quack-kernels, pinned at >=0.6.4 in both pyproject.toml and requirements.txt, and why there is a cu13 optional dependency group that pulls the CUDA 13 build of the same package.

The stack is not something you can swap out. SonicMoE is built on a specific kernel library targeting specific architectures, and the dependency chain is the design.

## The MoE class takes a routing count and a kernel backend

The usage is compact. You build a MoE layer with the expert counts and dimensions, move it to the GPU in bfloat16, and pass a kernel backend choice on the forward call:

```python
import torch
from sonicmoe import MoE, KernelBackendMoE
from sonicmoe.enums import ActivationType

# Create MoE layer
moe = MoE(
    num_experts=128,                           # Number of experts
    num_experts_per_tok=8,                     # Top-k experts per token
    hidden_size=4096,                          # Hidden dimension
    intermediate_size=1536,                    # Expert intermediate size
    activation_function=ActivationType.SWIGLU, # SwiGLU activation
    add_bias=False,                            # Add bias to linear layers
    std=0.02,                                  # Weight initialization std
).to(device="cuda", dtype=torch.bfloat16)

# Forward pass
x = torch.randn(32768, 4096, device="cuda", dtype=torch.bfloat16)
output, aux_loss = moe(x, kernel_backend_moe=KernelBackendMoE.sonicmoe)
```

Two things are worth noting in that snippet. The forward pass returns two values, the output and an auxiliary loss, which is the load balancing term a mixture-of-experts model needs and that most simple implementations leave out. And the kernel backend is chosen per call through the KernelBackendMoE enum, with sonicmoe as the value used here, so the fast path is opt-in rather than the only path.

The layer covers the usual configuration surface: expert count, top-k per token, hidden size, expert intermediate size, activation type, bias on the linear layers and the standard deviation for weight initialisation.

## Three routing strategies, selected on the command line

The benchmark scripts expose the routing variants, and the differences between them are the interesting part because routing strategy changes what the kernels have to do.

The first is TC top-K routing, described as softmax-over-topk, or softmax(topk(logits)):

```bash
python benchmarks/moe-cute.py --thiek 32768,4096,1024,128,8 --activation swiglu
```

The second is Qwen3-style routing, topk-over-softmax, or topk(softmax(logits)), with topk probability renormalization:

```bash
python benchmarks/moe-cute.py --thiek 32768,4096,1024,128,8 --topk_over_softmax --norm_topk_probs
```

The third is token rounding routing, with the SwiGLU activation:

```bash
python benchmarks/moe-token-rounding.py --routing nr --thiekq 16384,4096,1024,256,8,128
```

The difference between the first two is the order of the softmax and the top-k selection, and it is not cosmetic. Softmax over topk means the normalisation happens over the selected experts only, while topk over softmax normalises over all experts and then selects, which changes the probability values that the auxiliary loss sees. Models trained with one convention do not automatically transfer to the other, which is why the project ships both and why the Qwen3-style option comes with probability renormalization.

The naming of the flags is also inconsistent across the scripts, with the first using thiek and the third using thiekq. That is a small friction point worth knowing before you write automation around these commands.

## Interleaved by default, concatenated if your model is HuggingFace

The weight layout is the detail most likely to bite you, and the README is explicit about it.

By default, SonicMoE expects the gated up-projection weights in interleaved format, described as gate zero, up zero, gate one, up one and so on. HuggingFace models, with the README naming Qwen3, Mixtral and DeepSeek, store gate_up_proj in concatenated format, where all the gate weights come first and all the up weights follow. To use that layout there is a flag:

```bash
# Concatenated weight layout format with TC top-K routing
python benchmarks/moe-cute.py --thiek 32768,4096,1024,128,8 --concat_layout
```

So the default assumption is the opposite of what most open weight checkpoints on disk use. If you are loading a converted HuggingFace checkpoint, you either pass the concatenated layout or you convert the weights into the interleaved order yourself. The README does not document a conversion utility, so that is on you, and it is the kind of one-time cost that decides whether the package is convenient for your workflow.

The README only shows the flag on moe-cute.py, so treat the layout choice as per-command rather than assumed everywhere.

## Installing, testing and the cache the tuner needs

Installation from PyPI is one command:

```bash
pip install sonic-moe
```

Installing from source means cloning the repository, installing the requirements and installing the package in editable mode:

```bash
# Clone the repository
git clone https://github.com/Dao-AILab/sonic-moe.git
cd sonic-moe

# Install dependencies
pip install -r requirements.txt

# Install SonicMoE
pip install -e .
```

The Makefile has the same operations with the submodule step added, since install and install-dev both run git submodule update --init --recursive before pip install. The test target runs pytest tests, and the debug variant sets DEBUG_CUTOTUNE=1 and TRITON_PRINT_AUTOTUNING=1 so you can see the kernel autotuning decisions.

The cutotune-cache target is the one worth understanding. It runs a script that builds a cache of tuned kernel configurations with DEBUG_CUTOTUNE=1 and LOAD_CUTOTUNE_CACHE=1 and TORCH_CUDA_ARCH_LIST=9.0, which pins the cache to the Hopper architecture. Since the cache is architecture specific, a Blackwell user needs a different list, and that is the sort of detail that turns into a slow first run if you skip it.

The project is licensed under Apache 2.0, the repository ships a pre-commit configuration with ruff and pre-commit in the dev extras, and the citation block in the README points at arXiv 2512.14080.

## Conclusion

sonic-moe fits a team training mixture-of-experts models on H100, H200, GB200, B200, B300 or RTX 5090 hardware that wants the grouped GEMM path and the IO-aware activation memory handling rather than a generic PyTorch MoE layer. It is the wrong choice if you need CPU or older GPU support, if you need a mixture-of-experts implementation that runs on an architecture the package does not target, and it is the wrong choice if your weights are stored in the concatenated HuggingFace layout and you cannot afford the conversion cost, since the default expects interleaved gate and up weights. Verify first that your GPU and CUDA versions are on the supported list, that PyTorch is 2.11 or newer, and whether you need the cu13 extra for a CUDA 13 build.

## FAQ

### What hardware does sonic-moe need?

It targets NVIDIA Hopper GPUs such as the H100 and H200, Blackwell datacenter GPUs such as the GB200, B200 and B300, and Blackwell consumer cards such as the RTX 5090. The prerequisites are CUDA 12.9 or newer, 13.0 for B300, Python 3.12 or newer and PyTorch 2.11 or newer.

### How do I install sonic-moe?

Install the release with pip install sonic-moe, or clone the repository and run pip install -r requirements.txt followed by pip install -e . . The Makefile targets install and install-dev do the same and also run git submodule update --init --recursive.

### Which routing strategies does sonic-moe support?

The benchmarks cover TC top-K routing, described as softmax over topk, Qwen3-style routing, described as topk over softmax with probability renormalization, and token rounding routing. Each is selected through flags on the scripts in the benchmarks directory.

### What weight layout does sonic-moe expect for the expert projections?

By default it expects the gated up-projection weights in interleaved format, gate zero, up zero, gate one, up one. HuggingFace models such as Qwen3, Mixtral and DeepSeek store gate_up_proj in concatenated format, and you select that with the --concat_layout flag.

### How do I run the sonic-moe tests?

The Makefile test target runs pytest tests, and the test-debug variant sets DEBUG_CUTOTUNE=1 and TRITON_PRINT_AUTOTUNING=1 so the kernel autotuning output is visible. There is also a cutotune-cache target that prebuilds tuned kernel configurations for a given TORCH_CUDA_ARCH_LIST.

## Sources

- [Dao-AILab/sonic-moe on GitHub](https://github.com/Dao-AILab/sonic-moe)
- [Issues](https://github.com/Dao-AILab/sonic-moe/issues)
- [License: Apache-2.0](https://github.com/Dao-AILab/sonic-moe/blob/main/LICENSE)
- [README](https://github.com/Dao-AILab/sonic-moe/blob/main/README.md)
- [Releases](https://github.com/Dao-AILab/sonic-moe/releases)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/dao-ailab-sonic-moe