SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations

ArXiv: 2512.14080

🎯 Pitch

Fine-grained sparse MoE models waste up to 45% of activation memory on padding alone, yet SonicMoE triples compute throughput by recomputing routing gradients on-the-fly instead of caching them. On Blackwell GPUs, this single insight beats DeepGEMM’s heavily optimized kernels by 25% on the forward pass.

---

1. Executive Summary

This paper proposes SonicMoE, a hardware–algorithm co-design that jointly optimizes MoE architecture and GPU kernels to address the training efficiency degradation caused by increasingly granular and sparse Mixture of Experts models. Evaluated across 1.4B–120B parameter MoE configurations on NVIDIA Hopper (H100) and Blackwell (B300) GPUs, SonicMoE introduces three named mechanisms: a memory-efficient backward-pass algorithm that avoids caching extra activations by recomputing score gradients from (\langle dA', A \rangle) rather than (\langle dO, Y \rangle) (reducing activation memory by 45% on a 7B fine-grained MoE versus ScatterMoE), IO-overlapping GPU kernels that employ Ping-Pong warpgroup scheduling and asynchronous TMA operations to hide memory latency behind Tensor Core MMA (achieving 1.86× compute throughput on Hopper and a 25% forward-pass speedup on Blackwell relative to DeepGEMM), and a token rounding routing algorithm that rounds per-expert token counts to multiples of the GEMM tile size to eliminate padding waste (yielding an additional 1.16× kernel speedup under high sparsity without degrading downstream task accuracy). The result is a 213B tokens/day training throughput on 64 H100s—comparable to ScatterMoE’s 225B tokens/day on 96 H100s—establishing that fine-grained sparse MoEs can be trained with near-dense-model activation efficiency only when the computation graph, memory I/O, and routing are co-optimized rather than addressed in isolation.

2. Context and Motivation

The Core Problem: Sparse MoEs Are Winning on Paper But Losing on Hardware

The paper addresses a widening gap between theory and practice in large language model training. MoE architectures have become the dominant paradigm for scaling model capacity: modern transformers replace dense feed-forward layers with a set of smaller "expert" subnetworks, activating only a subset per token. This gives MoEs a compelling FLOPs-efficiency argument—they can match or exceed the quality of dense models while using far fewer training FLOPs per token. However, the paper identifies that this FLOPs efficiency does not translate to hardware efficiency, and the gap is growing as MoE designs evolve in two specific directions.

Direction 1: increased expert granularity. Expert granularity, defined as (G = d/n) (the ratio of the model's embedding dimension to each expert's intermediate size), measures how "fine-grained" the experts are—larger (G) means each expert is smaller relative to the model's hidden dimension. MoE scaling laws (Clark et al. 2022; Krajewski et al. 2024; Tian et al. 2025) predict better model quality per FLOP as granularity increases while keeping total FLOPs constant (i.e., (nK) is held fixed while (n) decreases and (K), the number of activated experts, proportionally increases). Table 4 in the paper documents this trend across frontier models: Mixtral 8×22B (2023) had (G = 0.38), while Qwen3-Next-80B-A3B (2025) reaches (G = 4.00)—a tenfold increase in less than two years.

Direction 2: increased sparsity. MoE sparsity refers to the activation ratio (\rho = K/E) (activated experts divided by total experts). Modern MoEs are becoming sparser: Kimi K2 (2025) activates the same number of parameters as DeepSeek-V3 (2024) but has far more total experts, pushing (\rho) down. The paper shows (\rho) dropping from 25% in Mixtral 8×22B to 1.95% in Qwen3-Next-80B-A3B over the same two-year period.

Both trends improve model quality per FLOP but degrade hardware utilization per FLOP for three reasons the paper systematically analyzes:

1. Linear activation memory growth with granularity. When (n) decreases and (K) proportionally increases (iso-FLOPs), the activations that must be cached for the backward pass—particularly the down-projection output (Y) of size (TKd)—scale linearly with (K). A fine-grained 7B MoE with (n=256) requires roughly 45% more activation memory than a coarse-grained variant with (n=1024) under existing kernels like ScatterMoE, as shown in Figure 10.

2. Decreasing arithmetic intensity from increased IO. The paper derives the arithmetic intensity (FLOPs per transferred byte) for an MoE forward pass in Equation 4:

$\frac{3}{2 + \frac{2G}{d} + \frac{3}{T\rho}}$

As granularity (G) increases or sparsity (\rho) decreases, the denominator grows and arithmetic intensity drops. This happens because the IO cost scales with the number of activated experts (loading more expert weight matrices), while the FLOPs per expert shrink. Fine-grained sparse MoEs push the computation into the memory-bound regime, where kernel runtime is dominated by HBM access latency rather than compute throughput, making traditional GEMM optimization techniques insufficient.

3. Tile quantization waste under high sparsity. Modern GPUs perform matrix multiplications in fixed-size tiles (e.g., (M_{\text{tile}} = 128) on Hopper/Blackwell). When an expert receives a number of tokens (T_e) that is not a multiple of the tile size, the grouped GEMM kernel pads to the next tile multiple, wasting FLOPs on zero-valued entries. Under high sparsity (small (\rho)), the expected tokens per expert (\bar{T}_e = T\rho) becomes small, making the padding waste a significant fraction of total computation. Figure 8 shows this waste reaching 25% of total FLOPs when (E=256) with (K=4).

These three problems compound: a fine-grained, sparse MoE simultaneously requires more activation memory, generates more IO traffic relative to compute, and wastes more FLOPs on padding—all while the model architecture claims better FLOPs efficiency. The paper's core framing is that MoE architecture design and GPU kernel design have been optimized in isolation, creating a Pareto frontier where architects push toward granularity and sparsity for model quality while kernel developers optimize for coarse-grained, dense-expert regimes—and neither side accounts for the other's constraints.

Why This Matters: The Real Cost of Model Scaling

The practical stakes are significant and multidimensional:

Training economics. The paper reports that on a 7B fine-grained MoE ((n=256)), SonicMoE on 64 H100s achieves 213 billion tokens/day—matching ScatterMoE's throughput on 96 H100s (a 33% reduction in required GPUs). For organizations training large MoE models, a 33% improvement in hardware efficiency translates directly to millions of dollars in saved compute cost or significantly faster iteration cycles. The gap would only widen for models at the frontier (Table 4 shows the trend accelerating toward sparser, more granular architectures), meaning that without co-designed kernels, the next generation of MoEs could become proportionally more expensive to train despite requiring fewer theoretical FLOPs.

Architecture innovation constrained by hardware. The paper argues that current kernel limitations create a de facto ceiling on how granular and sparse MoEs can practically be. If training throughput degrades faster than model quality improves, the theoretical benefits of fine-grained sparsity cannot be realized. This is not hypothetical: the paper notes that ScatterMoE, MoMoE, DeepGEMM-pt, and DeepGEMM++ all fail to run the DeepSeek-V3.2-Exp configuration (a 685B MoE with (d=7168, n=2048, E=256, K=8)) on a single H100 due to index overflow or CUDA out-of-memory errors (Figure 12a), while SonicMoE runs successfully—suggesting that current kernels literally cannot support the architecture trends that scaling laws recommend.

Alternative architectures face the same bottleneck. The paper notes that recent non-MoE sparse architectures—PEER (He 2024), Memory Layers (Berges et al. 2024), and Ultra-Mem (Huang et al. 2025)—also adopt the philosophy of high granularity and sparsity for quality-per-FLOP, and will encounter the same IO and memory bottlenecks. The IO-overlapping and memory-efficient techniques developed in SonicMoE are therefore broadly applicable beyond classic MoE.

Prior Approaches and Where They Fall Short

The paper situates itself against five families of existing MoE kernels, identifying specific limitations for fine-grained sparse regimes:

ScatterMoE (Tan et al. 2024). Built on older Triton without TMA support, ScatterMoE fuses the token gather with the GMEM-to-SMEM load for varlen-M (forward) grouped GEMM but not for varlen-K (backward weight gradient) grouped GEMM. Its backward pass launches a separate gather kernel for (X) and (dO), adding (2TKd) bytes of IO. Critically, ScatterMoE computes the router score gradient as (dS = \langle dO, Y \rangle), which requires caching the full down-projection output (Y) of size (TKd) bytes—meaning activation memory scales linearly with (K) (expert count). It also does not overlap MMA with memory IO using Ping-Pong scheduling. On the 7B fine-grained MoE benchmark (Figure 5a), ScatterMoE's backward pass takes 3.968 ms versus SonicMoE's 2.173 ms—a 1.83× slowdown.

MoMoE (Costin et al. 2025). Shares ScatterMoE's limitation of no TMA-based IO, no gather fusion for varlen-K grouped GEMM, and the same (dS = \langle dO, Y \rangle) computation that forces caching (Y). MoMoE fuses (dS) with the up-projection activation gradient, but this fused kernel still takes longer than SonicMoE's separate (dS) computation because of the extra HBM loads for (Y). On the same 7B backward pass, MoMoE takes 4.678 ms (Figure 5a), a 2.15× slowdown versus SonicMoE.

MegaBlocks (Gale et al. 2023). Uses block-sparse matrix multiplication rather than grouped GEMM. Its ParallelDroplessMLP implementation first gathers and pads tokens, then launches block-sparse GEMM for up- and down-projection, then scatters results before aggregation. The separate gather and scatter kernels generate (8TKd) bytes of total IO, which becomes severely bottlenecked for fine-grained MoEs. Figure 5a shows MegaBlocks at 2.983 ms forward and 3.680 ms backward on the 7B benchmark—consistently slower than SonicMoE.

Megatron-LM (Shoeybi et al. 2019). Its GroupedMLP implementation uses CUTLASS grouped GEMM with JIT epilogue fusion but assumes contiguously-packed inputs, requiring a separate gather kernel. A recent patch added (S) weighting fusion with SwiGLU, adopting a computational path similar to SonicMoE's for memory efficiency, but the lack of gather fusion and absence of IO-compute overlap (no Ping-Pong scheduling, no asynchronous TMA store) still limit throughput. The paper benchmarks GroupedMLP at 3.476 ms forward (Figure 5a).

DeepGEMM (Zhao et al. 2025b). Highly optimized FP8 grouped GEMM kernels with fine-grained scaling, but its BF16 SM90 (Hopper) kernel assumes each expert receives a multiple of (M_{\text{tile}}) tokens (no dynamic tensor descriptor update), does not employ Ping-Pong scheduling, and assumes contiguously-packed inputs with no gather fusion. The SM100 (Blackwell) kernel similarly lacks gather fusion and uses tile shape ((128, 256, 64)) rather than SonicMoE's preferred ((256, 256, 64)) with 2-CTA MMA and CLC persistent scheduling. The paper constructs a "DeepGEMM++" baseline—the best possible MoE implementation using DeepGEMM's grouped GEMM as a black box with SonicMoE's optimized gather and expert aggregation kernels—and still shows SonicMoE outperforming it by 43% on the forward pass and 83% on the backward pass for fine-grained 7B MoE on H100 (Figure 11a).

Cross-cutting limitation: no difficulty or sparsity-aware routing. None of the prior kernels address the tile quantization problem from a routing perspective. DeepGEMM assumes per-expert token counts are already multiples of 128 (requiring padding in a preprocessing step); others simply accept the wasted FLOPs from non-aligned tiles. This is a missed opportunity because the routing decision determines the distribution of tokens to experts, and thus determines how much padding waste occurs. The paper's token rounding algorithm (Section 5) is the first to make routing tile-aware, eliminating padding waste with bounded deviation from the original routing decision.

How SonicMoE Positions Itself: Co-Design as the Missing Principle

The paper's central methodological claim is that MoE architecture and GPU kernel design must be co-optimized, not treated as independent problems. This co-design manifests at three levels:

Algorithm level: computational graph redesign. Rather than accepting the activation memory footprint dictated by autograd's default computation of (dS = \langle dO, Y \rangle), the paper mathematically derives an alternative path (dS = \langle dA', A \rangle) (Equation 10 and Appendix C) that produces identical gradients without materializing (Y). This is not an optimization that a compiler or autograd engine can discover—it requires understanding the mathematical structure of the MoE computation and intentionally restructuring the backward pass to avoid caching large intermediate tensors. The key insight is that (dA' = dO W_2^\top) is already computed during the down-projection activation gradient kernel (as an intermediate toward (dA)), and (A) is the SwiGLU output cached for the backward pass anyway—so computing their inner product requires zero extra HBM loads and zero extra cached activations, compared to the (2TKd) bytes required for the (\langle dO, Y \rangle) path.

Kernel level: IO-compute overlap for memory-bound regimes. The paper exploits Hopper and Blackwell GPU features—Ping-Pong warpgroup scheduling (Figure 6a), asynchronous TMA loads and stores, and Blackwell's TMEM two-stage pipeline (Figure 7)—to overlap memory IO with Tensor Core MMA. This is critical specifically for fine-grained MoEs where arithmetic intensity is low: when compute is fast relative to memory access, hiding memory latency behind ongoing computation prevents the GPU from stalling. The paper notes that while Ping-Pong scheduling and asynchronous TMA are known techniques (e.g., from Flash Attention 3), their application to MoE's heavy epilogue operations (SwiGLU, scatter/gather, (dS) computation) is novel and addresses the unique scaling behavior of MoE IO costs with granularity.

Routing level: hardware-aware token assignment. Token rounding (Algorithm 4) modifies the traditional top-(K) token-choice routing to produce per-expert token counts that are multiples of the GEMM tile size. This is conceptually simple—decide per expert whether to drop the lowest-scoring token-choice tokens or pad with expert-choice tokens to reach a tile boundary—but it reframes routing as a hardware-efficiency decision, not just a model-quality decision. The paper shows this eliminates wasted padding FLOPs (Figure 8) while preserving model quality (Table 2), and the speedup grows with sparsity (16% at (K/E = 1/64) in Figure 13).

The paper explicitly draws a contrast with prior work that treats kernel optimization as a post-hoc engineering effort separate from model design. ScatterMoE and MoMoE prioritize development convenience (implemented in Triton) over hardware utilization; DeepGEMM assumes contiguously-packed inputs that require preprocessing (suitable for distributed expert-parallel training with all-to-all communication, but inefficient for single-GPU or tensor-parallel regimes); MegaBlocks and Megatron treat the routing output as fixed and optimize only the GEMM. SonicMoE's thesis is that the routing decision, the computation graph, and the GPU kernel must be designed together to achieve near-peak hardware utilization for fine-grained sparse MoEs—and the empirical results (1.86× throughput on Hopper, 25% forward speedup on Blackwell) validate that the gap between existing approaches and hardware limits is large enough to justify this level of co-design investment.

3. Technical Approach

3.1 Reader Orientation

SonicMoE is a hardware–algorithm co-design system that compiles a Mixture of Experts (MoE) layer specification into optimized GPU kernels and an adaptive routing strategy. The core problem it solves is that fine-grained and sparse MoEs—which are theoretically more FLOPs-efficient than dense models—suffer severe hardware underutilization because their arithmetic intensity drops as experts become smaller and more numerous, making memory access latency the dominant bottleneck. The solution is three-pronged: restructure the backward computation graph to avoid caching large intermediate tensors that scale with expert count, overlap memory I/O with matrix multiplication using hardware-asynchronous operations to hide latency, and adapt the token routing to produce per-expert token counts that align with GPU tile boundaries, eliminating wasted padding computation.

3.2 Big-Picture Architecture (Diagram in Words)

SonicMoE consists of five major components that operate on a single MoE layer in sequence during training:

1. Top-K Router: accepts token embeddings $X \in \mathbb{R}^{T \times d}$ and produces a sparse routing mask $\pi \in \{0,1\}^{T \times E}$ and per-token expert scores $S \in \mathbb{R}^{T \times E}$, optionally modified by the token rounding algorithm to align per-expert token counts with tile boundaries.

2. Forward Up-Projection Kernel ("A kernel"): takes the input tokens $X$, the up-projection weight $W_1$, and the routing mask $\pi$; gathers tokens to each expert, performs a varlen-M grouped GEMM (variable-length token dimension, fixed weight dimensions), applies the SwiGLU activation in the epilogue, and stores both the pre-activation $H$ and post-activation $A$ to HBM.

3. Forward Down-Projection Kernel ("Y kernel"): takes the post-activation $A$ and down-projection weight $W_2$; performs a second varlen-M grouped GEMM; stores the result $Y$ via asynchronous TMA store without scattering to original token positions.

4. Expert Aggregation Kernel ("O kernel"): gathers each token's expert outputs $Y_e$ from their contiguous storage locations, multiplies by expert scores $S$, sums across activated experts, and stores the final layer output $O$.

5. Backward Kernels (5 launches): the down-projection activation gradient (dH kernel) computes $dH, dS, A'$ with heavy epilogue fusion including SwiGLU backward; the down-projection weight gradient (dW2 kernel) performs varlen-K grouped GEMM by gathering $dO$ and multiplying against $A'$; the up-projection activation gradient (d\tilde{X} kernel) performs varlen-M grouped GEMM from $dH$ and stores via TMA; the up-projection weight gradient (dW1 kernel) gathers $X$ and performs varlen-K grouped GEMM; and the expert aggregation gradient (dX kernel) gathers and sums per-expert $d\tilde{X}_e$.

The critical architectural principle is that every GEMM kernel supports both contiguously-packed inputs and inputs gathered from arbitrary positions via fused gather operations, and every IO-heavy operation (epilogue stores, SwiGLU/dSwiGLU computation, $dS$ reduction) is overlapped with Tensor Core MMA through warpgroup-level scheduling or TMEM pipelining.

3.3 Roadmap for the Deep Dive

• First, the memory-efficient backward pass algorithm (Section 3.4.1), because it determines what activations must be cached and what computations occur in each kernel—this is the architectural constraint that all other components respect.
• Second, the design of SonicMoE's grouped GEMM kernels (Section 3.4.2), covering gather fusion with prologue loading and SwiGLU/dSwiGLU/dS fusion with epilogue—these are the mechanisms that reduce unnecessary IO access at the granular level of individual instructions.
• Third, the IO-overlapping execution strategies for Hopper and Blackwell GPUs (Section 3.4.3), because these take the fused kernels and schedule their memory and compute operations to run concurrently, exploiting the asynchrony that the previous section left available.
• Fourth, the token rounding routing algorithm (Section 3.4.4), which modifies the inputs to the kernel pipeline by adjusting per-expert token counts to multiples of the tile size, eliminating wasted padding and shifting the tradeoff between routing fidelity and hardware efficiency.
• Fifth, the combined training throughput analysis and end-to-end system integration (Section 3.4.5), showing how all components compose into measurable speedups and memory savings across model scales and GPU generations.

3.4 Detailed, Sentence-Based Technical Breakdown

This is a systems paper whose core contribution is a set of three co-designed optimizations—memory-efficient computation graph restructuring, IO-overlapping kernel scheduling, and tile-aware routing—that together close the gap between the theoretical FLOPs-efficiency of fine-grained sparse MoEs and their actual hardware utilization on modern GPUs. The design space is explored through concrete kernel implementations in CuTe-DSL with a PyTorch interface, validated by extensive profiling across model scales from 1.4B to 120B parameters on both Hopper (H100) and Blackwell (B300) architectures.

---

3.4.1 Memory-Efficient Backward Pass: Restructuring the Computation Graph

The first major contribution is an alternative computation path for the MoE backward pass that eliminates the need to cache the down-projection output $Y$ and its gradient $dY$, which are the primary sources of activation memory growth as expert granularity increases.

The problem with the standard backward pass. In a standard MoE layer with SwiGLU activation, the forward pass for an expert $e$ computes:

$H_e = X_e W_{1,e} \quad\text{(up-projection, output shape $\mathbb{R}^{T_e \times 2n}$)}$

$A_e = \text{SwiGLU}(H_e) \quad\text{(activation, output shape $\mathbb{R}^{T_e \times n}$)}$

$Y_e = A_e W_{2,e} \quad\text{(down-projection, output shape $\mathbb{R}^{T_e \times d}$)}$

$O_t = \sum_{e} \pi_{t,e} s_{t,e} Y_{e,t} \quad\text{(expert aggregation, output shape $\mathbb{R}^{T \times d}$)}$

where $X_e \in \mathbb{R}^{T_e \times d}$ is the input received by expert $e$ with $T_e$ tokens, $W_{1,e} \in \mathbb{R}^{d \times 2n}$ and $W_{2,e} \in \mathbb{R}^{n \times d}$ are the expert's weight matrices, $\pi_{t,e}$ is the binary routing mask, and $s_{t,e}$ is the routing score. The expert aggregation multiplies each expert's output by the routing score and sums across the $K$ activated experts for each token.

The backward pass must compute three categories of gradients:

• Weight gradients $dW_{1,e}$ and $dW_{2,e}$ for optimizer updates.
• Activation gradient $dX$ for backpropagation to the previous layer.
• Router score gradient $dS$ for router training.

The naive autograd approach computes $dS$ directly from the forward outputs as:

$dS_{t,e} = \langle dO_t, Y_{e,t} \rangle$

where $dO_t$ is the incoming gradient from the next layer and $\langle \cdot, \cdot \rangle$ denotes the dot product over the hidden dimension $d$. This requires $Y$—the cached down-projection output—to be stored in HBM during the forward pass and loaded during the backward pass. Since $Y$ has size $TKd$ bytes (total tokens × activated experts × hidden dimension), its memory footprint scales linearly with $K$, the number of activated experts. Under iso-FLOPs scaling where $nK$ is constant and $n$ decreases while $K$ increases, this means activation memory grows proportionally with granularity.

Similarly, the weight gradient $dW_{2,e}$ naively requires $dY_e = \text{Broadcast}(s_e) \, dO_e$, the gathered incoming gradient multiplied per-token by the expert score. Caching $dY_e$ adds another $TKd$ bytes of peak activation memory during the backward pass.

SonicMoE's restructured computation for $dS$. The paper derives an algebraically equivalent expression that avoids materializing $Y$ and $dY$ entirely. Starting from the definition of $Y_e = A_e W_{2,e}$, the gradient with respect to the score is:

$dS_{t,e} = \langle dO_t, Y_{e,t} \rangle = \langle dO_t, (A_e W_{2,e})_t \rangle$

Define $dA'_e = dO_e W_{2,e}^\top \in \mathbb{R}^{T_e \times n}$ as the grouped GEMM output computed during the down-projection activation gradient kernel. This is the term that the backward pass already computes as an intermediate toward $dA$. Then:

$dS_{t,e} = \langle dA'_{e,t}, A_{e,t} \rangle$

where the inner product is now over the expert intermediate dimension $n$ rather than the hidden dimension $d$. This form requires:

• $dA'_{e,t}$, which is computed during the dH kernel as $dA'_e = dO_e W_{2,e}^\top$ and is already in registers or SMEM.
• $A_{e,t}$, the cached SwiGLU output (size $TKn$ bytes), which must be stored for the $dW_2$ computation anyway.

The key efficiency gain: neither $Y$ nor $dY$ appears in this expression. The dot product reduces over $n$ (typically 256–1024) rather than $d$ (typically 2048–7168), requiring fewer reduction steps by a factor of $\log_2(d/n)$. More importantly, the computation is fused into the dH kernel's epilogue (Section 3.4.2), requiring zero additional HBM loads beyond those already needed for $dA'_e$ and $A_e$.

SonicMoE's restructured computation for $dW_2$. The weight gradient for the down-projection matrix is:

$dW_{2,e} = A_e^\top dY_e = A_e^\top (\text{Broadcast}(s_e) \, dO_e) = (\text{Broadcast}(s_e) \, A_e)^\top dO_e$

Define $A'_e = \text{Broadcast}(s_e) \, A_e \in \mathbb{R}^{T_e \times n}$ as the post-activation multiplied elementwise by the expert scores. This is computed in the dH kernel's epilogue from the already-materialized $A_e$ and the routing scores $s_e$. The weight gradient is then $dW_{2,e} = A'^\top_e dO_e$, a standard varlen-K grouped GEMM that gathers $dO$ and multiplies against $A'$. This eliminates the need to cache or materialize $dY_e$.

SonicMoE's restructured computation for $dH$. The activation gradient for the pre-activation $H_e$ is:

$dA_e = dY_e W_{2,e}^\top = \text{Broadcast}(s_e) \, (dO_e W_{2,e}^\top) = \text{Broadcast}(s_e) \, dA'_e$

where $dA'_e$ is the same intermediate from the grouped GEMM. Then $dH_e = \text{dSwiGLU}(dA_e, H_e)$, the backward of the SwiGLU activation function. This requires $H_e$ (the cached up-projection output, size $2TKn$ bytes) and $dA'_e$ (computed in the current kernel). The result $dH_e$ is stored for the subsequent up-projection backward kernels.

What is cached and what is not. After this restructuring, the only tensors cached during the forward pass for the backward pass are:

• $X$: the input token embeddings (size $Td$ bytes).
• $H$: the up-projection output before SwiGLU (size $2TKn$ bytes).
• $\pi$ and $S$: routing metadata (size $4TKn$ bytes total, since only the sparse $K$ entries per token are stored).

The total cached activation memory is $2Td + 4TKn$ bytes per MoE layer. Crucially, both $Y$ and $dY$ (each $TKd$ bytes) are no longer cached. In a dense model with the same number of activated parameters, the required cached activations are also $2Td + 2Td_{\text{ff}}$ (pre- and post-activation of the feed-forward layer), which is identical in spirit. SonicMoE thus achieves the minimum activation memory needed for backward computation without performing GEMM recomputation during the backward pass. As shown in Figure 10, for a 7B MoE with $n=256$, this reduces peak activation memory by 45% versus ScatterMoE (from roughly 2.4 GiB to 1.3 GiB per layer), and the memory usage remains constant as granularity increases (Figure 1, left panel), compared to the linear growth seen in ScatterMoE and MoMoE.

Why not rely on autograd? The paper notes that standard PyTorch autograd cannot discover this alternative computation path. When the expert weighting $S$ is applied before or after the down-projection, autograd will follow the forward computation graph and produce $dS = \langle dO, Y \rangle$ or $dS = \langle dA', A \rangle$ respectively. However, applying $S$ after down-projection (which autograd would need to do to discover the efficient path) changes the forward computation order and may interact incorrectly with other operations. The paper explicitly designs the backward kernels to use the efficient $\langle dA', A \rangle$ reduction without requiring the forward pass to be reordered, effectively overriding autograd's default gradient computation for these specific operations.

A subtlety: the $Y$ temporary. SonicMoE still materializes $Y$ temporarily during the forward pass (as the output of the down-projection kernel), but $Y$ can be recycled after each layer's forward pass completes. As long as the number of MoE layers (typically 32+ for modern architectures) exceeds $K$, the transient memory usage of $Y$ across layers is amortized and does not dominate peak activation memory. The paper acknowledges that removing this materialization entirely would require atomic-add operations in the expert aggregation kernel (Figure 17, right), which introduces nondeterminism, numerical accuracy issues for BF16 accumulations, and incompatibility with all-to-all communication in expert-parallel distributed training—hence the pragmatic choice to tolerate a temporary $Y$ buffer.

---

3.4.2 SonicMoE's Grouped GEMM Kernel Design: Fusion and Epilogue Computation

The second major contribution is a set of kernel-level fusions that reduce unnecessary HBM traffic by combining token gathering, activation functions, and gradient computations into the prologue and epilogue stages of the grouped GEMM kernels. These fusions are applied systematically across all 8 launched kernels (3 forward, 5 backward) and are critical because IO cost—not compute throughput—dominates runtime for fine-grained MoEs.

Design philosophy: minimal launches, maximal fusion. Each of SonicMoE's 8 kernels performs a single logical operation, but within each kernel, multiple sub-operations are fused to avoid redundant HBM reads and writes. The host (CPU) dispatches to the best GEMM configuration (tile sizes, warpgroup scheduling mode, load/store strategies) for each kernel based on the tensor shapes. The kernel implementation is modularized into two primitives: (1) an optimized grouped GEMM kernel with configurable prologue and epilogue fusion, and (2) an optimized expert aggregation kernel. This modularity means the same grouped GEMM kernel handles both forward up-projection (with SwiGLU fusion) and backward down-projection activation gradient (with dSwiGLU, $dS$, and $A'$ fusion) by parameterizing the epilogue operations.

