Can LLMs Write Fast Multi-GPU Kernels? — Simran Arora, Together AI

Simran Arora discusses the critical bottleneck shift in large AI workloads from GPU compute to inter-GPU communication. Her team's solution, ParallelKittens, offers a set of primitives to optimize multi-GPU kernels by leveraging fundamental transfer mechanisms and compute-communication overlapping. They introduce ParallelKernelBench, a benchmark to evaluate AI models' ability to generate such kernels, revealing that while models can handle syntax, they struggle with deeper reasoning about communication patterns and hardware trade-offs.

August 27, 2026 6 min read

Can LLMs Write Fast Multi-GPU Kernels? — Simran Arora, Together AI — source: https://www.youtube.com/watch?v=pOvWgX7IJsc

The Shifting Bottleneck in AI Workloads: From GPU Compute to Inter-GPU Communication

Simran Arora, from Together AI, discusses the evolving performance bottlenecks in modern AI workloads, moving from single-GPU compute limitations to multi-GPU communication. This shift is primarily due to the disparate pace of hardware advancements: NVIDIA’s BF16 tensor core throughput improved 7.2x between A100 (2020) and B200 (2024), while intra-node communication only improved 3x and inter-node communication a mere 2x. This widening gap means that communication now dictates overall performance, often leading to less than 50% utilization of communication-aware roofline bounds with standard PyTorch and NCCL baselines.

GPU Hardware Fundamentals and Interconnect Hierarchy

To understand multi-GPU kernel optimization, it’s crucial to grasp GPU hardware. An NVIDIA H100 GPU features numerous Streaming Multiprocessors (SMs) for computation, supported by a memory hierarchy:

Registers: Fastest (130 TB/s on H100), closest to compute, but very limited capacity.
L2 Cache: Slower, larger capacity.
High Bandwidth Memory (HBM): Slowest (e.g., 80GB), largest capacity. Data access speed is directly tied to its proximity to the compute units.

Multi-GPU systems introduce an interconnect hierarchy:

PCIe: For CPU-GPU communication.
InfiniBand/TCP: For multi-node (inter-node) communication.
NVLink: High-bandwidth point-to-point connections between GPUs within a node.
NVSwitch: Connects NVLink endpoints into a non-blocking fabric, enabling full GPU-GPU communication and providing in-network acceleration for primitives like multicast and reductions.

The Imperative for GPU Networking Optimization

Several factors underscore the critical need for optimizing GPU networking:

Workload Scale: Modern AI models are massive, requiring kernels that span multiple GPUs, where communication increasingly dominates runtime.
Hardware Diversity: The rapid evolution and diversification of networking stacks across vendors (NVIDIA’s NVLink/NVSwitch, AMD’s XGMI, Google’s TPU with 3D torus) present both challenges and opportunities.
Workload Disaggregation: Advanced inference systems disaggregate tasks (e.g., KV cache across CPU/GPU/disk, speculative decoding on different hardware) and utilize larger scale-up domains (e.g., 576 GPUs in future systems), demanding fine-grained control over communication.
Limited Flexibility of Existing Solutions:
- NCCL/RCLL: Tuned for bulk transfers, they struggle with fine-grained communication and fused, non-trivial collectives, leading to significant performance gaps.
- Compilers/DSLs (e.g., Triton distributed, TileLink): Often struggle to keep pace with rapid networking hardware improvements and exhibit architecture-specific performance issues.
- Hand-tuned Operators: While achieving peak performance, this approach is not scalable, requiring significant effort for each precision or architecture change.

ParallelKittens: A Principled Approach to Multi-GPU Kernels

Simran Arora’s team developed ParallelKittens, a small set of minimal primitives and patterns for multi-GPU kernel development, aiming to simplify kernel writing while achieving peak performance. It focuses on understanding fundamental trade-offs:

1. Intra-GPU Data Transfer Mechanisms

Copy Engine: Host/CPU-initiated, excellent for large message transfers, high bandwidth, doesn’t consume GPU registers or SMs.
Tensor Memory Accelerator (TMA): Device/GPU-initiated, saturates NVLink with relatively small messages, low register/SM consumption, ideal for fine-grained overlapping. Limitations: cannot leverage in-network computations (e.g., NVSwitch reductions).
Register-level Instructions (LDST.red.multime): Device/GPU-initiated, best for leveraging in-network computations (NVSwitch reductions), but consumes precious registers.

2. Overlapping Compute and Communication

Intra-SM Overlapping: Different warps within a single SM specialize in compute or communication concurrently. Effective when compute and communication patterns align on data.
Inter-SM Overlapping: Different SMs specialize in compute, communication, or memory. Useful when resources (register file, shared memory) are misaligned, or when maximizing NVLink traversal with intra-SM is difficult, especially for leveraging in-network reductions (e.g., all_reduce).

ParallelKittens encapsulates these principles, allowing developers to add roughly a dozen lines of code to a single GPU kernel to enable multi-GPU operations. It’s actively used in production at Together AI and Cursor, demonstrating state-of-the-art performance across various parallelism schemes (data, sequence, expert).

Evaluating Frontier AI Models with ParallelKernelBench

The research extended to investigate whether frontier AI models could reason about these multi-GPU kernel design principles. For this, ParallelKernelBench was developed – a benchmark comprising 87 multi-GPU problems drawn from real-world AI workloads and GitHub repositories. Each task provides an unoptimized PyTorch reference implementation (with torch.distributed and NCCL) and a system topology, challenging the model to rewrite it into a performant CUDA kernel using unified virtual addressing.

Model Performance and Limitations

Zero-shot: The best frontier model solved 28 out of 87 problems correctly, with 22 of those outperforming the PyTorch+NCCL baseline.
Multiple Samples: Increasing samples improved correctness to 36 solutions, but the percentage of solutions that were both correct and faster plateaued at around 31%.
Agent Harness: Wrapping a model in a multi-turn agent harness with a local bash environment improved correct solutions to 35, with 26 achieving a 1x speedup. However, performance plateaued, indicating a need for more advanced techniques.

Analysis of Model Failures

Crucially, model failures were not primarily due to CUDA syntax errors (these often compiled after retries). Instead, models struggled with deeper reasoning about:

Collective ordering.
Data partitioning.
Choice between intra-SM vs. inter-SM scheduling.
Selecting appropriate transfer mechanisms (often failing to utilize TMA or register-level instructions). Model successes were concentrated in patterns heavily represented on the internet (e.g., collective primitives, tensor parallel GEMMs), suggesting a lack of generalized reasoning for novel multi-GPU patterns.

Conclusion

The work highlights that while a small set of fundamental patterns governs effective multi-GPU kernel writing, current frontier AI models lack the reasoning capabilities to effectively apply these principles, even when provided in context. ParallelKittens offers a practical, production-ready solution for optimizing these kernels, and ParallelKernelBench serves as a vital benchmark for advancing AI’s capability in this complex domain. The research points towards the need for architectures and methods that can better adapt to evolving networking stacks and larger scale-up domains.