Can LLMs Write Fast Multi-GPU Kernels? — Simran Arora, Together AI
Summary
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.
Technical details
-
GPU Interconnect Hierarchy
340s
The communication stack includes PCIe (CPU-GPU), NVLink (point-to-point GPU-GPU), and NVSwitch, which connects all NVLink endpoints into a non-blocking fabric. The hierarchy dictates that data transfer speed decreases as physical distance increases.
-
Multi-GPU Kernel Trade-offs
180s
Three main intra-GPU data transfer methods exist: the Copy Engine (host/CPU initiated, good for large messages), Tensor Memory Accelerator (TMA) (device initiated, excellent for fine-grain communication using few registers), and register-level instructions (best for leveraging in-network reductions via NVSwitch). Each method involves different trade-offs regarding bandwidth saturation and resource usage.
-
Overlapping Strategies
240s
Kernel scheduling can be optimized using IntraSM (within a single processor) or InterSM (across multiple processors). The choice depends on whether the computation and communication patterns align, especially when maximizing NVLink traversal is required.
Mentioned resources
- ParallelKittens
- ParallelKernelBench
Channel & topics
Watch on YouTube · Back to latest
This independent, AI-assisted summary is provided for commentary and informational purposes. It may contain errors or omit important context. Please watch the original video for the creator's complete presentation. Video, thumbnail, and related copyrights belong to their respective owners.