CUDA Agent 训练大语言模型编写比 torch.compile 更快的 GPU 内核

引言

字节跳动 Seed 团队与清华大学人工智能产业研究院（AIR）联合推出了 CUDA Agent，这是一个大规模智能体强化学习系统，旨在教会语言模型生成高性能 CUDA 内核。

这项工作针对的是 AI 生成底层代码中长期存在的弱点。前沿语言模型往往能生成可编译且运行正确的 CUDA 代码，但对于生产级 GPU 工作负载而言，仅正确性是不够的。生成的内核还必须与高度优化的编译器生成实现相竞争。

CUDA Agent 正是围绕这第二个问题构建的：不仅生成有效的 CUDA 代码，还要学习如何迭代地进行性能分析、验证和优化内核，直到实现有意义的运行时改进。

为什么仅生成正确的 CUDA 代码不够

在引入智能体强化学习之前，Seed1.6 基础模型在 KernelBench 上已经展现出相当强的 CUDA 编码能力。

在研究人员使用的 250 个 KernelBench Level 1–3 任务中，Seed1.6 实现了 74.0% 的通过率。换句话说，它能为基础模型的大部分基准任务生成功能正确的解决方案。

但性能表现则是另一回事。

基础模型生成的内核中，仅有 27.2% 比 torch.compile 更快，其相对于编译器基线的几何平均速度仅为 0.69×。

模型	通过率	比 torch.compile 更快	相对 torch.compile 的几何平均速度
Seed1.6 基础模型	74.0%	27.2%	0.69×
CUDA Agent	98.8%	96.8%	2.11×

这一差距凸显了自动 GPU 内核生成的主要挑战。

语言模型可以理解 CUDA 语法，但仍可能产生低效的内存访问、过多的内核启动、不佳的分块选择、不必要的中

性能剖析

• 针对 CUDA 的优化指导
• 受限的沙箱环境
• 与实际执行结果挂钩的奖励信号

SKILL.md 文件还限制了可能使基准测试结果失真的捷径。例如，代理不能简单地在自定义内核实现中回退到任意的 PyTorch 操作。

这一点很重要，因为强化学习系统可能会找到最大化奖励的方式，而无需解决预期的优化问题。

PPO 训练将代理扩展到长优化轨迹

研究人员使用**近端策略优化（PPO）**训练 CUDA Agent。

一个主要挑战是 GPU 优化本质上是一项长时程任务。模型可能需要经历多轮代码编辑、编译、调试、剖析和进一步优化，才能达到理想结果。

因此，论文采用了不同的上下文长度和训练阶段来稳定学习。

基础模型是 Seed1.6，一个混合专家模型，拥有 2300 亿总参数和 230 亿激活参数。

对于智能体强化学习：

• 上下文窗口为 131,072 个 token。
• 训练回放允许最多 150 轮代理交互。
• 评估允许最多 200 轮代理交互。
• 模型训练 150 步 RL。
• Actor 和 critic 模型在完整长时程优化之前采用多阶段预热策略。

这比简单地在一对 CUDA 提示和答案上微调模型更为复杂。

目标是通过与执行反馈的反复交互，教会模型如何改进自身的内核。

沙箱将编译与 GPU 剖析分离

可靠的性能测量至关重要，因为运行速度是奖励的一部分。

因此，研究人员构建了一种 CPU–GPU 资源解耦的沙箱架构。

基于 Docker 的终端沙箱处理 CPU 导向的任务（如编译），而验证和性能剖析则被分派到包含 128 块 NVIDIA H20 GPU 的专用 GPU 沙箱池中。

这种分离有双重目的。

第一，它将编译和代理交互与 GPU 基准测试隔离。第二，专有的 GPU 分配减少了并发工作负载之间的干扰，使延迟测量更加稳定。

这对强化学习很重要：如果时序测量受到噪声影响，模型会收到不可靠的奖励信号，并可能学到错误的优化行为。

