梦晨 发自 凹非寺
量子位 | 公众号 QbitAI

让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

— 欢迎AI产品从业者共建 —


📚「AI产品知识库」是量子位智库基于长期产品库追踪和用户行为数据推出的飞书知识库，旨在成为AI行业从业者、投资者、研究者的核心信息枢纽与决策支持平台。


一键关注 👇 点亮星标

科技前沿进展每日见

内容中包含的图片若涉及版权问题，请及时与我们联系删除