近日，字节跳动 Seed 团队与清华大学 AIR 的研究人员联合发布了一项名为 CUDA Agent 的新研究，在 AI 代码生成领域引起了广泛关注。

该研究训练了一个能够编写高性能 CUDA 内核的智能体模型。其目标不仅是生成语法正确的代码，更是产出经过深度优化、执行效率显著提升的内核。

在基准测试中，CUDA Agent 展现出卓越的性能：在简单/中等复杂度内核上，其性能比 torch.compile 高出 2 倍；在复杂内核上，性能提升约 92%；即使在最具挑战性的任务设置下，其性能也领先于 Claude Opus 4.5 和 Gemini 3 Pro 等强大的专有模型约 40%。

• 论文链接：https://arxiv.org/abs/2602.24286
• 项目主页：https://cuda-agent.github.io/
• Github 链接：https://github.com/BytedTsinghua-SIA/CUDA-Agent
• 数据集链接：https://huggingface.co/datasets/BytedTsinghua-SIA/CUDA-Agent-Ops-6K

当前，以 GPT、Claude 为代表的大模型已能生成语法“正确”的 CUDA 代码，并得到了一定应用。然而，“能运行”与“运行得快”之间存在巨大鸿沟。GPU 内核优化是深度学习性能的关键，但这项工作高度专业化，需要深厚的硬件知识。

现有 AI 生成 CUDA 代码的方法多依赖于无训练的提示工程或多轮执行反馈微调。这类方法通常只能修正表面的语法错误，难以深入理解底层硬件逻辑，从而限制了其进行深度优化的内在能力。

真正极致的 CUDA 优化需要关注性能分析器才能观测到的硬件级指标，业界一直期待能出现像人类 CUDA 专家一样“思考”的 AI。

针对这一挑战，CUDA Agent 提出了一个核心理念：CUDA 性能的关键不在于代码正确性，而在于对硬件特性的深刻理解与利用。线程束调度、内存带宽利用、内存访问冲突——这些只有在性能分析器中才能洞察的细节，才是优化的核心。

为此，研究团队摒弃了传统的“编译成功即奖励”的模式，转而采用基于真实 GPU 执行速度的强化学习训练。智能体直接从性能分析数据中学习，奖励机制直接与内核的实际运行速度挂钩。

这一方法取得了显著成效。在 KernelBench 基准测试中，CUDA Agent 取得了领先水平（SOTA）的成绩：在 Level-1、Level-2 和 Level-3 三个难度划分上，相比 torch.compile 实现的平均加速比例（faster rate）分别达到 100%、100% 和 92%。

CUDA Agent 与 torch.compile 及强大专有模型在 KernelBench 上的性能对比。

简而言之，CUDA Agent 是一个大规模智能体强化学习系统，其核心包含三个部分：一个可扩展的数据合成机制、一个集成了技能增强与可靠验证/性能分析能力的 CUDA 开发环境，以及一套用于稳定长上下文训练的强化学习算法。

此外，研究团队同步发布了 CUDA-Agent-Ops-6K，这是一个经过严格筛选与数据污染控制的高质量合成训练数据集，旨在支持基于强化学习的 CUDA 内核优化研究的复现。

系统管线设计

数据合成

研究团队通过一个三阶段管线来构建训练任务：种子问题爬取、基于大语言模型（LLM）的组合式合成，以及基于执行结果的筛选。

• 种子挖掘：从 PyTorch 和 Transformers 库中挖掘种子算子，每个算子以 Python 类的形式表示，包含初始化和前向传播方法。
• 组合合成：在合成阶段，最多采样 5 个算子，并将它们按顺序组合，构造出融合计算任务。
• 执行筛选：仅保留在 Eager 模式和 Compile 模式下均能正常运行的任务，并移除包含随机性的算子。
• 防作弊处理：为防止模型“投机取巧”，剔除在不同输入下输出为常数或无法区分的任务。
• 负载控制：将 Eager 模式下的运行时间限制在 1ms–100ms 区间内，并移除与最终评估基准 KernelBench 高度相似的样本。

