Mixture-of-Kittens: MoE Megakernel for NVL72s

Stuart H. Sul; Nash Brown; Henry Wildermuth; William Lin; Federico Cassano; Christopher Ré

cs.DCSep 28, 2026

Mixture-of-Kittens: MoE Megakernel for NVL72s

Authors: Stuart H. Sul, Nash Brown, Henry Wildermuth, William Lin, Federico Cassano, Christopher Ré

Organizations: Department of Computer Science, Stanford University · Cursor Research

Abstract

AI accelerator systems are rapidly consolidating into scale-up architectures, where tens to thousands of GPUs communicate over high-bandwidth, single-hop fabrics. We find that existing Mixture-of-Experts (MoE) training systems, optimized for conventional scale-out networks, transfer poorly to this setting, often running slower than a naive baseline built with PyTorch and NCCL. With industry roadmaps pointing toward even larger scale-up domains, understanding the performance tradeoffs of this hardware regime is increasingly important. We present Mixture-of-Kittens (MoK), an MoE training system designed for Nvidia NVL72. MoK builds on three insights that unlock performance on scale-up domains: (1) choosing push- or pull-based communication per operator, (2) restructuring the computation-communication overlap, and (3) fully eliminating CPU-GPU synchronization. MoK distills these insights into a single deterministic training megakernel that fuses token dispatch, shared and routed expert FFNs, and token combine. Across MoE layer shapes from four widely used open-weight models, MoK delivers up to $2.37\times$ the throughput of the strongest publicly available baseline. In a production run on 512 GPUs spanning multiple GB300 NVL72 racks, MoK improves end-to-end training throughput by $1.41\times$ .

Figures & tables

Figure 1: Overview of MoK. Top left: communication SMs pull tokens from peer GPUs into a fixed-size macrobatch ring buffer, expert FFNs process each minibatch, and communication SMs push the outputs back. Bottom left: MoK partitions SMs between computation and communication, and pipelines minibatches (M0, M1, …) to overlap expert computation with dispatch and combine. Right: before the megakernel launch, MoK performs GPU-side router projection, top- $k$ selection and all-gather, and dispatch scheduling.

Figure 2: Expert parallelism example with an EP degree of 3 and top-1 routing.

Figure 3: Cost model terms on a naive overlapped MoE execution. Unlabeled gray boxes denote synchronization. Arrows mark execution dependencies.

Figure 4: Signaling under push and pull dispatch. Left: with push, GPU 0’s computation SMs wait for a completion signal from every peer; with pull, they only need a GPU-local signal from their own communication SMs. Right: dispatch latency and signaling overhead on NVL72 (7,168 BF16 tokens of width 7,168 per GPU, 4 local experts, 4 uniformly random routes per token). Signal overhead is the added latency of in-kernel completion signaling relative to the same kernel without it.

Panel of Figure 5 — Figure 5: Left: MoK MXFP8 forward-pass throughput on 16 GPUs as a function of minibatch size. Colors denote model shapes, symbols denote the number $N$ of input tokens per GPU, and enlarged symbols mark the throughput-maximizing minibatch size for each curve. For a fixed model and token count, the best minibatch size achieves $1.38\times$ – $3.53\times$ the throughput of the worst. Right: maximally deferred synchronization with four minibatches (M0–M3). Communication (yellow) and computation (red) traverse the work in different orders and meet only at the dashed arrows.

Figure 6: Left: host-side invocation time of common PyTorch operations on a GB300 NVL72, normalized to a DGX B300. The x-axis is the operation and the y-axis is the normalized CPU time. Right: latency overhead of MoK’s macrobatch ring buffer. We compare a ring buffer that can hold half of the routed tokens against an over-allocated buffer that can hold all of them, and report the percentage increase in latency.

Figure 7: MoE layer throughput per GPU on a GB300 NVL72. The x-axis is the MoE shape and the y-axis is the effective expert-FFN throughput in TFLOP/s. We use an EP degree of 64 and 2,048 input tokens per GPU before routing. Higher is better.

			Forward (ms)		Backward (ms)		Speedup
EP	MoE Shape	Tokens	MoK	Best	MoK	Best	Fwd	Bwd
16	Kimi K2.7 Code	1,024	0.901	2.334 N	1.786	2.413 N	2.59 $\times$	1.35 $\times$
		2,048	1.332	2.581 N	2.425	3.329 N	1.94 $\times$	1.37 $\times$
		4,096	2.555	3.308 H	4.320	5.053 H	1.29 $\times$	1.17 $\times$
	GLM-5.2	1,024	0.849	2.321 N	1.284	2.070 N	2.73 $\times$	1.61 $\times$
		2,048	1.180	2.367 N	1.965	2.762 N	2.01 $\times$	1.41 $\times$

