用多轮强化学习生成更高效、正确的CUDA内核
Kevin: Multi-Turn RL for Generating CUDA Kernels
- 设计多轮强化学习框架，模拟专家迭代优化过程
- 生成代码正确率从56%提升至82%，速度比基线快1.10倍
- 适合需要高性能GPU计算的开发者和研究者
编写GPU内核是一项挑战性任务，对人工智能系统的效率至关重要。该过程具有高度迭代性：领域专家编写代码并根据执行反馈持续优化。同时，其奖励可验证，如正确性和加速比，使其成为应用强化学习的理想场景。为显式融入这一迭代特性，我们提出一种灵活的多轮强化学习方法，解决了真实场景中的独特挑战，如长轨迹学习与跨轮次的有效奖励分配。我们构建了Kevin——首个基于多轮强化学习训练的CUDA内核生成与优化模型。在评估中，Kevin相比基线模型QwQ-32B，纯CUDA生成的代码正确率从56%提升至82%，平均加速比从0.53x增至1.10x（相对于PyTorch Eager），超越前沿模型o4-mini（0.78x）。进一步分析表明，串行迭代优化优于并行采样；给予更多优化轮次时，Kevin的性能提升速率更高。
原文摘要 · Abstract (English)
Writing GPU kernels is a challenging task and critical for AI systems' efficiency. It is also highly iterative: domain experts write code and improve performance through execution feedback. Moreover, it presents verifiable rewards like correctness and speedup, making it a natural environment to apply Reinforcement Learning (RL). To explicitly incorporate the iterative nature of this process into training, we develop a flexible multi-turn RL recipe that addresses unique challenges encountered in real-world settings, such as learning from long trajectories and effective reward attribution across turns. We present Kevin - K(ernel D)evin, the first model trained with multi-turn RL for CUDA kernel generation and optimization. In our evaluation setup, Kevin shows significant gains over its base model (QwQ-32B), improving correctness of generated kernels (in pure CUDA) from 56% to 82% and mean speedup from 0.53x to 1.10x of baseline (PyTorch Eager), and surpassing frontier models like o4-mini (0.78x). Finally, we study its behavior across test-time scaling axes: we found scaling serial refinement more beneficial than parallel sampling. In particular, when given more refinement turns, Kevin shows a higher rate of improvement.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。