Paper 2511.13940
ParallelKittens: Systematic and Practical Simplification of Multi-GPU AI Kernels
- Published
- Nov 2025
- Research lab
- Together AI
- Citations
- 9
- GitHub
- Not linked
01 In brief
Summary
ParallelKittens (PK) is a minimal CUDA framework that simplifies the development of overlapped multi-GPU kernels by distilling three key principles: transfer mechanisms, scheduling strategies, and design overheads.
It extends ThunderKittens with eight core primitives and a unified programming template, enabling efficient compute-communication overlap.
PK achieves up to 2.33× speedup for data- and tensor-parallel workloads, 4.08× for sequence-parallel, and 1.22× for expert-parallel, with fewer than 50 lines of device code.
It matches or surpasses hand-optimized kernels (Flux, Comet, CUTLASS), outperforms compiler-based approaches (Triton Distributed) by 1.07–5.63×, and communication library-based approaches (xDiT, YunChang) by 1.01–4.08×.
PK is validated on Hopper and Blackwell architectures and is open-sourced, with adoption at Cursor for large-scale training.
The framework addresses the growing bottleneck of inter-GPU communication, which can occupy over 50% of execution time in LLM workloads, by providing device-initiated communication via TMA and register-level instructions, supporting both intra-SM and inter-SM overlapping, and minimizing synchronization and buffering overheads compared to libraries like NCCL and NVSHMEM.
PK's design enables explicit control over memory allocation and synchronization, reducing communication latency by up to 4.5× and improving pure communication kernel performance by up to 1.79×.
02 From the paper
Abstract
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.