Can LLMs Write Fast Multi-GPU Kernels? — Simran Arora, Together AI
Multi-GPU AI workloads are increasingly bottlenecked by inter-node and intra-node communication links rather than individual GPU compute power. The speaker introduces ParallelKittens, a set of minimal primitives designed to simplify writing high-performance multi-GPU kernels. A benchmark called ParallelKernelBench tests whether frontier LLMs can generate optimized CUDA kernels for complex real-world patterns (e.g., data parallelism across tensor/expert dimensions). While models show promise in generating correct kernels, their ability to reason about critical architectural trade-offs—such as collective ordering, data partitioning, or choosing between transfer mechanisms (Copy Engine vs. TMA)—remains limited.
Key takeaways
-
The Bottleneck Shift
17:04
Improvements in compute (e.g., A100 to B200) have outpaced improvements in communication links (Intra-node: 3x; Inter-node: 2x). This forces the bottleneck off the individual GPU and onto the interconnects, causing standard PyTorch/NCCL baselines to fall below 50% of their communication-aware roofline.
-
Kernel Development Simplification
2:20
The speaker's team developed ParallelKittens, a small set of primitives that adds minimal lines to a single GPU kernel but enables state-of-the-art multi-GPU performance in production environments (e.g., Together AI).
-
LLM Performance on Kernel Generation
20:40
On the ParallelKernelBench, the best frontier models solved 28 out of 87 problems zero-shot. While scaling up attempts increased correctness to 36/87, the performance gain (speedup) plateaued near 31%. Failures are not syntax but stem from inability to reason about complex trade-offs.