用强化学习训练的智能体,自动优化CUDA内核,速度超越现有工具。
CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation
- 构建可扩展数据合成与验证环境,实现自动化奖励反馈
- 在KernelBench上比torch.compile快100%、92%,最难点超40%领先
- 适合需要高性能计算优化的研究者和工业开发者
GPU内核优化是现代深度学习的核心,但高度依赖硬件专业知识。尽管大语言模型在通用编程中表现良好,但在CUDA内核生成方面仍逊于编译器系统如torch.compile。现有方法或依赖无训练微调,或在固定多轮执行反馈中调整模型,均无法根本提升模型的内核优化能力,导致性能提升有限。本文提出CUDA Agent,一种大规模代理强化学习系统,通过三个组件建立CUDA内核专长:可扩展的数据合成管道、带有自动验证与性能分析的技能增强开发环境,以及支持稳定训练的强化学习算法。在KernelBench测试中,CUDA Agent在Level-1、Level-2和Level-3分组上分别比torch.compile快100%、100%和92%,在最难的Level-3设置上优于Claude Opus 4.5和Gemini 3 Pro等最强私有模型约40%。
原文摘要 · Abstract (English)
GPU kernel optimization is fundamental to modern deep learning but remains a highly specialized task requiring deep hardware expertise. Despite strong performance in general programming, large language models (LLMs) remain uncompetitive with compiler-based systems such as torch.compile for CUDA kernel generation. Existing CUDA code generation approaches either rely on training-free refinement or fine-tune models within fixed multi-turn execution-feedback loops, but both paradigms fail to fundamentally improve the model's intrinsic CUDA optimization ability, resulting in limited performance gains. We present CUDA Agent, a large-scale agentic reinforcement learning system that develops CUDA kernel expertise through three components: a scalable data synthesis pipeline, a skill-augmented CUDA development environment with automated verification and profiling to provide reliable reward signals, and reinforcement learning algorithmic techniques enabling stable training. CUDA Agent achieves state-of-the-art results on KernelBench, delivering 100\%, 100\%, and 92\% faster rate over torch.compile on KernelBench Level-1, Level-2, and Level-3 splits, outperforming the strongest proprietary models such as Claude Opus 4.5 and Gemini 3 Pro by about 40\% on the hardest Level-3 setting.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。