CUDA Agent 达到 98.8% 通过率，并在 96.8% 任务上超越编译器

最终系统在 KernelBench 第 1 至第 3 级的 250 个任务上进行了评估。

CUDA Agent 实现了：

• 98.8% 的总体通过率
• 比 PyTorch 即时执行快 98.4%
• 比 torch.compile 快 96.8%
• 相对即时执行有 2.60× 的几何平均加速
• 相对 torch.compile 有 2.11× 的几何平均加速

这些结果相对于 Seed1.6 起点有了大幅提升。

模型不仅在生成通过正确性测试的代码方面变得更好，其生成的内核也更有可能超越编译器基线。

按 KernelBench 难度划分的结果

在 KernelBench 的三个难度等级中，性能提升均保持强劲。

KernelBench 难度	通过率	快于 torch.compile 的比例	相比 torch.compile 的几何平均加速
等级 1	100.0%	97.0%	1.87×
等级 2	100.0%	100.0%	2.80×
等级 3	94.0%	90.0%	1.52×

等级 2 的表现尤为突出。

这些任务涉及算子序列，其中融合与内存移动带来的优化机会，是静态编译器启发式策略未必总能有效利用的。

论文作者认为，迭代式学习优化器能够搜索更广泛的实现策略空间，包括面向特定硬件的内存访问方式、分块以及融合选择。

训练数据即将公开

目前，完整的训练模型权重并未包含在该项目公开释出的内容中。

不过，研究人员已经发布了训练栈中几个重要的组成部分。

CUDA-Agent-Ops-6K

CUDA-Agent-Ops-6K 数据集包含 6,000 个合成的算子级训练任务。

其构建流程包括：

1. 从 torch、transformers 等库中收集种子算子。
2. 使用 LLM 将多个算子组合成更复杂的融合任务。
3. 通过基于执行结果的检查来过滤生成的任务。

过滤阶段会移除具有随机性、过于简单、无效或与评测任务过于相似的情况。

该数据集已在 Hugging Face 上发布，并采用 CC BY 4.0 许可协议。

SKILL.md 与智能体环境

GitHub 仓库中还包含了面向 CUDA 的 SKILL.md 指令文件以及智能体工作环境。

这些材料记录了优化工作流程、可用工具、限制条件以及用于引导智能体的执行环境设置。

奖励机制与训练方案

论文和仓库中描述了奖励设计、预热阶段、评论家初始化以及基于 PPO 的训练方案，用于稳定长周期优化过程。

这对于那些关注系统编程智能体训练的研究人员尤为有用，而不仅仅是复现最终基准分数。

为什么这超越 KernelBench 本身具有重要意义

CUDA 内核优化支撑着现代 AI 基础设施的很大一部分。

更快的内核可以提升：

• 模型训练吞吐量
• LLM 推理延迟表现
• GPU 利用率
• 服务成本效率
• 实时 AI 流水线
• 高性能数值计算负载

AIBase 的原始文章指出了其在 AI 基础设施、大模型推理服务、自动驾驶以及量化交易等领域的潜在应用——在这些领域中，微小的延迟差异都可能至关重要。

CUDA Agent 的成果并不意味着传统编译器已经过时。

torch.compile 仍然是一个强大的自动基线，而且 CUDA Agent 的评测本身依赖于受控的基准测试环境和特定的 GPU 环境。

这项研究展示的结论虽然范围更窄，但仍然重要：在充分的执行反馈和专门的强化学习支持下，语言模型能够学习到在绝大多数测试 GPU 内核任务上超越静态编译器基线的优化策略。

这意味着，将底层性能工程转化为一种

代理式学习的合理领域，而不仅仅是单次代码生成。

常见问题

什么是 CUDA Agent？

CUDA Agent 是一个由字节跳动 Seed、清华大学智能产业研究院（AIR）及其联合 SIA-Lab 开发的大规模智能体强化学习系统。它训练语言模型以迭代方式编写、验证、分析和优化 CUDA 内核。

