CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation

Weinan Dai1,2,3∗, Hanlin Wu1,2,3∗, Qiying Yu1,2,3, Huan-ang Gao1,2,3, Jiahao Li1, Chengquan Jiang1, Weiqiang Lou1, Yufan Song1, Hongli Yu1,2,3, Jiaze Chen1,3, Wei-Ying Ma2,3, Ya-Qin Zhang2,3, Jingjing Liu2,3, Mingxuan Wang1,3, Xin Liu1, Hao Zhou2,3†

1ByteDance Seed 2Institute for AI Industry Research (AIR), Tsinghua University 3SIA-Lab of Tsinghua AIR and ByteDance Seed

1. 主要贡献

本文提出了 CUDA Agent，这是一个大规模的代理强化学习系统，旨在解决现有大型语言模型（LLM）在生成高性能 CUDA 内核代码方面竞争力不足的问题。尽管 LLM 在通用编程中表现出色，但在 CUDA 内核生成方面仍落后于 torch.compile 等基于编译器的系统。现有的方法要么依赖无训练的微调，要么在固定的多轮执行反馈循环中微调模型，都未能从根本上提升模型内在的 CUDA 优化能力。

为了克服这些限制，本文的主要贡献集中在以下三个方面：

1. 可扩展的数据合成管道（Data Synthesis Pipeline）：设计了一个三阶段的数据收集管道，通过从 PyTorch 和 Transformer 库中抓取种子算子，利用 LLM 进行组合合成，并实施严格的基于执行的过滤，从而构建了一个涵盖广泛难度级别的训练数据集。
2. 增强技能的 CUDA 开发环境（Skill-Augmented Environment）：采用代理技能（Agent Skills）范式，为模型配备了结构化的规范（http://SKILL.md），形式化了编写、验证和优化 CUDA 内核的标准工作流程。环境集成了自动化的验证和性能分析脚本，并实施了严格的正确性和性能测试及权限隔离，以提供可靠的奖励信号并防止奖励欺诈（Reward Hacking）。
3. 支持稳定训练的 RL 算法技术：分析了 RL 优化不稳定的根源，并针对 Actor 和 Critic 模型提出了多阶段预热（Warm-up）策略，实现了大语言模型在长上下文（128k token）和多轮交互（200 turns）下的稳定训练。

实验结果表明，CUDA Agent 在 KernelBench 上取得了最先进（SOTA）的性能。在 KernelBench 的 Level-1、Level-2 和 Level-3 划分上，相较于 torch.compile 分别实现了 100%、100% 和 92% 的更优率（Faster Rate）。在难度最高的 Level-3 设置中，其表现超越了 Claude Opus 4.5 和 Gemini 3 Pro 等最强专有模型约 40%。

2. 方法细节

3.1 可扩展的训练数据合成管道

由于高性能 CUDA 内核的稀缺性，手动实现专家级参考代码成本过高，限制了监督微调的效果。为了支持强化学习（RL），需要大量且多样的基于 PyTorch 实现的参考算子作为训练任务。本文开发了一个可扩展的数据收集管道（如图 1 所示），包含以下三个步骤：

种子问题抓取（Seed Problem Crawling）
首先，从 torch 和 transformers 库中挖掘基于 PyTorch 实现的参考算子，建立全面的种子问题集。每个算子被表示为一个包含初始化和前向传播方法的 Python 类。选择这些库是因为它们被广泛使用且维护良好，从而排除了代码质量不足的个人仓库。

组合问题构建（Combinatorial Problem Construction）
为了扩展数据集并引入更高的复杂性，利用 LLM 合成聚合算子。具体而言，提示 LLM 从 torch 库中采样不超过 5 个算子类。这些采样的算子类被顺序堆叠成一个单一的计算层。此处不从 transformers 库中采样算子类，因为这些算子通常已经是封装了多个原语操作的高级模块。

基于执行的过滤（Execution-Based Filtering）
最后，实施严格的数据选择过程，根据执行反馈过滤掉太容易或太难的问题。对每个算子进行四项标准的验证：
1. 可执行性：算子必须在 Eager 模式和 Compile 模式下都能成功执行。
2. 确定性：为了确保可复现性，排除具有固有随机性的算子。
3. 防欺诈（Anti-hacking）：验证不同输入的输出既不是恒定值，在数值上也是可区分的。
4. 工作负载合理性：为了过滤掉微不足道或过重的任务，将 Eager 模式下的执行时间限制在 1 毫秒到 100 毫秒之间。