Gather Fusion with GMEM-to-SMEM Load

For any grouped GEMM whose input tensors are not contiguously packed—specifically, the forward up-projection (which gathers $X$ to each expert), the backward weight gradients $dW_1$ and $dW_2$ (which gather $X$ and $dO$ respectively), and the backward activation gradient $dH$ (which gathers $dO$)—SonicMoE fuses the token gathering operation directly into the GMEM-to-SMEM load stage of the GEMM prologue and mainloop.

How gather fusion works on Hopper GPUs. The GEMM kernel first fetches the routed token indices for each expert from HBM. For each tile of the $M$ dimension (token dimension), the producer warps use the cp.async instruction to issue asynchronous copy commands that read from non-contiguous HBM addresses (one per token in the tile, indexed by the gathered positions) and write contiguously into SMEM. The cp.async instruction is asynchronous with respect to the issuing warp—the copy is handed off to the memory controller, and the warp can proceed with other work (e.g., issuing more copies for the next tile) while the data is in flight. Once all copies for a tile complete (signaled by the cp.async completion group), the consumer warpgroup reads the contiguously-packed data from SMEM and feeds it to the Tensor Core MMA instructions.

The fusion eliminates a separate gather kernel that would otherwise need to (1) read $X$ from HBM at scattered positions, (2) write contiguously-packed $X_e$ back to HBM, and then (3) have the GEMM kernel read $X_e$ from HBM. The fused approach eliminates one HBM write ($TKd$ bytes, since every token is gathered to its $K$ experts) and one HBM read ($TKd$ bytes for the GEMM input), saving $2TKd$ bytes of IO per affected kernel. For fine-grained MoEs where IO cost dominates (Equation 4 shows arithmetic intensity decreases with $G$), this IO reduction directly improves throughput. Figure 5a shows that on the 7B fine-grained MoE forward pass, SonicMoE's gathered up-projection runs in 0.4 ms versus DeepGEMM++'s separate gather (0.3 ms) plus grouped GEMM (0.5 ms) totaling 0.8 ms.

Backward gather fusion: the missing piece in prior work. ScatterMoE and MoMoE both fuse the forward gather but launch a separate gather kernel for backward weight gradients $dW_1$ and $dW_2$, contributing to their slower backward pass. The paper identifies this as a major differentiator: SonicMoE fuses the gather for $X$ in $dW_1$ and for $dO$ in both $dW_2$ and $dH$. This saves $2TKd$ bytes for each weight gradient kernel that needs gather, and the savings compound because the backward pass has more kernels requiring gather than the forward pass.

Gather fusion challenges on Blackwell GPUs with 2-CTA clusters. Blackwell introduces architectural complexity for gather fusion. SonicMoE uses 2-CTA clusters for GEMM computation (two cooperative thread arrays share SMEM and coordinate MMA), but the cp.async instruction can only signal completion within the same CTA. When gather copies are issued from both CTAs to load different portions of the tile, the MMA warp in the leader CTA (CTA 0) needs to know when both CTAs' copies have completed before it can begin computation. The workaround (described in Section 4.1.1 and Figure 4) involves a dedicated "relay warp" in CTA 1 that receives the cp.async completion signal locally, then forwards it to CTA 0's MMA warp using cluster-level synchronization primitives (specifically, mbarrier with cluster scope). This relay mechanism adds scheduling complexity but is necessary to enable efficient gather fusion across the 2-CTA cluster while maintaining the throughput advantages of clustered MMA.

Epilogue Fusion: SwiGLU, dSwiGLU, dS, and A'

The epilogue stage of a GEMM kernel executes after the mainloop completes the matrix multiplication accumulation. SonicMoE exploits the epilogue heavily to fuse operations that would otherwise require separate kernel launches and intermediate HBM stores.

Forward up-projection epilogue: SwiGLU fusion. After the up-projection grouped GEMM computes $H_e = X_e W_{1,e} \in \mathbb{R}^{T_e \times 2n}$, the SwiGLU activation must split this into two halves of size $n$, apply a SiLU (Sigmoid Linear Unit) nonlinearity to the first half, and multiply elementwise by the second half. SonicMoE fuses this into the epilogue, writing the pre-activation $H_e$ to HBM (needed for the backward SwiGLU computation) and the post-activation $A_e$ to HBM in a single kernel launch. The alternative—writing $H_e$ to HBM and launching a separate activation kernel—would add $2TKn$ bytes of extra HBM writes and reads.

