字节+清华AIR发布CUDA Agent：让AI生成的GPU内核超越编译器

ByteDance Seed 与清华大学 AIR 联合发布了 CUDA Agent，这是一个智能强化学习系统，能够训练大语言模型编写出超越编译器性能的 GPU kernel。它所针对的差距细微却顽固：前沿模型已经能生成正确的 CUDA 代码，只是生成的 CUDA 执行速度慢。在 KernelBench 上，基座模型 Seed1.6 通过了 74.0% 的任务，但在速度上只有 27.2% 能超过 torch.compile，其几何平均加速比仅为 0.69×，这意味着它的 kernel 平均来看比编译器生成的还慢。CUDA Agent 通过将模型置于真实的 CUDA 开发环境中来弥补这一差距——环境包含性能分析、正确性检查和权限锁定的沙箱，随后使用 PPO 在 131,072 token 的上下文长度下训练 150 步。最终在 250 任务的基准测试上达到了 98.8% 的通过率和 96.8% 的超越 torch.compile 比率，几何平均加速达 2.11×，在最具难度的 Level-3 分组上领先 Claude Opus 4.5 和 Gemini 3 Pro 约 40 个百分点。

Partly，但训练后的 agent 并未开源。它基于 Seed1.6 构建，这是一款拥有 23B 活跃参数和 230B 总参数的私有 MoE 模型，论文也没有发布权重。公开的内容包括：CUDA-Agent-Ops-6K 数据集、SKILL.md 规范以及 reward 和 warm-up 配方。

Which companies：仅性能分析沙箱就使用了 128 块 NVIDIA H20 GPU，这使得完整复现仅限于前沿实验室、GPU 云和大型基础设施团队。中等规模的团队仍可在开源基座模型之上采用其中部分组件——数据集、里程碑式 reward、反 reward-hacking 约束、技能规范等。

Industries and applications：AI 基础设施与推理服务、GPU 云、自动驾驶、量化交易、医疗影像和推荐系统——任何在延迟关键路径上有融合 kernel 的场景。用途包括融合 torch.compile 处理不好的算子序列、降低每 token 成本、以及跨 GPU 代际重调 kernel。

研究团队从 torch 和 transformers 库中爬取参考算子。随后由 LLM 采样最多五个 torch 算子类，并将它们堆叠成一个融合层。过滤器仅保留满足以下条件的算子：在 eager 和 compile 模式下均可执行、是确定性的、产生非常量输出、且在 eager 模式下运行时间在 1 ms 到 100 ms 之间。与任何 KernelBench 任务 AST 相似度超过 0.9 的样本被移除。最终得到 CUDA-Agent-Ops-6K：6,000 个样本，其中 83.77% 为双算子组合。

Environment and reward

Agent 循环复刻了 OpenHands 的工具链——Bash、Read/Write、Edit/MultiEdit、Glob、Grep、NotebookEdit、BashOutput、KillBash——并在 ReAct 模式下运行。CUDA 指令以 Agent Skills 格式交付。SKILL.md 指示模型对 PyTorch 模型进行性能分析、用自定义 kernel 重写 model_new.py、在 GPU 沙箱中编译、持续迭代直至 kernel 至少比 torch.compile 快 5%（atol=1e-2, rtol=1e-2）。

Reward hacking 有五重应对措施：权限锁定的验证和性能分析脚本、禁止 torch.nn.functional 回退的上下文管理器、对五个随机输入的检查、使用设备同步和 warm-up 的性能分析、以及不提供 web search 工具。

Reward 采用离散值而非原始加速比。r ∈ {−1, 1, 2, 3}：正确性失败得 −1，kernel 同时超越 eager 和 torch.compile 超过 5% 得 3，仅超越 eager 得 2，其余情况得 1。

https://arxiv.org/pdf/2602.24286v1

表 1 总体结果：98.8% 通过率、98.4% 超越 eager 的比率、96.8% 超越 torch.compile 的比率，分别对应 2.60× 和 2.11× 的几何平均加速。Level 2（算子序列）是最强分组：100% 通过率、100% 超越比率、较 torch.compile 快 2.80×。Level 3 达到 94.0% 通过率、90.0% 超越比率和 1.52× 加速，在超越 torch.compile 的比率上领先 Claude Opus 4.5（50.0%）和 Gemini 3 Pro（52.0%）约 40 个百分点。

One inconsistency：摘要和引言中声明 Level 1–3 的超越比率分别为 100%/100%/92%，而表 1 报告为 97.0%/100.0%/90.0%。表 1 才是主要结果表。

消融实验非常直接。移除 agent 循环后，超越 torch.compile 的比率从 96.8% 降至 14.1%。使用原始加速比 reward 得到 60.4%，不使用 RFT 得到 49.8% 并伴随 reward 崩溃，不使用价值预训练得到 50.9% 并伴随轨迹失控。

案例研究展示了策略学到了什么。对角矩阵乘法重写为逐行缩放：较 torch.compile 快 73.31×。矩阵乘法-除法-求和-缩放链经重排和融合：24.04×。ResNet BasicBlock 将 BatchNorm 融入卷积并调用 cudnnConvolutionBiasActivationForward：3.59×。

CUDA Agent 在 KernelBench 上达到 98.8% 通过率和 96.8% 超越 torch.compile 的比率，几何平均加速 2.11×。

Level 2 融合最为突出：100% 超越比率，较 torch.compile 快 2.80×。

离散里程碑式 reward 在超越比率上比原始加速比 reward 高出 36.4 个百分点。

RFT 加上价值预训练将 17 步崩溃转变为 150 步稳定训练。

权重未公开；6,000 样本数据集、SKILL.md 和配方已公开。

Check out the Paper, the Project Page and the Dataset. Also, feel free to follow us on Twitter and don't forget to join our 150k+ ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.

Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us

Asif Razzaq is the CEO of Marktechpost Media Inc.. As a visionary entrepreneur and engineer, Asif is committed to harnessing the potential of Artificial Intelligence for social good. His most recent endeavor is the launch of an Artificial Intelligence Media Platform, Marktechpost, which stands out for its in-depth coverage of machine learning and deep learning news that is both technically sound and easily understandable by a wide audience. The platform boasts of over 2 million monthly views, illustrating its popularity among audiences.