摘要
在多个GPU架构上优化CUDA代码具有挑战性,因为要达到峰值性能需要对日益复杂的硬件特定优化空间进行广泛探索。传统编译器受限于固定启发式方法,而微调大语言模型(LLMs)成本较高。现有的CUDA代码优化代理工作流难以聚合先前探索的知识,导致采样偏差和次优解。本文提出KernelBlaster,一种基于记忆增强的上下文强化学习(MAIC-RL)框架,旨在提升基于LLM的GPU编码代理的CUDA优化搜索能力。该框架通过构建可检索的持久CUDA知识库,使代理能够从经验中学习,并对未来任务做出系统化的决策。我们提出了一种基于性能分析和文本梯度的新型代理流程,以实现在不同代GPU架构上的高性能CUDA生成与优化。与PyTorch基线相比,我们的方法在KernelBench Level 1、2和3上分别实现了1.43倍、2.50倍和1.50倍的几何平均加速。我们开源了KernelBlaster作为代理框架,并附带测试套件、验证组件和可复现的评估流程。
AI 推荐理由
论文核心提出基于记忆增强的强化学习框架,用于提升CUDA优化能力,直接涉及记忆机制的设计与应用。
论文信息