---
title: 'SonicMoE: Hardware-Aware MoE Optimization'
url: https://www.emergentmind.com/topics/sonicmoe
type: topic
---

# SonicMoE: Hardware-Aware MoE Optimization

SonicMoE is a hardware-aware optimization suite for Mixture of Experts (MoE) language models designed to address memory, I/O, and compute bottlenecks in fine-grained and highly sparse MoE layers on GPU accelerators. Combining a memory-efficient activation recomputation algorithm, GPU kernels that overlap high-bandwidth memory (HBM) I/O with tensor-core computation, and a novel tile-aware "token rounding" routing scheme, SonicMoE achieves large reductions in activation memory, hardware utilization improvements, and substantial end-to-end throughput gains—all while maintaining or slightly improving downstream model quality. The approach is implemented with open source kernels and demonstrated on Hopper GPUs with state-of-the-art large language models [2512.14080].

## 1. Motivating Challenges in Modern MoE Architectures

Recent MoE models increasingly feature high expert granularity—smaller intermediate expert dimensions ($n$) at fixed hidden size ($d$)—and higher sparsity, increasing the total number of experts ($E$) and keeping the per-token activated expert count ($K$) approximately constant. These design choices favor model quality per floating point operation (FLOP) but lead to several hardware and memory inefficiencies:

- **Activation Memory Scalability**: The activation memory footprint in a fine-grained MoE grows linearly with expert granularity ($G = d/n$); as $n$ decreases, retaining all necessary intermediates for the backward pass becomes increasingly costly.
- **Memory-Bound Arithmetic Intensity**: For SwiGLU-based MoE layers, the arithmetic intensity (AI, FLOPs/byte) sharply drops with increasing granularity or with reduced expert activation fraction ($\rho = K/E$), moving the computation into a memory bandwidth–bound regime.
- **Tile Quantization Overhead in Grouped GEMM**: Under high sparsity, the number of tokens per expert $T_e$ is often well below the GPU hardware-optimized matrix tile size ($M_{\rm tile}$), causing substantial wasted compute due to padding in grouped GEMM kernel dispatch.

These inefficiencies are not addressed by algorithmic or model-level innovations alone but demand architectural and kernel-level reengineering.

## 2. Memory-Efficient Forward and Backward Algorithm

SonicMoE reduces per-layer activation memory consumption to a constant in expert intermediate size $n$ by recomputing necessary quantities during backward without incurring additional large GEMM operations.

**Standard vs SonicMoE Activation Caching**:

- Naïve (ScatterMoE-style) caching stores $\{Y, X_e, H_e, A_e\}$ per expert, leading to $O(T K (d + n))$ bytes for $T$ tokens and $K$ experts activated per token.
- SonicMoE caches only $X$ ($T d$), $H$ ($2T K n$), and routing metadata, plus their gradients, resulting in a memory requirement of $M_{\rm sonic} = 2T d + 4T K n$ bytes and yielding a 45% reduction for a 7B MoE with $(d, n, K) = (1536, 256, 8)$.

**Forward/Backward Pass Pseudocode**:
```python
# Forward pass (pseudocode)
for expert e:
    X_e ← gather(X, π[:, e])
    H_e ← X_e @ W1_e
    A_e ← SwiGLU(H_e)
    Y_e ← A_e @ W2_e
O = sum over e of π[:, e] * S[:, e] * Y_e
cache X, H, π, S
return O

# Backward pass (pseudocode)
for expert e:
    dO_e ← gather(dO, π[:, e])
    dA'_e = dO_e @ W2_e.T
    A_e ← SwiGLU_forward(H_e)
    dH_e ← dSwiGLU(dA'_e, H_e)
    dS[:, e] = dot(dA'_e, A_e)
    A'_e = broadcast(S[:, e], A_e)
# ... further weight and input gradients as in original
```

This approach trades minimal recomputation for gross memory reduction, and ensures backward pass memory is constant in $n$.

## 3. IO–Computation Overlap and GPU Kernel Design

SonicMoE's GPU kernels are engineered to maximize utilization by overlapping memory I/O with compute using several concurrency and kernel fusion techniques:

- **Gather Fusion in Grouped GEMM Prologue:** The initial indexing and gathering stage is fused into the `cp.async` HBM load such that token data is packed in GPU shared memory (SMEM) in the correct order for matrix multiplication, efficiently staged by four producer warpgroups.
- **Epilogue Fusion for SwiGLU and dSwiGLU:** Both forward SwiGLU activation and its backward pass are fused into GEMM epilogues, obviating separate round-trips between HBM and SMEM for non-matrix activations.
- **Ping-Pong Scheduling:** On Hopper architecture, each GEMM kernel is managed by two consumer warpgroups in an alternating ping-pong cycle, with one executing the MMA main loop while the other handles epilogue I/O. This design realizes 20–40% TFLOPS recovery on memory-intensive workloads.

**Throughput Metrics (7B MoE, $d=1536, n=256, K=8, E=128$):**

| Kernel          | ScatterMoE | MoMoE | DeepGEMM++ | SonicMoE |
|-----------------|------------|-------|------------|----------|
| Forward up      |   380      |  365  |   420      |  450     |
| Forward down    |   290      |  275  |   340      |  485     |
| Backward $dH$   |   185      |  150  |   220      |  328     |
| Backward $dX$   |   270      |  240  |   300      |  355     |

SonicMoE achieves up to +17% forward and +50% backward throughput gains over state-of-the-art kernels.

## 4. Tile-Aware Token Rounding Router

To avoid waste in grouped GEMM tile quantization, SonicMoE introduces a two-stage "token rounding" router that guarantees each expert receives a token count divisible by the hardware tile size, minimizing padding overhead.

**Token Rounding Algorithm Outline**:
1. For each token, select top-$K$ expert scores.
2. Compute expert token counts $f_e$ and round to nearest or balanced multiple of $M_{\rm tile}$.
3. For each expert, retain the first $\text{round}[f_e]$ highest-scoring tokens.
4. Construct a new sparse score matrix enforcing tile-aligned counts.

Pseudocode:
```python
Input: S∈ℝ^{T×E}, K, tile
1. [S_topK, I_topK] = TopK(S, K)
2. f_e = count of tokens assigned to e
3. round[f_e] = floor(f_e/tile)*tile or ceil(f_e/tile)*tile
4. For each e, sort S[:, e] ⇨ π_e
5. Retain first round[f_e] tokens in π_e
Output: S', π' for GEMM dispatch
```

This routing strategy produces at most $M_{\rm tile}$ rounding error per expert and experimentally yields up to 1.16× (forward) kernel speedup at maximal sparsity ($K/E\approx1/128$), with gains of 1.20×/1.08× in Qwen3-80B class models.

## 5. Experimental Evaluation

SonicMoE's performance was evaluated on models ranging from 1.4B to 120B parameters with SwiGLU MoE layers, using Hopper H100 GPUs and the lm-engine codebase.

**Per-Layer Activation Memory (GiB):**

| Model   | ScatterMoE | MoMoE | SonicMoE |
|---------|------------|-------|----------|
| 1.4B    |   1.9      | 2.2   | 1.2      |
| 7B      |   2.6      | 3.0   | 1.4      |
| 30B     |   4.2      | 4.8   | 2.5      |
| 120B    |   6.0      | 7.3   | 3.1      |

**End-to-End Training Throughput** (7B MoE):
- ScatterMoE (96 $\times$ H100): 225B tokens/day
- SonicMoE (64 $\times$ H100): 213B tokens/day (1.86× compute throughput per GPU relative to ScatterMoE's BF16)

**Downstream Quality Metrics:**

| Method     | Val PPL | Avg Acc |
|------------|---------|---------|
| TC top-K   | 15.94   | 49.6%   |
| TR nearest | 15.92   | 50.0%   |
| EC (ft→TC) | 16.98   | 49.3%   |

Token rounding matches or exceeds top-K routing in perplexity and accuracy with no degradation in model performance for $K/E \leq 1/32$.

## 6. Architectural Generality and Broader Impacts

SonicMoE’s backward memory-saving and I/O-overlap kernels are applicable to any tiled sparse-compute architecture, such as block-sparse transformers and future GPU architectures leveraging large SMEM or TMEM (e.g., Blackwell). The token-rounding router is a direct substitute for top-K routing in MoE training and can potentially support hash-based or locality-sensitive expert selection at inference. No adverse effects on statistical efficiency or model quality are observed.

Planned extensions include expert parallelism across nodes with overlapping all-to-all and GEMM computation, FP8 support, and further cache optimizations for deep MoE architectures. A plausible implication is that as model sparsity increases or hardware tiling constraints become stricter, tile-aware routing and memory-efficient backward algorithms will be increasingly essential for scalable MoE deployment [2512.14080].

Source: https://www.emergentmind.com/topics/sonicmoe