此外，还排除了与 KernelBench 测试用例具有高相似度的算子（相似度分布见附录 A）。最终过滤后的合成训练数据集包含 6,000 个样本，构成了 CUDA-Agent-Ops-6K 数据集。

3.2 技能集成的代理循环（Skill-Integrated Agent Loop）

代理循环（Agent Loop）
鉴于 CUDA 内核编码是编码代理的一个子任务，代理循环的设计与广泛采用的 OpenHands 框架对齐以确保通用性。LLM 被提供了一套标准的 Shell 工具——BashTool、GlobTool、MultiEditTool 和 TodoWriteTool——以全面支持 CUDA 编码开发。在此基础上，代理循环（如图 2 所示）遵循 ReAct 风格范式，交错进行推理、动作执行和观察，以实现迭代式的编码、调试和性能优化。

CUDA 编码技能（CUDA Coding Skill）
受 Agent Skills 理念的启发，特意将 CUDA 编码特定的指令和工具（如用于比较生成内核与 torch.compile 性能的分析工具）格式化为代理技能。设计了专门的 CUDA 内核编码指令 SKILL.md，规范化了优化 CUDA 内核的标准流程：
1. 使用提供的 profile.py 脚本分析原生 PyTorch 实现的性能。此步骤识别性能瓶颈和优化机会，例如过多的内核启动和次优的内存访问模式。
2. 通过在 model_new.py 中重写模型并开发相应的 CUDA 内核源文件及绑定代码，实现自定义 CUDA 算子，针对分析阶段确定的性能关键算子进行优化。
3. 在提供的 GPU 沙盒环境中编译和评估优化后的模型，并迭代完善 CUDA 内核实现，直到数值正确性和性能要求都得到满足。
4. 从步骤 2 开始重复优化过程，直到最终实现通过所有数值正确性检查，并相较于 torch.compile 基线实现至少 5% 的加速。

鲁棒的奖励调度（Robust Reward Scheduling）
现有的 CUDA 生成 RL 方法通常使用相对于基线的加速比作为奖励信号。然而，这种方法容易受到异常值的影响，并偏向于简单的内核。为了解决这个问题，并联合优化正确性和执行延迟，提出了一种归一化的鲁棒奖励方案。

根据正确性和性能分配奖励分数 $r \in {-1, 1, 2, 3}$：

$$\begin{aligned} r = \begin{cases} -1 & \text{if correctness check fails} \\ 3 & \text{if } b(t, t_{\text{eager}}) \wedge b(t, t_{\text{compile}}) \\ 2 & \text{if } b(t, t_{\text{eager}}) \\ 1 & \text{otherwise} \end{cases} \end{aligned}$$

其中 $t$ 是生成内核的运行时间，$t_{\text{eager}}$ 和 $t_{\text{compile}}$ 分别是 PyTorch Eager 实现和 torch.compile 版本的运行时间，$b(t, t_0) = \mathbb{I} [(t_0 - t)/t_0 > 5\%]$ 表示相对于基线 $t_0$ 有显著加速。

防止奖励欺诈的措施（Efforts to Avoid Reward Hacking）
系统强制执行以下约束以避免奖励欺诈：
1. 权限控制：用于正确性验证和性能分析的 Python 脚本通过文件权限控制进行保护，防止代理修改或干扰评估逻辑。
2. 禁止回退：使用上下文管理器强制执行时间约束，明确禁止调用 torch.nn.functional 中的回退实现，确保性能提升仅源于生成的 CUDA 内核。
3. 多输入验证：对于每个问题，严格遵循 KernelBench 评估协议，针对五个随机采样的输入验证内核输出，以确保功能正确性。
4. 去噪分析：分析管道经过精心设计，包含适当的设备同步、预热迭代和重复测量平均，大幅减少了测量噪声和指标波动。
5. 无外部访问：不向代理提供任何网络搜索或外部信息检索工具，确保所有解决方案纯粹源自本地执行环境。

3.3 稳定 RL 训练的算法改进

针对初始 RL 尝试中模型在 17 步后即性能崩溃的问题，确定了根本原因并提出了多阶段预热策略。

