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

Multi-GPU kernel development is increasingly critical as compute efficiency improvements have shifted the bottleneck to communication, and while fundamental ...

By Sean Weldon

Abstract

As single-GPU optimization techniques have matured, the performance bottleneck in modern AI systems has shifted from intra-GPU memory access to multi-GPU communication. This synthesis examines the emerging challenge of multi-GPU kernel development, where communication bandwidth improvements significantly lag compute improvements - BF16 tensor core speeds improved 7.2x from A100 to B200, while intra-node communication improved only 3x and inter-node 2x. The analysis identifies fundamental principles governing effective multi-GPU kernels, including three distinct transfer mechanisms, two scheduling paradigms for compute-communication overlap, and critical synchronization patterns. These principles are encapsulated in the Parallel Kittens framework, which achieves state-of-the-art performance with minimal code overhead. Evaluation through Parallel Kernel Bench, an 87-problem benchmark, reveals that frontier AI models achieve only 32% success rates on multi-GPU kernel generation tasks, struggling with core architectural trade-offs despite strong reasoning capabilities in other domains.

1. Introduction

The landscape of GPU computing has undergone a fundamental transformation as optimization efforts have successfully addressed intra-GPU bottlenecks through innovations such as flash attention, sparse attention mechanisms, and memory-efficient architectures like Mamba. This progress has exposed a new constraint: multi-GPU communication has emerged as the primary performance limiter in distributed AI workloads. The asymmetric advancement of compute versus communication capabilities - with tensor core performance improving at more than twice the rate of interconnect bandwidth - necessitates a systematic approach to multi-GPU kernel development.

Multi-GPU kernels represent computational primitives that explicitly manage data movement and synchronization across multiple GPU devices while overlapping communication with computation. Unlike traditional approaches that treat communication as a separate phase handled by libraries such as NCCL (NVIDIA Collective Communications Library), effective multi-GPU kernels integrate communication primitives directly into computational logic. Naive implementations using PyTorch with NCCL achieve below 50% of the communication-aware roofline bound across representative problems, indicating substantial performance left unrealized.

This synthesis addresses three central questions: What fundamental principles govern effective multi-GPU kernel design? How can these principles be systematically encapsulated into reusable frameworks? To what extent can frontier AI models reason about multi-GPU optimization trade-offs? The analysis proceeds by establishing the technical context for multi-GPU optimization, identifying core design principles, examining the Parallel Kittens framework implementation, and evaluating model capabilities through Parallel Kernel Bench.

2. Background and Related Work

Modern GPU systems exhibit complex memory hierarchies with distinct performance characteristics. Register memory provides 130 terabytes per second bandwidth but limited capacity, L2 cache offers intermediate performance, and High Bandwidth Memory (HBM) delivers 80-192 gigabytes capacity at lower bandwidth. Inter-GPU communication occurs through specialized interconnects: NVIDIA's NVLink provides up to 900 gigabytes unidirectional bandwidth, AMD employs XGMI point-to-point links, and Google TPUs utilize 3D torus topologies with optical wraparound connections. The scale of multi-GPU systems continues expanding, with current architectures supporting 72 GPUs per node and NVIDIA planning 576-GPU single-system configurations by 2027.

Traditional multi-GPU programming relies on communication libraries such as NCCL and Rickle, which are optimized for bulk data transfers but exhibit inflexibility for fine-grained communication and non-trivial fused collectives. Domain-specific languages including Triton Distributed and TileLink attempt to address these limitations but struggle to adapt to rapid networking improvements and architectural diversity. Hand-tuned operator approaches such as DPP, Comet, Ring Attention, and FlashDoe achieve peak performance but lack scalability - adapting a single precision to another architecture requires 5-6 months of engineering effort. Furthermore, the multi-GPU problem space expands combinatorially: standard transformer architectures can parallelize across data, sequence, tensor, context, layer, pipeline, and expert dimensions, each with distinct communication patterns.

3. Core Analysis

3.1 Fundamental Principles of Multi-GPU Kernel Design

