ParallelKittens: Systematic and Practical Simplification of Multi-GPU AI Kernels
This addresses performance issues for AI practitioners scaling models on multi-GPU systems, though it is incremental as it builds on existing frameworks.
The paper tackles the bottleneck of inter-GPU communication in AI workloads by introducing ParallelKittens, a minimal CUDA framework that simplifies multi-GPU kernel development, achieving speedups of up to 2.33x for data-parallel, 4.08x for sequence-parallel, and 1.22x for expert-parallel workloads.
Inter-GPU communication has become a major bottleneck for modern AI workloads as models scale and improvements in hardware compute throughput outpace improvements in interconnect bandwidth. Existing systems mitigate this through compute-communication overlap but often fail to meet theoretical peak performance across heterogeneous workloads and new accelerators. Instead of operator-specific techniques, we ask whether a small set of simple, reusable principles can systematically guide the design of optimal multi-GPU kernels. We present ParallelKittens (PK), a minimal CUDA framework that drastically simplifies the development of overlapped multi-GPU kernels. PK extends the ThunderKittens framework and embodies the principles of multi-GPU kernel design through eight core primitives and a unified programming template, derived from a comprehensive analysis of the factors that govern multi-GPU performance$\unicode{x2014}$data-transfer mechanisms, resource scheduling, and design overheads. We validate PK on both Hopper and Blackwell architectures. With fewer than 50 lines of device code, PK achieves up to $2.33 \times$ speedup for data- and tensor-parallel workloads, $4.08 \times$ for sequence-parallel workloads, and $1.22 \times$ for expert-parallel workloads.