训练不稳定的根本原因
根源在于严重的领域分布不匹配。基础模型的学习先验与 CUDA 内核编码所需的数据分布偏差显著（CUDA 代码仅占预训练数据的 0.01%）。这种分布差距导致大量低概率的 CUDA 内核代码 token 被采样。此外，当训练和推理引擎使用不同的数值精度（如 BF16 vs FP16）时，这些低概率 token 会导致重要性采样比率方差极大，因为在精度下限附近的微小数值误差会导致重要性比率 $\rho_t(\theta)$ 剧烈波动或爆炸。

为了实现稳定的强化学习，提出了一种简单而有效的预热策略：使用基础模型在单轮 RL 后生成的代理轨迹来初始化 Actor 和 Critic 模型（如图 3 所示）。

单轮预热（Single-Turn Warm-up）
首先在基础模型上执行单轮 RL，以增强其生成 CUDA 内核的能力。使用 PPO 进行优化，其中基础模型既作为策略网络也作为价值网络。

Actor 初始化（Actor Initialization）
接着，在代理轨迹上采用拒绝微调（Rejection Fine-Tuning, RFT）阶段来初始化 Actor 模型 $\pi_\theta$。使用单轮 RL 得到的模型在 3.2 节所述的代理循环中运行，收集 CUDA 代理轨迹。然后，对收集到的轨迹应用 RFT。根据以下规则进行拒绝采样以保留高质量的 rollout：
1. 结果过滤：仅保留获得正奖励（$R > 0$）的轨迹。
2. 模式过滤：丢弃表现出低效或无效行为的轨迹，例如冗余的多轮循环或违反预定义工具调用模式的幻觉。

过滤后的轨迹用于通过标准监督微调优化 Actor，目标函数如下：

$$\mathcal{L}_{\mathrm{RFT}}(\theta)=-\mathbb{E}_{\tau \sim \mathcal{D}^{\prime}}\left[\sum_{t=1}^{T} \log \pi_{\theta}\left(a_{t} \mid s_{t}, a_{<t}\right)\right],$$ <p>其中 $\tau = (s_0, s_1, \dots, s_{T-1})$ 表示过滤后的 CUDA 代理轨迹，$\pi_\theta$ 是参数化为 $\theta$ 的策略，$D'$ 表示拒绝采样后的数据集。

Critic 初始化（Critic Initialization）
执行价值预训练（Value Pretraining）来初始化 Critic。具体而言，利用包含状态序列及其对应结果奖励的采样轨迹数据来预训练 Critic 网络。令 $\tau = (s_0, s_1, \dots, s_{T-1})$ 表示轨迹，其中 $s_t$ 表示 token 位置 $t$ 处的状态，$r$ 表示在最后一个 token 处分配的结果奖励。使用广义优势估计（Generalized Advantage Estimation）计算目标值：

$$V_{t}^{\text{targ}} = V_{\phi}(s_{t}) + \hat{A}_{t}, \quad \text{where} \quad \hat{A}_{t} = \sum_{l=0}^{T-1-t} (\gamma \lambda)^{l} \delta_{t+l},$$

其中 $\delta_t = r_t + \gamma V_\phi(s_{t+1}) - V_\phi(s_t)$ 是时间差分误差，且 $V_\phi(s_T) = 0$。实验中设置 $\gamma = 1$ 和 $\lambda = 0.95$。通过最小化均方误差来优化 Critic 参数 $\phi$：

$$\mathcal{L}_{\mathrm{VP}}(\phi)=\frac{1}{2} \mathbb{E}_{\tau \sim \mathcal{D}}\left[\frac{1}{T} \sum_{t=0}^{T-1}\left(V_{\phi}\left(s_{t}\right)-V_{t}^{\mathrm{targ}}\right)^{2}\right],$$

其中 $\mathcal{D}$ 表示代理轨迹的集合。

RL 算法（RL Algorithm）
采用 PPO 优化 Actor 模型 $\pi_\theta$。令 $\pi_{\theta_{\text{old}}}$ 表示用于轨迹采样的旧策略。使用剪切的代理目标函数最大化预期回报：

$$\begin{aligned} \begin{aligned} \mathcal{L}^{\text{CLIP}}(\theta) = \mathbb{E}_{\tau \sim \mathcal{D}} \bigg[ \frac{1}{T} \sum_{t=0}^{T-1} \min \big( & \rho_t(\theta) \hat{A}_t, \\ & \text{clip}(\rho_t(\theta), 1 - \epsilon_{\text{lower}}, 1 + \epsilon_{\text{higher}}) \hat{A}_t \big) \bigg] \end{aligned} \end{aligned}$$

