ParallelKittens: Systematic and Practical Simplification of Multi-GPU AI Kernels
Stuart H. Sul, Simran Arora, Benjamin F. Spector, Christopher Ré
cs.DC, cs.LG
2025-11-18
ParallelKittens distills overlapped multi-GPU kernels into 8 primitives and one template; under 50 lines of device code hit up to 4.08× speedup over tuned baselines on Hopper and Blackwell.
As models scale across GPUs, the data moving between cards becomes the bottleneck. From A100 to B200, BF16 tensor-core compute climbed 7.2× and HBM bandwidth 5.1×, but intra-node NVLink improved only 3× and inter-node PCIe/InfiniBand just 2×. In large-model workloads, communication eating over half of wall-clock time is normal, and the compute sits idle waiting.
The standard fix is to overlap communication with computation: while a GPU waits for data, it computes on data already on-chip. Existing overlap schemes each have a crack. Hand-written kernels (Flux, FlashDMoE, CUTLASS) push a single operator to peak, at the cost of bespoke code leaning on low-level primitives. Compiler approaches (Triton Distributed, TileLink) sometimes emit kernels slower than the non-overlapped baseline and do not survive a hardware generation. Off-the-shelf libraries (NCCL, NVSHMEM) impose two-way synchronization and intermediate buffering, landing up to 4.08× slower than hand-tuned code. ParallelKittens (PK) asks: instead of rewriting the plumbing for every operator, can a small set of reusable principles guide multi-GPU kernel design systematically?
PK sits on top of ThunderKittens (TK), a tile-based single-GPU CUDA DSL, and adds eight new primitives. All operate at tile granularity (16×16 to 256×256) and are device-initiated, bypassing the host. They split into three groups: peer-to-peer tile transfers (storeasync, storeaddasync), collectives that exploit in-fabric reduction (reduce, allreduce), and multicast-based barriers (signal, signalall, wait, barrier).
Three principles sit behind them. First, pick the transfer mechanism by message size. The copy engine reaches about 81% of peak bandwidth but only for messages of 256 MB and up; TMA (Hopper's Tensor Memory Accelerator) hits about 74% from 2 KB; register-level instructions stay efficient at 128 B but need 3 to 5× more SMs to saturate, and they are the only mechanism supporting in-network reduction.
Second, choose scheduling by granularity. Intra-SM overlap splits warps inside one SM into compute and comm pools, with sync latency near 64 ns; inter-SM overlap dedicates whole SMs, unlocking in-network reduction and better L2 reuse, but sync is near 832 ns. The right pick depends on whether compute and communication granularities line up.
Third, cut design overhead. NCCL-style libraries force two-way synchronization; PK uses one-way writes into pre-allocated buffers, no intermediate copies, and no superfluous syncthreads on peer access.
On top of the primitives sits one template (Load-Compute-Store-Communicate) with four workers: loader, storer, consumer, communicator. The template handles shared-memory layout, barrier setup, and SM/warp partitioning automatically. You describe what to move and what to compute; it wires the overlap.
Validated on 8×H100 (Hopper, 4th-gen NVLink 450 GB/s) and B200 (Blackwell, 5th-gen 900 GB/s), CUDA 12.6.
| Workload | Baseline | PK speedup |
| Data/tensor-parallel GEMM+comm | Flux, CUTLASS, Triton Distributed | up to 2.33× |
| Sequence-parallel Ring Attention | xDiT | 1.07–4.08× |
| DeepSpeed-Ulysses | YunChang | 1.01–1.39× |
| Expert-parallel MoE dispatch | Comet | 0.92–1.22× |
With a large enough reduction axis, non-overlapped communication on the data-parallel kernel falls to 1% of execution time; on sequence-parallel Ring Attention it falls to 9%. Each kernel's communication component was written in fewer than 50 lines of device code beyond the original single-GPU kernel (the MoE one took under 40). On non-contiguous layouts, PK also beats NCCL on fine-grained collectives, because NCCL has to reshape before copying.
The honest caveat: speedups are ranges, depending heavily on problem size and which baseline counts as fair. On MoE, PK is roughly on par with Comet (0.92–1.22×), and on some GEMM configurations it is slightly slower than Flux. This is a framework that lets a short program reach the Pareto frontier, not a uniform win.
Most teams running multi-GPU training or inference do not write their own collective kernels; they call NCCL and eat the overhead. PK argues that the layer above NCCL (the hand-written overlapped kernel) can be made cheap: a few dozen lines on a template that already handles the hard wiring. For the operators that dominate latency (attention, MoE dispatch, all-gather/reduce-scatter around GEMMs), that is where idle GPU time hides. The framework targets intra-node scaling (NVL72→NVL144→NVL576), which is exactly where Nvidia is concentrating density, so the relevance is current. It is a research framework extending ThunderKittens, not a drop-in library, so the practical takeaway is the three principles more than a line to paste into a Makefile.
The headline limitation, stated plainly: PK handles only intra-node communication. Multi-node clusters still need MPI/NCCL over the fabric, and inter-node extension is future work. Every speedup is a range, and the low ends (0.92× vs Comet, 1.01× vs YunChang) mean PK is not strictly better than existing hand-tuned kernels; on some configurations it ties or loses. Validation covers Hopper and Blackwell only; AMD MI-series and TPU are absent, and the principles lean on Nvidia-specific hardware (TMA, copy engine, NVSwitch multicast). PK also requires familiarity with ThunderKittens, itself a fairly low-level DSL; the "50 lines" counts device code, not the surrounding setup a practitioner still writes.