Reproduce “Mixture-of-Kittens: MoE Megakernel for NVL72s” on OpenResearch
Turn your favorite coding agents into research agents that review literature, develop hypotheses, and run experiments.
Universal app for Apple silicon and Intel. Requires macOS 11 or later.
About this paper
Stuart H. Sul, Nash Brown, Henry Wildermuth, William Lin, Federico Cassano, Christopher Ré
The paper presents MoK, a GPU-resident program for mixture-of-experts training that coordinates token communication and expert computation, with tunable minibatches, direction-specific transfers, and a device-managed ring buffer. On GB300 NVL72 workloads, MoK reduced a Qwen3.5-397B-A17B layer from 2.527 to 1.068 milliseconds, a 2.37× speedup; this is not end-to-end training, but rather a general hardware or end-to-end claim.
Key insights
- On NVL72, communication direction and completion signaling materially affect performance: push signaling consumed 8–18% of dispatch time in tested cases, while pull signaling stayed within 4 microseconds.
- There is no universal best minibatch size: the throughput-maximizing choice varied by workload, and the best tested setting delivered 1.38–3.53× the throughput of the worst.
- Co-optimizing scheduling, communication, and expert computation can outperform treating token exchange as a standalone communication problem.
Summary by alphaXiv.
Tip: change arxiv.org to autoarxiv.org in any arXiv abstract link to open the paper here.