CUDA Agent
论文: CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation 作者: ByteDance Seed & 清华大学 AIR 日期: 2026.02
1. 问题背景
GPU kernel 优化是现代深度学习基础设施的核心，但这项工作极度依赖对 GPU 微架构的深入理解和复杂的 profiling 工具链。尽管 LLM 在通用编程任务上已经接近人类水平，但在 CUDA kernel 生成上仍然不如 torch.compile 这样的编译器系统。
现有方法分两条路线，都有明显瓶颈：
Training-free 方法（STARK, ReGraphT, EvoEngineer, CudaForge）：设计手工搜索/优化启发式，依赖基座模型本身的 CUDA 编码能力。模型不会写 CUDA，再怎么搜索也没用。
Fine-tuning 方法（Kevin, CUDA-L1, ConCuR）：在固定的多轮执行反馈循环中微调模型。问题是把所有历史解决方案塞进 context，浪费 context length，同时限制了 agent 自主学习调试、搜索和 profiling 策略的能力。
CUDA Agent 的核心思路：不要给 agent 一个固定的优化流程让它照着走，而是让它在一个真实的 CUDA 开发环境中自主探索，通过 RL 从根本上提升模型的 CUDA 优化能力。
2. 系统架构总览
CUDA Agent 由三个核心组件构成：
- 可扩展的数据合成 Pipeline — 解决训练数据稀缺问题
- Skill-Integrated Agent Loop — 提供真实的 CUDA 开发环境和可靠的 reward 信号
- RL 算法改进 — 解决训练不稳定问题，支持 150 步稳定训练
3. 数据合成 Pipeline
高质量 CUDA kernel 数据极度稀缺，手动实现专家级参考代码成本过高。论文设计了三阶段数据合成流程：
3.1 种子算子爬取
从 torch 和 transformers 库中挖掘 PyTorch 实现的参考算子。每个算子表示为带 __init__ 和 forward 方法的 torch.nn.Module 子类。只选用维护良好的官方库，排除质量不可控的个人仓库。
3.2 组合式问题构造
用 LLM 从 torch 库中采样最多 5 个算子类，顺序组合成融合任务。
这里有一个关键洞察：组合多个算子产生的融合任务 ≠ 单独优化每个算子再串联。融合改变了优化格局：
- 避免中间结果写回全局内存（intermediate global-memory materialization）
- 耦合了 register/shared memory/occupancy 约束
- 需要统一的并行映射和数据布局
所以组合式问题天然就是好的 CUDA 训练数据。
3.3 严格过滤
四个过滤条件：
- 必须在 Eager 和 Compile 模式下都能成功执行
- 排除具有随机性的算子（确保可复现性）
- 反作弊验证：不同输入的输出不能是常量或数值不可区分
- Eager 模式执行时间限制在 1ms-100ms（过滤过简单或过重的任务）
额外做了 AST 相似度去重，排除与 KernelBench 测试集相似度 > 0.9 的样本。
最终数据集 CUDA-Agent-Ops-6K 包含 6000 个样本，组成如下：
| 类别 | 占比 |
|---|---|
| torch 2-op 组合 | 83.77% |
| torch 3-op 组合 | 7.62% |
| torch 单 op | 3.40% |
| torch 4-op 组合 | 2.80% |
| torch 5-op 组合 | 1.23% |
| transformers 算子 | 1.18% |
4. Skill-Integrated Agent Loop（核心）
这是论文最重要的设计。
4.1 Agent Loop 架构
基于 OpenHands 框架，采用 ReAct 范式（Reasoning → Action → Observation 交替循环），支持迭代式编码、调试和性能优化。
Agent 的工具集（Action Space）：
| 工具 | 功能 |
|---|---|
| Bash | 持久 shell 会话，执行编译、依赖管理、程序运行 |
| Read / Write | 文件读写（写操作有 read-before-write 保护） |
| Edit / MultiEdit | 确定性字符串级代码修改，支持多文件原子编辑 |
| Glob | 基于 glob 模式的文件发现 |
| Grep | 基于 ripgrep 的代码搜索 |
| BashOutput | 流式输出后台进程的增量输出（监控编译等长任务） |
| KillBash | 终止后台 shell 会话 |
注意：不提供任何 web 搜索或外部信息检索工具，确保所有解决方案纯粹来自本地执行环境。
4.2 CUDA Coding Skill
受 Anthropic 的 Agent Skills 思想启发，将 CUDA 编码特定的指令和工具以 Agent Skills 格式提供。核心是一个 SKILL.md 文件，规范了 CUDA kernel 优化的标准流程：
Step 1 — 分析：使用 profile.py 分析原生 PyTorch 实现的性能，识别瓶颈（过多 kernel 启动、次优内存访问模式等）
Step 2 — 实现：在 model_new.py 中重写模型，在 kernels/ 目录开发 CUDA kernel 源文件和绑定代码
Step 3 — 编译评估：在 GPU sandbox 环境中编译和评估，迭代优化直到数值正确性和性能要求都满足
Step 4 — 迭代：重复 Step 2-3，直到最终实现比 torch.compile 基线快至少 5%
SKILL.md 中还包含：
- 工作区结构规范（哪些文件可以修改，哪些不能）
- CUDA kernel 编写模板（
.cu文件 +_binding.cpp文件分离，加速编译） - 优化策略优先级：算法优化（>50% 收益）> 硬件利用率（20-50%）> 精细调优（<20%）
- 常见问题排查指南和完整的优化 checklist
这个设计的精妙之处在于：不是把 CUDA 知识硬编码到模型里，而是以”技能文档”的形式提供，让 agent 在交互过程中自主参考和运用。
4.3 Robust Reward Scheduling
论文发现直接用 speedup 作为 reward 有问题：不同算子优化难度差异巨大，原始 speedup 是代码质量的不可靠代理。一个简单算子 10x speedup 可能比复杂算子 1.5x speedup 容易得多。
提出归一化的离散 reward 方案：
r = -1 正确性检查失败
r = 3 比 eager 和 compile 都快 5% 以上
r = 2 只比 eager 快 5% 以上
r = 1 正确但不够快消融实验证实，这种基于里程碑的 reward 显著优于连续 speedup reward（Faster Rate: 96.8% vs 60.4%）。
4.4 防止 Reward Hacking
这是工程上非常重要的细节。之前的工作（如 Sakana AI 的 AI CUDA Engineer）就被发现存在 reward hacking 问题。CUDA Agent 采取了五重防护：
- 文件权限控制：正确性验证和性能 profiling 的 Python 脚本通过文件权限保护，agent 无法修改评估逻辑
- 禁止回退：通过 context manager 禁止调用
torch.nn.functional的回退实现，确保性能提升来自真正的 CUDA kernel - 多输入验证：每个问题用 5 个随机采样的输入验证 kernel 输出，遵循 KernelBench 评估协议
- 精确 profiling：proper device synchronization、warm-up iterations、重复测量取平均，大幅降低测量噪声
- 无外部信息：不提供 web 搜索工具，所有解决方案必须来自本地环境
4.5 Sandbox 环境
采用 CPU-GPU 资源解耦架构：
- Docker-based terminal sandbox 处理 CPU 密集任务（kernel 编译）
- Agent 通过预定义的 skill scripts 将验证和 profiling 任务分发到专用 GPU sandbox pool（128 块 NVIDIA H20 GPU）
- 进程级隔离和独占资源分配，消除进程间干扰，确保延迟测量稳定
5. RL 算法改进：解决训练不稳定
这是论文的另一个关键贡献。
5.1 训练不稳定的根因
初始 RL 试验只能稳定训练 17 步就崩溃。根因是严重的领域分布不匹配：
- CUDA 编码数据在预训练数据中占比 不到 0.01%
- 基座模型的先验分布与 CUDA kernel 编码所需的数据分布严重偏离
- 大量采样到低概率的 CUDA kernel 代码 token
- 训练和推理引擎使用不同数值精度（BF16 vs FP16）时，低概率 token（如 π(a|s) ≈ 10⁻⁹）导致 importance sampling ratio 方差爆炸
5.2 三阶段 Warm-up 策略
为了解决这个问题，论文提出了一个简单但有效的 warm-up 策略：
Stage 1: Single-Turn RL Warm-up
先在基座模型上进行单轮 RL（标准代码生成目标，模型一次性预测最终 kernel，没有执行反馈）。使用 PPO 优化，context window 32K tokens。这一步提升模型的基础 CUDA 生成能力。
Stage 2: Actor 初始化 — Rejection Fine-Tuning (RFT)
用 Stage 1 的模型在完整 agent loop 中收集 CUDA agent 轨迹，然后做 rejection sampling：
- 结果过滤：只保留正 reward (R > 0) 的轨迹
- 模式过滤：丢弃低效行为（冗余多轮循环、违反工具调用 schema 的幻觉）
用过滤后的轨迹做标准 SFT 初始化 actor。
RFT 的作用是提供一个强行为先验，约束 RL 训练过程中的策略熵增长。消融实验显示，去掉 RFT 会导致训练 reward 快速崩溃，策略熵急剧上升（策略分布变得越来越弥散，输出不连贯）。
Stage 3: Critic 初始化 — Value Pretraining
用采样的轨迹数据（状态序列 + outcome reward）预训练 critic network。使用 GAE（Generalized Advantage Estimation）计算目标值，γ=1, λ=0.95。
Value Pretraining 的作用是让 critic 能立即提供准确的价值估计。消融实验显示，去掉 Value Pretraining 会导致 critic 无法学到有意义的 value function（explained variance 极低），进而导致轨迹长度爆炸——未初始化的 critic 无法惩罚无效搜索路径，agent 陷入无限循环。
5.3 Agentic RL 训练
完成三阶段 warm-up 后，使用 PPO 进行 agentic RL 训练：
| 参数 | 值 |
|---|---|
| 基座模型 | Seed1.6 (23B active / 230B total, MoE) |
| Context window | 131072 tokens (128K) |
| 最大 agent turns | 150（训练）/ 200（评估） |
| Global batch size | 1024 |
| Actor LR | 3×10⁻⁶ |
| Critic LR | 6×10⁻⁶ |
| PPO clip (lower/higher) | 0.2 / 0.28（非对称） |
| 训练步数 | 150 |
经过这些改进，RL 训练可以稳定进行 200 步并持续 reward 增长（对比初始的 17 步崩溃）。
6. 实验结果
6.1 主要结果
在 KernelBench（250 个算子任务，Level 1-3）上的表现：
| 模型 | Pass Rate | Faster Rate (vs Compile) | Speed-up (vs Compile) |
|---|---|---|---|
| Seed1.6 (base) | 74.0% | 27.2% | 0.69× |
| GLM 4.6 | 75.6% | 19.2% | 0.57× |
| Kimi K2 | 66.8% | 22.8% | 0.66× |
| Gemini 3 Pro | 91.2% | 69.6% | 1.42× |
| Claude Opus 4.5 | 95.2% | 66.4% | 1.46× |
| CUDA Agent | 98.8% | 96.8% | 2.11× |
三个关键发现：
-
CUDA Agent 全面碾压顶级闭源模型。Claude Opus 4.5 和 Gemini 3 Pro 虽然 Pass Rate 不错（91-95%），但 Faster Rate 只有 66-69%，说明通用 LLM 经常生成”能跑但不快”的 naive kernel。CUDA Agent 的 96.8% Faster Rate 说明专门的 RL 训练让模型学会了持续生成正确且高度优化的实现。
-
在算子融合场景（Level 2）上优势最大。CUDA Agent 达到 100% Faster Rate 和 2.80× speedup。传统编译器依赖预定义的规则模式做 kernel fusion，遇到非平凡的算子组合就力不从心。CUDA Agent 通过迭代 agent loop 探索了更大的设计空间，发现了静态后端无法触及的硬件特定内存访问模式和 tiling 策略。
-
在最难的 Level 3 上比最强闭源模型高约 40%（Faster Rate: 90% vs 50-52%）。
有趣的细节：ChatGPT-5 系列（5, 5.1, 5.2）在评估中持续拒绝响应 CUDA 相关 prompt，无法参与评测。
6.2 消融实验
| 变体 | Faster Rate (vs Compile) | Speed-up (vs Compile) |
|---|---|---|
| w/o Agent Loop | 14.1% | 0.69× |
| w/o Robust Reward | 60.4% | 1.25× |
| w/o RFT | 49.8% | 1.05× |
| w/o Value Pretraining | 50.9% | 1.00× |
| CUDA Agent (full) | 96.8% | 2.11× |
每个组件都不可或缺：
- 去掉 Agent Loop 影响最大（96.8% → 14.1%）：没有编译错误、运行时失败和 profiler 反馈的迭代交互，模型根本无法有效优化
- 去掉 Robust Reward：Pass Rate 几乎不变，但优化质量大幅下降
- 去掉 RFT：训练 reward 快速崩溃，策略熵爆炸
- 去掉 Value Pretraining：critic 失效，轨迹长度爆炸
7. Case Studies：CUDA Agent 学到了什么
7.1 代数简化（Level 1）
任务：对角矩阵与稠密矩阵相乘。
原始实现显式构造对角矩阵再做 GEMM。CUDA Agent 识别出这等价于逐行缩放（diag(A) × B 等价于每行 i 乘以 A[i]），将复杂度从 O(N²M) 降到 O(NM)。
结果：73.31× speedup vs torch.compile。
7.2 Kernel 融合 + 代数重写（Level 2）
任务：矩阵乘法 → 除法 → 求和 → 缩放。
CUDA Agent 利用求和和矩阵乘法的线性性，将 Σⱼ(xᵢ·wⱼᵀ)/2 重写为 xᵢ·(Σⱼwⱼᵀ)/2，把矩阵乘法+归约变成列求和+点积。然后融合为两个高效 kernel，使用 float4 向量化加载和 shared memory tree reduction。
结果：24.04× speedup vs torch.compile。
7.3 硬件感知优化（Level 3）
任务：ResNet BasicBlock（卷积 + BatchNorm + ReLU + 残差连接）。
CUDA Agent 的优化策略：
- 将 BatchNorm 参数折叠进卷积权重和 bias，消除 BatchNorm 算子
- 使用
cudnnConvolutionBiasActivationForward融合 conv + bias + ReLU 为单个 cuDNN kernel - 启用 TF32 计算，利用 Tensor Core
- 融合残差加法和最终 ReLU 为单个自定义 CUDA kernel
- 尝试了 NCHW → NHWC 布局转换以更好利用 Tensor Core，但发现转换开销抵消了收益，最终保留 NCHW
结果：3.59× speedup vs torch.compile。
7.4 通用优化模式总结
跨所有难度级别，CUDA Agent 展现出几类反复出现的优化模式：
- 代数简化：识别高层张量表达式的数学结构，用更简单的等价计算替代
- Kernel 融合：合并逻辑相关的操作，消除中间张量物化，减少 kernel 启动开销
- 内存访问优化：coalesced global memory access、shared memory、vectorized loads (float4)
- 硬件感知优化：TF32 Tensor Core、精度选择
- 库感知优化：在合适时调用 cuDNN/cuBLAS 融合 API
8. 局限性
- 未与 TVM 等更复杂的编译器框架比较。TVM 的调优开销和部署复杂度难以集成到大规模 RL 训练循环（数千次 rollout）中，论文选择
torch.compile作为广泛采用的基线。 - 计算资源需求巨大。训练 pipeline 依赖 128 块 H20 GPU pool 和进程级隔离，计算和工程成本高，限制了更广泛研究社区的可及性。
9. Takeaways
这篇论文的核心贡献不只是一个更好的 CUDA kernel 生成器，而是验证了一个更广泛的范式：
给基础模型配备结构化环境和可靠的执行反馈 reward，可以将它们从被动的代码生成器转变为主动的系统优化器。
几个值得关注的设计选择：
- Agentic RL > 固定多轮循环：让 agent 自主决定何时编译、profiling、调试、重写，比预定义的固定交互协议更有效
- Agent Skills 范式：将领域知识结构化为 skill 文档而非硬编码到模型中，兼顾了灵活性和专业性
- 离散里程碑 reward > 连续 speedup reward：归一化的 reward 方案更稳定，更能引导策略发现有效优化
- 三阶段 warm-up 解决分布不匹配：RFT 初始化 actor + Value Pretraining 初始化 critic，是在数据稀缺领域做 agentic RL 的通用方案
- 防 reward hacking 的工程实践：文件权限隔离、禁止回退、多输入验证——这些细节决定了 RL 训练能否真正 work