核心问题:当评测允许每道题最多运行 200 轮,并从整条交互轨迹中挑出最快的正确候选时,96.8% 的 faster rate 说明训练后的策略稳定胜过 torch.compile,还是说明它在足够大的试错预算下最终能搜索到更快实现?

论文卡片

  • 原题: CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation
  • 作者: Weinan Dai、Hanlin Wu 等 16 位作者。
  • 机构: ByteDance Seed、清华大学智能产业研究院(AIR)及双方联合实验室。
  • 版本: 2026-02-27,arXiv v1;cs.LG、cs.AI。
  • 定位: 通过合成训练任务、带验证与 profiling 的 Agent 环境,以及多阶段 PPO 训练,学习 CUDA kernel 生成与优化策略。
  • 阅读入口: alphaXiv · arXiv · PDF 全文 · 项目页

论文讲了什么

CUDA Agent 的出发点是:只在推理时给通用模型增加编译、测试和性能反馈,仍会受限于模型原本的 CUDA 能力;固定的多轮修正模板又限制了 Agent 自己学习调试和搜索策略。因此,作者不只设计一个推理期工作流,而是把整个 CUDA 开发循环放进强化学习训练中。

训练任务来自三步数据管线。作者先从 PyTorch 与 Transformers 库抓取算子,再让 LLM 顺序组合最多五个 PyTorch 算子,最后过滤无法执行、带随机性、输出近似常量以及 eager 耗时不在 1–100 ms 范围内的样本,得到 6,000 个训练问题。为避免直接重合,训练样本若与 KernelBench 测试程序的 AST 相似度超过 0.9 就被删除。

Agent 在 ReAct 风格环境中读取代码、编辑 CUDA 与绑定文件、编译、验证并 profiling。验证和计时脚本由权限保护,禁止调用 PyTorch fallback;每个问题使用五组随机输入检查结果,并进行同步、预热和重复计时。奖励不是直接使用容易受离群值影响的 speedup,而是离散成四档:错误为 −1;正确但未跨过性能门槛为 1;比 eager 快 5% 以上为 2;同时比 eager 和 torch.compile 快 5% 以上为 3。

训练分为单轮 PPO 预热、从正奖励轨迹做 rejection fine-tuning(RFT)、用轨迹结果预训练 critic,最后再进行 agentic PPO。作者认为,这两个初始化阶段缓解了长上下文、多轮 Agent 强化学习中的分布错配与训练崩溃。

主要结果

实验基于 KernelBench Level 1–3,共 250 个问题。CUDA Agent 的 backbone 是 Seed1.6 MoE,具有 23B 激活参数和 230B 总参数;训练进行 150 步,使用 128 张 NVIDIA H20。训练 rollout 最多 150 轮,评测时放宽到 200 轮。Claude Opus 4.5、Gemini 3 Pro、GLM 4.6、Kimi K2 和未训练的 Seed1.6 都在相同 Agent 环境中评测。

模型Pass Rate快于 torch.compile 的任务比例相对 torch.compile 的几何平均加速
Seed1.6 base74.0%27.2%0.69×
Gemini 3 Pro91.2%69.6%1.42×
Claude Opus 4.595.2%66.4%1.46×
CUDA Agent98.8%96.8%2.11×

这里的结果不是某一次生成,也不是轨迹结束时的最后版本。对每个问题,评测从整条交互轨迹中取相对 torch.compile 加速最大的正确候选;几何平均加速也只在正确解上计算。按表 1,CUDA Agent 在 Level 1、2、3 上快于 torch.compile 的比例分别是 97%、100% 和 90%,相应加速为 1.87×、2.80× 和 1.52×。

论文的 leave-one-out 消融显示,去掉 Agent 循环后 faster rate 从 96.8% 降到 14.1%;把离散里程碑奖励换成原始 speedup 奖励后降到 60.4%。去掉 RFT 或 value pretraining 的训练最终都会崩溃,表中分别报告崩溃前最后一次验证的 49.8% 和 50.9%,不能把它们当成稳定训练完成后的最终模型。

证据边界

模型之间共享同一 Agent 循环,这使 CUDA Agent 与通用模型的比较控制了主要工具条件;但“Agent 与 torch.compile”比较的计算预算并不对称。编译器运行一次,Agent 最多进行 200 轮编写、编译、测试和 profiling,再从历史候选中选择最快者。论文没有报告随轮数变化的 pass/faster 曲线、达到最佳解所需的中位轮数、单题 wall-clock 或推理与 GPU 成本。因此,结果有力证明的是该策略与这份测试时搜索预算组成的系统能够找到更快候选,还不能单独说明一次运行的策略稳定性或搜索效率。