The analysis identifies three distinct intra-GPU transfer mechanisms, each suited to different communication patterns. First, the copy engine provides CPU-initiated transfers that achieve peak bandwidth for large messages without consuming GPU registers or processors. Second, the Tensor Memory Accelerator (TMA) enables device-initiated transfers that saturate NVLink bandwidth using relatively small message sizes, enabling fine-grain communication while consuming few registers. Third, register-level PTX instructions (LDST) facilitate in-network reductions via NVSwitch, leveraging compute capabilities integrated directly into the interconnect fabric.

Two scheduling paradigms govern the overlap of compute, memory, and communication operations. IntraSM scheduling assigns different warps within a single streaming multiprocessor to handle compute versus communication concurrently, requiring aligned computational patterns. This approach excels in scenarios such as matrix multiplication (GeMM) combined with AllReduce operations. Conversely, InterSM scheduling dedicates different processors to specialize in compute, communication, or memory operations, proving effective when patterns misalign or NVLink traversal complexity increases. This paradigm demonstrates superior performance in GeMM combined with ReduceScatter operations, particularly when leveraging in-network reductions.

3.2 The Parallel Kittens Framework

Parallel Kittens encapsulates the identified design principles into a minimal set of programming primitives and templates. The framework requires approximately a dozen additional lines of code beyond single-GPU kernel implementations to incorporate multi-GPU primitives. This abstraction provides developers with explicit control over buffering and synchronization between data senders and receivers - a critical requirement absent from higher-level communication libraries.

The framework achieves state-of-the-art performance across data parallelism, sequence parallelism, and expert parallelism compared to strong reference baselines. Deployment in production environments at Together AI, Cursor, and other organizations validates its practical applicability. The framework's design philosophy prioritizes developer flexibility over automated optimization, recognizing that effective multi-GPU kernels require explicit reasoning about hardware-specific trade-offs rather than one-size-fits-all abstractions.

3.3 Evaluating Model Capabilities on Multi-GPU Kernel Generation

Parallel Kernel Bench provides a systematic evaluation framework comprising 87 problems drawn from GitHub repositories, optimized library implementations, and DSL-based multi-GPU kernels. The benchmark presents models with unoptimized PyTorch reference implementations using torch.distributed and NCCL, along with system topology specifications detailing ranks and hardware configurations. Models must generate performance-optimized CUDA kernels using unified virtual addressing. The taxonomy covers representative parallelism patterns including data, sequence, tensor, context, layer, pipeline, and expert parallelism, along with their compositions.

Performance is measured through two metrics: Pass@K indicates correct kernels after K attempts, while Fast1@K measures correct kernels achieving at least 1x speedup over PyTorch + NCCL baselines. Zero-shot evaluation reveals that the best frontier model (GPT-4.5) solves 28 of 87 problems, with 22 achieving speedups over baseline implementations. Multiple sampling improves correctness to 36 solutions, but Fast1 performance plateaus at approximately 31%, indicating that additional sampling does not improve the quality of performance optimizations.

3.4 Model Reasoning Limitations

Analysis of model-generated solutions reveals success patterns concentrated in familiar, internet-represented patterns such as collective primitives, tensor parallel matrix multiplications, and Ulyses-style context parallelism. Models struggle with core architectural trade-offs including collective ordering, data partitioning strategies, IntraSM versus InterSM scheduling decisions, and transfer mechanism selection. Notably, models rarely utilize register transfer instructions or Tensor Memory Acceleration when writing kernels, despite these mechanisms being critical for achieving peak performance.

Evaluation of a coding agent (Claude with Gemini 3 Pro backend and bash environment access) demonstrates improvement from 24 to 35 of 87 problems solved, with 26 achieving speedups over baseline. However, performance plateaus with additional computational scaling, suggesting that current agentic approaches do not fundamentally address the underlying reasoning limitations. Models do not currently understand how to reason through the identified trade-offs even when provided relevant context, indicating a fundamental gap in architectural reasoning capabilities.

4. Technical Insights

The hardware characteristics of modern GPUs establish hard constraints on kernel design. H100 GPUs provide register memory bandwidth of 130 terabytes per second but limited capacity, while HBM offers 80-192 gigabytes capacity at substantially lower bandwidth. Modern AI GPUs incorporate 100-200 streaming multiprocessors, each representing an independent compute unit. NVSwitch integrates compute capabilities directly into the interconnect fabric, enabling in-network reductions and multicast operations that reduce communication overhead.

