LLMs Fail to Write Fast Multi-GPU Kernels

Simran Arora from Together AI discusses the challenges of multi-GPU kernel development and why current LLMs struggle to optimize them, despite ongoing research.

6 min read
Presentation slide showing the title 'Can LLMs Write Fast Multi-GPU Kernels?' with speaker Simran Arora of Together AI.
AI Engineer
Visual TL;DR
Multi-GPU BottleneckDriver
communication across multiple GPUs is a critical hurdle for AI performance
From the article 6 mentionsArora explained that while significant advancements have been made in optimizing single-GPU kernels and memory-efficient architectures, the bottleneck has shifted to multi-GPU communication.
Simran Arora's ResearchCore
From the articleSimran Arora, a principal scientist at Together AI and an incoming professor at Caltech, recently shed light on the challenges and opportunities in developing efficient multi-GPU AI kernels during a presentation titled 'Can LLMs Write Fast Multi-GPU Kernels?'.
LLMs StruggleDriver
current LLMs fail to write fast multi-GPU kernels for optimization
From the article 4 mentionsArora highlighted that while LLMs can often generate compilable CUDA code, they struggle with the nuanced reasoning required for optimizing multi-GPU kernels.
Kernel Optimization GapOutcome
LLM limitations prevent efficient orchestration of multi-GPU communication
From the articleShe concluded by emphasizing the ongoing challenges and the potential for future research to unlock more sophisticated AI-driven kernel optimization.
AI Hardware UtilizationContext
From the article 3 mentionsIn the ever-evolving landscape of AI hardware, optimizing performance for large language models and other complex AI workloads is paramount.
Simran Arora's ResearchCore
From the articleSimran Arora, a principal scientist at Together AI and an incoming professor at Caltech, recently shed light on the challenges and opportunities in developing efficient multi-GPU AI kernels during a presentation titled 'Can LLMs Write Fast Multi-GPU Kernels?'.
Single-GPU ProgressEffect
advancements in single-GPU kernels and memory-efficient architectures
From the article 2 mentionsThe pace of innovation in GPU compute and memory has outstripped the progress in communication hardware.
Communication Hardware LagDriver
From the articleThe pace of innovation in GPU compute and memory has outstripped the progress in communication hardware.
Multi-GPU BottleneckDriver
communication across multiple GPUs is a critical hurdle for AI performance
From the article 6 mentionsArora explained that while significant advancements have been made in optimizing single-GPU kernels and memory-efficient architectures, the bottleneck has shifted to multi-GPU communication.
LLMs StruggleDriver
current LLMs fail to write fast multi-GPU kernels for optimization
From the article 4 mentionsArora highlighted that while LLMs can often generate compilable CUDA code, they struggle with the nuanced reasoning required for optimizing multi-GPU kernels.
Kernel Optimization GapOutcome
LLM limitations prevent efficient orchestration of multi-GPU communication
From the articleShe concluded by emphasizing the ongoing challenges and the potential for future research to unlock more sophisticated AI-driven kernel optimization.
Contents(4)

In the ever-evolving landscape of AI hardware, optimizing performance for large language models and other complex AI workloads is paramount. Simran Arora, a principal scientist at Together AI and an incoming professor at Caltech, recently shed light on the challenges and opportunities in developing efficient multi-GPU AI kernels during a presentation titled 'Can LLMs Write Fast Multi-GPU Kernels?'.

LLMs Fail to Write Fast Multi-GPU Kernels - AI Engineer
LLMs Fail to Write Fast Multi-GPU Kernels, from AI Engineer

The Bottleneck in AI Hardware Utilization

Arora explained that while significant advancements have been made in optimizing single-GPU kernels and memory-efficient architectures, the bottleneck has shifted to multi-GPU communication. She highlighted that innovations like FlashAttention, sparse attention mechanisms, and efficient DSLs have pushed the boundaries of AI performance, but the efficient orchestration of communication across multiple GPUs remains a critical hurdle.

The pace of innovation in GPU compute and memory has outstripped the progress in communication hardware. Arora cited data showing that while NVIDIA A100s to B200s tensor core speeds improved by 7.2x, intra-node communication bandwidth only increased by 3x, and inter-node communication by a mere 2x.

Current Approaches and Their Limitations

Arora outlined three primary approaches for multi-GPU kernel development:

  • Off-the-shelf APIs: Libraries like NCCL and RCCL (for AMD) are popular but lack flexibility. They are optimized for bulk transfers of large, contiguous data, making them less suitable for fine-grained communication or fusing complex collectives.
  • Compilers: Frameworks like Triton Distributed and TileLink aim to simplify kernel mapping to hardware. However, Arora noted that these frameworks struggle to adapt to the rapid pace of networking improvements, as demonstrated by Triton Distributed's inability to efficiently adjust to newer architectures like H100s after being tuned for H800s.
  • Low-level primitives: Hand-tuning specific AI operators, such as those found in popular works like DPP, Comet, and FlashDoE, can achieve peak performance. However, this approach is time-consuming and difficult to scale, with a single precision implementation potentially taking months to adapt to another.

The Role of Foundation Models and LLMs

Arora then posed a critical question: Can foundation models, particularly LLMs, reason about these fundamental trade-offs to simplify the development of multi-GPU kernels? To investigate this, Together AI developed 'Parallel Kittens,' a set of minimal primitives and patterns for multi-GPU kernels, and 'Parallel Kernel Bench' (PKB), a benchmark designed to evaluate LLM performance in generating these kernels.

The PKB benchmark tasks models with rewriting unoptimized PyTorch kernels into performance-optimized CUDA kernels. The evaluation metrics include 'Pass@K' (correct kernels generated after K attempts) and 'Fast@1K' (solutions that are both correct and faster than a baseline). In a zero-shot setting, the best-performing frontier models solved only 28 out of 87 problems, with 22 of those solutions showing speedups over the baseline. Even with multiple samples, the 'Fast@1K' performance plateaued around 31%, indicating significant room for improvement.

LLM Limitations in Kernel Optimization

Arora highlighted that while LLMs can often generate compilable CUDA code, they struggle with the nuanced reasoning required for optimizing multi-GPU kernels. Specifically, models faltered in understanding trade-offs related to collective ordering, data partitioning, intra- vs. inter-SM scheduling, and choosing appropriate transfer mechanisms. They often failed to utilize advanced features like register transfer instructions or tensor memory acceleration.

Despite these limitations, Arora expressed optimism, noting that even the best models showed some success, and that agentic loops, which allow models to write, test, and refine their code, can improve performance. She concluded by emphasizing the ongoing challenges and the potential for future research to unlock more sophisticated AI-driven kernel optimization.

© 2026 StartupHub.ai. All rights reserved. Do not enter, scrape, copy, reproduce, or republish this article in whole or in part. Use as input to AI training, fine-tuning, retrieval-augmented generation, or any machine-learning system is prohibited without written license. Substantially-similar derivative works will be pursued to the fullest extent of applicable copyright, database, and computer-misuse laws. See our terms.