SonicMoE: Accelerating MoE with IO and Tile-aware Optimizations
Authors: Wentao Guo
Mayank Mishra
Xinle Cheng
Ion Stoica
Tri Dao
Venue: ICLR 2026 Poster (2026)
Keywords: Mixture of Experts
GPU
kernel
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: We propose an algorithm and system codesign for fine-grained and sparse MoE.

Abstract: Mixture of Experts (MoE) models have emerged as the de facto architecture for scaling up language models without significantly increasing the computational cost. Recent MoE models demonstrate a clear trend towards high expert granularity (smaller expert intermediate dimension) and higher sparsity (constant number of
activated experts with higher number of total experts), which improves model quality per FLOP. However, fine-grained MoEs suffer from increased activation memory footprint and reduced hardware efficiency due to higher IO costs, while sparser MoEs suffer from wasted computations due to padding in Grouped GEMM kernels. In response, we propose a memory-efficient algorithm to compute the forward and backward passes of MoEs with minimal activation caching for the backward pass. We also design GPU kernels that overlap memory IO with computation benefiting all MoE architectures. Finally, we propose a novel "token rounding" method that minimizes the wasted compute due to padding in Grouped GEMM kernels. As a result, our method SonicMoE reduces activation memory by 45% and achieves a 1.86x compute throughput improvement on Hopper GPUs compared to ScatterMoE's BF16 MoE kernel for a fine-grained 7B MoE. Concretely, SonicMoE on 64 H100s achieves a training throughput of 213 billion tokens per day comparable to ScatterMoE's 225 billion tokens per day on 96 H100s for a 7B MoE model training with FSDP-2 using the lm-engine codebase. On Blackwell GPUs, SonicMoE also achieves a 28.7% and 22.1% relative speedup on the forward and backward pass respectively compared to a highly optimized DeepGEMM baseline on OLMoE-sized 7B MoE models. Under high MoE sparsity settings, our tile-aware token rounding algorithm yields an additional 1.16x speedup on kernel execution time compared to vanilla top-K routing while maintaining similar downstream performance. We open-source all our kernels to enable faster MoE model training. An extended version of this paper can be found on arXiv.

