字节跳动&清华发布 CUDA Agent，通过强化学习编写高性能 CUDA 内核，性能超越 torch.compile。
原文标题：英伟达护城河被AI攻破，字节清华CUDA Agent，让人人能搓CUDA内核
原文作者：机器之心
冷月清谈：
CUDA Agent 的核心在于其数据合成机制、CUDA 开发环境和强化学习算法。研究团队还发布了 CUDA-Agent-Ops-6K 数据集，用于支持相关研究的复现。智能体循环管线采用 ReAct 风格的工作流，结合代码工具与 CUDA Skill 规范，支持迭代式的编码-编译-调试循环，以及基于性能分析器的优化过程。
训练过程采用分阶段设计，首先进行单轮 PPO 预热训练，随后分别初始化 actor 和 critic，最后进入完整的多轮智能体强化学习阶段。实验结果表明，CUDA Agent 在 KernelBench 上取得了 SOTA 的成绩，尤其在高难度设置下优势明显。尽管研究存在未与 TVM 等复杂编译器框架对比，以及训练成本较高等局限性，但CUDA Agent 预示着全自动、高度性能优化计算基础设施的未来。
怜星夜思：
2、文章中提到 CUDA Agent 通过奖励 GPU 速度而非代码正确性来进行强化学习，这种思路有什么优缺点？是否可以应用于其他编程领域的代码优化？
3、CUDA Agent 目前还存在一些局限性，比如未与 TVM 等复杂编译器框架进行对比，训练成本较高等。你认为未来 CUDA Agent 还有哪些潜在的改进方向？
原文内容
研究人员训练了一个能够编写快速 CUDA 内核的模型：不只是正确的内核，而是真正经过优化的内核。
在简单/中等内核上，它的性能比 torch.compile 高出 2 倍；在复杂内核上，它的性能比 torch.compile 高出约 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 代码，AI 生成的代码也已获得了一定程度的应用，但能跑通和跑得快完全是两码事。
GPU 内核优化是现代深度学习的基础，但它仍然是一项高度专业化的工作，需要深厚的硬件知识。现有的 AI 生成 CUDA 代码通常依赖无训练的提示工程（Prompting）或多轮执行反馈微调的机制。这导致模型只能解决表面上的语法错误，无法真正理解底层硬件逻辑，一定程度上限制了其内在的优化能力。
真正极致的 CUDA 优化需要处理的任务，是只有在性能分析器中才能看到的硬件级指标。人们一直期待能出现一个像人类 CUDA 专家一样思考的 AI。
针对这一矛盾，CUDA Agent 的核心理念简单而巧妙：CUDA 性能并非取决于正确性，而是取决于硬件。线程束、内存带宽、内存冲突——这些只有在性能分析器中才能看到的东西。
研究人员不再奖励「是否编译成功」，而是奖励实际的GPU速度。真实的性能分析数据。强化学习直接基于性能进行训练。
产生的效果出乎人们的预料。
在 KernelBench 基准测试上，CUDA Agent 取得了 SOTA 的成绩：在 Level-1、Level-2 和 Level-3 三个划分上，相比 torch.compile 分别实现了 100%、100% 和 92% 的加速比例（faster rate）。
CUDA Agent 与 torch.compile 和强大的专有模型在 KernelBench 上的对比。
简而言之， CUDA Agent 是一个大规模的智能体强化学习系统，包含三个核心组成部分：可扩展的数据合成机制、一个集成技能增强且具备可靠验证与性能分析能力的 CUDA 开发环境，以及用于稳定长上下文训练的强化学习算法技术。
此外，研究团队同时发布了 CUDA-Agent-Ops-6K，一个经过严格筛选与数据污染控制的高质量合成训练数据集，可支持基于强化学习的 CUDA 内核优化研究的复现。
系统管线设计
数据合成
研究团队通过一个三阶段的管线来构建训练任务：种子问题爬取、基于 LLM 的组合式合成，以及基于执行结果的筛选。
-
从 torch 和 transformers
中挖掘种子算子。每个算子都以一个 Python 类的形式表示，包含初始化和前向传播方法。 -
在组合式合成阶段，最多采样 5 个 torch 算子，并将它们按顺序组合，构造成融合任务。
-
筛选阶段仅保留那些在 eager 模式和 compile 模式下都能正常运行的任务，同时移除包含随机性的算子。
-
为防止「投机取巧」，剔除在不同输入下输出为常数或无法区分的任务。
-
在工作负载控制方面，将 eager 模式下的运行时间限制在 1ms–100ms 区间内，并移除与 KernelBench 高度相似的样本。
三阶段数据收集管线
最终整理得到 6000 条训练样本，构建了 CUDA-Agent-Ops-6K 数据集，该数据集专为可扩展的强化学习训练而设计，兼具广泛的任务多样性和较低的数据污染风险。
智能体环境
智能体循环管线遵循一种 ReAct 风格的工作流，结合代码工具与 CUDA Skill 规范（SKILL.md），支持迭代式的编码-编译-调试循环，以及基于性能分析器的优化过程。
-
标准工作流程：对原生 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 在相对于 Compile 的性能优化上展现出显著优势：在整体 KernelBench 基准测试中，其相对 Compile 的加速达成率达到 96.8%，几何平均加速比为 2.11 倍。
这一优势在高难度设置下尤为明显：在 Level-3 上，CUDA Agent 相对 compile 的加速达成率达到 90%，相比最强的专有基线高出约 40 个百分点；在 Level-2 的算子序列任务上，其加速达成率达到 100%，几何平均加速比达到 2.80 倍。
在 KernelBench 上的整体性能和加速指标。
本研究存在两个主要局限。
首先，此次研究未将 CUDA Agent 与更为复杂的编译器框架（如 TVM）进行对比。其次，训练流程依赖于大规模 GPU 资源池以及进程级隔离机制，这带来了相当可观的计算与工程成本。探索更加资源高效的训练策略，将是未来的重要研究方向。
看起来，CUDA Agent 等技术的出现即将打破传统编译器（如torch.compile 或 Triton）的优化瓶颈。它证明了：大语言模型不仅可以学习人类自然语言和高级编程语言，还可以通过基于硬件反馈的强化学习，内化出极高门槛的「硬件直觉」。
一条通向全自动、高度性能优化计算基础设施的道路正在出现。
© THE END
转载请联系本公众号获得授权
投稿或寻求报道：liyazhou@jiqizhixin.com