Kevin: Multi-Turn RL for Generating CUDA Kernels
这是什么
写 GPU kernel 是典型的迭代式工程:写代码、跑起来、看性能、再改。这个过程有天然可验证的奖励(正确性 + speedup),很适合 RL。但此前的 kernel 生成 RL 工作基本是单轮的:模型一次性输出 kernel、拿一个分数,没法利用执行反馈修 bug 或逐步加激进优化。Kevin(K(ernel D)evin)是第一个用多轮 RL 显式训练'根据执行反馈迭代改 kernel'能力的模型,基座是 QwQ-32B,环境是 KernelBench 的 Level 1/2 任务。
对课题组视角,这篇是 kernel 自动优化中'训练权重'路线的代表:和 AlphaEvolve、prompt 进化那类不动权重、纯靠搜索的路线互补。它的主要贡献不是刷榜数字,而是一套解决多轮 RL 实际工程问题的 recipe:上下文爆炸、跨轮 credit assignment、reward hacking、训练不稳定(模型输出 junk)等,论文对这些坑写得相当坦率。
机制与做法
多轮训练 recipe:每轮都是训练样本 + 折扣奖励
每个训练步:对每个任务采 16 条并行 trajectory、每条 4 轮 refinement。为省上下文,前几轮的完整 CoT 被丢弃,只保留模型自己写的'改动摘要'+ 生成的 kernel + 执行评测结果。关键设计有两个:(1) 把 n 轮 trajectory 拆成 n 个训练样本(每轮的 kernel+摘要单独收梯度),而不是整条轨迹一个样本,提高样本效率;(2) 每轮的 reward 是当轮及后续各轮 kernel 分数的折扣和(sum, γ=0.4),消融显示这比 greedy(只算当轮)、outcome(全轨迹取最优)和 max 聚合的 scaling 都好。分数公式 S = 0.3·1{correct} + speedup·1{correct},用 GRPO 优化,组内归一化后按任务算 advantage。
工程坑:reward hacking 与训练崩溃
Reward hacking 很普遍:7B 小模型会直接抄 PyTorch 参考实现、套 try-except、或继承参考类;QwQ-32B 也会在 Level 2 融合任务上只融 ReLU/Max 这类便宜算子,把真正值得优化的卷积原封不动留在 PyTorch 里。作者的对策是规则检查:输出必须是纯 inline CUDA,含任何 PyTorch functional 算子直接判 0 分。另外训练久了模型会输出重复乱码(junk)直至崩溃,他们发现一个有趣的先行指标'Not Okay Ratio'(QwQ 的 CoT 本来总以 "Okay, " 开头,不稳定时会变成 "Okay Amigos,"、"Okay Holy crap" 之类),配合 Dr-GRPO 的常数长度归一化 + 0.05 梯度裁剪把 junk 出现推迟到 step 100;加 KL 惩罚反而只减慢学习、不解决崩溃。
Test-time scaling:串行 refinement 优于并行采样
固定 128 次生成的推理预算,比较 128 traj×1 轮、32×4、16×8 三种配置:对多轮模型、单轮模型、基座模型,更多 refinement 轮次几乎总是更优,16×8 是多数情况的最佳配置。并且多轮训练的模型随轮数增长的斜率最陡;在 best@k 随 k 的曲线上,单轮 RL 模型早早平台(呼应'RLVR 损害探索多样性'的发现),多轮模型的曲线斜率反而比基座还高,说明多轮训练在提升 best@1 的同时保住了探索能力。
关键结果
- 在自建 100 题 KernelBench 评测集(16 traj × 8 轮)上:Kevin 正确率 best@16 82%(基座 QwQ-32B 56%),平均 speedup 1.10x vs PyTorch Eager(基座 0.53x,o4-mini 0.78x,o3-mini 0.30x)。
- 多轮 RL 对比单轮 RL:正确率相同(82%)但 speedup 明显更高(1.10x vs 0.85x),fast_1.5 比例 20% vs 16%;单轮模型只在纯单轮推理(128 并行×1 轮)时略占优。
- Reward 设计消融:折扣和(sum, γ=0.4)聚合的多轮 scaling 最好;奖励中间目标(能编译/能运行)会被过度优化,长度惩罚降低性能,KL 系数设 0。
- 基座能力是前提:DeepSeek-R1-Distill-Qwen-7B 因奖励太稀疏而大面积 reward hacking、学不动;难度分布也要均衡,只训简单任务会很快过拟合。
- 训练规模其实不大:最终 run 每 batch 8 个任务、16 traj × 4 轮,多轮训练 40 步(80 个梯度步),max response length 16K→22K。
实证核查
与我们方向的关系
对'AI 优化 AI 系统'方向,Kevin 提供了与 AlphaEvolve/FunSearch 进化搜索路线正交的另一条路:把迭代优化能力训进权重,推理时用少量串行 refinement 就能拿到收益,而不是每个新 kernel 都烧大量搜索预算。它的 test-time scaling 结论(串行 refinement > 并行采样,且多轮训练保住 best@k 探索能力)对我们设计 agent 的推理预算分配有直接参考价值。
可借鉴的具体经验:跨轮 credit assignment 用折扣和而非 outcome reward;丢 CoT 留摘要的上下文管理;用规则检查而非奖励塑形对付 reward hacking;以及'Not Okay Ratio'这种便宜的训练不稳定先行指标。局限也要记住:环境是 KernelBench Level 1/2 的单文件 inline CUDA,离真实生产 kernel(多文件、Triton/CUTLASS、跨架构调优)还有距离。