AI Research

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.

Daniel Singer
Daniel Singer

Editor, StartupHub.ai

3 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, 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. You may not republish this article in full without a license. Search engines and AI research tools may crawl and summarize for reference. Bulk reproduction or model training requires a license. See our terms.
#Simran Arora#Together AI#Nvidia#GPU#AI Research#LLM#CUDA#Parallel Computing#Deep Learning#Semiconductors
Daniel Singer

Written by

Daniel Singer

Editor, StartupHub.ai

Daniel Singer is the editor of StartupHub.ai, a technology expert and thought leader on AI and its applications across sectors, from fintech and healthcare to developer tooling and consumer software. He writes and tests the tools covered here thoroughly and regularly, and built StartupHub.ai to give founders, operators and buyers a clearer read on what they are actually being sold.

More from Daniel SingerLinkedIn