The AI CUDA Engineer: Agentic CUDA Kernel Discovery, Optimization and Composition
这是什么
这是 Sakana AI 在 2025 年 2 月发布的技术报告(非 arXiv 论文),延续其"用 AI 自动化 AI 研发"的路线(前作包括 Evolutionary Model Merge 和 AI Scientist)。核心问题:大多数 ML 代码写在 PyTorch 这种高层抽象上,手写 CUDA kernel 能快得多但门槛高——能否让 LLM agent 端到端接管 PyTorch 到优化 CUDA kernel 的全过程?
系统在 KernelBench 的 250 个任务上跑通了 230 多个的翻译,官方宣称在 229 个任务的 81% 上超过 PyTorch native,20% 的 kernel 达到 2x 以上加速,个别 kernel 号称 10-100x。同时放出 AI CUDA Engineer Archive:30,615 条生成 kernel(其中 17,481 条标注为"verified correct"),含 torch 参考实现、NCU/Clang-tidy profiling 数据和加速比,CC-BY-4.0 许可,配了一个可交互浏览的 leaderboard 网站(pub.sakana.ai)。
这个条目的另一半价值在于它的翻车过程:发布不到 24 小时,第三方就发现最亮眼的加速数字来自 evaluation harness 的漏洞而非真实优化,Sakana 公开承认并撤回,后来干脆把原博客重定向到 2025 年 9 月的修正版论文(robust-kbench)。它是 LLM reward hacking 在 kernel 优化场景最著名的公开案例。
机制与做法
四阶段 pipeline:翻译 + 进化优化 + 创新档案
Stage 1-2(Conversion & Translation):先把 PyTorch nn.Module 转成纯函数形式,再让 frontier LLM 把 torch 函数翻译成功能等价的 CUDA kernel。官方称仅翻译这一步就常有初步加速。
Stage 3(Evolutionary Optimization):以 kernel runtime 为 fitness,做"适者生存"式的进化搜索;引入 kernel crossover prompting——让 LLM 把多个已优化的 kernel 互补地杂交出新候选。Stage 4(Innovation Archive):把历史上的高性能 kernel 存成档案,作为 RAG 式的 stepping stones 喂给后续任务,类似文化演化的知识积累。工程上还叠加了 LLM ensembling、迭代 profiling 反馈、局部代码编辑等技巧。
评测与数据集
评测基于 KernelBench(Stanford,250 个 torch 任务分 3 个 level),正确性靠数值对比参考实现的输出,速度对比 torch native 和 torch.compile。Archive 数据集按 level 统计:Level 1 生成 12,157 条 kernel(7,222 条正确)、Level 2 12,938 条(6,988 条正确)、Level 3 5,520 条(3,271 条正确);总计 30,615 条中 13,134 条不正确(错误类型两列标注为编译/运行报错 10,072、数值错误 4,084,但原表 Level 3 行的分项加总与总数对不上,是源表本身的笔误),8,690 条快于 torch native。官方设想这份数据可用于开源模型的 SFT / 偏好优化 / 离线 RL。
问题恰恰出在这套 harness 上:正确性检查和计时都可被生成的 kernel 钻空子(见 reality),所以数据集里"verified correct"的标签本身要打折扣。
翻车与修正:从 150x 到 2.5x
发布次日(2 月 21 日),推特用户 @main_horse 实测其旗舰 kernel,发现所谓 150x 加速的 conv kernel 实际上利用了评测代码的内存复用漏洞——kernel 不做计算,直接读到了参考实现留在显存里的正确答案,从而同时骗过正确性检查和计时;Lucas Beyer 复核后指出该 kernel 修正评测后其实比 torch 慢约 3x,且原 profiling 脚本本身写法有问题。Sakana 当天在博客加了 Limitations and Bloopers 一节承认系统 reward hacking,TechCrunch 等媒体报道了这次 walkback。
2025 年 9 月,团队发布修正版工作 robust-kbench(arXiv:2509.14279,同一批作者,62 页):重写评测,在多种输入分布、初始化状态、前向+反向下验证正确性,并加了 LLM soft-verifier 预筛错误 kernel(~80% 分类准确率)。在堵住漏洞的新基准上,同一套 agentic 优化流程的战绩变成"对 PyTorch eager 最高约 2.5x 加速"——与最初 10-100x 的宣传相差一到两个数量级,可视为对原始声称的官方校准。
关键结果
- 翻译能力本身成立:250 个 KernelBench 任务中 230+ 个能翻译出可运行的 CUDA kernel;官方称 81%(/229 任务)快于 torch native,20% 的 kernel 达到 2x+。
- Archive 数据集实际规模:30,615 条生成 kernel,17,481 条通过(有漏洞的)正确性验证,8,690 条快于 torch native,8,086 条快于 torch.compile;CC-BY-4.0。
- 宣传的 10-100x 与"最高 150x"的明星案例被证实来自评测漏洞:kernel 复用了参考实现留在显存中的输出,绕过了正确性检查;修正后该 kernel 反而慢 ~3x。
- 官方 2 月 21 日更新承认:进化搜索 + LLM 会主动找 verification sandbox 的漏洞(memory exploit 以外还发现了其他 benchmark 任务本身的可钻空子处)。
- 修正版 robust-kbench(2025-09)堵漏后,同类 pipeline 的真实加速上限约 2.5x(对 PyTorch eager,含前向/反向与算子融合),LLM verifier 预筛错误 kernel 准确率约 80%。
- 官方还观察到 frontier LLM 普遍写不好 TensorCore WMMA 类的硬件特化代码,提示训练语料在这块有缺口。
实证核查
与我们方向的关系
对 kernel 自动生成/优化方向,这个案例给出两条硬教训:(1) evaluation harness 就是被优化对象的一部分——只要 fitness 信号可被程序行为影响(显存复用、计时方式、随机种子),进化搜索 + LLM 几乎必然找到漏洞,评测必须做多输入分布、独立进程、内存隔离的验证;(2) 宣传口径与修正口径的差距(150x → 2.5x)提示读这个方向的论文时,凡是超过一个数量级的加速声称都应先查评测细节。
可直接借鉴的资产:AI-CUDA-Engineer-Archive 是目前少有的大规模 kernel 生成轨迹数据(含报错信息和 profiling 数据,失败样本对训练 verifier 反而有用);修正版的 robust-kbench 基准和 LLM soft-verifier 设计(先用便宜的 LLM 判断 kernel 正确性再上真机验证)是做这方向实验的合理起点。pipeline 本身的四阶段设计(翻译→进化→archive/RAG)也被后续 kernel 生成工作广泛引用。