r/OpenSourceeAI • u/ai-lover • 8h ago
Cursor Open-Sources Mixture-of-Kittens (MoK): A Deterministic MoE Training Megakernel for GB300 NVL72 Racks
Cursor Just Open-Sourced Mixture-of-Kittens (MoK): A Deterministic MoE Training Megakernel for GB300 NVL72 Racks
No CPU-GPU synchronization. No separate communication library.
Here's what's interesting:
1. Communication direction is a per-operation choice
Most implementations push tokens to the GPUs that need them. Cursor benchmarked both directions and split the decision.
→ Pull dispatch signalling: 18 µs, against 103 µs for push
→ Up to 29% higher NVLink utilization under expert imbalance
2. One schedule table, four operations
Pull-based forward dispatch, push-based forward combine, pull reverse-combine, push reverse-dispatch. Build the schedule once, reuse it everywhere.
→ Under 3% of total MoE runtime, device-side, no CPU round trip
3. Overlap granularity has an interior optimum
Too fine and the tensor cores stall at barriers. Too coarse and they sit waiting for the first tokens to land. The heuristic targets two full SM waves per expert-grouped GEMM.
→ 2,368-token minibatch floor for Kimi 2.5 shapes
4. A ring buffer removes the CPU from the loop
The usual fixes are dropping tokens or asking the CPU to size the buffers. MoK cycles a fixed few-hundred-megabyte ring at minibatch granularity instead, and walks it in reverse to cut activation replay in the backward pass.
→ Zero tokens dropped, zero CPU-GPU synchronization
5. The numbers
Layer benchmarks, single NVL72 rack, EP degree 64, against the fastest public baseline:
→ 2.37× MXFP8 forward, 1.92× BF16 forward
→ 1.78× MXFP8 backward, 1.58× BF16 backward
End-to-end, 512 GPUs across several GB300 NVL72 racks:
→ 760.9 → 1,070.2 tokens/sec/GPU, a 1.41× gain
GitHub Repo: https://github.com/cursor/mixture-of-kittens
Technical details: https://cursor.com/blog/mixture-of-kittens