STARK: Strategic Team of Agents for Refining Kernels
Juncheng Dong, Yang Yang, Tao Liu, Yang Wang, Feng Qi, Vahid Tarokh, Kaushik Rangadurai, Shuang Yang
TL;DR
GPU kernel optimization is highly challenging due to irregular memory hierarchies and hardware-specific behavior. STARK proposes a novel agentic framework that decomposes the task into planning, coding, and debugging through a grounded-instruction workflow and dynamic context, guided by a strategic tree-based search. Across KernelBench, STARK achieves perfect success rates and substantial speedups over Torch baselines, with up to $16\times$ faster runtimes, and outperforms prior single-agent and reflexion-based approaches. The results demonstrate that coordinated LLM agents can effectively navigate complex kernel design spaces, enabling scalable, automated optimization with practical impact for AI workloads. The work also suggests directions for broader operator support, hardware architectures, and co-design of AI algorithms and infrastructure.
Abstract
The efficiency of GPU kernels is central to the progress of modern AI, yet optimizing them remains a difficult and labor-intensive task due to complex interactions between memory hierarchies, thread scheduling, and hardware-specific characteristics. While recent advances in large language models (LLMs) provide new opportunities for automated code generation, existing approaches largely treat LLMs as single-shot generators or naive refinement tools, limiting their effectiveness in navigating the irregular kernel optimization landscape. We introduce an LLM agentic framework for GPU kernel optimization that systematically explores the design space through multi-agent collaboration, grounded instruction, dynamic context management, and strategic search. This framework mimics the workflow of expert engineers, enabling LLMs to reason about hardware trade-offs, incorporate profiling feedback, and refine kernels iteratively. We evaluate our approach on KernelBench, a benchmark for LLM-based kernel optimization, and demonstrate substantial improvements over baseline agents: our system produces correct solutions where baselines often fail, and achieves kernels with up to 16x faster runtime performance. These results highlight the potential of agentic LLM frameworks to advance fully automated, scalable GPU kernel optimization.