Loads and Loads of Fluffy Kittens: Compute-Communication Kernels with Multi-GPU ThunderKittens

TL;DR

Part 2 of Hazy Research’s multi-GPU ThunderKittens series. Argues that writing efficient compute-communication kernels on modern interconnect-rich GPUs (HGX H100/B200, GB200 NVL72) reduces to three reusable principles: pick the right inter-GPU transfer mechanism (copy engine vs TMA vs register/multimem) for the message granularity, pick the right overlap scope (intra-SM warp specialization when comm aligns with compute, inter-SM otherwise), and keep tiling — fewer than 50 extra lines of device code on top of a single-GPU kernel are enough. Demonstrates the recipe on AG+GEMM, GEMM+RS, GEMM+AR (tensor parallelism), Ring Attention and DeepSpeed-Ulysses (sequence parallelism), and token-dispatch+GEMM (expert parallelism for MoE), matching or beating Flux, CUTLASS, Triton-Distributed, Comet, xDiT, and YunChang baselines on 8× H100. Companion arxiv paper at 2511.13940; code at github.com/HazyResearch/ThunderKittens/tree/main/kernels/parallel.

Key claims

Three inter-GPU transfer mechanisms exist with different sweet spots: copy engine needs ~256 MB chunks to saturate NVLink (no SMs, but fine-grained-hostile); TMA saturates NVLink at 2 KB messages using only ~15/148 SMs on B200 and zero SM under intra-SM overlap, but no in-network reduction; register instructions (ld/st/red/multimem) handle in-network reduction and element-granularity but cost SM resources [§“Using the Right Transfer Mechanism”, Fig. 4, Fig. 5].
Off-the-shelf NVSHMEM put/put_nbi defaults are suboptimal because they emit volatile st instructions, enforce __ldg for peer addresses, and add thread-level synchronization — all of which increase latency and cut NVLink bandwidth utilization [§“Using the Right Transfer Mechanism”].
Intra-SM overlapping (warp specialization within an SM, one or more warps dedicated to comm) yields zero loss in the best case — tensor cores stay busy and NVLink is fully hidden. In their fused GEMM+reduce-scatter the non-overlapped comm fraction drops below 1% [§“Using the Right Overlapping Schedule”].
Inter-SM overlapping (entire SMs dedicated to comm or compute) is required when the comm pattern can’t align with computation, when in-network reduction (NVSwitch multimem) is needed, and when many compute thread blocks would otherwise hit the same remote L2 cache and serialize on NVLink — dedicating a few SMs to remote-to-local prefetch lets compute SMs run from local L2 instead [§“Using the Right Overlapping Schedule”, Fig. 7].
Tiling does not reduce inter-GPU bandwidth utilization in practice as long as TMA is used or remote HBM accesses are coalesced; every released kernel adds fewer than 50 lines of device code on top of the single-GPU GEMM/attention kernel [§“Using the Tiles”].
“ParallelKittens” (PK) matches or beats hand-optimized baselines (Flux, CUTLASS, Triton Distributed) on AG+GEMM, GEMM+RS, GEMM+AR across 8× H100 sweeping N from small to 16K; cuBLAS+NCCL non-overlapped is the floor [Fig. 10, Fig. 11, Fig. 12].
For sequence parallelism, PK Ring Attention and DeepSpeed-Ulysses implementations beat xDiT and YunChang reference open-source baselines across sequence lengths on 8× H100 (B=16, H=16/128, D=128) [Fig. 14, Fig. 15].
For expert parallelism, PK overlapped token-dispatch+GEMM matches or beats Comet (the prior SOTA fine-grained overlap kernel) on H100 with top_k=8, num_experts=256, D=7168, D_expert=2048 [Fig. 17].

Method

The post is a design-principles writeup, not a single kernel. The unifying observation is that NVLink-era GPU systems expose heterogeneous resources (CUDA cores, tensor cores, SFUs, TMA, copy engine, NVSwitch in-fabric reduction, memory controller, NVLink bandwidth, register file, SMEM) that must be co-scheduled — naively saturating one or two leaves the rest idle. The three principles instantiate as: (1) pick the inter-GPU transfer mechanism per workload (copy engine for bulk pre-stage, TMA for tile-granularity comm, register/multimem for element-wise or in-network reduction); (2) choose intra-SM warp specialization when the comm tile maps 1:1 to a compute output tile (e.g., GEMM+RS: the output tile of a tiled GEMM is exactly what gets sent and atomically added at the destination), and inter-SM otherwise (e.g., attention’s repeated remote KV access, or NVSwitch in-network all-reduce that increases register pressure); (3) tile, as in single-GPU ThunderKittens, so the kernel composes from familiar primitives. The companion arxiv paper (2511.13940) presents the full taxonomy and proofs.

Results

Numbers are reported in figures rather than tables, all on 8× H100 NVLink:

AG+GEMM (local shape N × N/8 × N): PK matches or exceeds Flux, CUTLASS, and Triton Distributed across the swept N range; large lead over cuBLAS+NCCL [Fig. 10].
GEMM+RS (local shape N × N × N/8): PK competitive with Flux and CUTLASS; consistently above Triton Distributed and cuBLAS+NCCL [Fig. 11].
GEMM+AR: only PK and Triton Distributed compete (Flux and CUTLASS don’t have a GEMM+AR), with PK ahead [Fig. 12].
Ring Attention + Ulysses: PK ahead of xDiT and YunChang across context lengths [Fig. 14, Fig. 15].
MoE token-dispatch+GEMM: PK ahead of Comet [Fig. 17].
Their GEMM+RS configuration drives the non-overlapped comm portion below 1% — i.e., NVLink transfer + atomic-add are fully hidden behind tensor-core GEMM [§“Using the Right Overlapping Schedule”].

Why it’s interesting

This is the systems-side companion to FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling — same Hazy Research / Dao-AILab lineage, same “co-schedule heterogeneous hardware” thesis, but extended from single-GPU softmax-MMA-MUFU co-scheduling onto multi-GPU NVLink-MMA-TMA co-scheduling. The principles are also the explicit kernel-side analog of Distributed training parallelism‘s “non-uniformity as a first-class primitive” story: ParallelKittens generalizes one level below the framework layer (veScale-FSDP, Fireworks 4D, KnapFormer compute bags all assume the kernels they call are well-built), and the GEMM+RS sub-1%-overhead result is concrete evidence that the comm/compute mismatch Scaling and Optimizing Frontier Model Training (Fireworks AI) flags can in fact be closed in <50 lines if the right transfer mechanism is picked. Contrasts with Kittens virtual machine fuses entire training runs into a single GPU kernel (Ben Spector / Hazy Research), which sketched the megakernel for training endpoint of the same line — this blog is the building-block layer underneath that endpoint. Open question: do the same intra-SM / inter-SM choices hold on Blackwell B200 with NVSwitch in-network reduction as a first-class scheduling target, and what new failure modes appear at NVL72 scale where the topology becomes a hierarchical fabric rather than a fully-connected 8-GPU group.