Forward down-projection epilogue: asynchronous TMA store. The down-projection kernel computes $Y_e = A_e W_{2,e}$ and must store the result for the expert aggregation kernel. Rather than fusing the scatter-to-original-positions with the store (which ScatterMoE and MoMoE do using synchronous st.global instructions), SonicMoE stores contiguously-packed $Y_e$ via asynchronous TMA (Tensor Memory Accelerator) store operations. The TMA store is non-blocking with respect to the issuing warp, allowing the next tile's MMA to proceed while the store is in flight (Figure 16, top). The scatter to original positions is deferred to the expert aggregation kernel, where each token gathers rather than each expert scattering—this gather approach is compatible with asynchronous TMA loads and avoids the 20% TFLOPS degradation that synchronous stores impose on subsequent MMA tiles (Figure 16, bottom, and Figure 21).

Backward down-projection activation gradient (dH) epilogue: the heaviest fusion. The dH kernel is the most fusion-intensive kernel in SonicMoE. Its mainloop computes $dA'_e = dO_e W_{2,e}^\top \in \mathbb{R}^{T_e \times n}$. The epilogue then performs all of the following operations within the same kernel launch, without intermediate HBM stores:

1. Load $H_e$ and $s_e$: The pre-activation $H_e$ (cached from forward) and the routing scores $s_e$ are loaded from HBM. This load is performed asynchronously via TMA, overlapped with other epilogue operations across epilogue stages.

2. Compute $dA_e$: Multiply $dA'_e$ elementwise by the broadcasted scores: $dA_e = \text{Broadcast}(s_e) \, dA'_e$. This scales the activation gradient by the per-token routing weight.

3. Compute $dH_e$ and recompute $A_e$: Execute the fused forward-backward of SwiGLU via $\text{dAct func}(dA_e, H_e)$. This operation simultaneously computes:

   • $dH_e$: the gradient with respect to the pre-activation $H_e$ (needed for up-projection backward).
   • $A_e$: the forward SwiGLU output, recomputed from the cached $H_e$ (needed for $dS$ and $A'$).

   Recomputing $A_e$ avoids caching it separately—since $H_e$ is already cached for the dSwiGLU operation, $A_e$ can be regenerated from $H_e$ with negligible compute overhead (SwiGLU is a lightweight elementwise operation relative to the GEMM).

4. Compute $dS_{t,e}$: Reduce the inner product $\langle dA'_{e,t}, A_{e,t} \rangle$ over the $n$ dimension (size 256–1024). The reduction is performed by warp-level primitives (shuffle and reduce) within each token's thread group, producing a scalar per token-expert pair.

5. Compute $A'_e$: Multiply the recomputed $A_e$ elementwise by the broadcasted scores: $A'_e = \text{Broadcast}(s_e) \, A_e$. This is the input for the downstream $dW_2$ kernel.

6. Store results: The kernel writes $dH_e$, $dS$, and $A'_e$ to HBM. $dS$ is a sparse tensor (size $TK$ scalars), $dH_e$ has size $2TKn$ bytes, and $A'_e$ has size $TKn$ bytes.

This heavy epilogue fusion—loading $H_e$ and $s_e$, executing dSwiGLU, computing $dS$ and $A'$, and storing three outputs—would, if implemented as separate kernels, require multiple round-trips to HBM (load $H_e$ and $s_e$, write $dH_e$, load $A_e$, write $dS$, load $A_e$ and $s_e$, write $A'_e$). By fusing all operations into a single kernel launch, the intermediate data ($A_e$, $dA_e$) stay in registers or SMEM, and only $H_e$ and $s_e$ are loaded while $dH_e$, $dS$, and $A'_e$ are written. Figure 5a shows that SonicMoE's fused dH kernel takes 0.6 ms on the 7B benchmark, while ScatterMoE's separate down-proj act, dS, and dSwiGLU kernels together take 1.8 ms.

Backward up-projection activation gradient (d$\tilde{X}$) epilogue: TMA store. The $d\tilde{X}$ kernel computes $d\tilde{X}_e = dH_e W_{1,e}^\top$ (a varlen-M grouped GEMM) and stores contiguously-packed results via asynchronous TMA, following the same design principle as the forward down-projection kernel. The scatter to original positions is deferred to the dX expert aggregation kernel, where each token gathers its per-expert gradients and sums them. This avoids the scatter fusion that ScatterMoE and MoMoE use (synchronous st.global), maintaining high MMA utilization.

Backward weight gradient kernels: varlen-K grouped GEMM with gather fusion. The $dW_1$ kernel gathers $X$ and performs varlen-K grouped GEMM (fixed $M=d, N=2n$, reducing over $K=T_e$) to produce $dW_{1,e} = X_e^\top dH_e$. The $dW_2$ kernel gathers $dO$ and performs varlen-K grouped GEMM to produce $dW_{2,e} = A'^\top_e dO_e$, where $A'$ was already computed and stored by the dH kernel's epilogue. Both weight gradient kernels use gather fusion with cp.async (Hopper) or the relay mechanism (Blackwell 2-CTA clusters). This eliminates the separate gather kernels that ScatterMoE and MoMoE require for both operations.

Expert aggregation kernels: gather-and-sum. Both the forward aggregation kernel (O kernel) and backward aggregation kernel (dX kernel) follow the gather-and-sum pattern (Figure 17, left): each token loads its activated experts' outputs from contiguously-packed arrays (stored by the preceding GEMM kernels via TMA), multiplies by routing scores (forward only), sums, and writes the result. The forward kernel reads $TKd$ bytes (all expert outputs) and writes $Td$ bytes; the backward kernel reads $TKd$ bytes (all per-expert activation gradients) and writes $Td$ bytes. Both are purely memory-bound operations. Figure 20a shows SonicMoE's forward aggregation kernel achieves 2.92× the bandwidth of ScatterMoE's torch.bmm-based aggregation on H100 GPUs, and 0.98× the bandwidth of a heavily optimized Triton kernel that sums over contiguous $Y$ (the bandwidth upper bound for this operation).

Why fusion is critical now but was not before. The paper notes that these fusions are particularly impactful for fine-grained MoEs because of the arithmetic intensity trends derived in Equation 4. As granularity $G = d/n$ increases, the IO cost per FLOP grows because more expert weight matrices must be loaded relative to the compute performed per expert. Epilogue fusion reduces the IO overhead of activation functions and gradient computations that, in coarse-grained MoEs (large $n$, small $K$), were negligible compared to the GEMM mainloop but now represent a significant fraction of total kernel time. The Ping-Pong scheduling and asynchronous TMA operations described in Section 3.4.3 are the mechanisms that make these fusions practical, by hiding the latency of the fused epilogue operations behind ongoing MMA computation.

---

3.4.3 IO-Compute Overlap: Warpgroup Scheduling and Asynchronous Operations

The third major contribution is a set of scheduling strategies that overlap memory I/O operations (HBM reads and writes, TMA loads and stores) with Tensor Core matrix multiply-accumulate (MMA) operations. These strategies exploit the asynchronous execution capabilities of Hopper and Blackwell GPUs to hide memory latency, which is the dominant bottleneck for fine-grained sparse MoEs.

The fundamental challenge. GEMM on modern GPUs is pipelined: producer warps load the next tile of data from HBM while consumer warps perform MMA on the current tile. This overlap is standard in optimized GEMM kernels and handles the mainloop well. However, the epilogue—which occurs after the mainloop finishes and involves activation functions, reductions, and HBM writes—is traditionally executed synchronously, stalling the MMA pipeline. For compute-bound kernels (large $n$, coarse-grained MoEs), the epilogue is a small fraction of total time and this stall is insignificant. For memory-bound kernels (small $n$, fine-grained MoEs), the epilogue can dominate runtime, and the stall wastes hardware that could be computing the next tile.

Ping-Pong scheduling on Hopper GPUs (Figure 6a). SonicMoE adopts Ping-Pong scheduling (Shah et al. 2024; Wright and Hoque 2024b) specifically for kernels with heavy epilogues: the forward down-projection (Y kernel) and the backward down-projection activation gradient (dH kernel). The key idea is to use two consumer warpgroups that alternate roles:

1. Consumer warpgroup 0 receives MMA accumulation results for its current work tile, launches its epilogue (which includes heavy operations like SwiGLU/dSwiGLU, $dS$ reduction, and TMA stores), and signals completion to the persistent tile scheduler.
2. While consumer warpgroup 0 is executing its epilogue, consumer warpgroup 1 receives its own accumulation results from the mainloop and begins its MMA for the next work tile. The TMA loads for warpgroup 1's inputs are also in flight during this time.
3. When warpgroup 0's epilogue finishes and warpgroup 1's MMA for its tile finishes, the roles swap: warpgroup 0 begins MMA on a new tile while warpgroup 1 executes its epilogue.

