Yuzhe's Blog

yuzhes

CUDA Agent 论文解读:用 RL 训 LLM 写 CUDA Kernel

CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation
ByteDance Seed + Tsinghua AIR (SIA-Lab),2026
cuda-agent.github.io


写快 GPU kernel 是真正意义上的难。你需要理解内存层次、warp 调度、bank conflict、tensor core 布局,以及大约五十个在不同 GPU 代际之间会变化的微架构细节。大多数工程师——包括大多数 ML 工程师——没有这种知识。他们用库(cuBLAS、cuDNN、FlashAttention),然后祈祷一切都好。

CUDA Agent 的目标是通过强化学习教 LLM 做这件事。头条结果:在 KernelBench 上平均比 torch.compile 快 2.11x,96.8% 的生成 kernel 比编译器基准快。这是个吓人的数字。让我解释他们实际上做了什么,然后解释为什么我认为这个比较部分是误导性的——同时论证核心结论仍然成立。

他们做了什么

三大核心组件

1. 可扩展数据合成(6K 训练 Op)

RL 之前,你需要训练数据。团队合成了大约 6000 个 CUDA 算子实现,覆盖不同类别(逐元素算子、归约、矩阵运算、注意力变体等)。这些和参考实现以及性能分析结果配对,构建了一个从简单单 kernel 任务(Level 1)到复杂多 kernel 序列(Level 3)的课程体系。

合成流水线很关键:他们不是直接用现有的开源 kernel,而是生成了多样化的算子变体,确保训练分布覆盖模型在测试时需要优化的操作类型。这是无聊但关键的基础,所有其他东西都建在上面。

2. 技能增强的执行环境(ReAct + GPU 沙箱)

Agent 不是生成一次代码就祈祷。它在 ReAct 循环中运行——推理、行动、观察——在真实 GPU 执行环境里访问一组工具:

“技能增强”意味着模型还能访问结构化的 CUDA 优化知识——本质上是一份关于 tiling 策略、内存访问模式和优化技巧的精心整理的参考资料,在推理步骤中可以查阅。

这是真正新颖的东西。之前的大多数工作给模型固定次数的尝试加上黑盒反馈(“对/错”、“快/慢”)。CUDA Agent 给了模型真实的性能分析信号——就是人类专家会用的那种。

3. RL 训练流水线(分阶段)

训练部分是论文里最有意思的地方。他们用了分阶段的方式:

  1. 单轮 PPO 预热:从标准 PPO 单轮 kernel 生成开始,得到一个至少能产生有效 CUDA 的基础模型。
  2. Rejection Fine-Tuning(RFT):过滤成功的 rollout 并在上面微调,积累一个有效优化的知识库。
  3. Critic 预训练:训练一个能估计期望加速比的价值函数,这是稳定多轮 RL 所需要的。
  4. 多轮 RL:最后,完整的 agent 循环——多步推理和工具使用,由 critic 提供奖励塑形。

分阶段方式很重要,因为带稀疏奖励的多轮 RL 出了名地不稳定。等到进入完整 agent 循环时,模型已经有了扎实的基础。

结果

在 KernelBench 上:

对比强力闭源模型在 Level-3(最难的基准)上:CUDA Agent 比 Claude 3.5 Sonnet 和 Gemini 1.5 Pro(标准 prompting 方式)高约 40 分。

我的质疑:torch.compile 的比较并不公平

这里我要推回去说几句。

torch.compile 是一个通用优化编译器。它设计用于:

CUDA Agent 生成的 kernel 是针对特定单一 GPU(他们测试用的 H100)优化的。针对 H100 Hopper 架构优化的 kernel 在 A100 Ampere 架构上会更慢。Agent 没有可移植性的概念。

把针对单 GPU 专项手调的 kernel 和通用编译器比较,然后声称有 2.11x 加速,就像把职业赛车手的单圈时间和私家车巡航定速比较。专项的东西当然更快——这就是专项化的全部意义。

更公平的比较应该是:

关于正确性验证:论文用 5 个随机输入检查正确性。对很多算子来说这够了,但数值精度充满了 corner case——非规格化数、NaN/Inf 传播、归约中的累加顺序、半精度饱和。从标准分布中随机抽 5 个样本覆盖不了这些。更严格的正确性评估对生产使用很重要。

为什么核心结论依然成立

尽管我对基准比较有意见,我认为论文的核心主张是有效的:在窄、定义良好的任务上,RL 能训练 LLM 发展出真实的、可迁移的优化能力

最让我信服的证据是 Level-3 和闭源模型的比较。Claude 和 Gemini 在遵循指令和生成代码方面很强,但它们没有专门经过训练去理解 GPU 性能分析信号并迭代优化。Level-3 上约 40 分的差距不是因为一个模型更聪明——而是因为一个模型通过训练内化了特定技能。

分阶段训练流水线也是真正的贡献。带工具使用的多轮 RL 很难稳定;PPO 预热 → RFT → Critic 预训练 → 多轮 RL 的序列是一个具体的、有效的方案,对任何试图把类似技术应用到其他领域的人来说都是有价值的信息。

什么会让我相信它完全泛化

大图景

关于”LLM 写 CUDA”作为已解决问题,存在很多兴奋的情绪。它还没解决。但 CUDA Agent 是有意义的一步,展示了认真对待这个问题是什么样的:适当的数据合成、带真实性能分析的真实执行环境、以及针对任务特性设计的训练流水线。

结果是一个模型学到了真实的东西:如何读性能分析结果、识别瓶颈、应用正确的优化技术。这不是小事。对于一个从语言模型预训练出发的系统来说,这其实相当令人印象深刻。

我只是不认为”比 torch.compile 快 2.11x”是解释它为什么令人印象深刻的正确框架。