Kevin：用于生成 CUDA 内核的多回合强化学习
机器学习
2025-07-17 v1 人工智能
性能
软件工程
摘要
编写 GPU 内核是一项具有挑战性的任务，对于 AI 系统的效率至关重要。它也是高度迭代的过程：专家编写代码并通过执行反馈来提高性能。此外，它具有可验证的奖励，如正确性和加速，使其成为应用强化学习（RL）的自然环境。为将这一迭代性质明确地纳入训练，我们开发了一种灵活的多回合 RL 配方，解决在实际场景中遇到的独特挑战，如来自长轨迹的学习以及跨回合有效的奖励归因。我们提出了 Kevin - K(ernel D)evin，即第一个在多回合 RL 下训练的模型，用于 CUDA 内核的生成和优化。在我们的评估设置中，Kevin 相对于其基础模型（QwQ-32B）取得了显著的改进，将生成的 CUDA 内核正确性从 56% 提升到 82%，平均加速从基线（PyTorch Eager）的 0.53x 提升到 1.10x，并超过了前沿模型如 o4-mini（0.78x）。最后，我们研究了其在测试时间扩展轴上的行为：我们发现串行细化的扩展比并行抽样更有益。特别是，当给定更多的细化回合时，Kevin 显示出更高的改进率。
引用
@article{arxiv.2507.11948,
title = {Kevin: Multi-Turn RL for Generating CUDA Kernels},
author = {Carlo Baronio and Pietro Marsella and Ben Pan and Simon Guo and Silas Alberti},
journal= {arXiv preprint arXiv:2507.11948},
year = {2025}
}