CUDA Agent 基于哪个模型？

该研究使用 Seed1.6 作为基础模型。论文将其描述为一个混合专家模型，总参数量为 2300 亿，激活参数量为 230 亿。

什么是 KernelBench？

KernelBench 是一个开放基准测试，用于评估语言模型能否为 PyTorch 程序生成正确且高效的 GPU 内核。CUDA Agent 在涵盖 KernelBench 第 1 级至第 3 级的 250 个任务上进行了评估。

CUDA Agent 比 torch.compile 快多少？

在整个基准测试中，CUDA Agent 相对于 torch.compile 实现了 2.11 倍的几何平均加速。其生成的内核在 96.8% 的评估任务上快于编译器基线。

CUDA Agent 使用强化学习吗？

是的。该系统使用 PPO，并结合多阶段预热、价值模型预训练、稳健的奖励设计以及长时间多轮智能体轨迹。

一次 CUDA Agent 优化轨迹可以运行多长时间？

训练 rollout 允许最多 150 个智能体轮次，而评估允许最多 200 个轮次。智能体训练上下文窗口为 131,072 个 token。

CUDA Agent 是开源的吗？

该项目已发布其 GitHub 仓库、智能体环境、SKILL.md、训练配方以及 CUDA-Agent-Ops-6K 数据集。完整训练好的模型权重未列入当前公开发布的内容中。

什么是 CUDA-Agent-Ops-6K？

CUDA-Agent-Ops-6K 是一个包含 6,000 个合成运算符级 CUDA 训练任务的数据集。它专为大规模智能体强化学习而设计，并包含过滤步骤，以确保任务可执行、确定性、非平凡，且与 KernelBench 评估集分离。

相关工具

• CUDA Agent：字节跳动 Seed 和清华大学智能产业研究院 CUDA 内核生成系统的官方项目页面。
• CUDA Agent GitHub：包含智能体环境、SKILL.md 和已发布项目材料的官方仓库。
• CUDA-Agent-Ops-6K：随项目发布的官方 6,000 样本训练数据集。
• KernelBench：用于评估大语言模型生成的 GPU 内核的开放基准测试。
• PyTorch：深度学习框架，其 eager 模式和 torch.compile 执行是研究中的关键基线。
• NVIDIA CUDA：NVIDIA 官方的 CUDA 编程和 GPU 内核开发文档。

相关链接

• CUDA Agent 项目页面：官方概述、基准测试结果、项目图表、论文、代码和数据集链接。
• CUDA Agent 论文（arXiv）：完整研究论文，“大规模智能体强化学习用于

高性能CUDA内核生成。”

• CUDA Agent GitHub仓库：已发布环境和项目文件的源代码仓库。
• CUDA-Agent-Ops-6K数据集：官方Hugging Face数据集，包含6，000个综合算子任务。
• KernelBench GitHub仓库：用于高效GPU内核生成的基准测试和评估环境。
• PyTorchtorch.compile：研究中使用的编译器基线的官方PyTorch文档。
• NVIDIA CUDA C++编程指南：NVIDIA官方CUDA编程与优化指南。

摘要

CUDA Agent解决了大语言模型生成的系统代码中的一个具体弱点：生成的CUDA代码不仅要正确，还要真正比编译器生成的替代方案更快。

从Seed1.6开始，研究人员使用CUDA专属智能体环境、执行反馈、稳健的奖励设计以及长视界PPO训练，将通过率从74.0%提升至98.8%，将快于torch.compile的比例从27.2%提升至96.8%。

该项目还发布了包含6，000个任务的训练数据集、SKILL.md、智能体环境材料和训练配方，为研究人员在学到的GPU优化方向的进一步工作提供了具体基础。

CUDA Agent的主要贡献在于证明了性能工程本身可以通过迭代式智能体循环来学习——而不仅仅是依靠单次代码生成来近似实现。