Published as a conference paper at ICLR 2026
SONICMOE: ACCELERATING MoE WITH 1O AND TILE-AWARE OPTIMIZATIONS
Wentao Mayankk Mishra Xinle Cheng', Ion Stoica', Tri Dao Princeton University, Uniiversity of California, Berkeley, TToetherr Al ug0420@princet.omm edu, trietr tri@tridas me
ABSTRACT
Mixture of Experts (MoE) models have emerged as the de facto architecture for scaling up language models without significantly increasing the computational cost. Recent MoE models demonstrate a clear trend towards high expert granularity (smaller expert intermediate dimension) and higher sparsity (constant number of activated experts with higher number of total experts), which improves model quality per FLOP. However, fine-grained MoEs suffer from increased activation memory footprint and reduced hardware efficiency due to higher IO costs, while sparser MoEs suffer from wasted computations due to padding in Grouped GEMM kernels. In response, we propose a memory memory-efficient algorithm to compute the forward and backward passes of MoEs with minimal activation caching for the backward pass. We also design GPU kernels that overlap memory 10 with computation benefiting all MoE architectures. Finally, we propose novel "token rounding" method that minimizes the wasted compute due to padding in Grouped GEMM kernels. As a result, our method SonicMoE reduces activation memory by 45% and achieves 1.86x compute throughput improvement on Hopper GPUs compared Moo''s BF16 MoE kernel for fine-grained 7B MoF. Concretely, SonicMoE on 64 H100s achieves a training throughput of 213 billion tokens per day comparable to ScatterMoE's 225 billion tokens per day on 96 H100s for 7B MoE model training with FSDP-2 using the Im-engine codebase. On Blackwell GPUs, SonicMoE also achieves 28.7% and 22.1 % relative speedup on the forward and backward pass respectively compared to highly optimized DeePGEMM baseline on OLMoE-sized 7B MoE models. Under high MoE sparsity settings, tile-aware token rounding algorithm yields additional 1.16x speedup on kernel execution time compared to vanilla top- K routing while maintaining similar downstream performance. We open-source all our kernel to enable faster MoE model training, An extended version of this paper can be found on arXiv
INTRODUCTION
Mixture of Experts (MoE) Shazeer et al. 2017) models have emerged key technique for scaling up parameters (Zhao et al. 2025a Kimi et al. (025) without increasing the training computational require- ments. Modern transformers often have layers comprised of a sequence mixer block (e.g. Multi-head Attention (Vaswani et al. 2017)) followed by a channel mixer block (c.g. dense MLPs) where MoEs are an excellent substitute for dense MLPs for FLOPs efficiency. A MoE block is typically composed of a to- ken router and multiple smaller and often equal-sizza subnetworks, called "experts"" MoEs can reduce FLOPs consumption during training by only activating subset of all experts per token. However, re- ducing FLOPs does not directly translate to better hardware utilization since MoE computation features more dynamic I0 accesses when each expert needs to gather token embeddings from different positions, and also scatter the results back to the original positions. Moreover, such hardware- cettwarr-nfrriedlinnee becomes worse experts become more granular (experts have smaller intermediate sizes) and sparser (experts are increased while keeping the number of activated experts constant), shown in Table 3
https://githui com/bee AT1
https:///ree wwwgg//t..oogg///// 4080

---

Published as a conference paper at ICLR 2026
MoE scaling laws (Clark et al. 2022 Krajewski al. 2024) predict better model quality per FLOP with
increasing expert granularity (ratio between the model embedding dimension and each expert's 's inter- mediate size) and sparsity. Recent MoE models like DeepSeek V3 (DeepSeck-Al et al. 2024), Qwen3 120b
MoF. (QwenI M 2025) and gpt-oss- 120b (OpenA] 2025). have demonstrated superior performance of "fine-grained MoEs over "coarse-grained" MoEs scale. Besides granularity, the pursuit of MoEs with better model quality while keeping computational requirements constant has also led to modern MoEs becoming sparser. For example, Kimi K2 (Kimi et al. (2255) has the same amount of activated parameters as DeepSeek V3 (DeepSeek-Al et al. 2024) but much larger total parameter count. Overall, granularity and sparsity for MoEs have only increased over time as shown in Table We also note that the pursuit of granularity and sparsity is also adopted by recent alternative architectures to MoE such as PEER (He 2024), Memory Layers (Berges et al. 2024). and Ultra-Mem (Huang et al. 2022
Though more granular and sparser MoEs increase model quality per FLOPP, they suffer from hardware inefficiency due to: (1) larger activation memory footprint for granular MoE. models as activation size typically scales linearly with the number of activated experts, (2) lower arithmetic intensity and increased IO cost due to granular experts and (3) wasted computations due to tile quantization effects of grouped GEMM for highly sparse MoEs. The high granularity and sparsity both push MoE training towards the memory-boung regime requiring carefully designed MoE kernels to hide the increased I0 costs. Existing ofate-of-the-ar MoE kernels such as ScatterMoE (Tan et al. (0244] and MoMoF (Costin et al. (025) are not designed to handle these high 10 costs and they suffer significant training throughput degradation.
Activation Memory Usage Expert Cracularity
Forward Pass T Throughpat
500
E00
Artharan
200-
6006/2348
4006/0024
a21//112
4290/299
4005/2348
400670024
J20//112
4200/210
Engert Graaiarit (dim)
Expert Graadariyy (6/m)
SonicMoE's per-layer activation memory footprint (left) stays constant even when expert granularity (d/n where d is the embedding dimension and n is the expert intermediate dimension) increases, and is 0.20-1.59x more memory-effice cemorr-efficient than other baselines. SonicMoE's ' forward computation throughput (right) reaches an average of 88% (max 91% %. min 86%) of the upper bound (cuBLAS BMM + activation + cuBLAS BMM + 30B MoE configuration with microbatch size of 32768 tokens, and we vary the activated experts total number of experts as 2/32, 4/64, 8/128, and 16/256 from left to right.
Figure 1:
aggregation) H100 GPUs. Here
Summary of contributions. We propose SonicMoE, a hardware and model architecture co-design solution to address MoE training efficiency problems, making the following contributions:
MoE training with minimum possible activation memory footprint without increasing FLOPs: We analyze the impact of MoE granularity on the MoE layer's 's forward and backward passes and observe that increasing MoE granularity while maintaining constant FLOPSs leads to a linear increase in activation memory required by the backward pass. Leveraging this observation, we carefully redesign the computation graph to avoid caching the activations for the router gradient computation while maintaining the mathematical equivalence to the original MoE formulation. As a result, for a fine-grained 7B MoE, SonicMoE reduces activation memory usage per layer by up to 45%.
Efficient MoE kernel that overlaps 10 with computation yield SOTA training throughput: show that increasing both granularity and sparsity leads to MoEs becoming increasingly memory bandwidth bound. To alleviate this bottleneck, we exploit the asynchrony of the GEMM and IO operations by overlapping them to maximize throughput. For the same fine-grained 7B MoE model, our approach increases relative speedup by 43% on the forward pass compared to a highly optimized DeepGEMM bascline on H100 GPUs, and by 83% and 1 115% on the backward pass compared to the state-of- ttthhhart art MoE basclines ScatterMoE and MoMoE, respectively. On B300 GPUs, our approach also achieves 28.79 28.7% relative speedup on the forward pass and 22.10 1% on the backward pass compared to a highly optimized DeePGEMM baseline on OLMoE -ized 7B MoE models.
Token rounding routing that eliminates wasted FLOPs from sparse MoEs: We introduce drop-in routing algorithm that rounds the per-expert token counts to multiples of the tile size (c.g., 1 128) used by grouped GEMM in MoE kernels. This rounding reduces compute wasted on padding
2

---

Published as a conference paper at ICLR 2026
while preserving the original ensures that, for each expert, the maximum deviation from the original top- token-choice result is bounded by one tile. This method effectively eliminates padding waste in grouped GEMM while maintaining the same total number of tokens in expectation, and it delivers robust token-choice accuracy even under highly sparse MoE training regimes. We validate the performance of this token-rounding strategy in 1.4B-parameter sparse training setting, demonstrating that its compute throughput consistently exceeds that of the vanilla top-K token-choice routing. In highly sparse regimes, the improvement reaches up to 16% higher TFLOPS for end-to-end MoE computation.
assignment much as possible. The algorithm
We release SonicMoE mainly written in CuTe-DSL (NVIDIA 2025a) with a PyTorch interface and a permissive license to benefit researchers and practitioners.
2 BACKGROUND
We first provide an overview of the MoE architecture and standard MoE kernel employing grouped
GEMM in Section 2.1] In Section 2.2 we discuss how granularity and MoE sparsity will affect MoE's training efficiency. We then examine the impact of MoE routing methods on the MoE model quality and training efficiency in Section 2.3
Algorithm 1 MoR forward with Grouped GEMM
2.1 MoE USING GROUPED GEMM
Input
E x7'd, W, (Wi,e)
xdx20 W2 = (Wa,) E
Modern GPUs support Tensor Cores; specialized hardware units
gaxd meting scores 8 X
with high matrix multiplication throughput (NVIDIA 2022). A
(0.1)" X a binary-valued
GEMM (general matrix multiply) (Lawson et al.) 1979) kernel often has 3 stages: prologue (start input loading), mainloop (keep loading inputs and compute GEMM) and epilogue (miscellaneous IO/math operations on GEMM outputs). The kernel tiles computa- tions (dividing large matrices into small tiles), and optionally pads dimensions so computation aligns with hardware-permistie sizes. In this paper, we follow standard GEMM notations in most BLAS (Lawson et al. 1979) libraries: we have A IR M C E RMIN for C' AB with problem shape MMNNKK K). This notation is adopted by CUTLASS (NVIDIA
trix where represents whether token ronted to expert e,
Output routput activation O OERTx4
Parallel for e [E] da
X Gatter(X,,r.... (...t..)
up upppprrj. varien- M Grouped GEMM
H. 4 X2 L.,
Ag + act__fun(( (f))
tile
doww proj, varlen M Grouped GEMM
Ye + A222
Purallel ( E do
B E R K
/f expert aggrogation
O2 = Ee
2025ch which implements efficient GEMM CUDA.
A MoE block is typically composed subnetworks, called "experts". The router is responsible for dispatching tokens to the experts. The outputs from all experts in the layer are then aggregated and passed onto the next layer. MoE computa- tion be performed using Grouped GEMM (a list of GEMMs with possibly different (M. dimensions). Algorithm| Tillustrates running MoE forward with Grouped GEMM. Here we refer to a Grouped GEMM operation with fixed (N,K) dim but variable M as "varlen-M Grouped GEMM", while variable K but fixed M and N as "varlen-K Grouped GEMM".
token router and multiple smaller and often equal-sized
2.2 MOE COMPUTATION
Arithmetic intensity, defined as the ratio of FLOPS over the number of transferred bytes (IO), is a metric to quantify whether kernel is memory memory-bound -bound (kernel runtime dominated by memory I0 cost) or compute-bound (kernel runtime dominated by compute throughput).
The standard MoE computation for an expert e with SwiGLU activation can be broken down into the following components:
H. wppprooectioonXX.2. = W1,e
.R + x2n T22
A. = -..GGLU(H.)
T. + Toxn
(2)
Y. = dow---oojjcccion(( A. :IR T TT xd
Here, the up- projeeecttoon uses 4T,nd FLOPs and 4nd 4 HBM memory transfer bytes (we ignore the writes for H. here). Similarly, drooecrootttion uses nd FLOPs with 2 2T,d bytes. Assuming p= the MoE activation ratio, G = d as the granularity and uniform routing i.e
3

---

Published as a conference paper at ICLR 2026
T. =Tp, the arithmetic intensity (ignoring the writes for He) for the forward pass of an expert is
4T 2Tand
3
3
(4)
++
2+20 + Tp
d
For specific model size (constant d), it be seen that increasing granularity (increasing G) or increasing sparsity (decreasing p) leads to a decreasing arithmetic intensity. This is caused by the linear scaling of IO cost w.r.t expert granularity. Therefore for the case of fine ane-grained MoEs (high G becomes increasingly important to address the increased IO cost by maximally reducing IO access and hiding IO latency. We examine memory-efficiet MoE kernel design in Section 3 and discuss techniques to reduce I0 access and latency in Section
2.3 MoE ROUTING METHODS
In MoE, routing determines which experts to activate for each token. Token choice (TC) routing where each token independently selects the activated expert is often the default routing method for MoE models (Shazeer et al. (2017). We often have iop-K TC routing where the routing decision for token t Topkeej (St,e,K K) and St.e is the expert score for token f. Besides TC routing, expert choice (EC) routing is developed avoid load imbalance for expert parallelism (Zhou et al. 2022) by letting experts choose the tokens. However, EC routing is not directly usable for inference because it is incompatible with autoregressive decoding, and switching back to TC at inference time leads to a mismatch. In addition, EC breaks causality by future token information leakage (Wang et al. (022)) To address the inference issue of EC routing, Raposo et al. (024) introduce an auxiliary loss to promote the agreement between TC and EC routing results, or train an auxiliary router to explicitly predict the routing result of EC router and use this auxiliary router during inference.
In this paper, we propose novel Grouped GEMM tile-aware token rounding method that rounds the number of received tokens per expert (" ("expert frequency")) to nearby multiples of Grouped GEMM tile sizes and alters most one tile of tokens per expert. This approach effectively reduces wasted FLOPss caused by Grouped GEMM padding during sparse MoE training while preserving inference quality of trained MoE models. There are similar works that propose to drop and reroute tokens, including Rectify- Rootter (Zeng et al. 22024), but they do not focus on the tile structure of Grouped GEMM. Other works such as TMA-adaptive FP8 Grouped GEMM (Fu al.) 2025) focus on reducing padding-related load traffic but the FLOPSs wasted by non- non-aligned tile size in GEMM computation is not addressed.
3 MEMORY- MOE ALGORITHM
The MoE computation in SonicMoE launches 8 kernels: during forward, we have up-proj (A), down- proj (Y), and expert aggregation (O) kernels; during backward, we have activation gradient kernels for dHl (down-proj), dXX (up-proj), d.X (aggregating d.X), and weight gradient kernels dW, and dW2, The overall computation is shown in Algorithm
5 and the workflow is illustrated in Figure
Activation memory efficiency We now focus on the key olgorithmicohjeetive: minimizing activation memory during training. The FLOPs of MoE forward and backward computation is (6+ 12)TnKd d. For given T ,d, we need to keep nK constant for constant FLOPs. Therefore, increasing granularity requires decreasing n and proportionally increasing K. Hence, any activations with memory O Kd) should not be cached for backward computation avoid activation memory scaling with granularity. For current MoE kernels like ScatterMoE, activations scale linearly with expert granularity.
Activations Y (down-pro) proj output) and Xe (gathered X) have size TKd and avoiding caching them eliminates activation memory dependency on granularity. We avoid writing dY and dO, (gathered d0) to HBM as they increase the peak activation memory during the backward computation:
For X and dO, fusing the gather operation with the HBM load eliminates the need for materialization and activation caching in HBM.
We identify an alternative computation path that does not involve using Y and dY to compute dS and d.H. without increasing FLOPs, illustrated in Appendix D This eliminates the need for caching Y.
1 'eere we refer to "fine grained MoE" as MoE with large granularity ie.. is larger than n. We assume the setting of both iso FLOPss and isn-params.
4

---

Published as a conference paper at ICLR 2026
Algorithm 2 SonicMoF's MoF. kernel forward pass. Variables stored in HBM are colored blue. load and store means load from / store into HBM respectively.
Algorithm SonicMoE's MoE kernel backward pass of down projection.
Inpot :8, x; W2. 40.
Output:dlf, 2W2.d5.
Input :X, S. Ws. and Wa, same as Algorithm
Down-proj act dH kernel (dO, Wa. => (dM,, 4'):
Output :Moli layer outpet O
If Gather + varlen M Grouped GEMM + dSwiGLU d5
Up-peoj kernel (X, - (H:
# Appendii pelabraates this algorithm in more detail
If Gathher varien-M Grouped GEMM SWIGLU
Parallel for e de
Paralled for The (Xe, Wi,er miee)
do
d0., Wa,e: St ai, londddOO,,, S2,er..u..0
Xe, Wi.ei The
d0, + Gather(dOO s,,e)
X, -- Gathar(X,,..)
#f d.4' temp variable for computing d.A. dS, and 1/ dA,, Texn
H. Wi,
d.A. dd0, W2,+
Ag (H.
s. - Gather (.,e)
He As (He.
dA. + (roadcast((e)) (se))ttss)))))))))) dA
Dowj-proj kemel W::++Y:
Ae. dHz + -on.dddfcc.ffddd.e. He)
Rf varten- M Grouped GEMM
A, - -rooddaast(s))) An
11 Laput For dW2
Parallel for [B]
dS,,t + (.
// reduce over n dim
A.. 22,a load(Ae. W2,e)
dH,, 48, Ai adoo 48, A.)
Ye= Wa.e
stor
Down-proj weight dW2 kemel (4O, A', x) 2
Expert aparegation Kernel (Y. +)
I Gather + varlen-Ke K Grouped GEMM
7/ Gather and sum
Parallel for e [EE] do
Puralle for EE
d0., A'e. Ti.c load(do,, T.,e)
Yo, St, e +- load( 71,))
d0, Gat T..e)
O₂ [e
dW2,, - A. do_
(tore(
dW2,c store(dW2..)
As result, only cache X, H, and routing metadata for total size about 2T' 44'K bytes per layer. This activation memory usage the minimum required for backward computation without doing activation recomputation with GEMM, and is independent of expert granularity.
4 I0-AWARE KERNEL DESIGN
Ficvard pou with (T.d. (24076, 1606 256 (28.8)
(T. 224576, 256, 128.8)
The
Seald
Deep
(a) Runtime breakdown of 7B MoE training on H100 GPUs.
Forneed with DT A.A) (327788 2018, 1024, 64.8)
Dackored with (T. B,A) (3276O 2018 1084, (4.8)
30
(b) Runtime breakdown of OLMoE-sized 7B MMennighoff al 2025 MoE training B300 GPUs.
Figure 2: Runtime breakdown of different MoE kernels (ms 1) on H100 and B300 GPUs. We annotate the model memory bandwidth (TB/s t) for memory-boune bound kernels and compute throughput (TFLOPS 1, abbr TF/s in the figure) for grouped GEMM kernels. We use the GroupedMLP for Megatron, and ParalleerrollesMM... for McgaBlocks. DeepGEMM does not provide an efficient router implementation, gather and expert aggregation kernels during the forward pass, where optimized kernels ("DecpGEMM++") for them.
standard PyTorch inplementation ("DeepGEMM-pr)
5

---

Published as a conference paper at ICLR 2026
The expressivity of fine-grained MoE comes from the diversity of every token's 's expert selection, which in turn leads to linearly-scaled IO cost w.r.t. expert gramularity. To sustain high throughput, we need to maximally (1) reduce I0 access via fusion and (2) overlap the IO latency with compute. We first examine the token gather fusion with computation, and math and IO fusion with epilogue in Section 4.1
a.d respectively. We then describe the techniques to overlap MMA with 10 in Section 4.3 In
Appendix B we compare SonicMoE with other MoE kernel designs with summary in Table
4.1
GATHER FUSION WITH HBM LOAD
SonicMoE builds on efficient varlen-M and varlen-K Grouped GEMM kernels with two key forms of fusion: gather fusion during input loading from global memory (GMEM, often the HBM) to shared memory (SMEM), and epilogue fusion for elementwise operations. Both reduce unnecessary memory traffic and create additional opportunities to overlap IO with MMA. In particular, gather fusion is critical because many baseline implementations still require separate gather stage before GEMM, which introduces substantial IO overhead in fine-grained MoE training as illustrated in Figure 2a
4.2 EPILOGUE FUSION
We exploit the epilogue to reduce unnecessary accesses with the following design choices:
We fuse the SwiGLU U and backward of SwiGLU with the epilogue of forward up-proj and backward down-proj activation gradient kernel respectively (Costin al. 2025).
SwiGLU and dSwiGLU fusion:
Computing dH and d5 in backward down-proj activation gradient (dF) kernel's epilogue: in Appendix D.I] we show that SonicMoE's dS (dA, A') is the computationally and activation memory-efficient choice for fine-grained MoEs. Our d/H kernel produces the same output with far less total time than ScatterMoE's (lown-pro) act, dS, and dSwiGLU combined together in Figure 2a and 2b
4.3
GEMM MMA OVERLAPPING WITH ASYNCHRONOUS IO
Hopper GPUs. On NVIDIA Hopper GPUs, GEMM is performed asynchronously with producer- consumer paradigm (Shah et al. 2024). We can overlap the I0 I warpgroup with GEMM of another warpgroup with a smaller tile size. Once this is finished, we switch the roles of the warpgroups (effectively interleaving 10 and GEMM). This is often referred to as Ping-Pong scheduling (Wright & Hoque 2024b (Shah et al.) 2024) on Hopper GPUs in Figure Ping-Pong scheduling is particularly useful to maintain high Tensor Core throughput with heavy epilogue. For example, in the dH kernel's epilogue, we need load H and execute multiple activation and reduction operations. Besides Ping- Pong scheduling, SonicMoE also relies on asynchronous TMA operations to perform asynchronous 10 operations to overlap with the MMA operations:
Asynchronous TMA load during dHf kernel's epilogue: In the d'd kernel's 'ss epilogue, we create a dedicated pipeline to load H asynchronously, as illustrated in Figure 9
Asynchronous TMA store in forward Y and backward dXX kernel: SonicMoE does not fuse the scatter with HBM store for both kernels. This is primarily because the scatter fusion requires synchronous store instruction on Hopper GPUs.
Blackwell GPUs On NVIDIA Blackwell GPUs, GEMM kernels use the same "Ping-Pong" schedul- ing in spirit, but the implementation differs from Hopper. Blackwell introduces Tensor Memory (TMEM), a dedicated on-chip memory per SM (NVIDIA) 2025b). The MMA accumulator results are stored directly in TMEM which enables two-stage accumulator pipeline. While one stage performs MMA operations, the other stage executes the epilogue. The new Blackwell's MMA instruction, UMMA, is a threaded asynchronous operation that enables epilogue warps read and process results from one TMEM stage concurrently with MMA warps accumulating into the other stage.
TOKEN ROUNDING ROUTING
In this section, we analyze the hardware efficiency under sparse MoE training regime and identify that as MoEs become sparser, the wasted compute on padded GEMM tiles accumulate to nontrivial

---

Published as a conference paper at ICLR 2026
amount, known as "tile quantization" effects. In response, we propose novel routing method "token rounding" to eliminate tile quantization effects.
5.1 TRAINING EFFICIENCY OF SPARSE MOE
Wieled FLOPs (86)
Algorithm 4 Token rounding routing
FLLOR (50)
Tiden Clean
Input
number of experts E and expected activated number of experts
R 15
K per token; tile size Miste router scores S E [o, E round and spansify that determines rounding or down. -onn-dunded routed token list and scores (x (E(R
Output
12x
(L) Top-K token choice sorting (Supx, +- TopK(S, K)
Narchem esperts
Figure 3: Wasted FLOOPs by padding during MoEE forward & backward pass with T 16k, = 1k, n 1024, K = as illustrated in the bottom right 2 subfigures of Figure 7
(2) Calculate each expert's received taken frequencies and its Afttle. rounded multiples
(Je.Tal Meno--felMMun Myne)
(eCAARt (.-lSe/M/aaaa 1- Malee Maaaiie-aa/Mae M -Mettt)
EMde
(3) Build Top-K --pefeered S' for expert-wise ranking
I ensure non lop. K entries are smaller
S + S- Sc-1
for (EETTT] & [K] in parallel do
S Sy. (1,8) S.
144
(4) Token rounding per expert
for [F]
Figure 4: A demonstration tile quantization ef- fect for sparse MoE. The rounding subroutine in TR makes binary decision for discarding or padding tokens guarantee that each expert receives an Mull-multiple number of tokens.
// token onlering and sorted scores
Ter + sort(("))
It token rounding for expert
T. + and sww.realttttt.t........... -al.://-ee-...aaa--l-llaaaaaaa Mane)
Besides granularity, the arithmetic intensity of MoE also depends on the MoE activation ratio p as shown in Equation[4 When we scale down p. the expected number of received tokens per expert will also linearly decrease and the GEMM computation shifts towards a memory- -ooud regime.
Tile quantization effect. GEMM on modern GPUs is often computed in tiles (NVIDIA 2022) and always need pad the next tile-sized multiples if any dimensions of M,N. K are not fully divisible by tile sizes. Once the size of input (e.g. token dimension per expert) is small, the wasted TFLOPS by padding can be nontrivial. Therefore, we propose to token rounding to avoid launching such extra tiles, thereby leading to more efficient training. We also empirically show that our token rounding method does not affect model quality while achieving much higher training throughput.
5.2 TOKEN ROUNDING ROUTING
As such, we introduce token rounding (TR), a two-step sorting procedure shown in Algorithm 4 TR first computes the vanilla token-choice (TC) routing result and then sorts cach expert's 's assigned tokens by router score, similar to the EC sorting step. It then either discards some TC-selected tokens or pads additional tokens from the second sorting step. Between these two steps, we process the routing weight matrix so that TC-selected tokens are always preferred EC-selected tokens, ensuring that any discarding or padding affects only the last input tile of each expert.
TR relies on a round and sparsify subroutine to decide between discarding and padding. By default, we round each expert frequency to the nearest multiple of Mrile: we pad EC-selected tokens when [fe Ma.e - fe < Sc [fe] M.a.- T that TR is robust to the choice of rounding subroutine, and that this simple nearest-rounding rule already yields strong task performance; additional discussion of alternative rounding rules is provided in Appendix [E.2
TR guarantees that, for each expert, the maximum deviation from TC routing is at most one tile. Despite this small perturbation, it remains robust in sparse MoE training, serves as an effective substitute for TC (Table 1), and is generally stable when T /Maee 2 2 (Tables 7 [ At the same time. by eliminating dile-quantization waste, TR consistently improves training throughput over vanilla TC in highly sparse regimes, achieving up to 16% higher TFLOPS as E. increases while K remains fixed (Section 6.3.2
7

---

Published as a conference paper at ICLR 2026
1.4 Keears Thep
Mexer 360
200 Seean Thep 2m7
130 Merwey Time (T 20W 1190
Figure 5: Peak activation memory usage per layer across different model scales (1.4B-120B. MegaBlocks does not support small n. The henchmark configurations are listed in Table We only cache X, gathered X., Hr for each expert e and routing metadata which is the minimum required for backward computation without GEMM recomputation for DeepGEMM pt and DeepGEMM++. Their definitions are the same as Figure 2a
EXPERIMENTS
We now evaluate the three main components of SonicMoE. First, we compare activation memory usage against prior MoE implementations in Section 6. Next, we evaluate training throughput in Section Finally, we show the token rounding routing preserves downstream quality in Section 6.3. and improves training efficiency in sparse MoE training configurations in Section 6.3.2
E-M
A--2
150
106
(a) Forward & backward TFLOPS for different MoE kernels H100 GPUs
44B 1503722 . 20b) Bevnet
248 1000 285 Babnnii
BUGB 1530) Fermed
78 smaad 1981 Bedkned
32088 40E3- Badmmed
1200
1200 Reruod
1200 32708. 4080) Backrund
2314
(b) Forward & backward TFLOPS for different MoE kerness on B300 GPUs.
Figure 6: Forward & backward TFLOPS for different MoE kemelss H100 and B300 GPUs.
6.1 SONICMOE'S SS ACTIVATION MEMORY
We demonstrate that the peak activation memory for SonicMoE has the lowest activation memory footprint for single MoE layer as shown in Figure across all scales. For the 7B model with n5256, our approach reduces memory usage by 45% compared to ScatterMoE, and more significantly compared to MoMoE. For 30B and 120B models, the gap becomes our method saves more than 3GiB memory per layer compared to MoMoE. We also validate that SonicMoE's activation memory stays constant w.r.t. expert granularity as shown in Figure
wider: at 120B scale,
8

---

Published as a conference paper at ICLR 2026
6.2 SONICMOE'S TRAINING THROUGHPUT
reporss the compute throughput of forward and backward passes of one MoE layer on H100 GPUs. Across all model scales, our method consistently achieves the highest TFLOPS. In 1.4B and 7B settings, our approach improves TFLOPS by 40% compared to ScatterMoB and MoMoE. This throughput gap becomes wider for 30B and 120B when other baselines either fail to support certain sizes (MegaBlocks) or suffer from significant performance degradation (MoMoE).
We further measure the real training throughput of 7B MoE model with FSDP-2 on H100 GPUs: SonicMoeE on 64 H100s gets 213 billion tokens per day which achieves similar throughput as So on 96 H100s with 225 billion tokens per day. The throughput for this is measured using the Im-engine codebase (Mishra 2024). We shard the model using ZeRO-3 within single node (8x H100s) and replicate this sharded unit across nodes for this experiment.
Besides H100 GPUs, we also measure SonicMoF's compute throughput B300 GPUs in Figure 6b We mainly compare with DeepGEMM++ which powered by DeepGEMM SMI00 BF16 grouped GEMM kernels. SonicMoE still demonstrates overall speedup over DecpGEMM++, and such speedup is more pronounced when we increase expert granularity similar to the trend on H100 GPUs. For example, when increase the expert granularity d/n from 2 to 8 for 120B MoE training, SonicMoE achieves a greater relative speedup over DeepGEMM++ from 20.9% and 22. forward and backward to 35.2% and 30.9% respectively.
6.3 TOKEN ROUNDING
6.3.1 TOKEN ROUNDING' 'S GENERAL TASK EVALUATION
Table I: Comparison of different routing methods' task evaluation. "Train" and "Val" refer to the perplexity at the last 100 steps of training and at the validation set respectively. The next 11 columns are downstream tasks evaluated at the end of training and report the accuracy for each. "Avg" the downstream tasks. We use TC top- routing for TR. EC, and token dropping baselines when evaluating validation perplexity and task performance. Te means the average number of received tokens in cach microbatch per expert.
accuracy
these 11
(a) 0.5B params, 40B tokens, 2/64 activated (T, 512, Mule 128)
TR
16.22
15.92
51.4
41.6
78.4
65.4
31.6
38.1
65.0
31.0
61.1
57.4
29.1
50.0
to
16.34
15.94
51.0
41.9
78.5
64.8
33.0
38.1
67.0
30.8
54.7
55.8
30.1
49.6
TC (token drop)
16.44
16.10
51.1
41.4
78.7
64.9
31.6
38.0
62.0
61.9
58,9
30.8
50.2
FC
16.83
18.61
49.6
41.4
79.1
64.4
33.4
36.9
62.0
32.8
60.2
55.8
29.1
49.5
EC (sux couter)
16.80
21 21.80
50.0
40.9
75.2
63.7
28.2
35.2
61.0
31.5
57.2
53.3
24.7
47.4
EC router)
64.9
27.4
(b) 1.8B params, 40B tokens, 8/256 activated (T. 512, Mane 128)
TR
13.34
13.10
53.4
42.1
81.7
69.6
35.2
45.3
70.0
33.2
61.4
63.0
33.4
53.5
totop-K K
13.51
13.12
50.1
42,9
81.3
69.8
33.8
45.2
71.0
34.1
56.7
64.6
31.1
52.8
TC (token drop)
13.62
13.19
55.4
41.6
82.2
68.6
34.8
45.0
69.0
34.0
54.4
63.5
31.4
52.7
FC
65.8
39.8
28.4
49.9
BC (sux couter)
14.94
18.01
50.6
41.8
79.8
65.8
31.6
39,3
62.0
31.8
59.7
55.8
29.8
49.8
EC (1I TC rCuuter)
14.81
15.01
52.7
41.1
79.6
66.9
30.6
40.2
66.0
31.9
60.5
57.2
30.8
50.7
(c) 1.4B params, 100B tokens, 2/128 activated (T' 512, Mile 128)
TR
13.31
13.22
52.8
41.8
80.8
68.7
33.0
43.4
67.0
33.6
60.2
60.7
29.8
52.0
TC top- K
13.50
13.32
51.3
68.2
34.0
43.4
57.9
61.6
29.4
52.0
TC (token drop)
13.35
13.29
50.0
42.2
81.7
68.3
31.2
43.3
66.0
34.3
56.6
59.5
30.8
51.3
FC
14.08
24.79
51.5
41.7
81.0
66.
33.2
40.6
64.0
34.0
56.3
56.5
27.4
50.2
EC couter)
37.52
57.5
27.6
33.2
58.8
45.2
24.2
45.3
FC (ft muiter)
14.24
14.75
52.2
42.6
79.4
65.7
32.8
40.8
64.0
34.9
58.3
57.2
27.1
50.5
In this section, we assess the quality of trained MoEs using our token rounding algorithm. We use TR for training and during evaluation switch to token-choice top-K routing. This assesses whether TR can be replaced by TC after training. We use the OLMoE codebase and construct MoE models with OLMoB base architecture (Muennighoff al. 2025). We use deduplicated version of Fine Web-Edu (Ben Allal et al. 2024) for training all our models. More details are included in Appendix G
In Table T consistently use softmax renormalization for TR. We compare TR token-choice (TC) top- K routing and expert-choice (EC) routing (Zhou et al. 2022). We also consider MoD's

---

Published as a conference paper at ICLR 2026
approach (Raposo et al. 2024) that trains an auxiliary router to predict the EC router's selection during inference. This baseline is referred to as "EC (aux router)" in each subtable in [ We also adapt EC routing TC routing by finetuning learned TC top-K router and compare its task performance against TR's ' task performance without any adaptation. This is the "EC (ft TC router)" baseline in consider token dropping baseline in which set the capacity each expert as the largest multiple of Mrlle not exceeding its token frequencies and we discard the tokens with lowest scores. This the "TC (token drop)" baseline.
T Finally,
Forward
Bodooad
Recard
Backward
Tobra pradit
dilte tiput
15
15
expests
roperts
Figure 7: Forward & backward model TFLOPS for SonicMoE MoE kernels with different routing methods on H100 GPUs. We compare equipped with "nearest rounding to Maie-multiples via expert frequeney" subroutine against TC top- K routing. Configuration details are in Appendix
TR's train-test gap. Across all sparse MoE configurations in Table [ we consistently observe similar model quality between TRR TC. In fact, TR achieves slightly lower validation perplexity and higher average accuracy under the extremely sparse MoF settings for the and Tc There is noticeable discrepancy between EC and TC as the train and val PPL for EC can have 23 gap for Ib and compared to TC and TR's ss usual S 0.3 gap. TC finetuning is more effective than the auxiliary router to close this gap, but TR's 's task evaluation is still always better. In this case, TR can serve in-place substitute for TC during training.
6.3.2 TOKEN ROUNDING' 'S TRAINING THROUGHPUT
In Figure 7 we benchmark the token rounding's 's MoE main kernel runtime (without router) against top-K token choice routing. We focus on the iso-FLOPs setting by keeping T, n and K constant. We linearly increase the number of experts E while keeping K constant to increase MoE sparsity. As we linearly increase E, we observe drop in TFLOPS for token-choice routing. This is due to the (1) tile quantization effect as the wasted FLOPs spent on padding roughly linearly increases with the MoE sparsity shown in Figure and (2) the linearly increased IO due to more expert weights. We observe a drop in both TC and TR TFLOPS as we increase E, but the drop is more pronounced for TC.
For the 3rd and 4th columns in the top row in Figure an MoE model with 128 experts (K/E=1/64) and n = 1 with token rounding routing achieves 16.5% model TFLOPS improvement in forward and 6.1% in backward resulting in an end-to-end improvement of 9.4%. In general, we observe that as we move to larger intermediate sizes (more compute-bound) and higher MoE sparsity, the gap between TR and TC top-K becomes larger.
CONCLUSION
We present SonicMoE, a co-design solution that jointly optimizes MoE architecture and GPU kernels to address the training challenges of granular and sparse MoEs. Our contributions include: (1) memory-efficiest algorithm that minimizes activation size as MoEs become more fine-grained, (2) GPU kernels that overlap 10 with computation for throughput improvement, and (3) tile-aware token rounding that yields additional speedup without quality loss. Future directions include extending to low-precision and microscaling formats (FP8, MXFP8, MXFP4) for further memory savings, and overlapping communication with computation in distributed settings like expert parallelism. We envision future model architecture designs that optimize for quality per compute hour rather than just quality per FLOP jointly considering algorithmic and hardware efficiency.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENT
We gratefully acknowledge the support of Schmidt Sciences A12050 fellowship. the Google ML and Systems Junior Faculty Awards, and the Google Research Scholar program. We also thank the Princeton Language Intelligence program for the computing resources support. We thank Shawn Tan for his generous support on our experiments. We also thank Songlin Yang. Yilong Zhao, Bharat Runwal, Xinyu Yang, Andrew Sheinberg, Lijie Yang, Yongye Zhu, Zhuoqing Song and numerous anonymous reviewers for providing valuable feedback.
REFERENCES
Loubna Ben Allal, Anton Lozhkov, Guilherme Penedo, Thomas Wolf, and Leandro von Werra. Smollm-corpus, 2024.
HuggingFFccTB/ cotp:/cce.cooo/ttteeeeeeeeeeeeee /ttp:///o.oo.....a/.////oo//ooao
smollm 1m morppus
Vincent-Pierre Berges, Barlas Oğuz, Daniel Haziza, Wen-tau Yih, Luke Zettlemoyer, and Gargi Ghosh. Memory layers at scale. arXiv preprint arXiv:2412.09764,22 2024.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language. In Thirty-Fourth AAAI Conference on Artificial Intelligence, 2020.
Aidan Clark, Diego Las Casas, Aurelia Guy, Arthur Mensch, Michela Paganini, Jordan Hoffmann, Bogdan Damoc, Blake Hechtman, Trevor Cai, Sebastian Borgeaud, et al. Unified scaling laws for routed language models. In International conference on machine learning, pp. 4057-4086. PMLR, 2022.
Christopher Clark, Kenton Lee, Ming-Wes Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 2924-2936. 2019.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. 10X1::1803.05455v1.,000
A Feder Cooper, Wentao Guo, Duc Khiem Pham, Tiancheng Yuan, Charlie Ruan, Yucheng Lu, and Christopher M De Sa. Coordinating distributed example orders for provably accelerated training. Advances in Neural Information Processing Systems, 36:61198-56666 2023.
NVIDIA Corporation, Sep 2025. URL httpp:://oocs..ooaaaaaaaaoootttt Bocs/cpppcuttaase 3x httpr//sss___aalaaaeeeeroocoonna
Bobby Costin, Timor Averbuch, Dhruv Pai, Nathan Chen, and Ben Keigwin, Momoe: Memory opti- mized mixture of experts. Tilde Research Blog, 7 2025. URL
com/ homm/boog/momoe) Blog post.
DeepSeek-AI. Deepseck-v3.2-exp. Boosting long-context teffi-cc with deepseek sparse attention, 2025.
DeepSeek-AI, Aixin Liu, Bei Feng, Bing Xue, and et al. Deepseek-v? v33 technical report, 2024. URL https://arxiv 1ttp://wwwvv.or///////2/2122.1
abs
Raaz Dwivedi and Lester Mackey. Kernel thinning. Journal of Machine Learning Research, 25(152): 1-77, 2024.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, 2022. URL
03961
11

---

Published as a conference paper at ICLR 2026
Rong Fu, Weihan Cao, Jianfei Gao, Minxi Jin, Hui Wang, et al. Tma-adaptive fp8 grouped gemm: Eliminating padding requirements in low-prre oww-riiiioon training and inference on hopper. In ES- ES-FoMo III: 3rd Workshop Efficient Systems for Foundation Models, 2025.
Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. Megablocks: Efficient sparse training with mixture-of-experts. Proceedings of Machine Learning and Systems,
5:2888 5:288-300 304, 2023.
IBM Granite. Granite 3.1 language models.
bttpp:://iithuu.cmm.
Srantte-3.--larguuue---moo
2024. GitHub repository.
Xu Owen He. Mixture of a million experts. arXiv preprint 2r..1.207.......15..... 2024.
Zihao Huang, Qiyang Min, Hongzhi Huang, Yutao Zeng, Defa Zhu, Ran Guo, al. Ultra-sparse memory network. In The Thirteenth International Conference on Learning Representations, 2025.
Matt Gardner Johannes Welbi, Nelson F. Liu. Crowdsourcing multiple choice science questions. 2017.
Team Kimi, Yifan Bai, Yiping Bao, Guanduo Chen, and al. Kimi k2: Open agentic intelligence,
2025. RRLLhttss://axxttt org/ abs 20534
Jakub Krajewski, Jan Ludziejewski, Kamil Adamczewski, Maciej Pióro, Michal Krutul, Szymon Antoniak, Kamil Ciebiera, Krystian Król, Tomasz Odrzygózdž, Piotr Sankowski, al. Scaling laws for fine-grained mixture of experts. arXiv preprint urXiv:2402.07877 2024.
Chuck I. Lawson, Richard J. Hanson, David R Kincaid, and Fred T. Krogh. Basic linear algebra subprograms for fortran usage. ACM Transactions on Mathematical Software (TOMS), 5(3):308- 323, 1979.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Representations, 2024.
Yucheng Lu, Wentao Guo, and Christopher De Sa. Grab: Finding provably better data permutations than random reshuffling. Advances in Neural Information Processing Systems, 35:8969- 100--9981111 2022.
Microsoft. Announcing the availability of phi-3.5 moe in azure ai studio and github. : htt.ae/ceh....aalaaa.a.oo-----aa ai foundry http:/auurreri-ttoounb-bbbbboooo blog//
d-github/ thub/
4256278 2024. Microsoft Tech Community.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can a suit of armor conduct electricity? dataset for open book question answering. In EMNLP, 2018.
Mayank Mishra. LM Engine: A Oppei-Oe Library for Pretraining and Finetuning, June 2024.
URL httptttpp::ttttte com/ 1bm/ 1m- engine]
Mistral. Mixtral of A high quality sparse mixture-of-experts. https://mistraal
hhwws/miixzzal--f-expp...
2024. Mistral News.
Niklas Muennighoff, Luca Soldaini, Dirk Groeneveld, Kyle Lo, Jacob Morrison, Sewon Min, Weijia Shi, Evan Pete Walsh, Oyvind Tafjord, Nathan Lambert, et al. Olmoe: Open mixture-of-experts language models. In The Thirteenth International Conference Learning Representattons, 2025.
NVIDIA. Nvidia h100 tensor core gpu architecture: Exceptional performance, scal- ability, and security for the data center. Whitepaper V1.01. 2022. URL h wwwp::/wwwwaaaeeeseeeeeeeeeeeeee 2022/5//gte22221tttepppp--ogpppp pdf
NVIDIA, March
cwm / wwtp::/oowooootooteeeeoeeeeoot htw.onttteuullaass////
Grace Hopper "Hopper" Architecture.
NVIDIA. halaAAaWiiiaaalllllssoooouonnttt 2025a.
NVIDIA. Nvidia blackwell awwhitcctureecceeeee
aracdweellrchhiceetureeeeeciii techhhii brief. Wwt.eeeeeieppppei/IDi/..0....... URL https
//resources. //eeaa/www............aaaaaaaaea
Blackwell Architecture.
12

---

Published as a conference paper at ICLR 2026
NVIDIA. CUTLASS: CUDA Templates for Linear Algebra Subroutines.
https://githu..cm
cutlas 2025c. Version 4.2.0, Accessed: 2025-09-19.
OpenAL gpt-oss-120b & gpt-oss-20b model card, 2025, URL
10925
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang. Zach Vito,o, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjie Bai, and Soumith Chintala. Pytorch: An imperative style. high-performance deep learning library, 2019. URL http:////aaa.....rrr//////////00
Team Qwen. Qwen3 technical report, 2025. URL
QwenLM Qwen3: Think deeper, act faster. 2025. Official Blog.
David Raposo, Sam Ritter, Blake Richards, Timothy Lillicrap, Peter Conway Humphreys, and Adam Santoro. -fxture-of-depths: Dynamically allocating compute in transformer-based language models. arXiv preprint arXiv:2404.0224 02258, 2024.
Melissa Roemmele, Cosmin Adrian Bejan, and Andrew S Gordon. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 2011 AAAI Spring Symposium Series, 2011.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 8732 8740, 2020.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions. In Conference Empirical Methods in Natural Language Processing, 2019.
Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention- 3: Fast and accurate attention with asynchrony and low-precision, 2024. URL htttt://tssxiii org abs 2407 08608
Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gate mixture-of-experts layer, 2017.
http:tpp::/aaxi...om org/ abs/ //bs/11 1701. 06538
Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv: /tt.:/9990.0..5 2019.
Alon Talmor, Jonathan Herzig, Nicholas Lourie, and Jonathan Berant. Commonsense@A:: A question answering challenge targeting commonsense knowledge. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 4149-4158, Minneapolis, Minnesota, June 2019, Association for Computational Linguisties. doi: 10. 1111156////N1911211 (tthtt//
www.laathhoooy..rgggggooo 1421
Shawn Tan, Yikang Shen, Rameswar Panda, and Aaron Courville. Scattered mixture-of-experrs implementation. arXiv preprint arXiv:2403.08245 2024.
Databricks The Mosaic Research Team. Introducing dbrx: A new wfw.te-offthheeattopenn Ilm. https:/
www. dat alatabbicks....... com wtoo//booe http//lll///// www.ooo/innoouuuuoo htt.iinttoriinnnnnnnnen dbrx new state- art- 11m 2024. Databricks Blog, March 27, 2024.
Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled
neural network computations. Proceedings of the 3rd ACM SIGPLAN International Workshop on
Machine Learning and Programming Languages, MAPL 2019, pp. 10-19, New York, NY, USA. 2019, Association for Computing Machinery. ISBN 978 9781450367196 doi: 10.1145/3315508
13

---

Published as a conference paper at ICLR 2026
Ashish Vaswani, Noam Shazeer. Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need, 2017. wtl.attaaa//aii..ooooo//
1706.03786
Lean Wang, Huazuo Gao, Chenggang Zhao, Xu Sun, and Damai Dai. Auxsliary-loss--ree load balancing Strategy for nixture-of-experts. arXiv preprini arXiv:2408., 15664, 2024.
Less Wright and Adnan Hoque. Deep dive cutlass ping-pong gemm kemne,, Nov 2024a. URL
Less Wright and Adnan Hoque. Deep dive on cutlass pingn gemm Novebbe 2024b. URL
http:://wwwg.gpaaa...ooooooaaaar cessed: 2025-09-21.
Ac-
Rowan Zellers. Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? Proceedings of the 57th Anmual Meeting of the Association for Computational Linguistics, 2019.
Aohan Zeng. Xin Lv, Qinkai Zheng, Zhenyu Hou, Bin Chen, Chengxing Xie, Cunxiang Wang, Da Yin, Hao Zeng. Jiajie Zhang, et al. Glm-4.5: Agentic, reasoning, and coding (are) foundation models. arXiv preprint 250xv25008.0677,, 2025.
Zhiyuan Zeng, Qipeng Guo, Zhaoye Fei, Zhangyue Yin, Yunhua Zhou, Linyang Li, Tianxiang Sun, Hang Yan, Dahua Lin, and Xipeng Qiu. Turn waste into worth: Rectifying top-k router of moe. In
EMNLP, 2024.
Chenggang Zhao, Chengqi Deng, Chong Ruan, Damai Dai, Huazuo Gao, Jiashi Li, Liyue Zhang, Panpan Huang, Shangyan Zhou, Shirong Ma. Wenfeng Liang, Ying He, Yuqing Wang, Yuxuan Liu, and Y.X. Wei. Insights into deepseek-v3: Scaling challenges and reflections on hardware for ai architectures. In Proceedings of the 52nd Annual International Symposium on Computer Architecture. ISCA '25, pp. 1731-1745, New York, NY, USA, 2025a. Association for Computing Machinery, ISBN 97984007 12616. 10. I0.1145/365033.37331
doi: 10.1145/3699055.33314|11
BRLhhtts:://dot org/
Chenggang Zhao, Liang Zhao, Jiashi Li, and Zhean Xu. Deepgemm: clean and efficient fp8 gemm kernels with fine-grained scaling, 2025b. URL https::/githu...c wwtp:/twwu...cc/epoeeeeeee
ai/
DeepGEMM
Chenggang Zhao, Shangyan Zhou, Liyue Zhang, Chengqi Deng, Zhean Xu, Yuxuan Liu, Kuai Yu, Jiashi Li, and Liang Zhao. Deepep: an efficient expert-paaralee communication library. https //github. cww http://eepeeeeeeeeeeeeee ai/eeepp 2025c.
Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Zhao, Andrew M Dai, Quoe V Le, James Laudon, al. Mixture-of-experts with choice routing. Advances in Neural Information Processing Systems, 35:7103 2022.
Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models, arXiv preprint arXiv:2202.08906, 2022..
14

---

Published as a conference paper at ICLR 2026
APPENDIX
we summarize the notation used throughout the paper. In Section we compare SonicMoE with existing open-source MoE kernel designs. In Section tables, figures, and the up-projection backward algorithm cited in the main paper. In Section D we derive the gradient computation used by SonicMoE, and in Section D.I] we explain why our choice of computing dS is both activation-memory memory-efficient and computationally efficient. In Section E we present additional experiments, including the effect of expert granularity and ablations of token rounding with respect to the rounding subroutine, microbatch size, and tile size. In Section [F] we provide the benchmark configurations used for the activation- memooy and throughput results. Finally, in Section G we include the model-training and evaluation details.
In Section A
collect the referenced
15

---

Published as a conference paper at ICLR 2026
A NOTATIONS
In 2 we describe the notations used in this paper.
Table 2: Notations and their explanations
Notations
Explanation
number of tokens microbstch
model embedding dimension (hidden size)
expert 's intermediate dimension
total number of experts
K
number of activated experts
p= K/E represents MoE activation ratio
To
T. = Bes
represents the expected number of received tokens in each micmobttc by expert
G
G d/n represents the MoB expert gramularity. Greater G means a MoE is more fine-grainee
M.N.K
Dimensions for GEMM CUTLASS, We define A BMKK B BeRKXM and C CERMMN for AB=C
Male: Nille Kule
tile size of M, N, K dimension for single GEMM tile
R.
tile quantization residue := T, mod Mate
X
X KexTxe input token embeddings for an MoE layer
W1
Wi R Ex weight of up projection
W22
W2 R xnxd, weight of down projection
E (0, T.X. baney matrix where m,,, represents if tnken is routed to expert
S
SER Kouterscores In practice, only materialize the sparsified S insteud of the full
I1 RT output of up projection
SERT .ouuput SwiGLU
YER PK x4, output down projection
O Tx output of expert aggregation, also output of the entire MoE layer
d0
d0 d0 activation gradient for O
dA''
d.'' d0 w,J TRTx= GEMM output of dO and Intermediate result for computing and ds
dA (roadenst(s) dA' e E x activation gradient for
dY
dY = Broadcast d0 e R xx activation gradient for Y. dY is not used in SonicMoE.
dS
d3 x activation gradient for S
A'
Broadcast(s) A intermediate result and input computing dW2
dff
d#f e R Xmn, activation gradient
dx
d. é R XXxd activation gradient for X before aggregation
dX
d. e 7xx activation gradient for X after aggregution
dW1 Exxx2 weight gradient for Wi
dW,
dW2 Exxxd, weight gradient for W2
kernel
forward up-proj kemel
Y kemee
forward down-proj kernel
O kemel
forward expert aggregation kernel where each token aggregates all routed expert's result as the final forward output
dH kernel
backward down-pro) proj activation gradient kernel
aWas kemel
backward down- proj weight gradient kemnel
d. kernel
backward up-pro activation gradient kemnel
JW, kemel
backward up-proj weight gradient kemee
d.X kernel
backward expert aggregation kemel where each token aggregates the routed expert's d.&
[ Musec Life] [ Mine
rondedd multiples of expert frequency fc. [ [fe/Manl] [fe/ Manel Man.
Mile [fel Maat
[S] A Life Tadino
Ue Mate is Mut. rounded multiples of expert frequency fer Lfd Mune {[fe Mane -fee Masec and [S] Mast is the score after rounding in Algorithagii
16

---

Published as a conference paper at ICLR 2026
B SONICMOE'S COMPARISON WITH EXISTING MOE KERNEL DESIGN
Existing efficient MoE kernels also frame MoE computation as a Grouped GEMM, but their ingredients are different from SonicMoE. Here we provide an overview (but not complete list) of key differences:
ScatterMooE (Tan et al. 2024 implements gather fusion for varlen-M Grouped GEMM but not for varlen-K Grouped GEMM. ScatterMoE also does not overlap MMA computation with memory I0. Morcover, ScatterMoE is also built on older versions of Triton where TMA is not supported. ScatterMoE computes dS as ds (dO,Y) which requires caching Y. This results in large I0 cost and activation memory requirement. ScatterMoE's both forward and backward pass have limited fusion and hence is much slower than SonicMoF, especially for backward computation.
MoMor (Costin et al. 2025 also implements the gather fusion for varlen-M but not varlen-K Grouped GEMM similar to ScatterMoF. Although fused with up-proj activation gradient, the ds computation still utilizes d$ = (dO,Y). Similar to ScatterMoE, MoMoE does not use TMA for IO.
MegaBlocks (Gale et al., (023) has multiple MoE implementations and we focus on ParalleeDrroplsssMLLI which is built on top of block-sparse matrix multiplication ParallelDroplessMLPPP first gathers and pads the tokens and then launches block-sparse GEMM for up and down-proj. Then, it launches scatter kernel before reducing across the expert results. These sparse matrix multiplications usually take a longer time than the boghii--ptimized Grouped GEMM, as shown in Figure 2a] and the gather and scatter kernel have a total IO cost of 87' d bytes which can be a bottleneck for fine-grained MoEs. We consider MegaBlocks's ' ParallelDroples.MML. MLP block-sparse GEMM baseline in our benchmark and find that MoE implemented via Grouped GEMMs often have higher training throughput than MoEs implemented via block-sparse GEMM.
Megatron-LMM (Shoeybi et al. 2019) also has multiple MoE implementations and we focus on GroupedMLF which uses Grouped GEMM from the CUTLASS library (Corporation 2005) with JIT epilogue fusion as the GEMM backend. Similar to DeePGEMM, GroupedMLP does not fuse gather with the HBM load. A recent memory-efficient patct To fuses S weighting with SwiGLU computation during forward, and during backward which allows the PyTorch autograd engine (Paszke) et al. 2019) to follow similar computational path as SonicMoE.
Megatron-LM also implements TEGroupedMLI which launches 4 CUDA streams to execute a list of GEMM (without wotguuuussllppckkeeee inputs, and without a persistent tile scheduler). In this case, each expert independently launches new GEMM kernel leading to "bubbles" " on the CUDA streams. This leads to underutilization of the GPU resources. We empirically find that TEG runs slower than GroupedMLP and so we use GroupedMLP across all benchmarks.
DeePGEMM (Zhao al. 2025b) designs Grouped GEMM kernel for contiguously-packee inputs. They also don't 't implement any other fusion for both SM90 (Hopper) and SM100 (Blackwell) BF16 Grouped GEMM. DeePGEMM specializes more on distributed training with expert parallelism (Zhao et al. (025c), and it is common to launch separate all2all kernel (Lepikhin et al. (024) which is then followed by a contiguous Grouped GEMM. Both DeePGEMM SM90 and SM100 BF16 kernels assume that each expert receives a multiple of Mate tokens as it does not implement the TMA tensor descriptor online update during the Grouped GEMM computation. DeepGEMMM's BF16 GEMM on SM90 also does not employ fing-Pong scheduling.
Additionally, ScatterMoB and MoMoE are both implemented in Triton (Tillet et al. 2019) for the ease of development at the expense of losing full programmability of the asynchronous compute and memory 10 of Hopper GPUs (NVIDIA 2022). 2022) For example, they cannot implement -rained
aLtps: (tlps://qo0 //gttbbb.b/ com/ http:///hhoooooooocccccc//////en htw.a//wwwwaaaaaa//////////aaaia
scat srattermoe/m1l wwt.:eemeeeepppp.m syLLL
ltttps://itthus www/tt//////... wtw.mm//ileeeeeeeeeeeeeeeee
mamoo mamoe/ncmoce wttp///ccmccp...)
aLLps://cithu/ com/ dal http:///tt.....ct///////aa////// htteer////ttiaaccccc////c//rccee
megabl
megsblocks/
www.kk/aacyeee//eeeeeeeea dmoe. by
attps://qithul com/ dat http:///hw....aaoooo././////////
megabble
megablocks/
bbllkkk/ayyer//ngappp L348
http:///wwa..aaaoooooo..///////a gomm
magatron/comm
17

---

Published as a conference paper at ICLR 2026
control of asynchronous load and store during the GEMM's epilogue. They also cannot overlap MMA with heavy epilogue operations using aig-Pong scheduling. becomes increasingly important to overlap IO operations in epilogue when the GEMM computations small in size (as in the case of fine-grained MoEs) to achieve high GPU utilization.
C REFERENCED TABLES, FIGURES, AND ALGORITHMS
In Table 3 we provide a trending overview for open-source frontier MoE models. SonicMoE's compu- tational workflow is illustrated in Figure and high-level comparison between SonicMoE and related efficient MoE kernel designs is shown in Table Sesurepprovideer an example of SonicMoE employ- ing Ping-Pong warpgroup scheduling on Hopper GPUs. We present SonicMoE's backward pass of up-proj algorithm in Algorithm
Table 3:
MoE Scaling Trends: Here, we show the activation ratio experts activated token K / total experts
E' and expert granularity shown model embedding dimension (d) expert intermediate size (n) for frontier open source models. We do include the shared experts for the MoE sparsity calculation. The trend indicates new open-source MoE models tend be more granular and sparser.
Model
Release date
Parameters
Expert activation (K/E)
Expert granularity (d/n)
Mixtral 8x22B Mistral
2024
11/23
131B
25.0%-
6144/16358 0.38
Mosaic Research Team
2024
03/24
(4/16)
6144/ |144//17522 0.57
Phi-3.5-
Microsoft
2024
0W/24
428
12.5% (2/16)
4096/6400 0.64
OLMoB
Muennigholi al
2025
0W/24
7B
12.5% (8/64)
2048/1024 2.00
Granite 3.1 1oM (rranite 2024
12/24
20.0% (8/40)
1536/512 = 3.00
D3epSock-V3
DeepSeck-AA
2024
12/24
671B
3.13% (8/256)
7168/20:5 3.50
Qwen3 MoE
QwonL.
04/25
6.25 % (8/128)
1096/1588 2.67
QWen3-30B-A3U (Owen
2025
05/25
30.5B
6.25 % (8/128)
2048/768 2.67
Kimi K2
Kimi al.
2025
07/25
1.4T
2.08 % (8/384)
7168/2048 3.50
gpt
20b OpenAi 2025
08/25
120B
3.13% (4/128)
2880 (280/2280-0 1.00
GLM-4.5-Airr Zeng 2025
08/25
6.25% (8/128)
4096
Owen3-Next-S80B-o
(Qwen 2025
09/25
1.5% (10/512)
2018/512
V..pSeck--3..--xp DeepSock-Al 2025
10/25
685B
3.13% (8/256)
7168/ (2088/2018-111 2
Features Methods
SonicMoE ScatterMolt
MoMol!
MegaBlocks Megatron DeepGEMM
Gather fused with GMEM SMEM (HIB) load (Sec
4.1
fwd V. bwd X fwdd V, bwdx
X
SwiGLU and dSwiGLU fused with epilogue (Sec
4.2
x
NA
dS computed as (dA', Ac.e) (See 42
X
NA
Epilogue that computes d.H. d5 together (Se
NA
Overlap MMA with cpilogue/IOO (See (Scc
X
Do not need a separate scatter kemee Efficient top- K sorting
X
NA
NA
X
X
X
NA
Do need shape alignment efforts outside GEMM kernels
X
Table 4: Comparison between SonicMoE prior MoE kernels. means that kernel implements the feature or a functionality similar in semantics, and X means the feature is missing from the kernel. "NA" means that the feature is out of the expected scope. We use the GroupedMLP for Megatron and ParalleDrroplesMMLD for MegaBlocks. More discussion is included in Appendix
18

---

Published as a conference paper at ICLR 2026
Forward pass
Up-proj
Down-pprg sewn-ppoo]
Expert aggregation
Grrpped GEUS
Act
larienal Groupel GBVM
Wy
Backward pass
Expert aggrogation
Up Up-proj act grad
Down-prej act grad
minune exged aup
Watees
ddf
dBi
4A'
do
4x
4X
Grupst SONt!
Up-proj weight grad
Down proj weight grad
JWV
Vertens
dkG
JWS
Vatteal
JWY
Figure Computational workflow of SonicMoF's s 8 launched kernels, grouped by yellow boxes. 3 5 kernels are launched during forward and backward computation respectively. The incoming arrows to yellow circle indicate variable loaded from HBM to SRAM, and an outgoing arrow represents a variable stored to HBM. We color the boxes of all variables HBM, with purple boxes indicating the oulput of forward and backward while blue boxes indicate intermediate variables or weights (W1, W2). We color all cached activations X. H, m, the backward pass.
red. Algorithm Aigorith@@foomnn formally describes SonicMoE's forward pass, and Algorithm
Time
Prodk soer
Get next sie
Thik
TVA
warggroup
Comsumar
Churter Leunch
Ger tife
Taraco Car U/AA
agun rase
wargrou
Figure 9: SonicMoB's Ling-Pong warpgroup scheduling Hopper GPUs. The green arrows indicate that consumer warpgroup signals the of the epilogue and the other consumer warpgroup can proceed with the MMA. Once this step is complete, the roles of 2 consumer warpgroups switched. SonicMoE mainly uses Ping-Pong for forward down-proj Y kernel and backward down-proj activation gradient dHf kernel as they both have heavy epilogue. In dff kernel, SonicMoE has asynchronous TMA load during epilogue, and producer warps need to issue cp.async for gathering dO and load expert weights with TMA. This figure adapted from Wright Hoque (2024a)'s blog on Fing-Pong scheduling.
Algorithm SonicMoE's MoE kernel backward pass of up-proj.
Input xx Wi,dH dH
Output::.3 :dlX. dW..
Up-proj d.X kernel (4H7, Wi)
Parallel for e [E] do dH.. W., lda,(dH, W1.))
dx. + dH. W,T.
dx_ sttoro(dX.))
Up-proj weight dW, kemee (X, dH, x) dW1
Parallel for do
X, r.... d.H, + load(XX (x, n.,. dH.)
X. Gather(X.,.. s.,e)
dH.
dW, ttorooWW
Expert aggregation kernel (d.,, ax::
Parallel for [7] do
d. क load(dX me,))
d.X, F d8,.
dX, store(ddX))
19

---

Published as a conference paper at ICLR 2026
D
GRADIENT COMPUTATION
For an expert e, let
X. ER W1. ER W
(5)
The forward activation computation is given by:
He 1, Tux2n
A (
2.
(6)
The token aggregation with scores S= {st,e) is given by
O₂
Σ St,c Yats
dO, E GRIxd the gathered results from dO.
We know
dY 8t, d0: = dY,- = -Broadcast d0c-
(8)
Define the Grouped GEMM output as :- W ER Then from Equation
dA.=dY, W. Broadcast d.'
(9)
The activation gradient for score d5
d - (dO,, Yo,) (dO, 22, (dA'' Acat)
(10)
In addition, we can derive dHf from d.A. and Ae (recomputed from He) as:
dH, dSwiGLLAAA H.).
(11)
Using Equation g
dWze- A dY. AT (Sroadcas((s.) -Drooaasss.......ssssaaasssssss. Ae) dO
(12)
D.1 COMPUTATIONAL CHOICES FOR dS
we do not implement custom kernels and rely solely on PyTorch's autograd (AD) engine, we can add the expert weighting (S) either (I) before down-proj forward (2) after down-proj forward. Both yield identical results for forward and backward, but the computation for dS is different, For (1), we need compute (dA's Ae.:) which is used by SonicMoE and Megatron, ScatterMoF, and MegaBlocks compute (dO,, Ye,e) as required in (2).
Note that dS can be computed any of dSs,. (dA' A...) (dO,, Y.,)). however computing it as (dAA e, Ac,t) computationally and activation memory-efficient choice due to the following reasons:
Additional HBM traffic (0 2T bytes): (dA', Ac.i) requires d.A',, and A already computed during the dH kernel, so we can avoid extra unnecessary loads.
Extra cached activation memory vs. 27 Kd bytes): One of the reasons why the cached activation memory for ScatterMoE, MoMoE and MegaBlocks fails to stay constant w.r.t. expert granularity the required caching of for computing d.S.
Parallel reduction rounds (log, (n) vs. log2(d)): (dA',,> Ae,)) reduces over n while (dO,, Ye,)) reduces over d. This difference saves least log, (d/ m) rounds of reduction.
20

---

Published as a conference paper at ICLR 2026
E MORE EXPERIMENTS
In this section, we investigate the qualitative improvement from fine-grained MoE in Section E.1 We also investigate the effect of rounding subroutines lound_and_spp snasssfyy in Algorithm and the effects from microbatch size 7 and tile size Miite for token rounding in Section E.3
E.1 EFFECT OF EXPERT GRANULARITY
Here we validate the effectiveness adopting fine- tne-grained MoE. We fix the MoE activation ratio K/E for the 0.5B and 1.4B model and we proportionally scale up K and E' while linearly decreasing n from row to row in Table 5bb
In general, we observe a better performance for 256 than 1024 which is also consistent with the MoE scaling trends mentioned in Table In Figure right subfigure, we find both SonicMoE and cuBLAS still sustain the throughput from = 1024 to 256 under iso-FLOPs, but starting from n = 256 FLOPs will drop linearly w.r.t. granularity. Therefore, we choose n 256 for all experiments in T
Table 5: Evaluation MoE w.r.t granularity with iso -so-FLOPs (nK is constant) and iso- params En constant) settings. "PPI." refers to validation perplexity at end training. "Avg" the mean accuracy across 11 downstream tasks. The "dense, iso-FL.OPs" refers to dense model with n as the intermediate size, while the "dense, iso-params" refers to dense model with nF the intermediate size.
(a) 0.5B params, 20B tokens, 8/64 activated
(E.K.n))
PPL
Wino
SIQA
SciQ
PIQA
OBQA
HS
COPA
CSQA
BoolQ
ArcEE
AreC
Avg
16, 1024
16.23
53.0
41.3
79.8
65.0
32.6
37.8
66.0
32.2
55.8
53.9
29.1
49.7
6,,8, 256
16.01
51.0
41.4
79.2
63.5
31.6
38.4
66.0
31.5
60.2
57.5
25.7
49.8
32, 64
16.13
51.2
41.5
78.9
65.3
34.2
38.4
63.0
32.4
60.6
59.5
28.1
Dense, iso-PLOPs
19.90
48.9
41.4
74,9
62.2
32.6
62.0
61.7
27.1
Dense, iso-params
15:46
52.1
41.5
78.9
65.3
34.0
39.2
69.0
32.2
58.5
59.3
28.8
50.8
(b) 1.4B params, 50B tokens, 8/128 activated
22,
13.38
81.7
33.6
44.3
33.5
61.1
29.8
13.32
81.5
45.3
56.6
512,32,6 64
13.50
52.5
41.2
82.9
68.9
34.4
44.7
69.0
33.6
58.7
62.6
30.1
52.6
Dense, iso FLOPSs
17.90
52.2
41.0
79.2
63.4
31.0
34.7
61.0
30,5
60.3
51.8
25.1
48.2
Dense, iso- coo-params
12.74
52.2
42.6
83.3
70.1
34.8
46.8
67.0
35.1
61.7
63.5
31.8
53.5
E.2 ABLATION STUDY DIFFERENT ROUNDING SUBROUTINES FOR TOKEN ROUNDING
We conduct ablation studies to study the effect of the different routing subroutines on the trained MoE by TR. We compare token rounding with nearest rounding ("NR") with per-expert token counts with other rounding methods. Specifically, we compare against stochastic rounding with per-expert token count ("SR"), always round up ("UP"), ""p"),amm and always round down ("DOWN"). The results are shown in Tabd we find that our token rounding algorithm in general is robust to the specific rounding subroutines.
Following Algorithmp 4 for expert we denote the expert frequency from the TC sorting as fex and the last Mite-divisibll expert frequency as [Se] Mase and the next Miile-dvisiill expert frequency as [fee Mane We also denote the expert scores from TC sorting sg, the expert scores from the selected tokens in Rel: ] as [ Ma. and the scores for Te [e Maa. as [e Mane We note that all rounding algorithms only make binary decision between discarding TC tokens and padding EC tokens for each expert. Following are simple heuristics to perform rounding:
NR-f: nearest rounding to Mistemulltipls via expert frequency: We pad EC tokens [fel M.. fa < Sa
[Se] Mun. "NR-f" is our default choice of token rounding and we use
for Table 7 and Figures and
21

---

Published as a conference paper at ICLR 2026
Algorithm 6 Balanced rounding to Mste via expert frequency ("Balance-F" f" in Table
This algorithm satisfies eax Man- Se M and Σ -. [fe] Menee Ee M
Input zf (fa)eett (ej a list of expert frequeney with tCtop-K routing, {[fe) Mesta Jej a list expert frequency with TC totp- multiple of Mill tokens, (Ife Matt hejjy
routing and potemtial padding ensure expert
expert
trequency with IC top- routing and with potential token dropping to ensure each experve multiple of Matta tokens; We should eusure (8) Mute [/=] Mate/ S Meste-
of [[fa]Mane ee(jy as list of expent frequency that ensures each expert receives multiple of Mate tokens // accumulator that ensures the preservation of total expert frequency 0;
for E [E)
I/ calculate the residual error of both rounding choice
r [fe [fel fo:
r. + If I ~f::
if I1 +2| < | +2| then
// choose to pad with tokens
M + i
Metai
2
else
// choose to discard TC tokens
Mel -LelMeter
2 2 t
SR-f: stochastic rounding to -ilt-multtipes via expert frequency: We sample from Sa
[So Mase
Mane
Bernoulli
distribution for deciding whether to pad EC tokens for expert e.
Mile
NR-s: nearest rounding to Mulle- multiples via expert scores: score:: We sample from the following distribution for deciding whether to pad EC tokens for expert e:
Se E, [8... Mule
Bernoulli
(13)
E.[ M Ee [8et]] Mate
Balance-f: balanced rounding M -multipes via expert frequency: The Balance
algorithm (Dwivedi & Mackey] 2024 Lu et al.) [022] Cooper et al. (203) can be adapted to ensure the total number of routed tokens to all experts after tile-rounding is preserved regardless of the number of experts E. Algorithm is such example that ensures
E
E
max eE(E)
[f S
Σ lffe Menle
Σ
S
Maue/2
(14)
where the other rounding subroutine will have an expected deviation of O M vE for
[
UP: always round up expert frequency as M...: We always pad EC tokens chosen in the second step of sorting in Algorithm 4 This gives model TFLOPS bower-bound for Figure
DOWN: always round down expert frequency as [fe] M..:: We always discard TC top- K tokens chosen in the first step of sorting in Algorithm This gives a model TELOPS
bound for Figure 7
Always discarding TC tokens ("DOWN"). "DOWN" a baseline in which we always drop the last TC tile if the expert frequency is not divissbbe This idea is similar to the idea of token dropping in expert parallelism where the expert will sort and drop the token with the lowest scores when receives too many tokens (Fedus et al. 2022). We note that "DOWN" produces the shortest MoE kernel runtime
22

---

Published as a conference paper at ICLR 2026
Table 6: Evaluation of token rounding algorithms equipped with different round_and_sg sparsiff subroutines in 4 "PPL" refers the validation perplexity at the end of training. "Avg" is the mean accuracy across the downstream tasks.
(a) 0.5B params, 40B tokens, 2/64 activated (T. = 512, Mill 128)
Method
PPL
Wino
SIQA
SeiQ
PIQA
OBQA
HS
COPA
CSQA
BoolQ
ArcE
ArcC
Avg
TR (NR-F)
15.92
65.4
31.6
65.0
61.1
57.4
TR (SR-F)
15.93
50.8
40.9
77.4
66.9
33.0
38.4
64.0
31.1
60.7
55.8
28.1
49.7
TR (NR-s)
15.91
51.3
40.9
80,3
65.4
30.8
37.7
67.0
31,0
61.6
55.4
28.4
50.0
TR (Balance I)
15.93
51.9
41.8
78.8
65.9
32.6
38.4
66.0
31.6
60.3
56.8
27.1
50.1
TR (UP)
15.89
50.5
40.9
78.6
64.5
32.2
38.2
68.0
29.9
55.2
54.2
30.1
49.3
TR (DOWN)
16.10
51.1
41.4
78.7
64.9
31.6
38.0
62.0
32.8
61.9
58.9
30.8
50.2
TC top-
15:94
51.0
41.9
78.5
64.8
33.0
38.1
67.0
30.8
54.7
55.8
30.1
49.6
(b) 1.8B params, 40B tokens, 8/256 activated (T. 512, Mrile 1228
TR (NR-D)
42.1
81.7
35.2
45.3
33.2
61.4
33.4
53.5
TR (SR-F)
13.08
52.7
41.6
82.6
69.4
34.4
45.6
70.0
33.0
59.1
62.5
34.8
53.2
TR (NR-s)
13.09
54.1
42.3
82.8
69.3
33.8
45.7
70.0
34.1
59.0
64,6
32.4
53.5
TR (Balance-F)
13.08
52.5
42.0
82.7
70.0
33.2
45.3
68.0
34.6
59.4
63.3
33.4
53.1
TR (UP)
13.07
50.4
41.7
81.4
68.4
37.2
45.4
69.0
31.9
51.7
62.2
33.4
52.1
TR (DOWN)
13.19
55.4
41.6
82.2
68.6
34.8
45.0
69.0
34.0
54.4
63.5
31.4
52.7
TC lop-K
13.12
50.1
42.9
81.3
69.8
33.8
45.2
71.0
34.1
56.7
64.6
31.1
52.8
for any rounding algorithm. However, in Table 6 we observe that "DOWN" yields much higher validation perplexity than "NR-f", "SR-f" and "NR-s". Although we can expect shorter MoE kernel runtime by always discarding TC tokens, such quality degradation might not be acceptable in practice.
Always padding EC tokens ("UP"). "UP" is a baseline in which we always pad extra EC tokens to the last TC tile if the expert frequency is not Mijteddiviiile. Contrary to "DOWN", "UP" produces the longest MoE kernel runtime for any rounding algorithm. In T we find that "UP" often produces lower validation perplexity, but the average downstream task accuracy is not necessarily higher than other rounding algorithms. Given the longer MoE kernel runtime but not necessarily better trained MoE quality, we do not recommend the usage of always rounding up. We speculate this is due the train-test gap between TC and EC routing and "UP" reinforces the bias towards EC more strongly than other TR algorithms.
For a balance between training efficiency and trained MoE quality, neither always discarding TC tokens nor padding EC tokens is the right solution. In Table we pick "NR-f" as the round anddaandss sparsify subroutine for TR's 's main experiments.
E.3 ABLATION STUDY ON THE EFFECTS OF MICROBATCH SIZE T AND TILE SIZE Mile
Effect of microbatch size Since the token rounding is applied on the microbatch level, the choice of microbatch size will result in different qualitative results for TR. Note that this also holds true for EC routing. For example, EC over sequence will result in different model quality than EC over a text segment. In Table we vary the microbatch size while keeping the minibatch size (consumed tokens per optimization step) constant.
We find that TR will preserve its trained MoE quality when Te/Mule 2,, but Tf (the last row in both subtables), there is noticeable quality degradation for both validation perplexity and downstream task performance. However, the trained MoE quality with T, Ma/Mile is still better than training with EC and finetuning with TC top-K K routing.
Effect of the tile quantization size Mile- Similarly in Table X we can find that TR is generally robust w.r.t. Mik when T. MBe 2, and when Te/ TaeMa = there is noticeable degradation but the overall result is still better than EC baseline.
23

---

Published as a conference paper at ICLR 2026
Table 7: Evaluation of token rounding algorithms when vary microbatch size T change average number of
tokens per expert (T.). ) For each trial, vary the microbatch size from 4 (T, =512) (T, 128) and keep Mile always kept as 128. "PPL" refers the validation perplexity at the end of
minibatch size constant. training. "Avg" is the mean accuracy across the downstream tasks.
(a) 0.5B params, 40B tokens, 2/64 activated (Mile 128)
Method
PPL.
Wino
SIQA
SeiQ
PIQA
OBQA
BS
COPA
CSQA
BoolQ
AreE
AreC|
Arg
(R. 1024)
31.0
TR 512)
41.6
78.4
31.6
TR (T. 256)
15.98
52.2
41.4
77.7
66.1
32.2
37.9
66.0
31.0
59.6
57.2
30.1
50.1
TR (T. 128)
16.11
51.7
41.7
77.9
66.1
30.8
37.7
67.0
31.9
61.2
54.7
29.1
50.0
topo-
15.94
51.0
41.9
78.5
33.0
38.1
67.0
30.8
54.7
55.8
30.1
49.6
FC TC
16.98
50.0
41.7
79.7
64.9
31.6
36.8
63.0
32.1
60.7
54.6
27.4
49.3
(b) 1.8B params, 40B tokens, 8/256 activated (Mile 128)
TR (T. 1024)
13.08
59.5
32.1
TR 13:10
53.4
42.1
81.7
35.2
70.0
61.4
33.4
53.5
TR (T, = =26)
13.12
51.9
41.2
81.8
69.7
33.6
45.2
73.0
34.2
56.9
63.2
34.1
53.2
TR 128)
13.55
51.9
41.5
82.0
69.2
32.8
44.4
69.0
34,4
59.8
64.0
30.4
52.7
topp-K
13.12
50.1
42.9
81.3
33.8
45.2
71.0
56.7
64.6
31.1
52.8
(C
15.01
52.7
41.1
79.6
66.9
30.6
40.2
66.0
31.9
60.5
57.2
30.8
50.7
Table 8: Evaluation of token rounding algorithms when we vary the size of tile Mule for token rounding. "PPL" refers to the validation perplexity at the end of training. "Avg" is the mean accuracy across the 11 downstream tasks.
(a) 0.5B params, 40B tokens, 2/64 activated (I. 551)
Method
PPI.
Wino
SIQA
SciQ
PIQA
OBQA
HS
COPA
CSQA
BoolQ
AreF:
AreC)
Avg
TR (Miae 64)
15.90
51.3
41.7
78.
65.6
31.4
37.9
67.0
32.4
59.8
57.2
28.8
50.1
TR (Mdite 128)
15.92
51.4
41.6
78.4
65.4
31.6
38.1
65.0
31.0
61.1
57.4
29.1
50.0
TR (Mate 256)
16.00
51.7
41.4
78.7
66.3
32.4
37.7
67.0
31.3
60.1
58.2
29.1
50.4
TR (Mille 512)
16.17
52.5
41.2
65.2
32.0
37.9
31.0
59.4
57.2
30.4
49.9
TC Cotp-K
15.94
51.0
41.9
78.5
64,8
33.0
38.1
67.0
30.8
54.7
55.8
30.1
49.6
EC (A mmtte)))
16.98
50.0
41.7
79.7
64.9
31.6
36.8
63.0
32.1
60.7
54.6
27.4
49.3
(b) 1.8B params, 40B tokens, 8/256 activated (T, 512)
TR (Mine
13,07
52.3
42.9
82.7
69.4
35.4
45.6
70.0
32.4
56.6
64.4
31.4
53.0
TR (Mile 128)
13.10
53.4
42.1
81.7
69.6
35.2
45.3
70.0
33.2
61.4
63.0
33.4
53.5
TR (Mule 256)
13.13
52.0
41.6
82.1
69.2
35.4
45,3
69.0
34,2
$8.0
65.6
32.1
53.1
TR (Mate 512)
13.56
53.0
41.8
81.2
68.4
34.0
44.2
68.0
33.3
58.1
59.5
30.1
52.0
TC lop-
13.12
50.1
42.9
81.3
69.8
33.8
45.2
71.0
34.1
56.7
64.6
31.1
52.8
EC (ft router)
15.01
52.7
41.1
79.6
66.9
30.6
40.2
66.0
31.9
60.5
57.2
30.8
50.7
F ACTIVATION MEMORY AND TRAINING THROUGHPUT BENCHMARK CONFIGURATIONS
The configurations of Figure 5 ba and are
are included in Table
The configurations for the subfigures Figure are listed below. Notice that we consistently use Male 128 when we benchmark the TR's speed.
Top-left 2 subfigures: We use (T,d,n,)) (Td32_.))=(63334533322 (16384,1536,,56.69 and we vary E from 64 to 512.
Top-right subfigures: We use (T,d,n,K) (1634,,1111,,,,,,3,1 and vary E from 16
to 128.
Bottom-left 2 subfigures: We use (T, (16384,4096,5522
(1GG4,))))(G84440009008333 and we vary from 64
to 512.
24

---

Published as a conference paper at ICLR 2026
Table 9: Benchmark configurations for activation memory footprint and compute throughput.
(a) Benchmark configurations used by Figure
and
6a
(b) Benchmark configurations used by Figure 6b
Model Size
T
d
K
Model Size
K
40960
768
256
128
8
131072
768
256
128
8
1.48
40960
768
1.48
131072
768
40960
768
1024
32
2
131072
768
1024
32
2
24576
1536
256
128
8
81920
1536
256
128
8
24576
1536
512
64
81920
1536
512
24576
1536
1024
2
81920
1536
1024
32768
4096
256
32768
4096
256
256
30B
32768
4096
512
128
8
30B
32768
4096
512
128
8
32768
32768
4096
512
256
16
32768
4096
512
256
120B
32768
4096
1024
128
8
120B
32768
4096
1024
128
8
32768
4096
2048
64
4
32768
4096
2048
Bottom-right 2 subiigures: We use (T, (1600))66666644499112442,,, and we vary E from
32 to 256.
G HYPERPARAMETER DETAILS FOR LM TRAINING
We use the OoMoE codebase MMuennighoff et al. 2025 and its downstream tasks in the default
configuration 12] except for MMLU: WinoGrande ("wino") ") (Sakaguchi et al. 2020), Social IQA ("SIQA")) (Sap et al. 2019). SciQ (Johannes Welbl 2017). PIQA (Bisk et al. 2020), OpenBookQA ("OBQA") (Mihaylov et al. 20188) HellaSwag ("HS") (Zellers et al.11),)), COPA (Roemmele et al.
2011), QA ("CSQA") (Tammor al. 2019), BoolQ (Clark et. 2019), Are-Easy and
Are-Char Pre-Challenge ("ArcE" and "AreC") CClar et al. 2018) datasets. We use a deduplicated version of
FineWeb-FFdu (Ben Allal et al. 2024 for pretraining corpus, and train all models with context length of 4096 tokens.
We always MoE with SwiGLU for the MoE layers and
auxiliary load balancing loss (Shazeer et al. 2017) with coefficient 0.01 but we do not use the router Z loss (Zoph et al. (2022). Our attention block architecture identical to OLMoE's attention block. We always tie the weight of the LM head with the weight of the token embedding matrix.
Table 10: Common configurations for MoE pretraining experiment
Coufig name in Tablss I a
# layers # #atto beads
T
E
K minibatch tokeu size
LR WD
LR scheduler
0.5B params, 40B tokens, 2/64
12
12
768
256
64
2
IM
6e-4 0.01 cos w/. warmup (10% steps)
1.88 puramss 40B tokens, p 8/256
12
12
768
256
256
IM
6e-4 0.01 vos w/. warmup (10% steps)
1.4B params, 100B tokens, 2/128
18
12
768
256
128
2M
4e-4 0.01 warmup (10% steps)
For all EC with finetuned TC router experiments in Table use additional 4B tokens and only finetune the router weights with TC top-K routing (all other parameters are frozen). We always use a learning rate of 2c-4, weight decay of 0.01 and cosine learning rate scheduler with 10% warmup steps. The number of tokens per minibatch during finetuning is 1M. We disable auxiliary load balancing loss during TC finetuning.
For all EC with auxiliary router experiments, we use 2-layer MLP (each linear layer has size E x E with SiLU activation) which takes input the raw router logits and makes E independent binary predictions for all experts. We compute the averaged binary cross entropy loss over E labels using the multi-label prediction loss, and scale the loss by 0.01. During the evaluation, we will let EC router
12 http://chhhhotcccoooooooooooaeae
com/
O_Mog IB- 78- 3924 vmd
httpps//rugg ing wwt..a//a...cacooooo.aoo.oo.. aseatttsssssss/eee htt.://ggg..g......aaaaoeeeeee/a In-corpus
25

---

Published as a conference paper at ICLR 2026
compute the raw logits and raw scores and let the auxiliary router mask the token token-expert pair with its own confidence score.
We implement "TC (token drop)" by discarding tokens selected from the TC top-K sorting, always round down.
26


Reviews:
[{"review_id": "zdGVNXQm3I", "reviewer": "Reviewer_fRZc", "summary": "This paper introduces SNaX, a co-design solution for training fine-grained and sparse MoE models efficiently. The authors identify that increasing MoE granularity and sparsity shifts computation from compute-bound to memory-bound. They propose three contributions: (1) a memory-efficient backward pass maintaining constant activation memory regardless of expert granularity (45% reduction), (2) IO-aware kernels with asynchronous operations achieving 1.87x throughput improvement, and (3) token rounding that aligns token counts to GPU tile sizes, reducing wasted computation by up to 18%.", "questions": "The paper focuses on Hopper GPU optimizations (TMA, async operations). What is the performance degradation on A100 or V100 GPUs? Is SNaX still beneficial without these hardware features, or does it become worse than baselines?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper correctly identifies that modern MoE architectures face a critical shift from compute-bound to memory-bound operations as granularity and sparsity increase, supported by compelling arithmetic intensity analysis. The memory-efficient backward pass cleverly avoids materializing Y₂ through computation reordering, achieving 45% memory reduction while maintaining mathematical equivalence. Token rounding addresses the overlooked tile quantization problem where padding wastes significant compute in sparse MoEs, yielding 18% additional speedup with maintained accuracy. Comprehensive experiments across 1.4B-120B scales show consistent improvements over ScatterMoE (1.87x throughput) and MoMoE, with the planned open-source release providing valuable production-quality kernels to the community.", "weaknesses": "The fundamental incompatibility between token rounding (training) and token-choice (inference) is acknowledged but poorly explained, with authors only stating it \"creates difficulty\" without technical details or solutions. The training-inference mismatch may impact generation quality beyond perplexity (diversity, hallucinations) but remains uninvestigated. Heavy Hopper-specific dependence (TMA, tile sizes) limits portability with no analysis on A100/TPU performance degradation. While presenting ping-pong scheduling prominently, these are known techniques - the contribution is adaptation to MoE epilogues, not algorithmic novelty. Critical details like the \"state-of-the-art BF16 MoE kernel\" baseline and GEMM \"K dimension\" remain undefined, hindering reproducibility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Yu8iowMufT", "reviewer": "Reviewer_LFjW", "summary": "This work presents a new efficient implementation of MoE operator for reducing the peaking GPU memory of MoE's training process. This work also introduces a grouped-GEMM–aware token rounding method for tackling the wasted tile computation, as the required computation is lower than the minimum value that a hardware unit can compute in one pass. This phenomenon often occurs when training the highly sparse MoEs since each expert is small and some of them may only receive a few tokens in one pass. This work conducts experiments on a broad scale of MoEs (1.4B-120B), and comparisons against ScatterMoE, MoMoE, and MegaBlocks demonstrate its effectiveness.", "questions": "see weakness", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The writing is clear, and this work targets the efficiency bottlenecks of training highly sparse MoEs, which is an important research direction for the large models community. More and more recent LLMs and VLMs adopt the MoE structure, and as the trending analysis proposed in the appendix shows, recent MoEs tend to increase the sparsity of activated experts.", "weaknesses": "Weakness:\n1. Although the motivation of token rounding aims to reduce tile padding waste in grouped GEMM, the proposed method is similar to load balancing and token dispatching strategies. There are many prior works on token-choice and expert-choice. The proposed method's novelty is limited. The related work 2.4 briefly discusses the routing method, but it's better to discuss the hybrid methods for tackling imbalance issues in detail. \n2. May token rounding introduce hardware sensitivity? The strategy depends on grouped GEMM tile sizes and implementation details (which vary across GPUs/libraries/precisions). This could lead to unexpected behavior when open-sourcing the MoE models since the user may fine-tune it on different devices. \n3. The risk of training–inference consistency and behavioral bias. This work trains with token rounding but evaluates with token-choice. Recent work suggests routing consistency matters in RL/align phases; train–infer routing mismatch may increase the chance of mode collapse and bias amplification.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "A2YzkoFdVo", "reviewer": "Reviewer_tULt", "summary": "This paper addresses the memory bottlenecks and hardware inefficiencies posed by the trend towards increasingly fine-grained and sparse Mixture of Experts (MoE) models. It proposes SNaX, a system and algorithm co-design solution. SNaX tackles these challenges via three core contributions: (1) A memory-efficient MoE algorithm that minimizes activation memory by modifying the backward pass computation graph; (2) An IO-aware GPU kernel design that overlaps memory IO latency with computation ; and (3) A novel 'Token Rounding' routing method designed to reduce wasted compute resulting from tile quantization. Experimental results show that SNaX significantly reduces memory usage (e.g., a 45% reduction for a fine-grained 7B MoE model ) while substantially improving compute throughput (a 1.87x improvement on Hopper GPUs ).", "questions": "The paper avoids caching Y_2 by reordering the computation graph. Does this complex backward pass introduce any numerical instabilities, particularly when training in mixed precision (BF16)?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "● The paper astutely points out that as MoE models trend towards being more fine-grained and sparse, the training bottleneck is shifting from computation to memory bandwidth.\n● Consistent performance improvements over baselines like ScatterMoE and MoMoE across various model scales (1.4B to 120B) strongly demonstrate the practical value of the proposed method.\n● 'Token Rounding' is a novel and clever idea that addresses the often-overlooked problem of wasted computation in GEMM operations due to padding, which the paper refers to as 'tile quantization'.", "weaknesses": "● The paper states in Section 6.3 that Token Rounding (TR) is used during training, but it switches back to Token Choice (TC) for evaluation. The authors also admit that \"Token rounding is not a token-choice routing method which creates difficulty for autoregressive inference.\" This is a major concern. The authors do not clearly explain the root cause of this 'difficulty'.\n● Furthermore, this training-inference mismatch constitutes a theoretical weakness. Although the results in Table 1 show acceptable PPL and Avg accuracy, it is unclear if this mismatch could negatively impact other aspects of the model, such as generation diversity or the prevalence of 'hallucinations'.\n● Moreover, the paper does not discuss the differences in loss convergence between the two methods (TR vs. TC) during the training phase.\n● The paper highlights 'Ping-Pong' scheduling and asynchronous TMA as key to its IO-aware kernel design. While their application in SNaX is effective, these scheduling techniques themselves are not new. The authors should more clearly articulate that their core contribution lies in the specific application and adaptation of these techniques to the heavy-epilogue nature of MoE kernels, rather than claiming novelty for the techniques themselves.\n● The paper fails to clarify certain technical terms. For instance, what is the \"state-of-the-art BF16 MoE kernel\" that SNaX is compared against? Additionally, what does the \"K dimension\" (mentioned in Section 2.1) refer to in the context of GEMM?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "SFqy8yXRwo", "reviewer": "Reviewer_Y7dy", "summary": "This paper, SNaX, presents three improvements to current MoE models to allow them to better utilize modern GPU hardware. Specifically, they propose a new implementation that reduces the memory used to store activations during training and allow this activation size to be independent of the expert granularity. Secondly, they improve throughput by overlapping IO with computation. Finally, they propose to constrain the routing strategy such that the number of tokens routed to each expert is a multiple of the tile size of GPUs, reducing the amount of wasted computation resulted from padding. Empirical results show that the propose implementation and algorithm is much more efficient in terms of throughput and memory usage while maintaining performance. It also fosters the development of high-granularity MoE models since it provides better support for such models compared to existing open-source solutions.", "questions": "- Figure 5 and Section 6.2: Floating-point operations per second is usually abbreviated as FLOPS instead of FLOPs.\n- I think the citation in Line 264 should be moved to the first reference of Ping-Pong in Line 256. Moreover, \"Ping-Pong\" and \"pingpong\" both appears in the paper. The authors, should unify the capitalization for this term.\n- Figure 6 is never referred to in the paper.\n- Can this method be used for MoE models without a top-K function, such as ReMoE [1]?\n\n[1] ReMoE: Fully Differentiable Mixture-of-Experts with ReLU Routing", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper is well-written, the motivation is clear, and the description of the proposed algorithm is easy to understand.\n- The experimental results are convincing and the proposed method provides considerable efficiency improvements. The proposed method, which will be open-source as described in Line 107, will be useful for the community at large since it will greatly lower the cost of training large MoE models. Moreover, high-granularity MoE models hold great promise and is the focus of many research works. The method proposed by this paper will be helpful for many researchers.", "weaknesses": "- It would be helpful to see efficiency tests on more hardware other than H100.\n- It seems that the focus of this paper lies in training efficiency. It is currently unclear whether parts of the improvements can be used to improve inference efficiency. Nonetheless, a discussion on inference efficiency seems to be important.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)