其中 $\rho_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{\text{old}}}(a_t|s_t)}$ 是当前策略与旧策略之间的重要性采样比率，$a_t$ 表示位置 $t$ 处采取的动作（即 token）。遵循相关文献，设置 $\epsilon_{\text{lower}} = 0.2, \epsilon_{\text{higher}} = 0.28$。

3. 设计原则与关键观察

在方法细节之前，本文提出了一些关键的设计原则和观察，这些原则指导了系统的构建。

RL 不稳定性的根源分析：如 3.3 节所述，训练不稳定主要源于基础模型在 CUDA 领域的先验知识极度匮乏（<0.01% 数据占比）。这种分布失配不仅导致采样出低质量代码，还会在使用不同数值精度（如训练用 BF16，推理用 FP16）时，导致低概率 token 的概率计算出现微小误差，进而引发重要性采样比率（Importance Sampling Ratio）的剧烈波动甚至爆炸。

组合融合的重要性：在数据合成中（3.1 节），关键观察是组合多个算子形成融合任务能产生有价值的训练问题。这是因为组合问题通常不等同于简单地优化每个独立算子然后串联。融合改变了优化格局，避免了中间全局内存的物化，并通过共享寄存器/SMEM/占用率约束耦合了各个阶段，需要统一的并行映射和数据布局。

4. 实验环境

• 基础模型：Seed1.6，一个混合专家（MoE）模型，拥有 23B 激活参数和 230B 总参数。
• 训练设置：全局批量大小为 1024。Actor 和 Critic 的学习率分别为 $3 \times 10^{-6}$ 和 $6 \times 10^{-6}$。单轮 RL 上下文窗口为 32768，代理 RL 为 131072。训练 rollouts 最多 150 轮交互，评估时放宽至 200 轮。
• 沙盒环境：设计了 CPU-GPU 资源解耦的沙盒架构。Docker 终端沙盒处理 CPU 任务（如编译），并通过预定义的技能脚本将验证和分析任务分发到专用的 GPU 沙盒池（128 张 NVIDIA H20 GPU）。这种进程级隔离和独占资源分配消除了干扰，确保了稳定的延迟测量。
• 数据集：KernelBench (Level 1 至 Level 3 子集)，共 250 个不同算子任务。
• 基线模型：Claude Opus 4.5, Gemini 3 Pro, GLM 4.6, Kimi K2 以及 torch.compile。

5. 实验结果

主结果分析
Table 1 总结了 CUDA Agent 与强专有模型基线在 KernelBench 上的性能。

1. CUDA 编码能力：与 Claude Opus 4.5 和 Gemini 3 Pro 相比，CUDA Agent 展现了更强的编码能力。虽然专有模型通过率较高（91.2%–95.2%），但其 Faster Rate 较低（66%–69%），说明它们生成的内核往往无法超越 torch.compile。相比之下，CUDA Agent 实现了 98.8% 的通过率和 96.8% 的 Faster Rate。
2. 超越静态编译器：CUDA Agent 证明了学习到的优化策略可以持续超越静态编译器启发式算法，特别是在算子融合等复杂场景中。在 Level 2 任务（算子序列）中，CUDA Agent 实现了完美的 100% Faster Rate 和相对于 torch.compile 的 2.80 倍加速。

消融研究
Table 2 展示了不同组件的影响。

1. 代理循环的影响：移除代理循环（即单轮模型）导致正确性和优化质量大幅下降。单轮模型无法获得编译错误和分析器反馈，导致生成的内核性能甚至不如基线。
2. 奖励设计的影响：将鲁棒奖励替换为原始加速比奖励（w/o Robust Reward）后，虽然功能正确性相当，但优化性能显著减弱。鲁棒奖励通过设定明确的性能里程碑，引导策略发现真正的加速。
3. 多阶段训练的影响：移除 RFT 或 Value Pretraining 都会导致性能下降，且表现出训练不稳定和最终崩溃（Collapse）。
   • RFT：移除 RFT 导致训练奖励迅速崩溃（图 4a），且策略熵急剧增加（图 4b），表明策略分布变得发散且结构混乱。RFT 提供了关键的行为先验。
   • Value Pretraining：没有初始化的 Critic 无法捕捉多轮交互状态的价值（图 5a），导致轨迹长度爆炸（图 5b），因为未初始化的 Critic 无法惩罚无效的搜索路径。

6. 结论

