Mixture-of-Kittens: MoE Megakernel for NVL72s

1 Introduction

• Scale-up domains connect many accelerators through high-bandwidth fabrics, but prior MoE systems often transfer poorly, with PyTorch/NCCL outperforming evaluated alternatives in nearly half the tested configurations.MoE can consume more than half of end-to-end execution time during large-scale training, according to prior work.Source: §1 · p. 1
• 2.37× MXFP8 forward speedup over the strongest publicly available baseline across four open-weight models, plus 1.41× end-to-end tokens/s/GPU over the previous DeepEP-based implementation in production.The production result was measured on 512 GPUs across multiple GB300 NVL72 racks.Source: §1 · pp. 2–3
• MoK fuses dispatch, shared and routed expert FFNs, and combine into one deterministic megakernel, with production features including MXFP8 support and fused router-weight gradients.It also supports RDMA overlap for FSDP and tunable SM partitioning.Source: §1 · pp. 2–3
• MoK selects push or pull communication per operator, tunes overlap granularity, and uses device-side ring buffering to eliminate CPU-GPU synchronization.Throughput-maximizing granularity ranges from 512 to 32,768 tokens across workloads; tuning it achieves up to 3.53× throughput over a design fixed to a suboptimal granularity.Source: §1 · pp. 2–3

2.1 MoE and Expert Parallelism

• MoE applies shared and selected routed-expert FFNs to each token, then sums their outputs using normalized router weights.Expert parallelism partitions routed experts across ranks, requiring token transfers before expert computation (dispatch) and afterward (combine).Source: §2.1 · p. 3

2.2 GPUs and Scale-Up Architectures

• Scale-up architectures provide uniform, single-hop connectivity across many GPUs, allowing parallelism to remain within the domain and enabling more aggressive overlap.GPU performance depends on overlapping execution units with different saturation rates.Source: §2.2 · pp. 3–4

2.3 Related Work

• Prior overlap work spans host-triggered transfers, device-initiated communication, flexible primitives, operator-specific methods, and whole-model techniques.MoK targets distributed MoE training on scale-up platforms.Source: §2.3 · p. 4
• MoE systems optimize local expert computation and communication overlap through techniques including fused permutation, low-precision grouped GEMM, adaptive pipelining, and fine-grained device-initiated communication.MegaMoE is a fused EP megakernel for inference-only DeepSeek-V4 serving in FP8×FP4, whereas MoK supports full training on scale-up architectures.Source: §2.3 · p. 4
• End-to-end frameworks can use MoK as an MoE backend, while the megakernel approach combines persistent operators but does not itself prescribe fusion or overlap design.MoK is already a backend for Nvidia NeMo AutoModel.Source: §2.3 · p. 4

3 Mixture-of-Kittens

• MoK minimizes end-to-end latency by overlapping compute, memory, and communication while accounting for exposed work, synchronization, and host overhead.Peak utilization still matters: tensor cores must stay fed, and network lanes need outstanding transfers with evenly spread traffic.Source: §3 · pp. 4–5, Figure 3 · p. 4

3.1 Choosing the Right Communication Direction

• Push-based dispatch spends 8–18% of dispatch time signaling, while pull stays under 4 µs; MoK uses pull dispatch, push combine, and a prebuilt schedule table.Source: §3.1 · pp. 5–6

3.2 Restructuring the Computation-Communication Overlap

• Small minibatches underutilize tensor cores while large ones expose dispatch and combine; MoK tunes granularity from 512–32,768 tokens and defers synchronization within a fused megakernel.Source: §3.2 · pp. 6–7, Figure 5 · p. 7

3.3 Eliminating CPU-GPU Synchronization

• A device-side macrobatch ring buffer avoids host synchronization and token dropping, adds 1.7% geometric-mean latency versus full over-allocation, and ring-aware backward replay recovers overwritten intermediates.Source: §3.3 · pp. 7–8

3.4 Additional Features and Optimizations

• MoK fixes atomic ordering or uses deterministic temporary-buffer reductions; Cluster Launch Control yields SMs for higher-priority inter-rack FSDP overlap.Source: §3.4 · p. 8
• MXFP8 keeps the shared expert FFN in BF16 and fuses activation quantization; backward fuses router-weight gradients from gated activations and down-projection input gradients.Source: §3.4 · pp. 8–9
• The shared-expert FFN runs in an otherwise-idle first dispatch window; communication-SM allocation is tunable per workload and separately for forward and backward.Source: §3.4 · p. 9

4 Experiments

• Experiments benchmark standalone MoE forward and backward passes and 512-GPU production-stack training on GB300 NVL72 racks against publicly available NVL72-compatible baselines.The layer evaluation covers shapes from four widely used open-weight models; the production comparison varies only the MoE implementation.Source: §4 · p. 9, §4.1 · p. 9, §4.2 · p. 10

4.1 MoE Layer Benchmarks

• Figure 7 shows MoK speedups over the fastest baseline for each combination, reaching 2.37× MXFP8 forward, 1.78× MXFP8 backward, 1.92× BF16 forward, and 1.58× BF16 backward at EP 64 and 2,048 tokens per GPU before dispatch.Source: §4.1 · p. 10
• MoK’s relative gains are larger for MXFP8 than BF16 and for forward than backward, where a smaller communication share limits the opportunity for overlap or elimination.The paper attributes these patterns to lower peak BF16 tensor-core throughput and roughly twice as many floating-point operations in backward as in forward.Source: §4.1 · p. 10
• MoK achieves larger speedups at smaller local token counts, where communication and scheduling overheads make up a larger fraction of execution time.The paper relates this regime to scaling across more GPUs, which reduces tokens processed per GPU when global batch size cannot keep increasing indefinitely.Source: §4.1 · p. 10

4.2 End-to-End Benchmarks

• Replacing the DeepEP-based implementation with MoK increased production end-to-end training throughput by 41% on 512 GB300 GPUs.Both Table 2 runs use EP degree 32, MXFP8 routed experts, and a BF16 shared expert.Source: §4.2 · p. 10, Table 2 · p. 11

• MoK combines per-operator communication-direction choices, tunable computation-communication overlap granularity, and device-side ring buffering to eliminate CPU-GPU synchronization.The design is fused into a single deterministic megakernel for MoE computation and communication.Source: §5 · p. 11
• The megakernel includes production-critical features such as hardware-accelerated MXFP8 execution, RDMA overlap for FSDP, and fused router weight-gradient computation.Source: §5 · p. 11
• MoK substantially outperforms the strongest publicly available baselines across widely used open-weight model shapes and delivers end-to-end throughput gains in production-scale LLM training.Source: §5 · p. 11

A Scheduling Overhead

• MoK scheduling builds a dispatch schedule after gathering local top-k routes, while Comet also computes per-expert counts and scatter indices, sorts routes, and schedules GEMM tiles on the host.Source: A Scheduling Overhead · p. 19
• Scheduling overhead is compared as a relative share of total time, since MoK requires Blackwell GPUs and Comet supports GPUs only up to Hopper.The BF16 forward measurements use EP degree 4 and the workloads, inputs, and local token counts from Section 4.1.Source: A Scheduling Overhead · p. 19
• Table 3 reports scheduling time before dispatch, total time including forward latency, and scheduling’s share of that total for BF16 forward passes at EP degree 4.Source: Table 3 · p. 19