Transfer mechanism selection critically impacts performance. The copy engine reaches peak bandwidth for large message sizes without consuming GPU resources, making it suitable for bulk transfers between computation phases. TMA saturates NVLink bandwidth with smaller messages, enabling fine-grain communication patterns where compute and communication must be tightly interleaved. Register-level instructions enable the lowest-latency communication but require careful orchestration with NVSwitch capabilities to achieve in-network reductions.

The choice between IntraSM and InterSM scheduling depends on computational pattern characteristics. IntraSM scheduling minimizes synchronization overhead when compute and communication patterns align naturally, allowing warps within a single processor to execute concurrently. InterSM scheduling provides greater flexibility when patterns misalign or when NVLink topology complexity makes communication routing non-trivial, at the cost of inter-processor synchronization overhead.

Parallel Kittens demonstrates that a small set of carefully designed primitives can capture the essential trade-offs in multi-GPU kernel design. The framework's requirement of approximately a dozen additional lines of code suggests that the abstraction level appropriately balances expressiveness with simplicity. Production deployments have already generated novel kernels including vocabulary parallel filtering for the Nemo architecture, context parallelism for the Hyena architecture, and IOU suppression for the SAM 3 video segmentation model.

5. Discussion

The findings reveal a fundamental tension in multi-GPU kernel development: while a small set of principles governs effective design, current frontier AI models cannot reliably reason through these trade-offs. The 32% success rate on Parallel Kernel Bench stands in stark contrast to model performance on other technical domains, suggesting that multi-GPU optimization requires forms of architectural reasoning not captured by current training distributions. The concentration of model successes in familiar internet-represented patterns indicates that models primarily succeed through pattern matching rather than principled reasoning about hardware characteristics.

The asymmetric improvement of compute versus communication capabilities - 7.2x for tensor cores versus 3x for intra-node and 2x for inter-node communication - suggests that multi-GPU optimization will increase in importance. As scaleup domains expand to 576 GPUs per system and workloads increasingly disaggregate across heterogeneous hardware backends, the combinatorial explosion of parallelization strategies will make systematic approaches essential. The diversity of networking stacks across vendors (AMD XGMI, TPU 3D Torus, NVIDIA NVSwitch) further complicates the development of portable optimization strategies.

Future research directions include developing architectures that grow with evolving networking stacks, addressing larger scaleup domains, and adapting to the shift from scaleout to massive on-chip memory structures. The plateau in coding agent performance suggests that incremental improvements to sampling and environmental feedback will not suffice; fundamentally new approaches to architectural reasoning may be required. The success of Parallel Kittens in production environments indicates that human-designed abstractions currently outperform automated approaches, but the scalability limitations of hand-tuning motivate continued investigation of AI-assisted kernel development.

6. Conclusion

This synthesis establishes that multi-GPU kernel development has emerged as a critical bottleneck in modern AI systems, driven by asymmetric improvements in compute versus communication capabilities. Three transfer mechanisms (copy engine, TMA, register-level instructions) and two scheduling paradigms (IntraSM, InterSM) constitute the fundamental design space for effective multi-GPU kernels. The Parallel Kittens framework demonstrates that these principles can be encapsulated in minimal primitives requiring approximately a dozen additional lines of code, achieving state-of-the-art performance across diverse parallelization strategies.

Evaluation through Parallel Kernel Bench reveals that frontier AI models achieve only 32% success rates on multi-GPU kernel generation, with performance plateaus indicating fundamental reasoning limitations rather than sampling insufficiency. Models succeed primarily on familiar patterns rather than through principled architectural reasoning, rarely utilizing critical mechanisms such as TMA or register-level transfers. These findings suggest that multi-GPU optimization represents a domain where human expertise, properly abstracted into frameworks like Parallel Kittens, currently surpasses automated approaches. As systems scale to hundreds of GPUs and workloads diversify across heterogeneous backends, developing AI systems capable of architectural reasoning about communication-computation trade-offs remains an important open challenge.

Sources

About the Author

Sean Weldon is an AI engineer and systems architect specializing in autonomous systems, agentic workflows, and applied machine learning. He builds production AI systems that automate complex business operations.

LinkedIn | Website | GitHub