正确性门槛也有明确范围。五组随机输入、受保护脚本与 fallback 禁令能拦住许多投机行为,却不等于所有输入上的语义等价;附录中的容差为 atol=1e-2, rtol=1e-2。最快候选的反复选择还会增加遇到测试漏检实现的机会,论文没有报告额外隐藏测试、形式等价检查或候选数量增长时的误通过分析。

数据去重使用 AST 结构相似度阈值,能排除近似拷贝,却无法排除语义相同但语法重写的训练任务,也没有检验由同一 PyTorch 算子库生成的任务与 KernelBench 在算子分布上的接近程度。它支持“没有高 AST 相似样本”,但不足以证明更强意义上的任务分布独立。

论文自己承认没有与 TVM 等更强编译器框架比较,并指出 128 张 H20 与进程级隔离带来的成本会限制可复现性。另一个未解释的不一致是:摘要声称 Level 1–3 的 faster rate 为 100%、100%、92%,而表 1 相对 torch.compile 的数字是 97%、100%、90%;本导读采用定义更明确的表格结果。

重点读哪里

先读 §3.1–3.2,理解 6K 合成任务、Agent 工具闭环和四档奖励;再读 §3.3 与图 3–5,看 RFT 与 value pretraining 如何用于稳定多轮 PPO。随后细读 §4.1、表 1 与表 2,特别留意 200 轮上限、best-in-trajectory 选择和仅对正确解计算加速的口径;最后读 附录 A 与 E,核对 AST 去重方法、硬件规模和作者承认的 baseline 限制。

思想启发

以下是导读者由论文引出的理解角度,不是作者额外验证过的结论。

CUDA Agent 把“模型能力”拆出了一个经常被忽略的维度:模型不只要会写一个 kernel,还要会决定何时检查、怎样根据错误改写、何时继续搜索。表 2 说明执行闭环对最终成绩非常重要;但正因为闭环如此强,最终指标衡量的是策略、工具、迭代上限与候选选择共同形成的系统能力。

这使“比编译器快”至少有两种不同含义。一种是给定相近时间或资源,Agent 更有效地找到实现;另一种是不限制搜索成本时,Agent 的搜索空间里存在比编译器输出更快的实现。论文为第二种含义提供了强证据,对第一种含义则缺少预算曲线。两者都重要,却不能由同一个 best-of-trajectory 数字相互替代。

三个思考问题

1. Best-of-trajectory 何时会把搜索能力误读成稳定能力?

若一个策略在 200 轮中大多能找到一次很快的正确解,但不同运行到达它的时间和概率波动很大,96.8% 的 faster rate还能支持“稳定优化”的结论吗?

展开思考线索

Best-of-trajectory 回答的是“预算内是否曾经找到过”,不是“第几轮通常能找到”或“重复运行是否得到相近结果”。如果缺少按轮数的成功曲线、多个随机种子的方差和达到最优候选的成本分布,搜索空间中存在好解与策略可靠地抵达好解仍是两个判断。

2. 同一 Agent 环境足以构成公平比较吗?

通用模型与 CUDA Agent 使用相同的工具和 200 轮上限,但 torch.compile 只是被调用一次;这种评测公平地回答了什么,又没有回答什么?

展开思考线索

相同 Agent 环境有助于比较不同模型在同一工具条件下的搜索效果。与编译器的比较则回答“投入整套 Agent 搜索后,最终结果能否超过编译器”,并未控制生成成本、墙钟时间或能源。若结论是最终可达性能,这种比较成立;若结论涉及优化效率或可替代性,就还需要成本相称的指标。

3. 奖励的 5% 门槛塑造了怎样的“CUDA 能力”?

离散奖励抑制了噪声和极端 speedup,但把 4.9% 与 0% 放在同一档、5.1% 与更大提升放入另一档;训练出的策略会因此偏向哪些搜索行为?

展开思考线索

里程碑奖励降低了不同任务尺度和计时离群值对 PPO 的影响,这是表 2 支持的优势。与此同时,阈值会制造不连续的信用分配:策略可能更偏好快速跨过门槛,而不是持续探索更高上限;也可能在噪声边缘反复测量。论文展示了最终性能改善,但没有分离“更稳定跨线”与“学到更深硬件规律”各自贡献多少。