字节清华智能体自动写CUDA内核，比torch.compile加速2.11倍

让AI自己写高性能GPU代码，字节Seed与清华AIR团队做到了。

最新开源CUDA Agent，在GPU内核优化基准KernelBench上拿下当前最优成绩：

Level-1和Level-2的faster rate（相对torch.compile的加速通过率）双双达到100%，最难的Level-3同样达到90%。

并且全面超越单用大模型的方案。

通讯作者为清华AIR副研究员周浩，共同一作Weinan Dai，Hanlin Wu。

GPU内核优化，强化学习上场了

GPU内核优化门槛一直很高，要求开发者对GPU硬件架构、内存层次、线程调度有深入理解。

同时，模型训练和推理的速度，很大程度上取决于底层CUDA内核写得好不好。

现有的AI辅助方案大体分两类：

一类是无训练的迭代优化，依靠提示词引导模型反复修改代码；

另一类是固定的执行-反馈循环，模型在编译运行结果的指导下做调整。这两条路线都没有从根本上提升模型的内核优化能力。

CUDA Agent走出第三条路。

它是一个完整的大规模智能体强化学习系统，包含三个核心组件：

• 一套可扩展的数据合成流水线，
• 一个配备技能规范文档的CUDA开发环境（内置可靠的验证和性能分析工具）
• 一个专门为长上下文场景设计的RL训练算法。

目标只有一个——通过强化学习，让模型真正学会生成和优化高性能CUDA内核。

6000道合成训练题，和滴水不漏的防作弊机制

CUDA Agent训练数据的构建经历了三个阶段。

第一步是种子算子挖掘，来源是torch和transformers，每个算子被表示为一个包含初始化方法和前向方法的Python类。

第二步是组合合成，随机采样最多5个torch算子，按顺序拼接成融合任务。第三步是执行驱动过滤，只保留在eager模式和compile模式下都能正常运行的任务，同时去除含有随机算子的样本。

第三部加入反作弊检查，排除那些对不同输入产生常数输出或无法区分输出的任务。

eager运行时间被控制在1毫秒到100毫秒的范围内，与KernelBench高度相似的样本也被移除以降低数据污染风险。

经过层层筛选，最终留下6000个训练样本，命名为CUDA-Agent-Ops-6K，已对外开源发布。

智能体的工作环境采用ReAct风格的交互循环，配合编码工具和一份名为SKILL.md的CUDA技能规范文档。

标准流程如下：

先对原生PyTorch代码做性能分析，然后实现CUDA内核和绑定代码，在GPU沙箱中编译运行，再根据结果迭代优化。通过验证的标准是：正确性检查全部通过，同时速度超过torch.compile至少5%。

为了防止模型走捷径，团队设置了一整套约束：

验证脚本和性能分析脚本被设为受保护状态，模型无法篡改；回退函数调用被禁止；正确性检查使用5组不同输入而非单一样本；性能分析引入同步预热步骤以消除测量噪声；网络检索功能完全关闭。

奖励机制采用基于里程碑的离散奖励，分别对应正确性达标和速度提升两个阶段。所有这些约束确保奖励信号反映的是真实的内核质量。

分阶段训练：从单轮热身到128K上下文

长上下文场景下的强化学习很容易不稳定。为此，训练流水线被拆成多个阶段。

第一阶段是单轮PPO热身，模型先在非交互式的设置下学习基础的CUDA代码生成能力，为后续多轮交互打基础。

进入多轮智能体RL之前，actor和critic需要分别初始化。

Actor的初始化使用拒绝微调（Rejection Fine-Tuning，RFT）：

从大量采样轨迹中筛选出取得正向结果的样本进行训练。筛选过程还会进一步去除包含低效循环和无效工具调用模式的轨迹，以降低策略崩溃的风险。Critic则通过价值预训练完成初始化，使优势估计在训练早期就保持可靠。

这套分阶段设计支撑了极大的上下文规模。训练时上下文窗口达到128K token，单条轨迹最多150轮交互。评估时交互轮数进一步放宽到200轮。在整个训练过程中，奖励信号持续增长，没有出现常见的长序列RL训练崩溃现象。

全面超越商业模型

在KernelBench的整体评测上，CUDA Agent的通过率为98.8%。相对torch.compile，faster rate达到96.8%，几何平均加速比为2.11倍。相对eager模式，faster rate为98.4%，加速比为2.60倍。

分级来看，Level-1和Level-2的faster rate（相对compile）均为100%。其中Level-2的算子序列任务上，几何平均加速比达到了2.80倍。最难的Level-3，通过率94%，faster rate（相对compile）为90%，几何平均加速比1.52倍。

与强商业模型的对比拉开了明显差距。Claude Opus 4.5和Gemini 3 Pro在KernelBench整体上的faster rate（相对compile）约为66.4%到69.6%，几何平均加速比在1.42倍到1.46倍之间。而CUDA Agent的整体faster rate为96.8%，加速比2.11倍。差距在难题上更为显著——Level-3上CUDA Agent的90% faster rate领先最强商业基线约40个百分点。

团队同步开源了训练数据集CUDA-Agent-Ops-6K，包含完整的过滤流程和污染控制方案，供后续基于强化学习的CUDA内核优化研究使用。

论文地址：

https://cuda-agent.github.io

文章来自于“量子位”，作者 “梦晨”。