用大模型自动优化AI加速器核函数,越用越快。
AccelOpt: A Self-Improving LLM Agentic System for AI Accelerator Kernel Optimization
- 基于记忆的迭代生成,让大模型自主学习优化策略。
- 在Trainium上将核函数性能从49%提升至61%,效率显著提高。
- 成本仅为商用模型的1/26,适合资源有限的研究者使用。
我们提出AccelOpt,一个自进化的大型语言模型代理系统,可自主优化新兴AI加速器的核函数,无需依赖专家提供的硬件特定优化知识。AccelOpt通过迭代生成探索优化空间,并利用优化记忆库整合过往慢-快核函数对的经验与洞察。我们构建了NKIBench,一个基于AWS Trainium加速器、从真实大模型工作负载中提取的具有不同复杂度的核函数基准测试套件,用于评估AccelOpt的有效性。评估结果表明,AccelOpt的能力随时间持续提升,在Trainium 1上将平均峰值吞吐率从49%提升至61%,在Trainium 2上从45%提升至59%。此外,采用开源模型的AccelOpt在性能上达到Claude Sonnet 4水平,但成本仅为后者的1/26。代码已开源:https://github.com/zhang677/AccelOpt。
原文摘要 · Abstract (English)
We present AccelOpt, a self-improving large language model (LLM) agentic system that autonomously optimizes kernels for emerging AI acclerators, eliminating the need for expert-provided hardware-specific optimization knowledge. AccelOpt explores the kernel optimization space through iterative generation, informed by an optimization memory that curates experiences and insights from previously encountered slow-fast kernel pairs. We build NKIBench, a new benchmark suite of AWS Trainium accelerator kernels with varying complexity extracted from real-world LLM workloads to evaluate the effectiveness of AccelOpt. Our evaluation confirms that AccelOpt's capability improves over time, boosting the average percentage of peak throughput from $49\%$ to $61\%$ on Trainium 1 and from $45\%$ to $59\%$ on Trainium 2 for NKIBench kernels. Moreover, AccelOpt is highly cost-effective: using open-source models, it matches the kernel improvements of Claude Sonnet 4 while being $26\times$ cheaper. The code is open-sourced at https://github.com/zhang677/AccelOpt.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。