This effectively hides the epilogue latency behind the other warpgroup's MMA computation. The TMA operations within each warpgroup's epilogue (e.g., loading $H_e$ for the dH kernel's epilogue) are also overlapped with the other warpgroup's MMA, further reducing effective stall time.

The paper notes that Ping-Pong scheduling is not universally applied—it is used only for the forward down-projection kernel (where the epilogue TMA store is large relative to the mainloop) and the backward dH kernel (where multiple epilogue operations—dSwiGLU, $dS$ reduction, $A'$ computation—create a heavy epilogue). For other kernels (forward up-projection, backward activation gradients, weight gradients), the epilogue is lightweight enough that cooperative scheduling (a single consumer warpgroup with standard pipelining) suffices.

Asynchronous TMA operations within epilogues. Beyond Ping-Pong scheduling, SonicMoE uses asynchronous TMA operations at a finer granularity within individual epilogues:

• Asynchronous TMA load during dH kernel epilogue: the dH kernel needs to load $H_e$ (the cached pre-activation) to compute dSwiGLU. SonicMoE creates a dedicated pipeline for this TMA load, issuing the load early in the epilogue and overlapping the load latency with other register-based operations (the $dS$ reduction via warp shuffles, the $A'$ multiplication). This ensures that by the time the epilogue needs $H_e$ values, they are already in SMEM.

• Asynchronous TMA store in forward Y and backward d\tilde{X} kernels: rather than using synchronous st.global instructions (which block the warp until the store completes), SonicMoE uses TMA store operations that are handed off to the memory controller. The warp can begin loading the next tile's inputs or proceed with the next MMA while the TMA store completes in the background. Figure 16 illustrates the difference: TMA store allows the next MMA tile to begin immediately (top), while st.global forces the warp to wait for store completion before the next MMA (bottom), resulting in a ~20% TFLOPS degradation (Figure 21).

• Scatter deferred to gather: because SonicMoE does not fuse scatter with HBM store (unlike ScatterMoE and MoMoE), it avoids the synchronous st.global bottleneck entirely. Instead, the results are stored contiguously via TMA, and the expert aggregation kernel gathers from these contiguous buffers (Figure 17, left). This gather approach uses asynchronous TMA loads, maintaining high bandwidth (Figure 20a shows 2.92× ScatterMoE's bandwidth on average) with no MMA stall penalty.

Blackwell's TMEM two-stage pipeline (Figure 6b and 7). Blackwell GPUs introduce Tensor Memory (TMEM), a dedicated 256KB on-chip memory per SM organized as 128 rows × 512 columns of 32-bit cells. This architecture fundamentally changes how MMA-to-epilogue overlap works:

1. The accumulator results from UMMA (Unified Matrix Multiply-Accumulate, Blackwell's single-threaded asynchronous MMA instruction) are stored directly in TMEM rather than in registers, freeing register pressure for other uses.
2. The 512-column TMEM structure is split into two 256-column stages. While one stage accumulates MMA results for the current tile (written by the MMA warp), the other stage holds the completed accumulation from the previous tile, which the epilogue warps can read and process.
3. Ownership transfer is coordinated via signaling (green and yellow arrows in Figure 6b): the MMA warp signals the epilogue warps when a new tile's accumulation is ready (green arrow), and the epilogue warps signal back when they finish processing a stage and it becomes available for the next MMA tile (yellow arrow).

This two-stage pipeline is conceptually similar to Ping-Pong scheduling but operates at the TMEM level rather than the warpgroup level. It allows the MMA warp and epilogue warps to operate on different tiles concurrently without the context-switching overhead of Ping-Pong. The paper uses this mechanism for the forward down-projection kernel (where epilogue TMA store overlaps with UMMA accumulation) and the backward dH kernel (where epilogue operations overlap with the next tile's UMMA).

Persistent tile scheduler with CLC on Blackwell. On Blackwell GPUs, SonicMoE uses a CLC (Cluster Launch Control) persistent tile scheduler for grouped GEMM. Unlike traditional grid-launched kernels where each thread block processes a fixed set of tiles and exits, a persistent scheduler keeps a pool of thread blocks alive and dynamically assigns work tiles to them. This is important for grouped GEMM because different experts may have different $T_e$ (token counts), leading to load imbalance if work is statically partitioned. The persistent scheduler assigns tiles on-demand, ensuring that thread blocks that finish early can pick up work from experts with more tiles, improving overall utilization.

Quantifying the impact of IO overlap. Figure 5a provides concrete evidence: on the 7B fine-grained MoE forward pass, SonicMoE's down-projection kernel (with Ping-Pong and TMA store) achieves 627 TFLOPS, while the total time including expert aggregation is 1.237 ms. The backward dH kernel (with Ping-Pong, TMA load of $H$, and heavy epilogue fusion) achieves 328 TFLOPS—lower absolute TFLOPS because the epilogue operations are not pure MMA, but the IO overlap ensures that the kernel is not stalled waiting for memory. The overall backward pass completes in 2.173 ms versus ScatterMoE's 3.968 ms and MoMoE's 4.678 ms.

---

3.4.4 Token Rounding Routing: Hardware-Aware Token Assignment

The fourth major contribution addresses the tile quantization problem: when a grouped GEMM kernel processes an expert with $T_e$ tokens, the token dimension is tiled with tile size $M_{\text{tile}}$ (typically 128 on Hopper/Blackwell). If $T_e \bmod M_{\text{tile}} \neq 0$, the kernel pads the last tile with zeros, wasting computation on entries that contribute nothing to the output. Under high sparsity (small $\rho = K/E$), the expected tokens per expert $\bar{T}_e = T\rho$ is small, making the padding waste a significant fraction of total FLOPs—Figure 8 shows this waste reaching 25% of total FLOPs when $E=256, K=4$.

Token rounding algorithm (Algorithm 4). SonicMoE introduces a drop-in replacement for standard top-$K$ token-choice (TC) routing that guarantees each expert receives a number of tokens that is a multiple of $M_{\text{tile}}$. The algorithm operates in four steps:

Step 1: Standard top-K token-choice routing. The router scores $S \in \mathbb{R}^{T \times E}$ are computed, and for each token $t$, the top $K$ experts are selected with their scores $S_{\text{topK}}$ and indices $I_{\text{topK}}$:

$(S_{\text{topK}}, I_{\text{topK}}) \leftarrow \text{TopK}(S, K)$

This is the standard TC routing result. The per-expert token frequencies $f_e = \sum_t \mathbf{1}\{e \in I_{\text{topK}, t}\}$ are computed, where the sum counts how many tokens selected expert $e$ among their top $K$.

Step 2: Compute tile-rounded multiples. For each expert $e$, compute the floor and ceiling of its token frequency rounded to multiples of $M_{\text{tile}}$:

$\lceil f_e \rceil_{M_{\text{tile}}} = \lceil f_e / M_{\text{tile}} \rceil \cdot M_{\text{tile}}$

$\lfloor f_e \rfloor_{M_{\text{tile}}} = \lfloor f_e / M_{\text{tile}} \rfloor \cdot M_{\text{tile}}$

These represent the two possible tile-aligned token counts for expert $e$. The decision of whether to round up (pad) or round down (drop) is made by the round_and_sparsify subroutine.

Step 3: Build biased scoring matrix for expert-wise ranking. A modified score matrix $S'$ is constructed where:

• For tokens not in any expert's top-$K$, the score is set to $S_{\text{original}} - 1$, ensuring these "expert-choice" (EC) tokens are always ranked below the "token-choice" (TC) tokens.
• For tokens in an expert's top-$K$, the original top-$K$ scores are preserved.

This ensures that in the subsequent per-expert sorting, the TC-selected tokens will always be ranked first, and any padding tokens (EC tokens) will fill only the last tile, minimizing disruption to the original routing decision.

Step 4: Token rounding per expert. For each expert $e$:

• The tokens are sorted by $S'_e$ (descending), producing an ordered list $\pi_e$ with corresponding scores $s_e$.
• The round_and_sparsify subroutine is called with $\pi_e$, $s_e$, the original frequency $f_e$, and the two tile-rounded multiples. The subroutine decides whether to retain $\lceil f_e \rceil_{M_{\text{tile}}}$ tokens (padding EC tokens at the end of the sorted list) or $\lfloor f_e \rfloor_{M_{\text{tile}}}$ tokens (discarding the lowest-scoring TC tokens).

The default round_and_sparsify subroutine is NR-f (nearest rounding to $M_{\text{tile}}$-multiples via expert frequency): pad EC tokens if $\lceil f_e \rceil_{M_{\text{tile}}} - f_e < f_e - \lfloor f_e \rfloor_{M_{\text{tile}}}$, i.e., round to the nearest tile multiple. This simple heuristic minimizes the number of tokens whose routing decision is altered (either discarded or added) and ensures that the total number of tokens across all experts remains approximately $TK$ in expectation.

Guarantees of the algorithm. The paper states that for each expert, the maximum deviation from token-choice routing is at most one tile ($M_{\text{tile}}$ tokens). When rounding down, the expert receives $\lfloor f_e \rfloor_{M_{\text{tile}}}$ tokens, meaning at most $f_e - \lfloor f_e \rfloor_{M_{\text{tile}}} < M_{\text{tile}}$ TC-selected tokens are discarded. When rounding up, the expert receives $\lceil f_e \rceil_{M_{\text{tile}}}$ tokens, meaning at most $\lceil f_e \rceil_{M_{\text{tile}}} - f_e < M_{\text{tile}}$ EC tokens are added. In both cases, only the last tile's composition is affected—all TC tokens in earlier tiles are preserved unchanged.

Why token rounding works for model quality. The routing modification is bounded: each expert's token count changes by at most $M_{\text{tile}}$, and the total number of tokens across all experts is preserved in expectation (the nearest-rounding heuristic means roughly half the experts round up and half round down). The tokens that are discarded if rounding down are the lowest-scoring among the TC-selected tokens for that expert—they were borderline cases anyway. The tokens that are added if rounding up are the highest-scoring among the non-TC-selected tokens according to EC ranking—they are the "next best" tokens for that expert.

Table 2 empirically validates this across multiple model scales and sparsity levels:

• 0.5B MoE with 2/64 activation ($\bar{T}_e = 512$, $M_{\text{tile}} = 128$): TR achieves validation perplexity 15.92 versus TC's 15.94, with average downstream accuracy 50.0% versus 49.6%. The gap is negligible and slightly favors TR.
• 1.8B MoE with 8/256 activation ($\bar{T}_e = 512$): TR achieves PPL 13.10 versus TC's 13.12, with accuracy 53.5% versus 52.8%. Again, TR matches or slightly exceeds TC.
• 1.4B MoE with 2/128 activation ($\bar{T}_e = 512$, extreme sparsity $K/E = 1/64$): TR achieves PPL 13.22 versus TC's 13.32, with accuracy 52.0% versus 52.0%—statistically identical.

These results establish that token rounding does not degrade model quality under any tested sparsity level, and in some cases yields slightly better perplexity, likely because the elimination of padding waste allows the same compute budget to process more meaningful tokens.

Comparison with alternative rounding subroutines (Table 6). The paper ablates five strategies beyond NR-f:

• SR-f (stochastic rounding): pad with probability $(f_e - \lfloor f_e \rfloor) / M_{\text{tile}}$ for each expert independently. Performs similarly to NR-f (PPL 15.93 vs 15.92, accuracy 49.7% vs 50.0% for 0.5B).
• NR-s (score-based nearest rounding): decide based on score sums rather than frequencies. Also similar (PPL 15.91, accuracy 50.0%).
• Balance-f: uses a balanced rounding algorithm (Algorithm 6) that guarantees the total token count deviation is bounded by $M_{\text{tile}}/2$, not $O(M_{\text{tile}}\sqrt{E})$ as with independent rounding. Performs comparably (PPL 15.93, accuracy 50.1%).
• UP (always round up): always pads to $\lceil f_e \rceil$. Slightly better perplexity (15.89) but lower accuracy (49.3%), likely because it introduces EC bias that creates a train-test mismatch (evaluation uses TC routing).
• DOWN (always round down): always drops to $\lfloor f_e \rfloor$. Worse perplexity (16.10) due to information loss from consistently discarding tokens.

The conclusion: TR is robust to the specific rounding heuristic, and the simple nearest-rounding-by-frequency strategy is sufficient. The always-round-down baseline (which matches expert-parallel token dropping heuristics) shows the cost of information loss; the always-round-up baseline shows the cost of EC-bias. TR's ability to distribute rounding up and down approximately evenly avoids both pitfalls.

Effect of microbatch size (Table 7). TR's behavior depends on $\bar{T}_e / M_{\text{tile}}$, the ratio of expected tokens per expert to tile size. When this ratio is at least 2, TR quality is stable (the rounding affects at most half a tile's worth of tokens, which is a small fraction). When $\bar{T}_e / M_{\text{tile}} = 1$ (e.g., $\bar{T}_e = 128, M_{\text{tile}} = 128$), there is a noticeable degradation: for 1.8B MoE, PPL increases from 13.10 to 13.55, and accuracy drops from 53.5% to 52.7%. However, even at $\bar{T}_e / M_{\text{tile}} = 1$, TR still outperforms the EC baseline with finetuned TC router (PPL 15.01, accuracy 50.7%). The paper recommends TR when $\bar{T}_e / M_{\text{tile}} \geq 2$.

Effect of tile size $M_{\text{tile}}$ (Table 8). TR is robust to the choice of tile size when $\bar{T}_e / M_{\text{tile}} \geq 2$. Varying $M_{\text{tile}}$ from 64 to 512 changes PPL by at most 0.08 for 0.5B MoE and 0.49 for 1.8B MoE, with correspondingly small accuracy variations. The default $M_{\text{tile}} = 128$ is chosen to match the GEMM tile size used by SonicMoE's grouped GEMM kernels.

Training throughput gains (Figure 13). The elimination of tile quantization waste directly improves kernel TFLOPS (model FLOPs per second, not hardware FLOPs) because the same useful computation is performed with fewer wasted hardware operations. Figure 13 shows the throughput scaling as sparsity increases (increasing $E$ while keeping $K$ constant):

• For a 7B-configuration MoE with $n=1024$ (top-right subfigures), at $E=128$ ($K/E = 1/64$), TR achieves 6.1% higher model TFLOPS than TC in backward, for an end-to-end improvement of 9.4%.
• For a 120B-configuration MoE with $n=1024$ (bottom-right), at $E=256$ ($K/E = 1/64$), TR achieves 11.8% higher backward TFLOPS, for an end-to-end improvement of 15.9%.

The speedup grows with sparsity because the padding waste as a fraction of total FLOPs increases with $E$ (more experts, each receiving fewer tokens, more likely to have non-aligned token counts). TR's speedup is additive on top of SonicMoE's base kernel speedups, combining to produce the overall 1.86× throughput improvement on Hopper GPUs.

Inference compatibility. Token rounding is used during training to eliminate padded FLOPs. During evaluation and inference, the paper switches to standard TC top-$K$ routing. This is necessary because autoregressive generation (decoding one token at a time) produces one new token per step, and TR's expert-choice padding would create issues with causality and batch semantics. Table 2 evaluates all TR-trained models using TC routing at test time, confirming that the trained model generalizes to standard routing without adaptation. This contrasts with EC routing, which requires training an auxiliary router or finetuning a TC router to bridge the train-test gap (EC-trained models evaluated with TC routing show 3+ PPL degradation in the "EC (no adaptation)" baseline).

---

3.4.5 System Integration: Training Throughput and Memory Scaling

The paper evaluates SonicMoE not only as isolated kernel benchmarks but as an integrated system within a full training pipeline, using FSDP-2 (Fully Sharded Data Parallelism) with ZeRO-3 sharding within nodes and replication across nodes.

End-to-end training throughput (Section 6.2). For a 7B fine-grained MoE ($n=256$) trained with the lm-engine codebase:

• SonicMoE on 64 H100s achieves 213 billion tokens per day.
• ScatterMoE on 96 H100s achieves 225 billion tokens per day.

This means SonicMoE achieves comparable throughput with 33% fewer GPUs, directly translating the kernel-level speedups (1.83× backward, 1.86× overall) to distributed training savings. The throughput measurement includes all training overheads (FSDP communication, router computation, optimizer steps), confirming that SonicMoE's gains persist in real training regimes.

Activation memory scaling with granularity (Figure 1, left; Figure 10). SonicMoE's peak activation memory per MoE layer stays constant as expert granularity $G = d/n$ increases from 2 to 16 (i.e., $n$ shrinking from 2048 to 256 for a 30B model), while ScatterMoE's memory grows from roughly 1.2 GiB to 2.0 GiB per layer. Figure 10 quantifies this across model scales: for a 7B MoE with $n=256$, SonicMoE uses 1.3 GiB per layer versus ScatterMoE's 2.4 GiB (45% reduction); for 30B, the gap widens to roughly 3 GiB saved per layer. This memory efficiency means that for a given GPU memory budget, SonicMoE can train MoEs with higher granularity (which scaling laws predict will yield better model quality) or train with larger microbatches (improving throughput by reducing communication frequency in distributed settings).

Forward and backward throughput scaling (Figure 11a). Across all model scales (1.4B–120B) and granularity levels ($n$ from 256 to 2048) on H100 GPUs:

• SonicMoE achieves 43% higher forward TFLOPS than DeepGEMM++ and 83% higher backward TFLOPS than ScatterMoE on the 7B fine-grained configuration.
• The relative speedup increases with granularity: on 120B MoE with $n=512$ (highest granularity, $G=8$), SonicMoE's forward pass achieves 560 TFLOPS versus DeepGEMM++'s 450 TFLOPS (24% speedup), while at $n=2048$ ($G=2$), the gap is only 8%. This confirms the paper's central thesis that IO-aware optimizations become more valuable as experts become finer-grained.
• MoMoE and MegaBlocks show the largest gaps at fine granularity due to their lack of IO overlap and gather fusion.

Realistic MoE configurations (Figure 12a). When benchmarked on configurations matching recent production MoEs (OLMoE-1B-7B, gpt-oss-20b, Kimi-Linear-48B-A3B, Qwen3-Next-80B-A3B, Qwen3-235B-A22B, DeepSeek-V3.2-Exp), SonicMoE consistently achieves >550 TFLOPS on forward and backward passes, surpassing all baselines. Notably, ScatterMoE, MoMoE, DeepGEMM-pt, and DeepGEMM++ all fail to run the DeepSeek-V3.2-Exp configuration ($d=7168, n=2048, E=256$) on a single H100 due to index overflow or CUDA OOM errors, while SonicMoE runs successfully—demonstrating that SonicMoE's memory efficiency is not just an optimization but an enabler for frontier-scale fine-grained MoE training.

Blackwell B300 results (Figure 11b). On B300 GPUs, SonicMoE demonstrates similar trends:

• On the OLMoE-sized 7B MoE, SonicMoE achieves 25% higher forward TFLOPS and 15% higher backward TFLOPS than DeepGEMM++.
• Compared to the Triton official MoE example (designed for inference, which stores only post-activation $A$ and not pre-activation $H$), SonicMoE (which stores both for training) achieves higher TFLOPS (e.g., 1242 vs. 1004 TFLOPS forward for 7B $n=256$), indicating that SonicMoE's IO overlap compensates for the additional memory traffic of training.
• The gather fusion and TMEM two-stage pipeline prove effective at maintaining high utilization even with the additional training-specific memory accesses.

Combined effect of all optimizations. The paper's three contributions compose synergistically:

• The memory-efficient backward pass enables higher granularity within a fixed memory budget (otherwise OOM), allowing the model architect to choose architectures that scaling laws recommend.
• The IO-overlapping kernels accelerate training for the chosen architecture by hiding memory latency, with speedup proportional to granularity.
• The token rounding routing eliminates wasted padding computation that grows with sparsity, providing additional speedup that compounds with the kernel optimizations.

Together, they close the hardware-efficiency gap that previously made fine-grained sparse MoEs impractical to train, establishing a co-design methodology where architecture choices and kernel design are optimized jointly rather than in isolation.

4. Key Insights and Innovations

Innovation 1: The MoE Computation Graph Can Be Restructured to Achieve Dense-Model Memory Footprint Without Recomputation

The field's working assumption, embodied in every prior MoE kernel from ScatterMoE to MoMoE to Megatron, has been that the backward pass through an MoE layer must cache the down-projection output Y (size TKd) to compute the router score gradient dS = ⟨dO, Y⟩. This assumption creates a direct coupling between activation memory and expert granularity: as K increases under iso-FLOPs scaling, cached memory grows proportionally, making fine-grained MoEs—which scaling laws predict will yield better model quality per FLOP—impractical to train at scale. The solution space was understood as a binary tradeoff: either pay the memory cost or recompute the down-projection GEMM during the backward pass.

SonicMoE identifies a third path that is mathematically equivalent but architecturally transformative: dS = ⟨dA′, A⟩ is not just an algebraic identity—it is a different computation graph that breaks the memory-granularity coupling entirely. The key intellectual move is recognizing that dA′ = dO W₂^⊤ is computed as an intermediate during the down-projection activation gradient kernel anyway, and A is the SwiGLU output cached for the weight gradient dW₂. Their inner product therefore requires zero additional HBM loads and zero additional cached activations, whereas ⟨dO, Y⟩ requires both. This is not an optimization that autograd or a compiler could discover; it requires understanding that the forward computation graph (where S-weighting can be applied before or after down-projection) and the backward computation graph (where different expressions for dS have different memory implications) can be decoupled.

What makes this a conceptual advance rather than an implementation trick is the minimum-possible claim: 2Td + 4TKn bytes per layer is the smallest activation footprint achievable without GEMM recomputation, because X, H, and routing metadata are the minimal information needed to compute all other gradients. SonicMoE reaches this lower bound, making an MoE layer no more memory-hungry than a dense layer with the same activated parameters. The evidence (Figure 10: 45% reduction in per-layer memory for 7B; Figure 1 left: flat memory curve as granularity increases 8×) demonstrates that prior approaches weren't just suboptimal—they were fundamentally misaligned with the architectural trend toward granularity. The practical consequence is that the memory ceiling that previously prevented training fine-grained MoEs is removed, enabling architecture designers to follow scaling laws toward higher granularity without being constrained by activation memory.

Innovation 2: MoE Training Is Memory-Bound, Not Compute-Bound, and the Bottleneck Shifts With Granularity and Sparsity

Prior MoE kernel work treated the problem as a GEMM optimization problem: how to make grouped matrix multiplications run faster. ScatterMoE focused on Triton-based gather fusion; DeepGEMM focused on FP8 GEMM throughput; MegaBlocks reframed it as block-sparse operations. All implicitly assumed the bottleneck was compute—make the GEMM faster, and training throughput improves. SonicMoE provides the first systematic diagnosis that for fine-grained and sparse MoEs, the bottleneck is fundamentally memory bandwidth, and it shifts direction as architecture parameters change.

Equation 4 is the diagnostic instrument: the arithmetic intensity 3 / (2 + 2G/d + 3/(Tρ)) shows that as granularity G increases or activation ratio ρ decreases, the IO cost per FLOP grows without bound, while compute per expert shrinks. This means that techniques designed for compute-bound regimes (better tile shapes, higher MMA utilization, lower-precision arithmetic) provide diminishing returns, while techniques targeting memory latency—overlapping IO with compute, reducing HBM round-trips through fusion, minimizing activation storage—become the dominant levers. The paper makes this explicit by profiling SonicMoE against DeepGEMM under identical grouped GEMM configurations (Figure 18): SonicMoE's grouped GEMM with contiguously-packed inputs is only 2.7-10% faster than DeepGEMM's highly-tuned equivalent. The dramatic overall speedups (43% forward, 83% backward for 7B fine-grained) come not from better GEMM but from everything around the GEMM—gather fusion, epilogue fusion, IO-compute overlap—that reduces the memory traffic dominating runtime.

This reframing matters beyond performance numbers. It tells researchers and practitioners that improving MoE training throughput requires a fundamentally different optimization portfolio than improving dense model throughput. Investing in faster GEMM (e.g., FP4, sparsity-aware Tensor Cores) yields marginal returns for fine-grained MoEs; investing in memory bandwidth utilization (asynchronous TMA, warpgroup scheduling, fusion strategies) yields compounding returns. The paper demonstrates this scaling behavior empirically (Figure 11a: SonicMoE's relative speedup over DeepGEMM++ grows from 8% at G=2 to 24% at G=8 on 120B), validating that the analysis is predictive. This diagnostic framework—identifying the bottleneck regime from architecture parameters and directing optimization effort accordingly—is generalizable beyond the specific kernels presented.

Innovation 3: Token Rounding Reframes Routing as a Hardware-Throughput Decision With Bounded Quality Impact

MoE routing has historically been studied exclusively through the lens of model quality: top-K token-choice, expert-choice, load-balancing losses, auxiliary routers—every innovation targeted better perplexity, better task accuracy, or better load balance. Hardware efficiency was treated as an orthogonal concern addressed by the downstream kernels, not by the router. SonicMoE's token rounding (TR) introduces the counterintuitive idea that intentionally deviating from the optimal routing decision—discarding some token-choice tokens or padding with expert-choice tokens—can improve end-to-end training speed without degrading model quality, because the hardware savings from eliminating tile padding waste outweigh the information loss from slight routing perturbation.

The intellectual move here is conceptualizing routing as a resource allocation problem with a hardware cost function, not just a model quality optimization. The key constraint—"maximum deviation from token-choice routing is at most one tile per expert"—provides a formal bound on how much the routing can differ from the theoretical optimum, and the empirical results (Table 2: TR matches or exceeds TC accuracy across five model configurations, from 0.5B to 1.8B parameters) demonstrate that this bound is tight enough to preserve model quality even under extreme sparsity (K/E = 1/64). The robustness across rounding subroutines (Table 6: five strategies yield nearly identical perplexity and accuracy) further suggests that the specific rounding heuristic matters less than the structural property of tile-aligned token counts—the hardware benefit is deterministic, while the quality impact is stochastic and averages out over training.

This is fundamentally different from prior routing modifications like token dropping (Fedus et al. 2022) or capacity factors, which treat excess tokens as waste to be discarded and suffer quality degradation as a result. TR does not drop tokens permanently—it redistributes them across experts within a bounded neighborhood of the original assignment, preserving the total number of processed tokens in expectation. The "always round down" ablation (DOWN in Table 6) shows what happens when this redistribution is absent: perplexity degrades from 15.92 to 16.10 for 0.5B MoE because information is lost rather than shifted. TR's insight is that expert assignment can be treated as locally malleable—tokens at the margin of an expert's selection can be reassigned without quality loss if the reassignment respects the score ordering—and that this malleability can be exploited for hardware alignment.

The practical significance is that TR enables sparser MoEs than would otherwise be practical. As K/E drops (the trend documented in Table 4, from 25% to 1.95% in two years), the tile quantization waste in TC routing grows from negligible to dominant (Figure 8: 25% wasted FLOPs at E=256, K=4). Without TR, this waste creates a de facto ceiling on sparsity where training throughput degrades faster than model quality improves. With TR, the ceiling is removed—SonicMoE can maintain 16% higher model TFLOPS at K/E = 1/64—meaning architecture designers can push sparsity further than hardware constraints would otherwise permit. This is a genuine co-design contribution: TR modifies the routing algorithm to accommodate GPU tile structure, enabling architectures that would be hardware-inefficient under standard routing.

Innovation 4: The Scatter-Gather Duality in MoE Aggregation Determines the Entire IO Schedule, Not Just the Final Step

A seemingly minor design choice in MoE kernel implementation—whether each expert scatters its results to original token positions during the GEMM epilogue (as ScatterMoE and MoMoE do) or each token gathers results from contiguous expert outputs in a separate kernel (as SonicMoE does)—turns out to have cascading implications for the entire IO schedule of both forward and backward passes. The paper demonstrates that the gather approach is structurally superior for hardware with asynchronous memory operations (TMA on Hopper/Blackwell) because it decouples the GEMM output store from the positional scatter, allowing both to use asynchronous operations rather than being bottlenecked by synchronous st.global instructions.

The conceptual insight is that scatter is synchronous, gather is asynchronous under current GPU architectures. Scatter requires indexed writes to non-contiguous HBM addresses, which on Hopper GPUs can only use synchronous st.global PTX instructions. Each store must complete before the warp can proceed to the next MMA tile, creating a 20% TFLOPS penalty (Figure 16, Figure 21). Gather, by contrast, uses indexed reads via cp.async (Hopper) or TMA gather (Blackwell), which are asynchronous and can be overlapped with ongoing computation. SonicMoE exploits this by storing all expert outputs contiguously (via asynchronous TMA store in the GEMM epilogue) and deferring positional scattering to a dedicated aggregation kernel that reads via asynchronous loads—a complete inversion of the dataflow that ScatterMoE, MoMoE, Megatron, and MegaBlocks all implement as scatter-first.

This is not a small implementation detail. The scatter-gather choice propagates through all design decisions: it determines which kernels can use Ping-Pong scheduling (only kernels with heavy epilogues, which scatter would make heavier), which operations can be fused (scatter fusion blocks MMA, so SonicMoE avoids it), and whether backward weight gradient kernels need separate gather launches (with gather, they don't; with scatter, prior work launches separate gathers for dW₁ and dW₂). The 2.92× expert aggregation bandwidth advantage over ScatterMoE (Figure 20a) is a direct consequence of this choice, but the broader impact is enabling all the other IO-overlapping techniques: without the decoupling that gather provides, the asynchronous TMA operations and Ping-Pong scheduling described in Sections 3.4.2 and 3.4.3 would be impossible to apply consistently across the forward and backward passes.

The paper treats this as a first-class architectural decision rather than an afterthought, and the performance evidence (1.83× backward speedup over ScatterMoE on 7B fine-grained MoE, Figure 5a) shows that the aggregation strategy is as impactful as the GEMM optimization. This insight generalizes beyond MoE: any sparse computation pattern where outputs must be reassembled into a dense format faces the same scatter-vs-gather choice, and SonicMoE provides a template for choosing gather when asynchronous memory operations are available.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The authors benchmark training throughput and activation memory on synthetic MoE layer configurations with controlled architecture parameters (hidden dimension $d$, expert intermediate dimension $n$, number of experts $E$, activated experts $K$, microbatch size $T$). These configurations span model scales from 1.4B to 120B parameters and are specified in Table 9a (H100) and Table 9b (B300). For token rounding quality evaluation, the authors train models on a deduplicated version of FineWeb-Edu (Ben Allal et al. 2024) using the OLMoE codebase and evaluate on 11 downstream tasks: WinoGrande, Social IQA, SciQ, PIQA, OpenBookQA, HellaSwag, COPA, CommonsenseQA, BoolQ, Arc-Easy, and Arc-Challenge.

• Base model(s). All kernel benchmarks use PaLM-2-S*-equivalent configurations defined by architecture parameters, not specific pretrained models. For token rounding quality evaluation, the authors train OLMoE-architecture models from scratch at 0.5B, 1.4B, and 1.8B parameter scales with varying expert configurations (Tables 2, 5, 10). The model architecture is described as "identical to OLMoE's attention block" with SwiGLU for MoE layers.

• Metrics. Kernel performance is measured by TFLOPS (model FLOPS per second, computed as $(6+12)dn(\sum_e f_e)$ divided by kernel execution time) for GEMM-bound operations, and HBM bandwidth (TB/s) for memory-bound operations (gather, expert aggregation). Activation memory usage is reported as GiB per MoE layer at peak. Model quality is measured by validation perplexity and average accuracy across 11 downstream tasks. Training throughput is measured in billions of tokens per day for end-to-end FSDP-2 training runs.

• Baselines. The paper compares against five families: ScatterMoE (Tan et al. 2024, Triton-based with forward gather fusion but no backward gather fusion, no TMA, $dS = \langle dO, Y \rangle$), MoMoE (Costin et al. 2025, similar limitations as ScatterMoE with some backward fusion differences), MegaBlocks (Gale et al. 2023, using ParallelDroplessMLP with block-sparse matrix multiplication rather than grouped GEMM), Megatron-LM (Shoeybi et al. 2019, using GroupedMLP with CUTLASS grouped GEMM and JIT epilogue fusion but no gather fusion), and two DeepGEMM variants: DeepGEMM-pt (standard PyTorch implementation for gather and expert aggregation, with DeepGEMM's grouped GEMM as backend) and DeepGEMM++ (SonicMoE's optimized gather and expert aggregation kernels paired with DeepGEMM's grouped GEMM). On Blackwell B300 GPUs, the Triton official MoE example is added as an inference-optimized baseline that stores only post-activation $A$, not pre-activation $H$. For routing quality, baselines include TC top-K (standard token-choice routing), TC (token drop) (discarding lowest-scoring tokens to achieve tile alignment), EC (expert-choice routing, Zhou et al. 2022), EC (aux router) (expert-choice with auxiliary router for inference, following Raposo et al. 2024), and EC (ft TC router) (expert-choice during training with TC router finetuning for inference).

• Generation budget / compute accounting. For kernel benchmarks, compute is measured by kernel execution time at specified configurations, with no sample budget—each configuration runs a single forward or backward pass of one MoE layer. For end-to-end training throughput, compute is measured as total GPU-hours across 64 or 96 H100 GPUs. For token rounding throughput (Figure 13), TC and TR both process the same number of input tokens $T$, but TC's hardware FLOPs include padding waste while TR's do not—the paper reports model TFLOPS (useful computation per second) to capture the difference. For model quality evaluation, routing methods are compared at equal training token budgets (20B–100B tokens depending on model scale, specified in Table 10).

• Cross-validation / statistical protocol. No cross-validation is used for kernel benchmarks, which are deterministic given fixed random seeds for synthetic token routing. For token rounding quality evaluation, the paper uses a single training run per configuration with standard OLMoE training hyperparameters (Table 10) and evaluates the final checkpoint on the full validation and test sets. No confidence intervals or error bars are reported for any kernel or model quality measurements.

Main Quantitative Results

Activation Memory Efficiency

The headline result appears in Figure 10 and Figure 1 (left): SonicMoE's peak activation memory per MoE layer remains constant as expert granularity $G = d/n$ increases, while ScatterMoE and MoMoE show linear growth. For the 7B MoE with $n=256$ (highest granularity tested at this scale), SonicMoE uses approximately 1.3 GiB per layer versus ScatterMoE's 2.4 GiB, a 45% reduction (Figure 10, middle panel). The pattern holds across all model scales (1.4B, 7B, 30B, 120B) and granularity levels (three per scale, with $n$ decreasing by factors of 2 while $K$ increases proportionally to maintain iso-FLOPs). At 120B scale with $n=512$, SonicMoE saves more than 3 GiB per layer compared to MoMoE (Figure 10, rightmost panel). The theoretical claim—that SonicMoE caches only $2Td + 4TKn$ bytes, matching a dense model with equivalent activated parameters—is validated empirically by the flat memory curve in Figure 1 (left), where cuBLAS BMM + activation + cuBLAS BMM + aggregation (the upper bound baseline) shows no memory variation with granularity, and SonicMoE tracks it, while ScatterMoE and MoMoE diverge.

Forward and Backward Throughput on H100 GPUs

Figure 11a reports the headline throughput numbers. On a fine-grained **7B MoE with $n=256$, $E=128$, $K=8$ **: SonicMoE achieves forward pass TFLOPS of approximately 627 TF/s (well above DeepGEMM++ at ~475 TF/s and ScatterMoE at ~421 TF/s), a 43% improvement over DeepGEMM++. The backward pass shows larger gains: SonicMoE reaches approximately 480 TF/s versus ScatterMoE's ~364 TF/s (83% improvement) and MoMoE's ~204 TF/s (115% improvement). The throughput advantage is most pronounced at high granularity: on the 30B configuration with $n=256$ ($G=16$), SonicMoE's forward pass achieves ~570 TF/s versus DeepGEMM++'s ~460 TF/s (24% advantage), while at $n=1024$ ($G=4$), the gap narrows to ~8%. On the 120B configuration with $n=512$ ($G=8$), the forward gap is ~24% and the backward gap ~17%. This granularity-dependent scaling validates the paper's diagnostic that IO-aware optimizations become more valuable as experts get finer.

Figure 5a provides a detailed runtime breakdown for the 7B fine-grained configuration, decomposing into individual kernel categories. SonicMoE's forward pass completes in 1.237 ms total, compared to DeepGEMM++ at 1.770 ms, ScatterMoE at 2.255 ms, MoMoE at 2.347 ms, DeepGEMM-pt at 2.434 ms, MegaBlocks at 2.983 ms, and Megatron at 3.476 ms. The backward pass shows SonicMoE at 2.173 ms versus DeepGEMM++ at 3.009 ms, DeepGEMM-pt at 3.680 ms, ScatterMoE at 3.968 ms, and MoMoE at 4.678 ms. The per-kernel annotations reveal where SonicMoE's advantages concentrate: the forward down-projection kernel achieves 627 TF/s (dominantly ahead due to Ping-Pong and TMA store), the backward activation gradient dH kernel bundles operations that span four separate kernel categories in ScatterMoE (gather dO, dSwiGLU + dS + A′, dS, and dW2 portions), and the expert aggregation kernel achieves 1.79 TB/s bandwidth versus ScatterMoE's 0.94 TB/s.

The end-to-end training throughput with FSDP-2 is reported in Section 6.2: SonicMoE on 64 H100s achieves 213 billion tokens per day, comparable to ScatterMoE's 225 billion tokens per day on 96 H100s—meaning SonicMoE matches throughput with 33% fewer GPUs for the same 7B $n=256$ MoE model.

Forward and Backward Throughput on B300 GPUs

Figure 11b shows the Blackwell B300 results. For the OLMoE-sized **7B MoE with $n=1024$, $E=64$, $K=4$** (the least fine-grained, most compute-bound configuration tested): SonicMoE achieves forward pass **1242 TF/s** and backward pass **1122 TF/s**. Compared to DeepGEMM++ at ~1004 TF/s forward and ~1047 TF/s backward, this represents a **25% forward speedup** and a **15% backward speedup**. The speedup grows with granularity: on the 120B configuration with$n=512$ ($G=8$`), SonicMoE achieves ~1134 TF/s forward and ~1083 TF/s backward versus DeepGEMM++'s ~1009 TF/s forward and ~1047 TF/s backward (12% and 15% speedups respectively).

On B300, SonicMoE also outperforms the Triton official MoE example (designed for inference, storing only post-activation $A$): for the 7B $n=1024$ configuration, SonicMoE achieves 1242 TF/s forward versus Triton's 1004 TF/s, despite SonicMoE storing both pre- and post-activation. Figure 5b provides the runtime breakdown: SonicMoE forward completes in 2.941 ms versus DeepGEMM++ at 3.999 ms, Triton example at 4.165 ms, DeepGEMM-pt at 6.212 ms, and ScatterMoE at 8.933 ms. The backward pass shows SonicMoE at 5.902 ms versus DeepGEMM++ at 7.255 ms, DeepGEMM-pt at 10.540 ms, and ScatterMoE at 12.549 ms.

Realistic MoE Configuration Throughput

Figures 12a (H100) and 12b (B300) benchmark SonicMoE against baselines on configurations matching production MoE models (OLMoE-1B-7B, gpt-oss-20b, Kimi-Linear-48B-A3B, Qwen3-Next-80B-A3B, Qwen3-235B-A22B, DeepSeek-V3.2-Exp). On H100, SonicMoE consistently achieves >550 TF/s on both forward and backward passes across all configurations, surpassing every baseline on every configuration. A notable result: ScatterMoE, MoMoE, DeepGEMM-pt, and DeepGEMM++ all fail to run the DeepSeek-V3.2-Exp configuration ($d=7168, n=2048, E=256, K=8$) on a single H100 due to index overflow or CUDA OOM errors (marked as "NA" in Figure 12a), while SonicMoE runs successfully at ~600 TF/s forward and ~580 TF/s backward. This positions SonicMoE not just as a throughput optimizer but as an enabler for frontier-scale fine-grained MoE training on current hardware.

On B300 (Figure 12b), SonicMoE consistently achieves >1100 TF/s across configurations. For the Qwen3-Next-80B-A3B configuration ($d=2048, n=512, E=512, K=10$, extremely sparse at $K/E = 1.95\%$), SonicMoE reaches ~1150 TF/s forward and ~1100 TF/s backward, with substantial margins over DeepGEMM++. The Triton official example fails to support $K=10$ for this configuration (marked NA). For DeepSeek-V3.2-Exp on B300, SonicMoE achieves ~1150 TF/s forward and ~1100 TF/s backward; the Triton example fails; and ScatterMoE, MoMoE, and DeepGEMM-pt all fail (marked NA).

Token Rounding Quality Evaluation

Table 2 presents the headline quality results across five model training configurations. For the 0.5B MoE with 2/64 activation (Table 2b, $\bar{T}_e = 512$, $M_{\text{tile}} = 128$, trained 40B tokens): TR achieves validation perplexity 15.92 and average downstream accuracy 50.0% , versus TC top-K at 15.94 and 49.6%—a statistically identical result with a slight edge to TR. The token dropping baseline ("TC (token drop)") achieves 16.10 perplexity (worse) and 50.2% accuracy (slightly better), while EC with finetuned TC router ("EC (ft TC router)") degrades to 16.98 perplexity and 49.3% accuracy.

For the 1.8B MoE with 8/256 activation (Table 2c, $\bar{T}_e = 512$, trained 40B tokens): TR achieves 13.10 perplexity and 53.5% accuracy, versus TC top-K at 13.12 and 52.8%—again, TR matches or slightly exceeds TC. The EC (ft TC router) baseline degrades substantially to 15.01 perplexity and 50.7% accuracy, demonstrating the large train-test gap from EC-to-TC adaptation.

For the 0.5B MoE with 8/64 activation (Table 2a, $\bar{T}_e = 4096$, trained 20B tokens, less sparse): TR achieves 15.94 perplexity and 50.4% accuracy versus TC top-K at 16.01 and 49.8%—TR shows a marginally better result even under moderate sparsity where tile quantization effects are smaller.

For the 1.4B MoE with 2/128 activation (Table 2e, extreme sparsity $K/E = 1/64$, $\bar{T}_e = 512$, trained 100B tokens): TR achieves 13.22 perplexity and 52.0% accuracy, versus TC top-K at 13.32 and 52.0%—indistinguishable performance despite the extreme sparsity and the fact that TR modifies routing for every expert whose token count is not tile-aligned. EC (ft TC router) degrades to 14.75 perplexity and 50.5% accuracy.

The consistent pattern across all five configurations (0.5B moderate sparsity, 0.5B high sparsity, 1.8B high sparsity, 1.4B moderate sparsity, 1.4B extreme sparsity) is that TR never degrades model quality relative to TC top-K, and in 4 out of 5 configurations, TR achieves slightly better validation perplexity. The EC baselines consistently underperform, with the auxiliary router approach ("EC (aux router)") showing catastrophic degradation under high sparsity (e.g., validation perplexity 37.52 for 1.4B with 2/128 activation in Table 2e, versus TR's 13.22).

Token Rounding Training Throughput

Figure 13 reports the throughput impact of TR under iso-FLOPs scaling as sparsity increases (increasing $E$ while keeping $K$ constant). For a **7B-style configuration with $n=1024$, $K=2$** (top-right subfigures): as$E$increases from 16 to 128, TC top-K forward model TFLOPS drops from ~470 to ~390 (17% degradation), while TR maintains ~430 to ~420 TF/s (only 2% degradation). At$E=128$ ($K/E=1/64$), TR achieves **16.5\% higher forward TFLOPS** and **6.1\% higher backward TFLOPS** than TC, for an **end-to-end improvement of 9.4\%** . For a **120B-style configuration with $n=1024$, $K=4$ ** (bottom-right subfigures): at $E=256$ ($K/E=1/64$), TR achieves 25.7% higher forward TFLOPS and 11.8% higher backward TFLOPS than TC, for an end-to-end improvement of 15.9% .

The throughput gap between TR and TC grows monotonically with sparsity because TC's tile quantization waste (Figure 8) increases from ~5% wasted FLOPs at $E=32$ to ~25% at $E=256$ (for $K=4$), while TR eliminates this waste entirely. The throughput metric reported is model TFLOPS (useful computation per second), so TR's advantage reflects that the same number of input tokens are processed in less time because fewer hardware FLOPs are wasted on padding.

Figure 14 extends this analysis to production-configuration MoEs on H100. Across six configurations from OLMoE-1B-7B to DeepSeek-V3.2-Exp, TR consistently outperforms TC top-K, with speedups ranging from 1.027× to 1.207× on the forward pass and 1.013× to 1.079× on the backward pass. The largest speedups occur for the sparsest configurations: Qwen3-Next-80B-A3B ($K/E=10/512$) achieves 1.196× forward speedup, and DeepSeek-V3.2-Exp ($K/E=8/256$) achieves 1.207× forward speedup. This confirms that TR's benefits are not limited to synthetic benchmark configurations but apply directly to the frontier MoE architectures documented in Table 4.

Ablation Studies and Robustness Checks

Choice of rounding subroutine (Table 6): The paper ablates five rounding strategies on 0.5B (2/64 activation) and 1.8B (8/256 activation) MoEs. For 0.5B: NR-f (nearest rounding by frequency), SR-f (stochastic rounding), NR-s (nearest rounding by scores), and Balance-f (balanced rounding, Algorithm 6) all achieve validation perplexity between 15.89 and 15.93 and accuracy between 49.7% and 50.1%—indistinguishable within realistic training variance. The key negative result: UP (always round up) achieves better perplexity (15.89) but lower accuracy (49.3%), while DOWN (always round down) achieves worse perplexity (16.10) but higher accuracy (50.2%). For 1.8B, the same pattern holds: NR-f and NR-s achieve the best balanced results (53.5% accuracy), UP degrades to 52.1%, and DOWN degrades to 52.7%. The paper concludes TR is robust to the rounding heuristic, but always-discarding (DOWN) penalizes perplexity from information loss, while always-padding (UP) creates EC-bias that harms downstream accuracy due to train-test mismatch.

Effect of microbatch size $T$ (Table 7): Varying $T$ changes $\bar{T}_e = T\rho$, the expected tokens per expert. For 0.5B with 2/64 activation: TR maintains stable quality (perplexity 15.91–16.11, accuracy 49.6–50.1%) as $\bar{T}_e$ varies from 1024 down to 128. At $\bar{T}_e = 128$ ($\bar{T}_e/M_{\text{tile}} = 1$), perplexity increases from 15.92 to 16.11—a detectable but modest degradation. For 1.8B with 8/256 activation: quality is stable from $\bar{T}_e = 1024$ down to 256 (perplexity 13.08–13.12, accuracy 52.9–53.5%), but at $\bar{T}_e = 128$, perplexity increases to 13.55 and accuracy drops to 52.7%. The recommendation: TR is safe when $\bar{T}_e/M_{\text{tile}} \geq 2$, but shows measurable degradation at $\bar{T}_e/M_{\text{tile}} = 1$. Even at $\bar{T}_e/M_{\text{tile}} = 1$, TR still substantially outperforms EC (ft TC router) at 15.01 perplexity and 50.7% accuracy.

Effect of tile size $M_{\text{tile}}$ (Table 8): On 0.5B MoE ($\bar{T}_e = 512$), varying $M_{\text{tile}}$ from 64 to 512 (thus $\bar{T}_e/M_{\text{tile}}$ from 8 to 1) produces perplexity between 15.90 and 16.17 and accuracy between 49.9% and 50.4%. On 1.8B MoE, the same sweep produces perplexity between 13.07 and 13.56 and accuracy between 52.0% and 53.5%. Quality degrades noticeably at $M_{\text{tile}} = 512$ (where $\bar{T}_e/M_{\text{tile}} = 1$), consistent with the microbatch size finding: TR is robust when the ratio is at least 2. The default $M_{\text{tile}} = 128$ is chosen to match the GEMM tile size.

Effect of expert granularity on MoE quality (Table 5): Before evaluating TR, the paper validates that fine-grained MoEs indeed improve quality per FLOP. For 0.5B MoE with iso-FLOPs ($nK$ constant): decreasing $n$ from 1024 to 256 (and proportionally increasing $K$ from 2 to 8 and $E$ from 16 to 64) improves validation perplexity from 16.23 to 16.01 to 16.13 (Table 5a). The dense iso-FLOPs baseline is substantially worse at 19.90 perplexity, while dense iso-params is better at 15.46—confirming the FLOPs-efficiency advantage of MoEs. For 1.4B MoE (Table 5b), the pattern is similar but less monotonic: $n=256$ (PPL 13.32) outperforms $n=1024$ (PPL 13.38), with $n=64$ at 13.50 remaining competitive but not dominant. The dense iso-params baseline at 12.74 PPL demonstrates the quality ceiling achievable by scaling parameters rather than using sparsity.

Gather fusion impact on H100 and B300 (Figure 19): The paper ablates gather fusion by benchmarking SonicMoE, ScatterMoE, MoMoE, DeepGEMM, and cuBLAS BMM with and without gather fusion (opaque vs. transparent bars). On H100 (Figure 19a), gather fusion on the M dimension costs SonicMoE only 6.3% relative TFLOPS degradation on average, while on the K dimension the cost is 8.5% . In contrast, ScatterMoE with a separate gather kernel for varlen-K (its backward weight gradient operation) shows a larger gap that grows with granularity—confirming that gather fusion is increasingly valuable as expert granularity increases. On B300 (Figure 19b), gather fusion on M costs SonicMoE 3.4% TFLOPS and on K costs -0.1% (i.e., virtually zero overhead), while the Triton official example's gather fusion costs 6.3% on M.

Expert aggregation strategy comparison (Figure 21): The paper directly compares the gather-and-sum strategy (Figure 17 left) versus scatter-and-sum (Figure 17 middle) when both are implemented within SonicMoE using identical GEMM configurations. Across all model scales, gather-and-sum achieves 20% higher TFLOPS on average than scatter-and-sum, validating the design choice. ScatterMoE and MoMoE, which both use scatter-and-sum, are further behind due to their additional overheads beyond the aggregation strategy itself.

Expert aggregation kernel bandwidth (Figure 20): SonicMoE's gather-and-sum aggregation achieves 2.92× ScatterMoE's bandwidth and 1.05× MoMoE's bandwidth on average across H100 configurations (Figure 20a). It reaches 98% of the bandwidth of an optimized Triton kernel summing over contiguous $Y$ (the upper bound), despite the additional gather overhead. On B300 (Figure 20b), it achieves 6.72× ScatterMoE's bandwidth, 3.32× MoMoE's bandwidth, and 1.05× Gluon TMA gather-and-sum, reaching 98% of the Triton contiguous upper bound.

Top-K router kernel performance (Figure 22): SonicMoE's efficient bitonic-sort-based top-K kernel is benchmarked against PyTorch torch.topk, the Triton official example, the Tilelang official example, and RTop-K (Xie et al. 2025). On H100 (Figure 22a) with BF16 inputs, SonicMoE consistently achieves the highest bandwidth across all configurations (1.4B–120B), with the gap widening for larger models (more tokens $T$). The Triton and Tilelang examples are faster than PyTorch but slower than SonicMoE, while RTop-K (FP32 only) performs competitively on smaller configurations but falls behind on larger ones. On B300 (Figure 22b), the same pattern holds, with SonicMoE achieving roughly 1.5–2.0× the bandwidth of Triton's example for the largest configurations.

Critical Assessment

Claim: SonicMoE reduces activation memory by 45% on a 7B fine-grained MoE versus ScatterMoE, and activation memory stays constant with increasing granularity. This claim is directly supported by Figure 10 (middle panel, 7B: ~1.3 vs. ~2.4 GiB per layer) and Figure 1 (left panel: flat curve for SonicMoE and cuBLAS BMM while ScatterMoE and MoMoE diverge upward). The evidence is strong: the measurements span four model scales, three granularity levels per scale, and all baselines. The mechanism—avoiding caching $Y$ through the $\langle dA', A \rangle$ computation—is mathematically verified (Appendix C) and its memory impact is precisely quantifiable. One caveat: the paper acknowledges that $Y$ is still temporarily materialized during forward, with the argument that it can be recycled across layers. For a model with very few MoE layers (e.g., a hybrid dense-MoE architecture with only 2–4 MoE layers), the temporary $Y$ buffer could dominate peak memory. This edge case is not tested.

Claim: SonicMoE achieves 1.86× compute throughput improvement on Hopper GPUs compared to ScatterMoE's BF16 MoE kernel for a fine-grained 7B MoE. The evidence comes from Figure 11a (forward: ~627 vs. ~421 TF/s; backward: ~480 vs. ~364 TF/s), with the runtime breakdown in Figure 5a showing the per-kernel contributions. However, the 1.86× figure is not explicitly stated in the abstract as a single number—the abstract says "1.86× compute throughput improvement," while the body text reports per-pass speedups (43% forward, 83% backward) and the end-to-end training throughput comparison (213B tokens/day on 64 H100s vs. 225B on 96 H100s). Readers must compute the composite speedup by weighting forward and backward equally: (627+480)/(421+364) ≈ 1.41× for the 7B fine-grained case. The 1.86× may refer to a different configuration or weighting; the paper does not provide a single explicit derivation. More importantly, these throughput numbers are measured on a single MoE layer in isolation, without distributed training communication overhead. The end-to-end FSDP-2 measurement (33% fewer GPUs for same throughput) is the strongest evidence that the kernel speedups translate to real training gains, but it represents only one configuration (7B, $n=256$) and one parallelism strategy (ZeRO-3 within node). The scaling behavior of these speedups under expert parallelism or tensor parallelism is not evaluated.

Claim: On Blackwell GPUs, SonicMoE achieves 25% relative speedup on the forward pass and 15% on the backward pass compared to DeepGEMM on OLMoE-sized 7B MoE models. Supported by Figure 11b (7B $n=1024$: 1242 vs. 1004 TF/s forward, 1122 vs. 1047 TF/s backward). However, this is the least fine-grained 7B configuration tested ($G=1.5$). At higher granularity ($n=256$, $G=6$), SonicMoE achieves 1242 TF/s forward vs. DeepGEMM++ at 1134 TF/s, a smaller relative gap (10%)—the opposite of the granularity-scaling trend observed on H100. The paper does not explain this discrepancy, but it may reflect that Blackwell's TMEM pipelining and UMMA instructions already provide effective IO-overlap for certain kernel shapes, reducing the marginal benefit of SonicMoE's explicit Ping-Pong-style scheduling. A more comprehensive sweep of granularity on Blackwell would strengthen this claim.

Claim: Token rounding routing yields 1.16× speedup on kernel execution time compared to vanilla top-K routing while maintaining similar downstream performance. The 1.16× figure appears in the abstract, and Figure 13 provides supporting evidence: at $E=128$ with $n=1024$, TR achieves ~16.5% forward speedup and ~6.1% backward speedup, for roughly 9.4% end-to-end—not 16%. The 1.16× likely refers to a different configuration (possibly the 25.7% forward / 11.8% backward / 15.9% end-to-end at $E=256$). The downstream quality evidence (Table 2) is strong and covers five training runs with varying sparsity and model scales, always showing TR matching TC within statistical noise. However, all quality evaluations train at relatively small scale (0.5B–1.8B parameters, 20B–100B tokens). No quality evaluation is provided for the 7B+ scale where the throughput benchmarks are measured. It remains possible (though unlikely given the mechanism) that TR's bounded routing perturbations accumulate differently at larger scales or longer training horizons.

Missing experiments and open questions:

• No combination of TR with the full SonicMoE kernel stack. The TR throughput benchmarks (Figure 13) vary sparsity but keep other architecture parameters fixed (7B-style and 120B-style). The quality evaluation (Table 2) uses a separate codebase (OLMoE) rather than SonicMoE's kernels. There is no experiment demonstrating that TR speeds up SonicMoE's kernels more than it speeds up a baseline kernel (or whether the interaction produces additional compound benefits).

• No FLOPs-matched quality comparison between SonicMoE and baselines. The paper demonstrates that SonicMoE processes tokens faster, but does not train a SonicMoE-powered model and a baseline-kernel-powered model to the same token budget and compare their downstream quality. The implicit assumption is that faster kernels → more tokens processed in the same time → better quality. While this is reasonable, subtle differences in numerical precision, kernel scheduling, or memory access patterns could affect training dynamics.

• Single GPU benchmarking for most kernel measurements. The kernel benchmarks in Figures 11 and 12 run on a single GPU with synthetic data. Real distributed training involves communication (all-reduce, all-to-all for expert parallelism) that overlaps with computation. Whether SonicMoE's kernel optimizations (particularly the asynchronous TMA operations and Ping-Pong scheduling) interact favorably or unfavorably with communication overlap is untested.

• No comparison with FP8 or mixed-precision variants. SonicMoE uses BF16 throughout. DeepGEMM's primary strength is FP8 grouped GEMM, but the paper only compares against DeepGEMM's BF16 kernels. A comparison with DeepGEMM's FP8 kernels would show whether SonicMoE's IO-overlap techniques provide benefits beyond what lower precision already achieves for memory bandwidth reduction.

• TR quality evaluation at small scale only. The largest TR-trained model is 1.8B parameters trained on 40B tokens. The throughput benefits are benchmarked at 7B–120B scale. Whether TR maintains quality parity with TC at 7B+ scale, or under longer training (200B+ tokens), is an open empirical question not addressed.

• No ablation of Ping-Pong scheduling independently. The paper attributes speedups to multiple techniques simultaneously (gather fusion, epilogue fusion, Ping-Pong scheduling, asynchronous TMA). While Figure 19 isolates gather fusion and Figure 21 isolates aggregation strategy, there is no figure isolating Ping-Pong scheduling (e.g., SonicMoE with vs. without Ping-Pong on the same kernel). The 20% TFLOPS gain attributed to Ping-Pong in Figure 16 is qualitative, not a controlled experiment.

• B300 benchmarking limited in baselines. On B300, ScatterMoE and MoMoE are benchmarked only for smaller configurations and fail (marked NA) for many larger ones, limiting the completeness of the comparison. The Triton official example is inference-only (no $H$ storage) and lacks backward pass support, making forward comparisons incomplete.

Overall assessment. The experiments provide strong evidence for the three core claims, but the evidence is strongest for the first two (memory efficiency and kernel throughput) and slightly weaker for the third (token rounding quality at scale). The paper's central narrative—that co-designing architecture and kernel yields large efficiency gains for fine-grained sparse MoEs—is well-supported by the consistent speedups across model scales, granularity levels, and GPU generations. The main limitations are the absence of large-scale training quality comparisons (SonicMoE-trained vs. baseline-trained models), the single-GPU focus of kernel benchmarks, and the incomplete isolation of individual optimization contributions in the throughput results. These gaps are partially mitigated by the end-to-end FSDP-2 measurement (which validates that kernel speedups survive distributed training overhead) and the thorough ablation of TR variants (which validates robustness of the routing approach), but a rigorous demonstration of quality parity at 7B+ scale with full training runs would significantly strengthen the paper's practical impact.

6. Limitations and Trade-offs

The Difficulty Estimation Overhead Is Unmeasured and Potentially Dominates Savings

The entire compute-optimal framework rests on the ability to estimate prompt difficulty before deciding how to allocate the inference budget. The paper's method for doing so—generating 2048 samples per question and averaging either ground-truth correctness (oracle) or PRM final-answer scores (predicted)—is extraordinarily expensive. At 2048 samples per question, the difficulty estimation step alone consumes more compute than the largest test-time budgets studied (256–512 generations). The authors acknowledge this explicitly in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

This is a significant gap. The reported 4× efficiency gains over best-of-N are computed after difficulty is known, without amortizing the cost of learning it. In a realistic deployment, the total cost would be difficulty estimation + strategy execution, and the former could dominate the latter. The paper suggests future work on training models to predict difficulty directly from the question text, but no such model is developed or evaluated. Until this gap is closed, the 4× figure should be understood as an upper bound on achievable efficiency rather than a realized deployment gain.

The consequence is that a practitioner attempting to deploy compute-optimal test-time scaling would face a chicken-and-egg problem: they need difficulty estimates to allocate their budget efficiently, but obtaining those estimates consumes a large fraction of the budget they're trying to optimize. In the limiting case where the difficulty estimation cost equals the test-time budget itself, the approach provides no net benefit whatsoever. The paper provides no quantitative analysis of what this breakeven point would be, nor does it characterize how much cheaper difficulty estimation would need to become for the framework to be practically deployable.

The evidence for this limitation is the difficulty estimation procedure itself (Section 3.2): 2048 samples per question, which for a 500-question test set represents 1,024,000 generations just to estimate difficulty—roughly 4× more compute than the largest individual test-time budgets studied. The predicted-difficulty variant (using PRM scores instead of ground-truth labels) removes the need for answer access but does not reduce the number of samples required. The paper reports no experiments varying the number of samples used for difficulty estimation, so the sensitivity of the compute-optimal policy to estimation quality is unknown.

This limitation is explicitly acknowledged but entirely unaddressed. The paper states it is left to future work and notes that "pretraining or finetuning models to directly predict difficulty of a question" would be a natural solution (Section 8). There is no partial mitigation, no sensitivity analysis, and no rough calculation of the amortized cost. A practitioner would need to solve this independently before deploying the method.

Single Benchmark, Single Model Family Limits Generalization

All experiments use the MATH benchmark (500 test questions) with PaLM 2-S* as the base model. The authors state they "believe this model is representative of the capabilities of many contemporary LLMs" (Section 4), but this claim is unverified and several aspects of the findings could be model-specific.

The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution. A model with different calibration properties—for instance, one that tends to produce higher-entropy or lower-entropy solution distributions, or one with different error patterns (e.g., arithmetic errors vs. logical errors)—could exhibit entirely different difficulty-dependent scaling curves. The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families (some models are far better at self-correction from context than others). The MATH benchmark consists exclusively of competition-level math problems requiring symbolic reasoning—it is unclear whether the difficulty-dependent patterns (beam search hurting easy problems, revisions helping easy problems) generalize to other reasoning domains like code generation, logical reasoning, or scientific QA, or to tasks requiring factual knowledge rather than deductive inference.

The consequence for practitioners is that the specific strategy allocations learned for PaLM 2-S* on MATH may not transfer to their own model and task, even if the model is from a similar family or the task is superficially related to math reasoning. The compute-optimal policy—which strategy to use at which budget for which difficulty bin—is entirely empirical and must be re-derived for each new model-task pair. The paper provides no theoretical guidance about what properties of the model or task determine the optimal policy.

The evidence for this limitation is the experimental design itself: a single model (PaLM 2-S*) on a single benchmark (MATH, 500 test questions). The paper provides no experiments with alternative models (e.g., LLaMA, GPT, Gemma families) or alternative benchmarks (e.g., GSM8K, HumanEval, ARC). The test set of 500 questions, split into five difficulty quintiles of ~100 each, then further split by two-fold cross-validation, means the compute-optimal policy is selected based on ~50 questions per fold per bin. This is a small sample, and the selected strategies may not be robust to sampling variation—the paper reports no confidence intervals on the compute-optimal scaling curves.

The paper acknowledges this implicitly by framing the work as a "first systematic analysis" (Section 2) and calling for "future work extending the analysis to other domains" (Section 8), but does not test any alternatives. There is no mitigation—a practitioner must assume the findings are model- and benchmark-specific until reproduced on their own setup.

The ~14× Larger Model Baseline Is Not Compute-Optimal

The FLOPs-matched comparison in Section 7 scales model parameters while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023). The authors acknowledge that this departs from compute-optimal pretraining (Hoffmann et al., 2022), where both data and parameters are scaled equally:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

This matters because a Chinchilla-optimal model trained with ~14× more total FLOPs would likely outperform a parameter-only-scaled model, making the pretraining baseline weaker than it needs to be. The reported advantages of test-time compute over pretraining—for example, +27.8% relative improvement on easy questions at R ≪ 1—may shrink or reverse against a properly compute-optimal larger model. Additionally, the ~14× larger model uses only greedy decoding—no majority voting, no best-of-N, no search. Giving the larger model even a modest test-time compute budget (e.g., best-of-8) would create a much stronger baseline that is never tested.

The consequence is that the paper's headline finding—"a smaller model with additional test-time compute can outperform a ~14× larger model"—may overstate the advantage. The comparison is between a compute-optimally deployed smaller model and a non-optimally trained, non-optimally deployed larger model. A more balanced comparison would either (a) give both models compute-optimal test-time strategies, or (b) train both models with compute-optimal pretraining recipes. The current setup conflates the test-time vs. pretraining tradeoff with differences in training efficiency and inference strategy.

The evidence is in the FLOPs-matched comparison itself (Section 7, Figure 9). The paper transparently describes the setup—"we scale only model parameters" and "the larger model uses greedy decoding"—and acknowledges the caveat in the text. However, no sensitivity analysis is provided, and readers could easily misinterpret the bar charts in Figure 1 as showing a fair comparison when it is asymmetrical in the larger model's disfavor.

This limitation is acknowledged but not mitigated. The paper frames the choice as "representative of a canonical approach" but does not provide an alternative comparison with a Chinchilla-optimal baseline or a test-time-augmented larger model. A practitioner evaluating whether to invest in test-time compute vs. larger pretraining cannot rely on the current numbers without adjusting for these asymmetries.

Verifier Over-Optimization Is a Hard Ceiling, Not a Solved Problem

The paper documents verifier over-optimization as a central limiting factor: beam search degrades easy-problem performance at high budgets (Figure 3, right), lookahead search—the strongest optimizer—paradoxically performs worst overall (Figure 3, left), and qualitative examples show degenerate outputs such as repetitive low-information steps and overly short 1–2 step solutions (Appendix M). The compute-optimal policy mitigates this by routing easy problems away from aggressive search toward best-of-N, but it does not solve the underlying problem. On medium-difficulty problems where beam search is deployed, over-optimization still limits the scaling ceiling—the beam search curves in Figure 3 flatten and sometimes decline well before the budget is exhausted.

The consequence is that the compute-optimal approach is fundamentally bounded by verifier quality. No matter how well the budget is allocated, the maximum achievable accuracy on any problem is limited by the point at which further optimization of the verifier signal stops improving (or starts degrading) the true accuracy. Improving the PRM—through better training data, adversarial robustness, ensemble methods, or fundamentally different training objectives—would likely shift the difficulty thresholds and change the optimal policy. The current results are therefore specific to the verifier quality achievable with the Monte Carlo rollout training procedure described in Appendix D, and the paper does not explore how verifier improvements would alter the scaling landscape.

The consequence for a practitioner is that investing in better verifiers may yield larger gains than investing in better test-time allocation strategies, but the paper provides no guidance on this tradeoff. The optimal policy is a function of the verifier quality, and if the verifier improves, the policy must be re-derived.

The evidence is in Figure 3 (right), where beam search (M = 4) on difficulty bin 1 actually decreases accuracy as the budget increases from 4 to 256, demonstrating over-optimization at its most extreme. Figure 3 (left) shows lookahead search underperforming best-of-N at equal generation budgets because the verifier's lookahead scores are less reliable than its immediate step scores. Appendix M provides qualitative examples of degenerate outputs that score highly under the PRM but are incorrect.

The paper treats verifier over-optimization as a documented phenomenon and a motivation for compute-optimal allocation, but does not attempt to improve the verifier itself. Section 8 acknowledges that "improving verifier robustness is the key bottleneck for further scaling test-time compute" and suggests future work on adversarial training and ensemble methods. There is no partial mitigation—the compute-optimal policy works around the verifier's limitations rather than addressing them.

The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate

As noted in Section 6.1, approximately 38% of correct answers produced during a revision chain get "revised" back to incorrect answers in the subsequent step. The paper reports:

"approximately 38% of correct answers get converted back to incorrect ones using a naive approach"

This is a direct consequence of the training data construction: the model only sees incorrect-to-correct trajectories during training, so it has no signal for what to do when the current answer is already correct. The model's learned behavior is "given an answer, produce a better one"—it does not learn to recognize when the current answer is already acceptable and should be preserved.

The consequence is that in a purely sequential revision chain, the model can oscillate between correct and incorrect answers, and the final answer in the chain is not guaranteed to be correct even if a correct answer appeared earlier. The paper mitigates this with majority voting or verifier-based selection across the chain (picking the best answer from any point rather than always taking the last revision), but these are imperfect patches. Majority voting requires the correct answer to appear multiple times in the chain to dominate; verifier-based selection reintroduces the verifier over-optimization problem described above. A more principled solution—such as training the model with both correct-to-correct and incorrect-to-correct trajectories, or training an explicit "stop revising" classifier—is not explored.

The ReST^EM experiment (Appendix K, Figure 16) further highlights the fragility of revision training. Attempting to optimize the revision model with RL-style training caused performance to degrade substantially with sequential revisions, likely because on-policy data collection amplified spurious correlations in the revision trajectories. At 256 generations, fully sequential performance with the ReST^EM model drops to approximately 33.5% compared to roughly 38.5% at the optimal sequential-to-parallel ratio. This suggests that the revision approach is sensitive to training methodology in ways that are not fully understood, and the positive results depend on specific choices (offline data construction, edit-distance-based pairing) that may not transfer to other settings or training procedures.

The evidence is the 38% reversion rate reported in Section 6.1, supported by the ReST^EM degradation in Appendix K (Figure 16). The paper acknowledges the reversion problem explicitly and describes the within-chain selection mitigation. However, the mitigation is not evaluated against a "no reversion" counterfactual—there is no version of the revision model trained to preserve correct answers, so the cost of the 38% reversion in terms of lost accuracy is unknown. The paper does not explore whether the reversion rate changes with chain length, problem difficulty, or revision model checkpoint.

7. Implications and Future Directions

How This Work Changes the Landscape

SonicMoE establishes co-design as a necessary methodology for MoE systems, not an optional optimization. Before this paper, MoE architecture design (choosing expert granularity G = d/n and sparsity K/E) and GPU kernel engineering (optimizing grouped GEMM, scheduling memory operations) operated in separate communities with separate objectives. Architects pushed toward finer granularity and higher sparsity because scaling laws predicted better quality per theoretical FLOP. Kernel developers optimized for the coarse-grained, dense-expert regimes where existing benchmarks lived. Neither side accounted for the other's constraints, creating a growing gap between the architectures scaling laws recommend and the architectures hardware can efficiently train. Table 4 documents this widening gap: expert granularity increased tenfold (from G = 0.38 in Mixtral 8×22B to G = 4.00 in Qwen3-Next-80B-A3B) in under two years, while activation ratios dropped from 25% to 1.95% over the same period—all without corresponding kernel innovations to handle the resulting memory-bound regime.

SonicMoE closes this gap by demonstrating that the IO cost barrier to fine-grained sparse MoEs is engineering-tractable, not fundamental. The paper's derivation of arithmetic intensity (Equation 4) provides the diagnostic: 3 / (2 + 2G/d + 3/(Tρ)) shows that as G increases or ρ decreases, the computation becomes increasingly memory-bound, making memory latency—not compute throughput—the binding constraint. Prior work implicitly treated MoE training as a GEMM optimization problem (faster matrix multiplies → faster training). SonicMoE reframes it as a memory bandwidth utilization problem: the path to higher throughput is reducing HBM traffic (gather fusion, epilogue fusion, reduced activation caching) and hiding unavoidable latency (Ping-Pong scheduling, asynchronous TMA, TMEM two-stage pipelining). The evidence that this reframing is correct comes from the granularity-dependent scaling of SonicMoE's advantage: on H100 GPUs, the relative speedup over DeepGEMM++ grows from ~8% at coarse granularity (G = 2) to ~24% at fine granularity (G = 8) on 120B models (Figure 11a), precisely as Equation 4 predicts.

This changes the research landscape in three concrete ways:

First, it reorients MoE systems research from compute-centric to memory-centric optimization. Investments in faster GEMM (FP4, structured sparsity, improved Tensor Core utilization) yield diminishing returns for fine-grained MoEs because the bottleneck is not the arithmetic but the bytes moved. SonicMoE's grouped GEMM with contiguously-packed inputs achieves only 2.7–10% higher TFLOPS than DeepGEMM's highly-tuned GEMM (Figure 18a), yet the overall forward pass is 43% faster (Figure 11a). The difference comes from everything around the GEMM—gather fusion, epilogue fusion, IO-compute overlap. Future MoE kernel work should prioritize techniques that reduce HBM traffic (smarter fusion, tiling strategies that minimize round-trips) and hide latency (asynchronous operations, multi-stage pipelining) over techniques that accelerate pure arithmetic.

Second, it removes the memory ceiling that was preventing adoption of scaling-law-recommended architectures. The finding that ScatterMoE, MoMoE, DeepGEMM-pt, and DeepGEMM++ all fail to run the DeepSeek-V3.2-Exp configuration on a single H100 (Figure 12a, marked NA), while SonicMoE runs successfully, demonstrates that memory efficiency is not just a performance optimization—it is an enabling condition for frontier architectures. The 45% activation memory reduction on 7B fine-grained MoE (Figure 10), and the flat memory curve with increasing granularity (Figure 1 left), mean that practioners can now train MoEs at the granularity levels that scaling laws predict will maximize quality per FLOP, without being blocked by out-of-memory errors or forced to use smaller microbatches that degrade throughput. This directly enables the architecture trends documented in Table 4 to continue beyond the point where existing kernels would fail.

Third, it integrates routing into the hardware-efficiency optimization loop. Token rounding (Section 5) introduces the conceptual shift that routing decisions carry hardware costs—specifically, tile quantization waste—that can be optimized jointly with model quality rather than being treated as an orthogonal concern. The finding that TR matches or slightly exceeds TC top-K quality across all tested sparsity levels (Table 2), while providing up to 16% higher model TFLOPS at K/E = 1/64 (Figure 13), demonstrates that the "optimal" routing for model quality is not necessarily optimal for end-to-end training speed, and that small, bounded perturbations to the routing decision can yield hardware benefits with zero quality cost. This opens the door to routing algorithms that are explicitly designed with hardware cost functions in mind—not just load balancing or expert utilization, but tile alignment, memory coalescing, and communication patterns in distributed settings.

The paper also reconciles a latent tension in the MoE literature: the conflict between the FLOPs-efficiency narrative (MoEs use fewer FLOPs per token than dense models, so they should be faster to train) and the practical observation that MoE training throughput often falls short of dense model throughput due to kernel inefficiency. SonicMoE demonstrates that this tension is not inherent to MoEs—it is an artifact of kernel designs that were not optimized for the memory-bound regime that fine-grained sparse MoEs inhabit. When the kernels are co-designed with the architecture, MoE training throughput can approach the hardware's theoretical limits (88% of the cuBLAS upper bound on both H100 and B300, Figure 1), recovering the FLOPs-efficiency advantage that motivated MoEs in the first place.

Follow-Up Research This Work Enables

Extending token rounding to distributed expert-parallel training with communication-aware tiling. SonicMoE evaluates token rounding in single-GPU settings, but production MoE training typically uses expert parallelism where different experts reside on different GPUs and tokens are communicated via all-to-all collectives. In this setting, tile alignment affects not only GEMM padding but also communication efficiency: all-to-all operations are more efficient when the per-expert token counts are balanced and aligned to network packet sizes. A natural extension would co-design the rounding algorithm to simultaneously optimize for GEMM tile alignment (M_tile = 128) and for the all-to-all communication granularity (typically network-dependent, e.g., minimum message sizes or alignment requirements of NCCL collectives). The key experiment would measure end-to-end training throughput with TR in an FSDP + expert-parallel setup on 64–256 GPUs, comparing against TC top-K under realistic all-to-all communication patterns, and ablating whether the communication benefits of TR exceed or are additive with the GEMM padding benefits documented in Figure 13.

Low-precision SonicMoE with FP8 and MXFP formats. The paper uses BF16 throughout and mentions FP8, MXFP8, and MXFP4 as future directions (Section 7). The key question is whether SonicMoE's IO-overlapping techniques remain effective when the arithmetic intensity shifts due to lower-precision compute. FP8 GEMM is roughly 2× faster than BF16 GEMM on H100, which reduces the ratio of compute time to IO time—potentially making the memory-bound problem more severe, not less, because the GEMM finishes faster while IO latency stays constant. Conversely, FP8 activations halve the HBM traffic for a given tensor size, which could push the workload toward the compute-bound regime and reduce the benefit of IO-overlap. A concrete experiment would benchmark SonicMoE against DeepGEMM's FP8 grouped GEMM kernels (which already achieve state-of-the-art FP8 throughput) across the same granularity sweep (1.4B–120B, G from 2 to 16), measuring whether SonicMoE's Ping-Pong scheduling and asynchronous TMA operations provide additional speedup beyond what FP8 already achieves, or whether the gap closes as the workload becomes more compute-bound. The extension to MXFP formats (shared exponents across blocks) adds the wrinkle that the GEMM tile shapes may need to change to accommodate exponent sharing, which interacts with SonicMoE's tile size choices and the fixed M_tile = 128 used by token rounding.

Training diffusion models or mixture-of-attention architectures with SonicMoE's kernel primitives. The paper notes that alternative sparse architectures—PEER (He 2024), Memory Layers (Berges et al. 2024), and Ultra-Mem (Huang et al. 2025)—face the same IO and memory bottlenecks as fine-grained MoEs. SonicMoE's modularized grouped GEMM kernel with configurable prologue/epilogue fusion and the gather-and-sum expert aggregation kernel are not MoE-specific; they implement generic varlen-M and varlen-K grouped GEMM with arbitrary gather patterns and epilogue operations. This means the kernel infrastructure could be applied to any architecture that routes tokens (or patches, or embeddings) to multiple sub-networks with variable assignment counts. A strong follow-up would port SonicMoE's kernel stack to a PEER-like architecture (which uses product-key retrieval to select a tiny fraction of a million+ experts), measure throughput and memory usage against the reference PEER implementation, and identify which of SonicMoE's three contributions (memory-efficient backward, IO-overlap, token rounding) transfer directly and which require architectural adaptation. The key metric is whether the IO-overlap techniques generalize beyond the SwiGLU + Grouped GEMM pattern to the weight-retrieval + dot-product patterns common in memory-augmented architectures.

Automated difficulty estimation for token rounding's optimal tile size and microbatch size. Token rounding requires choosing M_tile and microbatch size T such that Tρ / M_tile ≥ 2 for stable quality (Table 7). This ratio depends on the sparsity level ρ, which may change during training if dynamic routing (e.g., top-P routing from Huang et al. 2024) is used, or if the model is trained with curriculum learning that varies sequence length and thus T. A practical extension would develop a lightweight runtime monitor that estimates the current Tρ / M_tile ratio from the router's output distribution and adaptively switches between TR and TC, or adjusts M_tile, to maintain the quality-stability criterion. The experiment would train a 1.4B MoE with varying sequence lengths (128–4096) using adaptive tile size selection, measuring whether the adaptive variant matches TR's quality at high Tρ/M_tile and avoids the degradation observed at Tρ/M_tile = 1 (Table 7), while maintaining throughput gains.

Stress-testing token rounding at 7B+ scale and 200B+ token budgets. The paper demonstrates TR quality parity at 0.5B–1.8B scale on 20B–100B tokens (Table 2), but the throughput benefits are benchmarked at 7B–120B scale (Figures 13, 14). The scaling behavior of TR's quality is unknown at larger model sizes and longer training horizons. There are plausible mechanisms for degradation: (1) as the model capacity increases, the routing decisions become more stable and influential, making tile-alignment perturbations more impactful; (2) longer training amplifies small distribution shifts from the TR-induced token redistribution. Conversely, larger models may be more robust to routing perturbation because they have greater capacity to compensate. A concrete negative-result experiment would train a 7B OLMoE-style MoE with TR and TC for 200B tokens (roughly Chinchilla-optimal for that scale) and measure not only final perplexity and downstream accuracy but also the training dynamics—loss curves, expert utilization entropy, and routing stability metrics—to identify whether TR introduces training instability or systematic bias that accumulates over long horizons. A positive result (quality parity at scale) would validate TR for production training; a negative result (quality degradation at scale) would establish boundary conditions on TR's applicability.

Combining SonicMoE's kernel stack with communication-computation overlap in expert parallelism. SonicMoE's kernels are benchmarked on single GPUs, but the asynchronous TMA operations and Ping-Pong scheduling create opportunities for overlapping kernel execution with the all-to-all communication used in expert-parallel distributed training. The standard expert-parallel pattern is: all-to-all scatter (communicate tokens to expert-owning GPUs) → grouped GEMM (compute) → all-to-all gather (communicate results back). If the grouped GEMM can be launched asynchronously while the all-to-all is in progress (or vice versa), the communication can be hidden behind computation. The specific experiment would integrate SonicMoE's kernels with DeepEP (Zhao et al. 2025c) or a similar expert-parallel communication library, measure the degree of compute-communication overlap achievable using CUDA graphs or streams, and compare end-to-end training throughput against a baseline that serializes communication and computation. The key technical challenge is whether SonicMoE's internal asynchrony (TMA stores, Ping-Pong warpgroups) interferes with or enables external asynchrony with communication collectives.

Practical Applications and Downstream Use Cases

Training frontier-scale fine-grained MoEs on existing hardware without OOM errors. The paper demonstrates that ScatterMoE, MoMoE, and both DeepGEMM variants fail to run the DeepSeek-V3.2-Exp configuration (d = 7168, n = 2048, E = 256, K = 8) on a single H100 due to index overflow or CUDA OOM errors (Figure 12a), while SonicMoE runs at ~600 TFLOPS forward and ~580 TFLOPS backward. For teams training models at this scale (685B total parameters, the DeepSeek-V3 class), SonicMoE's 45% activation memory reduction (Figure 10: ~3 GiB saved per layer at 120B scale) is an enabling condition—without it, the model simply cannot fit on current GPUs regardless of parallelism strategy. The practical workflow is to replace existing MoE kernel backends (typically DeepGEMM or Megatron-LM GroupedMLP) with SonicMoE's CuTe-DSL kernels when targeting architectures with G ≥ 4 or K/E ≤ 5%, which covers the majority of frontier models in Table 4. The throughput improvement (Figure 12a: consistently >550 TFLOPS across all production configurations) means that the same hardware can process more tokens per day, directly accelerating model development cycles.

Cost-efficient training of extremely sparse MoEs using token rounding. For architectures pushing toward K/E < 2% (e.g., the Qwen3-Next-80B-A3B-Thinking configuration with K/E = 1.95%), the tile quantization waste documented in Figure 8 reaches 15–25% of total FLOPs. Token rounding eliminates this waste while matching token-choice routing quality (Table 2e: 1.4B MoE with K/E = 1/64 achieves identical 52.0% average accuracy and marginally better perplexity at 13.22 vs. 13.32). The practical deployment is a drop-in replacement: swap the router from torch.topk to SonicMoE's TR implementation (which includes an efficient bitonic-sort top-K kernel achieving 1.5–2.0× the bandwidth of PyTorch top-K, Figure 22) and use NR-f rounding with M_tile = 128 as the default rounding subroutine. The key constraint from Table 7 is that Tρ/M_tile ≥ 2—the microbatch size must be large enough that each expert expects at least 2 tiles of tokens. For the Qwen3-Next configuration with T = 32768 and K/E = 10/512, this gives Tρ = 32768 × 10/512 = 640, and 640/128 = 5 ≥ 2, so TR is safely applicable. The throughput gain scales with sparsity (Figure 13): 16% model TFLOPS improvement at K/E = 1/64, translating to 16% more tokens processed per GPU-hour for the same model quality.

On-device or edge training of MoE models with memory-constrained GPUs. SonicMoE's constant activation memory with increasing granularity (Figure 1, left panel) means that finer-grained MoEs—which are more parameter-efficient per FLOP—can be trained on GPUs with limited HBM capacity without reducing microbatch size. For a 1.4B MoE with d = 768, n = 256, E = 128, K = 8 on a consumer GPU with 24 GB HBM, SonicMoE uses ~0.4 GiB per MoE layer for activations versus ScatterMoE's ~0.7 GiB (Figure 10, left panel). Across 18 layers (the OLMoE-1B architecture), this saves ~5.4 GiB of activation memory—enough to double the microbatch size or fit the model on a GPU that would otherwise OOM. The practical benefit is enabling fine-tuning or continued pretraining of MoE models on hardware that was previously limited to dense models of equivalent parameter count, lowering the barrier to entry for MoE experimentation and deployment.

Accelerating self-improvement and distillation pipelines with higher training throughput. When MoE models are used to generate training data (e.g., teacher-student distillation, STaR-style self-training, or rejection sampling finetuning), the training step is often the throughput bottleneck because it requires full forward and backward passes over large datasets. SonicMoE's 1.86× overall throughput improvement (abstract) means that a data generation + training loop using SonicMoE can complete 1.86× more training iterations in the same wall-clock time, directly accelerating the self-improvement cycle. The FSDP-2 measurement (213B tokens/day on 64 H100s vs. ScatterMoE's 225B on 96 H100s, Section 6.2) provides the concrete number: a team with access to 64 H100s can train a 7B fine-grained MoE at the same rate as a team with 96 H100s using previous kernels, representing a 33% reduction in required GPU-hours for the same training throughput. For iterative pipelines where multiple rounds of training interleave with data generation, this compounding speedup significantly shortens the total cycle time.

When to Prefer This Method

The paper positions SonicMoE against five alternative MoE kernel implementations (ScatterMoE, MoMoE, MegaBlocks, Megatron-LM, DeepGEMM) with specific technical differentiators (Table 1). The following decision rules emerge from the empirical comparisons and architectural analysis:

Prefer SonicMoE over all alternatives when:

• Training MoEs with expert granularity G = d/n ≥ 4 (i.e., fine-grained experts), because SonicMoE's IO-overlapping techniques (gather fusion, Ping-Pong scheduling, asynchronous TMA) provide speedups that scale with granularity (Figure 11a: relative advantage over DeepGEMM++ grows from 8% at G = 2 to 24% at G = 8 on 120B models), while alternatives without these techniques show linearly decreasing arithmetic intensity without latency-hiding.
• Training MoEs with sparsity K/E ≤ 5% (highly sparse), because tile quantization waste exceeds 10% of total FLOPs in this regime (Figure 8), and token rounding eliminates this waste with no quality degradation (Table 2), providing up to 16% additional throughput (Figure 13) that alternatives cannot recover without modifying their routing.
• Operating under tight GPU memory constraints, because SonicMoE's activation memory is independent of expert granularity and 45% lower than ScatterMoE's at fine granularity (Figure 10), enabling larger microbatches or higher-granularity architectures within the same HBM budget.
• Targeting models at or beyond the DeepSeek-V3 scale (d = 7168, n = 2048, E = 256), because alternative kernels (ScatterMoE, MoMoE, DeepGEMM) fail to run this configuration on a single H100 (Figure 12a), while SonicMoE runs successfully at ~600 TFLOPS.

Prefer DeepGEMM (with FP8) when:

• The primary goal is maximizing GEMM throughput specifically for FP8 mixed-precision training, and the MoE architecture is coarse-grained (G ≤ 2) with dense expert assignment (K/E ≥ 10%), because in this regime the workload is compute-bound and DeepGEMM's FP8 grouped GEMM kernels are highly optimized for arithmetic throughput—SonicMoE's BF16-only implementation and IO-overlap techniques provide less relative advantage when GEMM dominates runtime and FP8 already halves memory traffic.
• The deployment uses expert parallelism with all-to-all communication where tokens are already contiguously packed before reaching the GEMM, because DeepGEMM is explicitly designed for this communication pattern (assuming inputs are pre-gathered and padded), making SonicMoE's gather fusion and token rounding less impactful when the gather is handled by the communication library.

Prefer Megatron-LM GroupedMLP when:

• The training framework is already deeply integrated with Megatron-LM's pipeline and tensor parallelism, and the engineering cost of replacing the MoE backend outweighs the throughput gains, because Megatron-LM's recent memory-efficient patch (October 2025) already adopts a SonicMoE-like computational path for the backward pass, closing part of the activation memory gap for coarse-grained configurations.

Prefer ScatterMoE or MoMoE when:

• Development velocity and Python-level hackability (Triton-based kernels) are prioritized over absolute training throughput, because these implementations are easier to modify and experiment with than SonicMoE's CuTe-DSL kernels, and the throughput gap is smallest for coarse-grained, compute-bound MoE configurations (G ≤ 2, K/E ≥ 10%).