Table 1: MoK sensitivity to EP degree and local token count. “Tokens” is local tokens per GPU before dispatch. “Best” is the lowest latency among the four baselines for the given workload, with the subscript identifying the best baseline (N for NCCL+PyTorch, D for DeepEP+PyTorch, H for HybridEP+Megatron). “Speedup” is “Best” divided by “MoK”. Lower latency and higher speedup are better.

	DeepEP-based	MoK
Tokens / second / GPU	761.0	1,070.2

Table 2: End-to-end production training stack throughput on 512 GB300 GPUs, comparing the DeepEP-based backend with MoK. Both runs use an EP degree of 32, MXFP8 routed experts, and a BF16 shared expert. We report the mean per-GPU throughput. Higher is better.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

		MoK (GB300)			Comet (H100)
MoE Shape	Tokens	Sched.	Total	Share	Sched.	Total	Share
Kimi K2.7 Code	1,024	0.068	2.250	3.04%	0.331	4.892	6.76%
	2,048	0.093	2.571	3.60%	0.389	5.640	6.89%
	4,096	0.133	4.507	2.96%	0.515	10.537	4.89%
GLM-5.2	1,024	0.064	1.411	4.50%	0.294	3.268	9.00%
	2,048	0.077	2.050	3.74%	0.349	4.774	7.30%

Appendix

Table 3: Scheduling overhead of MoK and Comet for BF16 forward passes at EP degree 4. “Tokens” is local tokens per GPU before dispatch. “Sched.” is the time spent building the schedule before dispatch begins, “Total” is “Sched.” plus the forward latency, both in ms, and “Share” is “Sched.” divided by “Total”.

Distributed Training

Distributed training coordinates computation across multiple devices or machines to train models that exceed single-device limits.

See all 133 papers

Scalable Training of Mixture-of-Experts Models with Megatron Core
Mar 8, 2026Zijie Yan, Hongxiao Bai, Xin Yao +42Low-Precision Training Mixture-of-Experts Models
NVIDIA
DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism
May 10, 2026Zhichen Zeng, Chi-Chih Chang, Jiayi Wang +10Expert Parallelism Communication-Efficient Distributed Training
ByteDance Seed · University of Washington · Cornell University
Fine-grained Computation-Communication Overlap via Tile-level Signaling and Scheduling for Mixture-of-Experts
Jul 21, 2026Minyu Cui, Anna Wingkvist, Morgan EricssonGPU Kernel Optimization Communication-Efficient Distributed Training
Linnaeus University · Växjö, Sweden

Mar 8, 2026J/K move · Enter open · S save

Scalable Training of Mixture-of-Experts Models with Megatron Core

Zijie Yan, Hongxiao Bai, Xin Yao, Dennis Liu+41

NVIDIA

Low-Precision Training Mixture-of-Experts Models Communication-Efficient Distributed Training Memory-Efficient Training

Open paper AskPDF arXiv

Scaling Mixture-of-Experts (MoE) training introduces systems challenges absent in dense models. Because each token activates only a subset of experts, this sparsity allows total parameters to grow much faster than per-token computation, creating coupled constraints across memory, communication, and computation. Optimizing one dimension often shifts pressure to another, demanding co-design across the full system stack. We address these challenges for MoE training through integrated optimizations spanning memory (fine-grained recomputation, offloading, etc.), communication (optimized dispatchers, overlapping, etc.), and computation (Grouped GEMM, fusions, CUDA Graphs, etc.). The framework also provides Parallel Folding for flexible multi-dimensional parallelism, low-precision training support for FP8 and NVFP4, and efficient long-context training. On NVIDIA GB300 and GB200, it achieves 1,233/1,048 TFLOPS/GPU for DeepSeek-V3-685B and 974/919 TFLOPS/GPU for Qwen3-235B. As a performant, scalable, and production-ready open-source solution, it has been used across academia and industry for training MoE models ranging from billions to trillions of parameters on clusters scaling up to thousands of GPUs. This report explains how these techniques work, their trade-offs, and their interactions at the systems level, providing practical guidance for scaling MoE models with Megatron Core.

Abstract

Distributed Training

Explore similar work

Scalable Training of Mixture-of-Experts Models with Megatron Core

DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism

Fine-grained Computation-Communication Overlap via Tile-level Signaling and Scheduling for Mixture-of-Experts