三阶段数据收集管线示意图。

最终，团队整理得到 6000 条训练样本，构建了 CUDA-Agent-Ops-6K 数据集。该数据集专为可扩展的强化学习训练设计，兼具任务多样性和较低的数据污染风险。

智能体环境

智能体循环管线遵循一种 ReAct 风格的工作流，结合代码工具与 CUDA 技能规范，支持迭代式的编码-编译-调试循环，以及基于性能分析器的优化过程。

• 标准流程：对原生 PyTorch 实现进行性能分析，编写 CUDA 内核及其绑定代码，在 GPU 沙盒环境中完成编译，并不断迭代优化。
• 优化目标：要求生成的代码通过正确性检查，并在性能上相对于 torch.compile 实现至少 5% 的加速。
• 奖励机制：采用基于里程碑的离散奖励设计，根据正确性达标情况和性能提升幅度分别给予奖励。
• 防作弊措施：包括对验证与性能分析脚本进行保护、禁止回退调用、采用 5 组不同输入进行正确性检查、在同步预热后进行性能分析，以及禁止网络检索。

智能体环境与工作流程示意图。

这些约束共同构建了可靠的、基于真实执行结果的反馈机制，使策略学习聚焦于内核质量的实质性提升，而非依赖取巧或捷径行为。

训练流程

训练过程采用分阶段设计，以稳定 CUDA 编码这一长时序强化学习任务。首先进行单轮 PPO 预热训练，随后分别初始化执行者（actor）和评判者（critic），最后进入完整的多轮智能体强化学习阶段。

• 单轮预热：旨在提升基础的 CUDA 代码生成能力，为后续交互式智能体训练打下基础。
• Actor 初始化：采用基于正向结果轨迹采样的拒绝式微调（RFT）。RFT 过滤机制会剔除低效循环以及无效的工具调用模式，从而降低策略崩溃的风险。
• Critic 初始化：通过价值函数预训练，使得从训练早期开始，优势估计就具备较高可靠性。

多阶段训练流程示意图。

借助这一设计，系统在长上下文设定下（最长 128k 上下文、训练阶段最多 150 轮、评估阶段最多 200 轮）依然保持稳定，从而实现持续的奖励增长。

核心实验结果

研究团队在 KernelBench 基准上报告了整体及 Level-3 难度拆分的完整指标，包括通过率、提速率（相对于 Eager/Compile 模式）以及几何平均加速比（相对于 Eager/Compile 模式）。

CUDA Agent 在 KernelBench 上的详细实验结果表格。

与强大的专有基线模型相比，CUDA Agent 在相对于 torch.compile 的性能优化上展现出显著优势：
* 在整个 KernelBench 基准测试中，其相对 torch.compile 的 加速达成率达到 96.8%，几何平均加速比为 2.11 倍。
* 这一优势在高难度设置下尤为明显：在 Level-3 任务上，其相对 torch.compile 的加速达成率达到 90%，相比最强的专有基线高出约 40 个百分点；在 Level-2 的算子序列任务上，其加速达成率达到 100%，几何平均加速比达到 2.80 倍。

CUDA Agent 在 KernelBench 上的整体性能和加速指标可视化。

本研究也存在一定的局限性。

首先，此次研究未将 CUDA Agent 与更为复杂的编译器框架（如 TVM）进行对比。其次，其训练流程依赖于大规模 GPU 资源池及进程级隔离机制，带来了显著的计算与工程成本。探索更高效的训练策略，是未来的重要研究方向。

CUDA Agent 等技术的出现，有望突破传统编译器（如 torch.compile 或 Triton）的优化瓶颈。这表明，大语言模型不仅能学习自然语言和高级编程语言，还能通过基于硬件反馈的强化学习，内化出高门槛的“硬件直觉”。

一条通向全自动、高性能优化计算基础设施的道路正在显现。

关注“鲸栖”小程序，掌握最新AI资讯