本文介绍了 CUDA Agent，这是一个大规模代理强化学习系统，赋予了大型语言模型在真实的、执行驱动的开发工作流中生成和优化 CUDA 内核的能力。通过联合扩展数据合成、代理环境和面向稳定性的 RL 训练，CUDA Agent 将 LLM 从被动的代码生成器转变为主动的系统优化器，在 KernelBench 上相对于 torch.compile 和强专有模型取得了一致的增益。这些结果表明，为基础模型配备结构化环境和可靠的执行反馈，是实现 GPU 计算中性能关键型软件开发自动化的有效途径。

7. 附录

A 训练数据详情

• 算子格式：每个训练样本是一个用 PyTorch 实现的 Python 类（torch.nn.Module 子类）。包含 __init__ 和 forward 方法，以及用于生成的辅助函数 get_init_inputs() 和 get_inputs()。
• 数据去污染：使用基于 AST 的代码相似度工具进行去污染。如果训练样本与任何评估程序的 AST 相似度超过 0.9，则将其移除。过滤后无高相似度样本（图 7）。
• 数据集构成：主要由通过顺序堆叠 1 到 5 个 torch 算子构成的复合算子类组成，以及部分 transformers 算子。

B 代理循环详情
* 工具列表：Bash（执行命令）、Read/Write（文件操作）、Edit/MultiEdit（代码修改）、Glob（文件查找）、Grep（代码搜索）、NotebookEdit（Notebook 操作）、BashOutput/KillBash（进程管理）。
* http://SKILL.md：定义了严格的限制（禁止在 C++ 中使用 torch 算子，禁止修改基础设施文件），统一的工作流（实现 -> 编译测试 -> 性能优化），以及优化清单（内存合并、内核融合、共享内存等）。

C 并行工作讨论
与其他工作的主要区别：
* STARK：多代理固定角色 vs 本文单代理自主交互。
* ReGraphT：侧重于蒸馏到小模型。
* EvoEngineer：仅在 KernelBench 的子集上评估。
* CudaForge：基于固定交互协议的双代理系统。
* Kevin & CUDA-L1 & ConCuR：存在使用 KernelBench 数据进行训练或微调的数据泄露问题，使得比较不公平。

8. 补充细节（案例研究）

常见优化模式
CUDA Agent 展示了多种常见优化模式：
1. 代数简化和算子规约：将复杂的张量表达式简化。例如将对角矩阵乘法简化为行缩放。
2. 内核融合：将多个逻辑相关的操作（如加法、乘法、激活）合并为一个内核，减少中间张量物化和启动开销。
3. 内存访问优化：强调合并全局内存访问，使用向量化加载（float4），以及利用共享内存进行块内归约。
4. 硬件感知优化：启用 TF32 计算以利用 Tensor Cores。
5. 库感知优化：识别并调用高度优化的库（如 cuDNN）来替代多个独立算子。

案例 D.2 (Level 1): 对角矩阵乘法
参考模型计算对角向量构造的矩阵与稠密矩阵的乘积。CUDA Agent 识别出这等价于行缩放，从而避免了构建中间对角矩阵和执行 GEMM，将复杂度从 $O(N^2M)$ 降低到 $O(NM)$。生成的内核使用网格跨步循环直接执行行缩放，实现了 73.31 倍的加速。

案例 D.3 (Level 2): 矩阵乘、除、求和及缩放
参考模型包含一系列操作。CUDA Agent 通过代数重排，将“矩阵乘法后归约”转换为“权重矩阵列归约后点积”：

$$\sum_{j} \frac{x_{i} \cdot w_{j}^{T}}{2}=x_{i} \cdot\left(\sum_{j} w_{j}^{T}\right) / 2,$$

生成的实现包含两个内核：一个计算权重列和，另一个执行点积并融合了除法和缩放。利用 float4 加载和共享内存归约，实现了 24.04 倍的加速。

案例 D.4 (Level 3): ResNet BasicBlock
对于包含卷积、BN、残差连接的 ResNet 块。CUDA Agent 采用了：
1. 在 Python 层面将 BatchNorm 参数折叠进卷积权重。
2. 使用 cudnnConvolutionBiasActivationForward API 在单个 cuDNN 内核中执行卷积、偏置加法和 ReLU。
3. 启用 TF32。
4. 将残差加法和最后的 ReLU 融合为一个自定义 CUDA 内核。
最终实现了 3.59 倍的加速。