用多个智能体协作优化GPU内核,速度提升最高16倍。
STARK: Strategic Team of Agents for Refining Kernels
- 设计多智能体协同框架,模拟专家工程师的优化流程。
- 在KernelBench上实现正确率提升,最快比基线快16倍。
- 适合对高性能计算和自动化代码优化感兴趣的开发者。
GPU内核的效率是现代AI发展的核心,但其优化因内存层次结构、线程调度和硬件特性之间的复杂交互而困难且耗时。尽管大语言模型(LLMs)为自动代码生成带来新机遇,现有方法大多将LLM视为一次性生成器或简单的改写工具,难以应对内核优化的非规则性。本文提出一种基于LLM的智能体框架,通过多智能体协作、指令引导、动态上下文管理和战略搜索,系统探索设计空间。该框架模仿专家工程师的工作流程,使LLM能够推理硬件权衡、融入性能分析反馈,并迭代优化内核。我们在KernelBench(一个面向LLM的内核优化基准)上评估,结果表明:相比基线智能体,本系统能生成基线常失败的正确解,且实现最高达16倍的运行速度提升。这些结果凸显了智能体式LLM框架在实现全自动、可扩展的GPU内核优化方面的潜力。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。