AI workloads have shifted their bottleneck from single-GPU computation to communication between GPUs.
2
ParallelKittens packages a small set of multi-GPU communication and scheduling patterns, adding roughly a dozen lines to a single-GPU kernel and running in production at Together AI and Cursor.
3
Frontier models can generate some correct CUDA kernels, but their performance stalls because they struggle with collective ordering, data partitioning, scheduling, and transfer-mechanism choices.
Summary
Simran Arora explains why multi-GPU communication has become a larger performance problem as GPU compute has advanced faster than the links between GPUs. She introduces the hierarchy from PCIe to NVLink and NVSwitch, then describes the choices kernel authors must make around copy engines, tensor memory acceleration, register-level transfers, and the placement of communication relative to computation. Together AI captured these patterns in ParallelKittens, a small set of primitives used in production and designed to add about a dozen lines to a single-GPU kernel. Arora then presents ParallelKernelBench, which contains 87 multi-GPU problems drawn from real repositories. The best zero-shot model solved 28, with 22 faster than the PyTorch and NCCL baseline. More samples raised correctness to 36, while correct-and-faster results stalled around 31 percent. An agent with a bash environment reached 35 solved problems, but further time did not keep improving performance. The failures involve multi-GPU reasoning rather than basic CUDA syntax.
AI performance bottlenecks have moved from single GPUs to communication between GPUs
Arora says earlier GPU utilization problems came from poor intra-GPU memory access and weak single-GPU kernels. Work such as FlashAttention, memory-efficient architectures, sparse attention, Mamba, and better DSLs has reduced those problems. Production training and inference workloads now span multiple GPUs, where communication can consume most of the runtime and reduce model FLOP utilization. The hardware is also becoming more varied, with AMD XGMI, TPU 3D torus networks, and NVIDIA NVLink and NVSwitch. KV cache data can move across GPUs, CPUs, disks, and remote machines, while inference stages such as speculative decoding, decoding, and prefill can run on different hardware.
Memory distance determines both the speed and capacity available to GPU computation
Using an H100 diagram, Arora grounds kernel design in the hardware. Threads run on many processor groups, retrieve data through registers, L2 cache, and high-bandwidth memory, and are scheduled in blocks. Registers are closest to computation and reach about 130 terabytes per second on an H100, but there is little register capacity. Farther memory has more capacity and takes longer to reach. This same distance principle extends beyond one GPU. PCIe handles CPU and GPU communication, InfiniBand or TCP can connect nodes, and NVLink provides point-to-point links inside a node. NVSwitch connects NVLink endpoints into a non-blocking fabric and can perform multicast and reductions inside the network.
Compute has advanced faster than the communication hardware connecting GPUs
Comparing NVIDIA's A100 from 2020 with the B200 from 2024, Arora reports that BF16 tensor-core throughput improved by 7.2 times. Intra-node communication improved by 3 times, while inter-node communication improved by 2 times. The gap makes communication increasingly important in large workloads. Arora also says networking stacks differ more across hardware vendors than tensor cores and memory hierarchies do. A standard PyTorch and NCCL-style baseline falls below 50 percent of the communication-aware roofline on most ParallelKernelBench problems. NCCL and similar libraries are tuned for large, contiguous transfers, so they struggle with fine-grained communication and fused, non-trivial collectives.
Multi-GPU kernel development sits between generic libraries, DSLs, and expensive hand tuning
Arora describes three existing paths. Libraries such as NCCL and RCCL make common bulk collectives easier, but require synchronization around transfers and offer limited flexibility. Compilers and DSLs such as Triton Distributed and TileLink can provide an abstraction, yet adapting them to changing network hardware is difficult. Arora cites cases where code tuned for H800 GPUs did not adapt efficiently to H100s. Hand-tuned operators can reach peak performance, with examples including DPL, Comet, Ring Attention, FLUX, FlashDecoding, and distributed GEMM kernels from CUTLASS. That approach does not scale well across operators and precisions, with some changes taking five or six months.
ParallelKittens reduces multi-GPU kernels to a small set of reusable choices
Together AI built ParallelKittens to study the trade-offs in multi-GPU kernels and package the resulting patterns. The primitives cover data movement, buffering, synchronization, and scheduling. Arora says developers usually add roughly a dozen lines over a single-GPU kernel to insert the multi-GPU operations. The library supports data, sequence, and expert parallelism, and is used in production at Together AI and Cursor. Its design gives developers control over how senders and receivers buffer and synchronize data instead of hiding those decisions inside a bulk-communication library.
Transfer mechanisms have different costs and fit different communication patterns
Arora describes the per-GPU copy engine as host- or CPU-initiated work that works well for large messages and can reach peak communication bandwidth without taking GPU registers or processor groups away from computation. Device-initiated transfers through tensor memory acceleration can saturate NVLink with smaller messages while using few registers and processors, which makes them useful for fine-grained overlap. Register-level instructions can take advantage of in-network reductions supported by NVSwitch. TMA cannot effectively use those in-network computations. Choosing among the copy engine, TMA, and register-level transfers is therefore part of the kernel's algorithm, rather than a minor implementation detail.
Communication and computation can be overlapped within an SM or across SMs
One schedule assigns different warps within the same processor group to computation and communication. This intra-SM approach works when both activities use compatible data and align well. Another schedule assigns separate processor groups to computation, communication, and memory. That inter-SM approach helps when the kernel would otherwise need incompatible allocations across registers or shared memory, or when intra-SM overlap cannot maximize NVLink traffic. Arora gives different examples for the two choices. Intra-SM overlap works well for a GEMM with reduce-scatter, while inter-SM scheduling is effective for a GEMM with all-reduce because it can use NVSwitch's in-network reductions.
Models generate some fast kernels, but their performance does not scale with more samples
ParallelKernelBench gives a model an unoptimized PyTorch reference using torch.distributed and NCCL operations, plus a topology describing ranks and intra-node hardware. The model must produce a CUDA kernel using unified virtual addressing. The benchmark has 87 problems drawn from real repositories and covers communication patterns from data, sequence, tensor, context, layer, pipeline, and expert parallelism. In zero-shot tests, the best frontier model solved 28 problems, with 22 faster than the PyTorch and NCCL baseline. More samples raised correctness to 36 problems, but the fraction that was both correct and faster plateaued around 31 percent. The wins clustered around familiar collective, tensor-parallel GEMM, and Ulysses-style context-parallel patterns.
The difficult failures involve multi-GPU decisions rather than CUDA syntax
Arora says repeated sampling and error correction often let models compile their kernels. The harder failures concern collective ordering, data partitioning, intra-SM versus inter-SM scheduling, and choosing between transfer mechanisms. Models often omit register-level transfer instructions and tensor memory acceleration. A multi-turn agent using Gemini 3 Pro and a local bash environment improved the result from 24 to 35 solved problems out of 87, with 26 faster than the reference. However, giving the agent more time caused performance to plateau. Arora's examples of new kernels include Nemo vocabulary-parallel filtering, Hyena context parallelism, and an IOU suppression kernel for the SAM 3 video segmentation model.
"BF16 tensor core speeds improved by 7.2x, while intra node communication by just 3x and inter node communication by just 2x."10:34
Who should watch
You are writing CUDA or GPU kernels that must move data across NVLink and want a compact set of design patterns.
Your distributed training or inference workload is limited by communication, especially fine-grained transfers or fused collectives that NCCL does not handle well.
You are evaluating coding models on systems work and need a benchmark that tests topology, collective ordering, data partitioning, and performance.
