计算机科学 > 机器学习
[提交于 2025年7月16日
]
标题: 凯文:生成CUDA内核的多轮强化学习
标题: Kevin: Multi-Turn RL for Generating CUDA Kernels
摘要: 编写GPU内核是一项具有挑战性的任务,对AI系统的效率至关重要。 它也是高度迭代的:领域专家编写代码并通过执行反馈来提升性能。 此外,它提供了可验证的奖励,如正确性和加速比,使其成为应用强化学习(RL)的自然环境。 为了在训练中显式地融入这一过程的迭代特性,我们开发了一个灵活的多轮RL方案,解决了现实环境中遇到的独特挑战,例如从长轨迹中学习以及在各轮中的有效奖励分配。 我们提出了Kevin - K(ernel D)evin,第一个使用多轮RL训练的CUDA内核生成和优化模型。 在我们的评估设置中,Kevin相比其基础模型(QwQ-32B)取得了显著提升,将生成的内核的正确性(纯CUDA)从56%提高到82%,平均加速比从0.53倍提升至1.10倍的基准(PyTorch Eager),并且超越了前沿模型如o4-mini(0.78倍)。 最后,我们研究了其在测试时扩展轴上的行为:我们发现扩展串行优化比并行采样更有益。 特别是,当提供更多的优化轮次时,Kevin表现出更高的改进率。
文献和引用工具
与本文相关的代码,数据和媒体
alphaXiv (什么是 alphaXiv?)
CatalyzeX 代码查找器 (什么是 CatalyzeX?)
DagsHub (什么是 DagsHub?)
Gotit.pub (什么是 GotitPub?)
Hugging Face (什么是 Huggingface?)
带有代码的论文 (什么是带有代码的论文?)
ScienceCast (什么是 ScienceCast?)
演示
推荐器和搜索工具
arXivLabs:与社区合作伙伴的实验项目
arXivLabs 是一个框架,允许合作伙伴直接在我们的网站上开发和分享新的 arXiv 特性。
与 arXivLabs 合作的个人和组织都接受了我们的价值观,即开放、社区、卓越和用户数据隐私。arXiv 承诺这些价值观,并且只与遵守这些价值观的合作伙伴合作。
有一个为 arXiv 社区增加价值的项目想法吗? 了解更多关于 arXivLabs 的信息.