CUDA Agent：面向高性能 CUDA 内核生成的大规模智能体强化学习系统 CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation
通过大规模智能体 RL 训练 LLM 生成超越 torch.compile 的高性能 CUDA 内核
前置知识
CUDA 内核优化
CUDA（Compute Unified Device Architecture）是 NVIDIA 提供的并行计算平台和编程模型。CUDA 内核（kernel）是运行在 GPU 上的并行函数，开发者需要手动管理线程层次（grid、block、thread）、共享内存、寄存器等底层硬件资源。高性能 CUDA 内核优化涉及多个层次：算法层面的内核融合（kernel fusion）减少中间结果的全局内存写回；内存层面的合并访问（coalesced access）和共享内存 tiling 提升带宽利用率；计算层面利用向量化加载（float4）和 warp 级原语（__shfl_sync）减少指令数。传统上这是一项高度专业的工作，需要对 GPU 微架构有深入理解。
本文的核心任务就是自动生成和优化 CUDA 内核，理解 CUDA 编程的复杂性是理解本文动机和方法设计的基础。
强化学习（Reinforcement Learning）用于代码生成
强化学习是一种通过与环境交互、获取奖励信号来优化策略的训练范式。在代码生成场景中，智能体（agent）通过执行代码获得编译反馈、正确性验证和性能测量作为奖励信号，从而学习生成更高质量的代码。相比监督学习，RL 的优势在于不需要高质量的人工标注数据，而是通过试错探索来发现更好的策略。本文使用的 PPO（Proximal Policy Optimization）是一种稳定的策略梯度算法，通过裁剪重要性采样比率来避免过大的策略更新。
本文的整个训练框架建立在大规模 RL 之上，理解 RL 的基本原理（特别是 PPO 的裁剪机制和奖励设计）是理解本文技术贡献的关键。
智能体循环（Agent Loop）与 ReAct 范式
智能体循环是指 LLM 在多轮交互中交替进行推理（Reasoning）、行动（Action）和观察（Observation）的过程。ReAct 范式让模型先思考当前状态，然后调用工具执行操作（如编写代码、编译、运行测试），观察执行结果后继续推理。这种多轮交互模式模拟了人类程序员的开发工作流——写代码、编译、调试、优化的迭代过程。OpenHands 是一个广泛采用的代码智能体框架，本文的 agent loop 设计与其保持一致以确保通用性。
本文的核心创新之一是将单轮代码生成扩展为多轮智能体交互，理解 agent loop 和 ReAct 范式有助于理解为什么多轮交互比单轮生成更有效。
Actor-Critic 架构与优势估计（Advantage Estimation）
Actor-Critic 是 RL 中的一种经典架构，Actor（策略网络）决定采取什么动作，Critic（价值网络）评估当前状态的好坏。广义优势估计（GAE）通过结合多步时序差分误差来计算优势函数 $\hat{A}_t$，平衡偏差和方差。在 PPO 训练中，优势函数用于衡量某个动作相对于平均表现的好坏，指导策略更新的方向。如果 Critic 初始化不当，优势估计不准确会导致策略更新方向错误，最终导致训练崩溃。
本文发现并解决了 Critic 初始化不当导致的训练不稳定问题，理解 Actor-Critic 架构和 GAE 的工作原理是理解本文技术贡献中多阶段训练策略的前提。
KernelBench 基准测试
KernelBench 是一个用于评估 LLM 生成 CUDA 内核能力的基准测试，包含三个难度级别：Level-1（单个算子，100 题）、Level-2（算子序列，100 题）和 Level-3（复杂模型组件如 ResNet BasicBlock，50 题）。评估指标包括 Pass Rate（编译和正确性通过率）、Faster Rate（比基线更快的比例）和 Speed-up（执行加速比的几何均值）。基线通常为 PyTorch Eager 模式和 torch.compile。
KernelBench 是本文评估的主要基准，理解其难度分级和评估指标是理解本文实验结果的前提。
研究动机
GPU 内核优化是现代深度学习基础设施的基础组件，NVIDIA 的 CUDA 架构目前主导着 AI 硬件生态。然而，开发和优化高性能 CUDA 内核仍然是一项极具挑战性的任务，需要对 GPU 微架构特性有深入理解，并熟练使用复杂的性能分析工具。尽管 LLM 在通用软件开发任务上已展现出接近人类的能力，但在 CUDA 内核生成方面，现有方法甚至无法与自动优化工具如 torch.compile 竞争，更不用说与人类专家相比。已有两条研究路径：一是设计免训练的工作流（如 STARK、EvoEngineer），依赖手工设计的启发式策略在执行反馈下迭代优化，但这些方法无法弥补基础模型在 CUDA 编码能力上的根本缺陷，性能提升受到模型内在能力的严重制约；二是在固定多轮反馈循环中微调模型（如 Kevin、CUDA-L1），但这类方法将所有历史方案塞入上下文窗口浪费了宝贵的上下文长度，同时限制了智能体自主学习调试、搜索和性能分析策略的能力。
本文的目标是本文的具体目标是构建一个大规模智能体强化学习系统 CUDA Agent，通过三个互补维度——数据合成、智能体环境和 RL 算法技术——系统性地提升基础模型的 CUDA 内核编码能力。在量化指标上，CUDA Agent 的目标是在 KernelBench 基准测试上实现超越 torch.compile 的性能，在三个难度级别上均达到更高的正确率和加速比，确立 LLM 驱动的内核生成作为传统编译器优化的有竞争力甚至更优的替代方案。
与已有工作不同的是，本文的独特切入角度在于它认识到现有方法的根本瓶颈：不是缺乏推理能力，而是模型在 CUDA 编码这一高度专业化领域的内在能力不足。CUDA 编码数据在预训练数据中占比不到 0.01%，这意味着基础模型在这个领域的知识几乎是空白的。本文的洞察是：与其通过推理时的搜索来弥补能力不足，不如通过大规模 RL 训练从根本上提升模型的 CUDA 编码能力。具体来说，本文抓住了三个被忽视的关键点：一是训练数据的多样性——通过组合式数据合成生成跨越不同难度级别的训练问题；二是训练环境的可靠性——通过严格的权限隔离和防奖励黑客机制确保训练信号的准确性；三是训练稳定性——通过多阶段热启动策略解决 RL 训练中的分布不匹配导致的崩溃问题。
核心方法
CUDA Agent 的整体思路可以用一个类比来理解：传统的 CUDA 内核优化方法就像是给一个从未学过编程的人一份操作手册让他写程序——他可能会照搬模板，但无法真正理解底层原理。CUDA Agent 则是通过大规模实践训练一个程序员，让他在不断试错中掌握 CUDA 优化的精髓。技术路线上，系统包含三个核心组件：第一，可扩展的数据合成管线，从 PyTorch 和 Transformers 库中挖掘种子算子，通过 LLM 组合生成 6000 个不同难度的训练任务；第二，技能增强的 CUDA 开发环境，通过 SKILL.md 文件将标准化的内核开发流程注入模型，配合自动化的正确性验证和性能分析工具提供可靠的奖励信号；第三，多阶段 RL 训练策略，包括单轮 RL 热启动、基于拒绝采样的 actor 初始化（RFT）和 critic 价值预训练，确保大规模训练的稳定性。整个系统支持 128k 上下文长度和最多 200 轮交互，基于 Seed1.6（23B 激活参数、230B 总参数的 MoE 模型）进行训练。
本文的核心创新点在于将 CUDA 内核生成从一个静态代码生成问题转化为一个动态的智能体优化问题，并通过大规模 RL 训练从根本上提升模型的内在能力。与已有方法最本质的区别有三点。首先，已有方法（如 Kevin）的 RL 训练仅在固定多轮循环中进行，将所有历史方案堆积在上下文中，而 CUDA Agent 采用 ReAct 风格的智能体循环，让模型自主决定何时调用工具、如何解释反馈、怎样迭代修改内核，这赋予了模型更大的自主性和更高效的上下文利用。其次，已有方法使用原始加速比作为奖励信号，但不同算子的优化难度差异极大，导致奖励信号噪声很大且偏向简单任务。本文提出离散化的鲁棒奖励方案 $r \in \{-1, 1, 2, 3\}$，通过明确的性能里程碑（比 eager 快 5%、比 compile 快 5%）来归一化奖励，显著提升了优化效果。第三，本文发现了 RL 训练崩溃的根本原因——CUDA 编码数据在预训练中占比不到 0.01% 导致的分布不匹配问题，并通过 RFT 和价值预训练的多阶段热启动策略解决了这一问题，使训练从只能稳定运行 17 步提升到 200 步。
方法步骤详情
CUDA Agent 的完整方法流程分为三个阶段。第一阶段是数据合成管线：首先从 PyTorch torch 库和 transformers 库中爬取种子算子，每个算子表示为包含 __init__ 和 forward 方法的 Python 类；然后利用 LLM 进行组合式合成，从 torch 库中采样不超过 5 个算子类按顺序堆叠为单个计算层，生成更复杂的融合任务；最后通过严格过滤（必须在 Eager 和 Compile 模式下成功执行、无随机性、输出非平凡、执行时间在 1-100ms 范围内），得到 6000 个精选训练样本。第二阶段是智能体环境构建：将 CUDA 编码流程编码为 SKILL.md 文件，指导模型按照分析性能瓶颈、实现自定义 CUDA 内核、编译验证和迭代优化的标准化流程工作；同时设计了 CPU-GPU 资源解耦的沙箱架构，Docker 终端沙箱处理编译等 CPU 任务，GPU 沙箱池（128 张 NVIDIA H20 GPU）处理验证和性能分析任务，通过进程级隔离消除干扰并确保稳定的延迟测量。第三阶段是多阶段 RL 训练：先进行单轮 RL 热启动（使用 PPO 优化基础模型）；然后用热启动模型收集智能体轨迹，通过 RFT（仅保留正奖励且无低效行为的轨迹）初始化 actor 模型；接着用收集到的轨迹进行价值预训练，通过 GAE 计算目标值并最小化均方误差初始化 critic 模型；最后进行完整的智能体 RL 训练，使用 PPO 的裁剪代理目标优化策略，上下文长度从 32768 扩展到 131072，训练 150 步。
技术新颖性
CUDA Agent 在技术新颖性上有多个层面的贡献。在数据层面，本文的关键观察是：将多个算子组合为融合任务并不是简单地串联各算子的优化结果——算子融合通过避免中间全局内存写回、通过共享寄存器/SMEM/占用率约束耦合各阶段、以及要求统一的并行映射和数据布局，从根本上重塑了优化景观。这种基于组合的数据合成方法系统地扩展了任务空间，而不仅仅是增加数据量。在奖励设计层面，离散化的鲁棒奖励方案 $r \in \{-1, 1, 2, 3\}$ 相比常用的连续加速比奖励有本质区别：它不直接回归噪声较大的运行时间比率，而是通过明确的性能里程碑分配信用，使策略更可靠地发现真正有效的优化转换。在训练稳定性层面，本文对训练崩溃的根因分析具有独到见解：CUDA 编码数据在预训练中占比不到 0.01% 导致的领域分布不匹配，使得模型采样到大量低概率的 CUDA 代码 token，当训练和推理引擎使用不同数值精度（如 BF16 vs FP16）时，这些低概率 token 导致重要性采样比率 $\rho_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$ 剧烈波动甚至爆炸。多阶段热启动策略（RFT + 价值预训练）是解决这一问题的简洁而有效的方案。
实验结果
CUDA Agent 在 KernelBench 基准测试上取得了全面领先的结果。在总体指标上，CUDA Agent 达到 98.8% 的 Pass Rate、98.4% 的 vs. Eager Faster Rate 和 96.8% 的 vs. Compile Faster Rate，几何均值加速比分别为 2.60x（vs. Eager）和 2.11x（vs. Compile）。对比最强的商用模型 Claude Opus 4.5（95.2% Pass Rate、66.4% vs. Compile Faster Rate）和 Gemini 3 Pro（91.2% Pass Rate、69.6% vs. Compile Faster Rate），CUDA Agent 在 vs. Compile Faster Rate 上提升了约 27-30 个百分点。在 Level-1（单算子）上，CUDA Agent 实现 100% Pass Rate 和 97% vs. Compile Faster Rate，几何均值加速比 1.87x，超越 torch.compile 基线。在 Level-2（算子序列）上表现最为突出，CUDA Agent 达到 100% Faster Rate 和 2.80x vs. Compile 加速比，这是因为传统编译器依赖预定义的规则模式进行内核融合，面对非平凡的算子组合往往力不从心，而 CUDA Agent 通过智能体循环探索更大的设计空间，发现了硬件特定的内存访问模式和 tiling 策略。在最具挑战性的 Level-3（复杂模型组件如 ResNet BasicBlock）上，CUDA Agent 达到 94% Pass Rate 和 90% vs. Compile Faster Rate，vs. Compile 加速比 1.52x，比 Claude Opus 4.5（1.10x）和 Gemini 3 Pro（1.17x）分别高出约 40%。消融实验进一步验证了各组件的贡献：移除智能体循环导致 vs. Compile Faster Rate 从 96.8% 骤降至 14.1%；移除鲁棒奖励导致 vs. Compile Faster Rate 从 96.8% 降至 60.4%；移除 RFT 导致训练在 30 步内崩溃；移除价值预训练导致交互轨迹长度爆炸。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| KernelBench Level-1（单算子，100题） | vs. Compile Faster Rate | 97.0% | Claude Opus 4.5: 72.0% | +25 个百分点 |
| KernelBench Level-1（单算子，100题） | vs. Compile Speed-up（几何均值） | 1.87x | Claude Opus 4.5: 1.54x | +21.4% |
| KernelBench Level-2（算子序列，100题） | vs. Compile Faster Rate | 100.0% | Gemini 3 Pro: 76.0% | +24 个百分点 |
| KernelBench Level-2（算子序列，100题） | vs. Compile Speed-up（几何均值） | 2.80x | Gemini 3 Pro: 1.46x | +91.8% |
| KernelBench Level-3（复杂模型组件，50题） | vs. Compile Faster Rate | 90.0% | Claude Opus 4.5: 50.0% | +40 个百分点 |
| KernelBench Level-3（复杂模型组件，50题） | vs. Compile Speed-up（几何均值） | 1.52x | Claude Opus 4.5: 1.10x | +38.2% |
| KernelBench Overall（250题加权） | Pass Rate | 98.8% | Claude Opus 4.5: 95.2% | +3.6 个百分点 |
| KernelBench Overall（250题加权） | vs. Compile Faster Rate | 96.8% | Claude Opus 4.5: 66.4% | +30.4 个百分点 |
局限与改进
本文在附录中坦诚地承认了两个主要局限。第一，CUDA Agent 没有与更复杂的编译器框架（如 TVM）进行对比。TVM 等系统虽然可能提供更强的基线，但由于其调优开销大、部署复杂，难以集成到需要数千次 rollouts 的大规模 RL 训练循环中，因此本文选择了 torch.compile 作为广泛采用且训练友好的基线。这意味着 CUDA Agent 的优势可能在面对更先进的编译器时会缩小。第二，训练管线依赖大规模 GPU 资源池（128 张 NVIDIA H20 GPU）和进程级隔离，带来相当大的计算和工程成本，这可能限制更广泛研究社区的可及性。此外，从我的独立观察来看：（1）论文中 ChatGPT-5 系列模型（5、5.1、5.2）因一致拒绝回应 CUDA 相关提示而未能纳入评估，这可能是模型安全策略导致的，但也意味着缺少了一组重要的对比基线；（2）CUDA Agent 的评估和训练虽然严格避免了 KernelBench 数据泄露，但训练数据中 83.77% 为二元算子组合，与 Level-3 的复杂模型组件之间仍存在分布差距；（3）论文未提供推理时间成本的对比，智能体最多 200 轮的交互可能带来显著的延迟开销。
独立分析的弱点
尽管 CUDA Agent 取得了出色的结果，但仍有几个值得关注的弱点。首先，在实际部署场景中，智能体的推理成本可能是一个瓶颈——模型最多进行 200 轮交互、上下文长度达 128k tokens，这意味着生成一个优化内核可能需要大量的 GPU 推理时间和算力。对于需要快速响应的开发场景（如交互式编程），这种延迟可能难以接受。改进方向可以探索更高效的推理策略，如早期终止、分层优化（先粗粒度探索再细粒度调优）或小模型蒸馏。其次，训练管线的资源需求极高（128 张 H20 GPU 进行 RL 训练），这使得大多数研究团队难以复现。可以探索更资源高效的训练策略，如离线 RL、课程学习的自适应调度或更小的基座模型适配。第三，数据合成管线虽然系统化，但 83.77% 的训练样本为二元算子组合，多样性仍有限。对于 Level-3 涉及的复杂神经网络模块（如 ResNet BasicBlock），训练数据和评估任务之间的分布差距可能限制了泛化能力。可以引入更多来自实际深度学习框架的复杂模块作为训练数据。第四，鲁棒奖励方案虽然有效，但离散化的 $r \in \{-1, 1, 2, 3\}$ 在奖励为 2 时仅要求比 eager 快，这可能导致模型在略微超过 eager 后就停止优化。更细粒度的奖励信号或课程式奖励调整可能进一步提升性能。
未来方向
本文作者在结论中暗示了更广泛的愿景：通过结构化环境和可靠的执行反馈，将基础模型从被动的代码生成器转变为主动的系统优化器。这开辟了一条通往 GPU 计算中性能关键型软件开发自动化的道路。基于本文成果，未来研究可以向多个方向延伸：（1）扩展到更多硬件平台——将 CUDA Agent 的方法论迁移到 AMD ROCm、Intel oneAPI 或 Apple Metal 等平台，构建跨平台的内核优化系统；（2）与编译器框架结合——将 CUDA Agent 与 TVM、Triton 等编译器结合，利用编译器的搜索能力作为测试时扩展策略，叠加在 RL 训练后的能力之上；（3）真实世界部署——将系统应用于实际的深度学习框架优化中，如 PyTorch 核心算子库的优化，验证在分布外任务上的泛化能力；（4）更高效的训练范式——探索离线 RL、人类反馈强化学习（RLHF）或基于执行轨迹的知识蒸馏来降低训练成本；（5）扩展优化目标——除了速度优化，还可以探索内存占用最小化、能耗优化等多目标优化场景。
复现评估
从复现评估角度来看，CUDA Agent 的复现面临较大挑战但有一定基础。有利因素方面：（1）论文开源了训练数据集 CUDA-Agent-Ops-6K（发布在 HuggingFace: BytedTsinghua-SIA/CUDA-Agent-Ops-6K），包含 6000 个精选训练样本；（2）SKILL.md 的完整内容在附录中公开，智能体的工具集和开发流程有详细描述；（3）训练超参数（batch size 1024、学习率 $3 \times 10^{-6}$ / $6 \times 10^{-6}$、上下文长度 131072、训练步数 150）均有明确记录；（4）基座模型 Seed1.6 为字节跳动的 MoE 模型（23B 激活/230B 总参数），属于公开可获取的模型。挑战方面：（1）训练需要 128 张 NVIDIA H20 GPU 的大规模资源池，这对大多数研究团队来说是巨大的硬件门槛；（2）CPU-GPU 资源解耦的沙箱架构设计复杂，涉及 Docker 容器化和进程级隔离；（3）Seed1.6 模型的具体权重版本和训练配置未完全公开，可能影响精确复现；（4）智能体的完整训练轨迹和 RL 训练的详细工程细节（如 PPO 的实现细节、经验缓冲区设计等）未在论文中充分披露。总体而言，复现难度较高，主要瓶颈在于算力资源。
论文图表
