CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation高性能CUDA内核生成的大规模智能体强化学习
《CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation》的核心内容是:设计并实现了一个名为 CUDA Agent 的大规模智能体强化学习系统,用于自动生成和优化高性能 CUDA 内核,显著超越现有方法(如 torch.compile 和顶级专有模型)。
以下是其主要研究内容的全面总结与概括:
1. 研究背景与问题
痛点:CUDA 内核优化对深度学习至关重要,但需要深厚的硬件知识,非常专业和复杂。
现状瓶颈:
现有 LLM 在生成 CUDA 内核方面甚至不如编译器的自动优化(如
torch.compile)。先前方法主要有两类,但均有明显缺陷:
免训练方法:依赖手工设计的反馈循环,无法从根本上提升模型本身的 CUDA 能力。
微调方法:在固定多轮反馈中微调,会浪费上下文长度,且限制了智能体自主学习调试和搜索策略的能力。
2. 提出的方法:CUDA Agent
CUDA Agent 是一个基于大规模智能体强化学习的系统,通过三个核心组件系统性地提升基础模型的 CUDA 编程能力:
组件一:可扩展的数据合成管道
目的:解决高质量 CUDA 训练数据稀缺的问题。
方法:
种子爬取:从 PyTorch 和 Transformers 库中收集基础算子。
组合合成:利用 LLM 将多个基础算子组合成更复杂的融合任务(这些任务比单独优化每个算子更具挑战性和价值)。
严格过滤:通过执行反馈过滤掉太简单、太难、随机或与测试集相似的任务,最终得到6,000 个高质量训练样本。
组件二:技能增强的智能体环境
目的:为智能体提供一个真实、不可破解的开发环境,并提供可靠的奖励信号。
方法:
智能体循环:遵循 ReAct 模式(推理-行动-观察),让智能体可以迭代地编码、调试和优化。
CUDA 编码技能:以 Agent Skills 格式提供标准化的 CUDA 优化工作流程(分析 -> 实现 -> 编译评估 -> 重复直到达到 >5% 加速)。
鲁棒奖励设计:设计了一个离散奖励方案(-1, 1, 2, 3),基于正确性和是否显著超越
torch.compile和 Eager 模式,解决了原始加速比奖励易受离群值和任务难度影响的问题。防黑客机制:通过文件权限隔离、禁止回退到 PyTorch 原生实现、多输入验证等手段,确保智能体无法通过捷径获得奖励。
组件三:稳定 RL 训练的算法改进
目的:解决基础模型在 CUDA 任务上训练容易崩溃的问题(原始训练仅 17 步就崩溃)。
根本原因:预训练数据中 CUDA 代码占比极低(<0.01%),导致严重的分布不匹配,以及训练/推理精度不一致导致的重要性采样比率不稳定。
解决方案(多阶段热身策略):
单轮 RL 热身:先用标准 PPO 对基础模型进行单轮训练,提升其基本 CUDA 能力。
Actor 初始化(拒绝采样微调 RFT):用单轮 RL 模型收集智能体轨迹,仅保留正奖励且行为良好的轨迹,再用这些轨迹进行监督微调,作为 Actor 的初始策略,防止熵爆炸和策略崩溃。
Critic 初始化(价值预训练):用收集的轨迹和结果奖励预训练 Critic 网络,使其能准确评估状态价值,避免无效的搜索行为。
3. 实验与结果
基线与数据集:在KernelBench(Level 1-3,共250个任务)上评估,与
torch.compile、Claude Opus 4.5、Gemini 3 Pro 等强基线对比。主要结果:
超越 torch.compile:在 Level-1、2、3 上分别实现100%、100%、92%的更快率(Faster Rate),即生成的核函数比
torch.compile更快。超越专有模型:在最难的 Level-3 上,性能比 Claude Opus 4.5 和 Gemini 3 Pro 高出约40%。
通过率与加速比:总体通过率 98.8%,几何平均加速比相对
torch.compile达到2.11 倍。
消融实验:
智能体循环:移除后,性能大幅下降,证明了交互式环境的重要性。
鲁棒奖励:替换为原始加速比奖励后,优化性能显著减弱。
RFT 初始化:移除后,训练奖励迅速崩溃,策略熵激增,证实了其防止崩溃的关键作用。
价值预训练:移除后,Critic 无法提供有效优势估计,导致轨迹长度爆炸式增长。
4. 核心贡献与意义
提出了 CUDA Agent:一个集成了数据、环境、算法三大组件的大规模智能体 RL 系统。
实现了 SOTA 性能:在 CUDA 内核生成任务上,首次让 LLM 能够稳定、可靠地生成超越传统编译器和顶级专有模型的高性能代码。
范式转变:将 LLM 从被动的“代码生成器”转变为主动的“系统优化器”,为自动化高性能计算软件开发开辟了新道路。
CUDA Agent 通过构建大规模训练数据、模拟真实开发环境的智能体以及一系列稳定强化学习训练的算法,成功训练出一个能够自动编写远超现有工具(包括torch.compile和顶级商业模型)性能的 CUDA 内核的 LLM。这里是自己的论文阅读记录,感兴趣的话可以参考一下,如果需要阅读原文的话可以看这里,如下所示:
官方项目主页地址在这里,如下所示:
项目地址在这里,如下所示:
数据集地址在这里,如下所示:
摘要
GPU内核优化是现代深度学习的基础,但仍然是一项高度专业化的任务,需要深厚的硬件专业知识。尽管大型语言模型(LLM)在通用编程方面表现出色,但在CUDA内核生成方面,它们仍然无法与基于编译器的系统(如torch.compile)相抗衡。现有的CUDA代码生成方法要么依赖于免训练的精化流程,要么在固定的多轮执行-反馈循环中微调模型,但这两种范式都未能从根本上提升模型内在的CUDA优化能力,导致性能提升有限。
我们提出了CUDA Agent,一个大规模的智能体强化学习系统,通过三个组成部分来培养CUDA内核的专业知识:
一个可扩展的数据合成管道。
一个技能增强的CUDA开发环境,具备自动化验证和分析功能,以提供可靠的奖励信号。
能够实现稳定训练的RL算法技术。
CUDA Agent在KernelBench上取得了最先进的结果。与torch.compile相比,在KernelBench的Level-1、Level-2和Level-3分割上,CUDA Agent分别实现了100%、100%和92%的“更快率”(Faster Rate)。在最难的Level-3设置上,它优于最强的专有模型(如Claude Opus 4.5和Gemini 3 Pro)约40%。
1. 引言
GPU内核是现代深度学习基础设施的基础组成部分[12, 17],而NVIDIA的CUDA架构目前在AI硬件生态系统中占据主导地位。尽管被广泛采用,但高性能CUDA内核的开发和优化仍然极具挑战性,这需要深入理解GPU微架构特性[16]和复杂的分析工具包[4]。
尽管大型语言模型(LLM)在通用软件开发任务中已展现出与人类相当的能力[23],但现有的CUDA内核生成方法[4, 12]在CUDA内核代码生成方面甚至无法与torch.compile等自动优化工具[17]竞争,更不用说人类专家了。
先前工作的局限性:
训练-免费工作流:一类工作聚焦于设计免训练的工作流[5, 6, 8],依赖由执行反馈引导的手工设计精化启发式方法。然而,这些方法并未弥补基础模型在CUDA编码能力上的根本缺失,导致性能提升受到模型自身能力的严重限制。
多轮微调:另一类研究[1, 4, 11, 14]尝试在由代码执行反馈驱动的固定多轮精化循环中对基础模型进行微调。但这种方法浪费了上下文长度(包含所有之前的解决方案),并限制了智能体学习调试、搜索和分析策略的自主性。
为了克服这些限制,我们引入了CUDA Agent,一个大规模CUDA智能体训练系统,通过在三个互补维度做出贡献,系统地增强了基础模型的CUDA内核编码能力:数据、智能体环境和强化学习(RL)算法技术。
数据:为支持大规模强化学习,我们设计了一个可扩展的数据合成管道,生成涵盖广泛难度级别的训练问题,从而实现有效的基于课程的RL训练。
智能体环境:我们采用了智能体技能(Agent Skills)范式[2],为模型配备了一个结构化的规范,该规范形式化了编写、验证和优化CUDA内核的标准工作流,以及用于基于执行反馈的自动化测试和分析脚本。值得注意的是,我们实施了严格的正确性和性能测试,以及系统级权限隔离,以防止奖励黑客(reward hacking)并确保准确的奖励信号。
RL算法:关于RL算法的改进,我们分析了RL优化中不稳定的根源,并为actor和critic模型提出了一个多阶段热身策略,使得大型语言模型能够稳定训练150步。
整合这些组件,CUDA Agent成功地将上下文长度扩展到128k tokens,并支持多达200个交互轮次,实现了最先进的性能。具体来说,在KernelBench的Level-1、Level-2和Level-3分割上[17],CUDA Agent比torch.compile分别实现了100%、100%和92%的加速比,在Level-3分割上比Claude Opus 4.5和Gemini 3 Pro等先进专有模型高出约40%。
我们的贡献总结如下:
我们提出了CUDA Agent,一个用于自动CUDA内核生成的大规模智能体强化学习系统。它通过可扩展的数据合成、技能增强的CUDA内核开发环境,以及为稳定的长上下文、多轮智能体训练设计的RL技术,系统地提升了基础模型的CUDA编码和优化能力。
CUDA Agent在KernelBench上取得了最先进的性能,在大多数测试用例中优于
torch.compile,并在多个难度级别上实现了最大的加速比。这些结果表明,基于LLM的内核生成正成为传统编译器驱动内核优化的一个具有竞争力——且通常更优——的替代方案。
2. 相关工作
2.1 用于内核生成的免训练系统
最近的一些工作探索了设计带有显式搜索的方法来处理内核优化的不规则问题。
STARK [5]:采用一个具有规划、编码和调试角色的智能体策略团队,在KernelBench上探索树状搜索空间,利用编译、正确性检查和时序反馈迭代精化内核。
ReGraphT [6]:提出了一个免训练、检索增强的框架,将大型语言模型中的CUDA优化轨迹提炼成一个推理图,然后通过蒙特卡洛图搜索来引导较小的模型达到有竞争力的性能。
EvoEngineer [8]:将CUDA内核优化形式化为一个受约束的代码演化问题,并应用LLM驱动的演化循环,迭代地编辑和验证内核,以在保持正确性的同时提高性能。
CudaForge [26]:引入了一个免训练的双智能体系统,其中一个Judge使用Nsight Compute分析信息和硬件规格来诊断瓶颈,并向Coder提供有针对性的优化反馈,在不同GPU上实现了一致的加速。
尽管这些工作展示了可观的性能提升,但它们严重依赖基础模型自身的CUDA编码能力。此外,这些测试时扩展方法是正交的,也可以应用于我们的模型。
2.2 为内核生成微调LLM
与此同时,另一条研究路线试图通过监督微调或强化学习来训练基础模型,以改进CUDA内核生成。
Kevin [4]:为CUDA内核引入了一个多轮强化学习框架,模拟了迭代式的开发者工作流,在KernelBench上,相对于其QwQ-32B基础模型,在正确性和性能方面都取得了显著提升。
CUDA-L1 [14]:提出了一个对比强化学习框架,使用基于执行的奖励来评估多个CUDA内核变体。然而,他们的训练和评估在同一个KernelBench数据集上进行,没有适当的训练-测试分割。这种数据泄露使得他们报告的结果无法与我们直接比较。
ConCuR [11]:合成并整理了带有推理轨迹的CUDA内核,并使用生成的数据集将QwQ-32B微调为KernelCoder,在开源模型中实现了性能提升。
总体而言,这些方法仍然受到高质量训练数据稀缺、训练规模有限以及手工设计的优化循环的根本限制,这些因素共同限制了它们的性能提升。对这些相关工作的更详细分析和比较见附录C。
3. 方法
我们的大规模智能体强化学习系统CUDA Agent建立在三个关键组件之上:
一个可扩展的数据收集管道。
一个技能集成且不可破解的训练环境,具有鲁棒的奖励调度设计。
提出的用于稳定训练的强化学习算法技术。
每个组件在实现最终强大的实证结果中都扮演着至关重要的角色。
3.1 可扩展的训练数据合成管道
高性能CUDA内核的稀缺性给监督微调带来了重大瓶颈,因为手动实现专家级参考代码的成本过高。为了克服这一点,我们采用强化学习(RL)来训练CUDA Agent。RL需要一个庞大且多样化的参考算子语料库(在PyTorch中实现)作为训练任务。由于现有的公共数据集缺乏所需的多样性和规模,我们开发了一个可扩展的数据收集管道,通过种子问题爬取、基于LLM的组合合成以及严格的基于执行的过滤,系统地扩展了任务空间(图1)。
我们基于组合的数据合成的关键观察是:将多个算子组合成融合任务能产生有价值的新的CUDA智能体训练问题。这是因为组合问题通常不等同于独立优化每个算子然后串联起来。融合通过避免中间全局内存物化、通过共享寄存器/共享内存/占用率约束耦合阶段,以及需要统一的并行映射和数据布局(可能有利于下游消费),重塑了优化格局。
种子问题爬取:首先,我们从torch和transformers库中挖掘在PyTorch中实现的参考算子,建立一个全面的种子问题集。每个算子表示为一个Python类,包含初始化(__init__)和前向(forward)方法。这些算子类被广泛使用且维护良好。因此,我们排除了那些代码质量不佳的个人维护的仓库。
组合问题构建:接下来,为了扩展数据集并引入更高的复杂性,我们利用LLM来合成聚合算子。具体地,提示LLM从torch库中采样不超过5个算子类。采样的算子类通过堆叠到一个计算层中进行顺序组合。我们不从transformers库中采样算子类,因为这些算子通常是更高级别的模块,已经封装了多个原始操作。
问题过滤:最后,我们实施了一个严格的数据选择过程,基于执行反馈过滤掉太简单或太难的问题。我们根据四个标准验证每个算子:
算子必须在Eager模式和Compile模式下都能成功执行。
为确保可重现性,我们排除具有固有随机性的算子。
为防止奖励黑客,我们验证对于不同输入,输出既不是常数值,也不是数值上无法区分的。
为了过滤掉琐碎或过于繁重的任务,我们将Eager模式的执行时间限制在1毫秒到100毫秒的范围内。
此外,我们排除了与KernelBench测试用例表现出高度相似性的算子;相似度分布见章节A。
总共,经过过滤的合成训练数据集包含6,000个样本,形成了CUDA-Agent-Ops-6K,一个用于训练具备CUDA能力的智能体的精选算子级数据集。关于数据格式、污染分析和数据集组成的更多细节见章节A。
3.2 技能集成的智能体循环
智能体循环:由于CUDA内核编码自然地作为编码智能体的一个子任务出现,我们将智能体循环设计为与广泛采用的OpenHands框架[22]保持一致,以确保泛化性。LLM配备了一套标准的shell工具——BashTool、GlobTool、MultiEditTool和TodoWriteTool——以全面支持CUDA编码开发(完整列表见章节B)。在此基础上,我们的智能体循环(如图2所示)遵循ReAct式范式[24],交错进行推理、动作执行和观察,以实现迭代编码、调试和性能优化。
CUDA编码技能:受智能体技能(Agent Skills)理念[2]的启发,我们特意将CUDA编码特定的指令和工具(例如,用于比较生成内核与torch.compile性能的分析工具)以Agent Skills格式呈现。我们还设计了一个CUDA内核编码特定的指令文件SKILL.md,阐述了优化CUDA内核的标准流程(原文见章节B):
分析:使用提供的
profile.py脚本分析原生PyTorch实现的性能。此步骤识别性能瓶颈和优化机会,如过多的内核启动和次优的内存访问模式。实现:通过重写
model_new.py中的模型并开发相应的CUDA内核源文件和绑定代码,实现自定义CUDA算子,针对分析阶段识别的性能关键算子。编译和评估:在提供的GPU沙盒环境中编译和评估优化后的模型,迭代精化CUDA内核实现,直到同时满足数值正确性和性能要求。
重复:从第2步开始重复优化过程,直到最终实现相比
torch.compile基线获得至少5%的加速,同时通过所有数值正确性检查。
鲁棒的奖励调度:现有的CUDA生成RL方法[1, 4, 14]使用相对于基线的加速比作为奖励信号。然而,我们观察到这种方法存在离群值和偏向于简单内核的问题。这是因为不同算子的优化难度差异很大,使得原始加速比成为代码质量的不可靠代理。为了解决这个问题,我们提出了一种归一化的、鲁棒的奖励方案来补救这一点,并联合优化正确性和执行延迟。
避免奖励黑客的努力:我们注意到先前的工作Lange等人[12]存在黑客问题[9]。具体来说,我们的系统强制执行以下约束以避免奖励黑客:
用于正确性验证和性能分析的Python脚本通过文件权限控制得到保护,防止智能体修改或干扰评估逻辑。
为避免琐碎的fallback,我们使用上下文管理器强制执行执行时间约束,明确禁止调用来自
torch.nn.functional的fallback实现,确保性能提升完全来源于生成的CUDA内核。对于每个问题,我们严格遵循KernelBench评估协议,针对五个随机采样的输入验证内核输出,以确保功能正确性。
分析流程经过精心设计,包含适当的设备同步、预热迭代和重复测量与平均,显著降低了测量噪声和指标波动。
智能体没有配备任何网络搜索或外部信息检索工具,确保所有解决方案纯粹来源于本地执行环境。
3.3 稳定RL训练的算法改进
我们观察到,最初的RL试验只能稳定训练17步,之后模型性能就崩溃了。我们找出了这种不稳定的根本原因,并提出了一个热身策略,使actor和critic模型都能适应模型的数据分布。通过这一修改,我们的RL试验可以稳定训练200步,并伴随着持续的奖励增长。
4. 实验
4.1 实验设置
RL训练设置:我们使用Seed1.6 [20]作为基础模型,这是一个混合专家(MoE)模型,拥有230亿活跃参数和2300亿总参数。我们将全局批量大小设为1024,与在线PPO更新的小批量大小相匹配。actor和critic的学习率分别为 3×10−6 和 6×10−6。这些超参数在单轮RL和智能体RL之间共享。单轮RL的上下文窗口长度为32768,智能体RL为131072。在训练滚动期间,我们设置最大智能体轮次为150,在评估期间放宽到200轮次。模型总共训练150步。
沙盒环境:为了确保加速比测量的准确性以获得可靠的奖励信号,并实现高效的GPU利用,我们设计了一个CPU-GPU资源解耦的沙盒架构。一个基于Docker的终端沙盒处理CPU密集型任务(如内核编译),而智能体利用预定义的CUDA Agent技能脚本将验证和分析作业分派到一个专用的GPU沙盒池(128块NVIDIA H20 GPU)。这种进程级隔离和独占资源分配消除了进程间的干扰,确保了稳定的延迟测量和保证的HBM容量。
基准测试集:我们在KernelBench [17]上进行评估,使用其Level 1到Level 3的子集,共计250个不同的算子任务。为了确保公平比较,我们将这些任务从其原始的单文件格式适配到我们的多文件开发环境(章节3.2)。
基线模型:我们将我们的方法与前沿的专有编码模型:Claude Opus 4.5 [3]和Gemini 3 Pro [7],以及两个强大的开源编码模型:GLM 4.6 [27]和Kimi K2 [21]进行比较。这些模型目前在主要编码智能体排行榜上处于领先地位,是通用推理和编码能力的强基线。为了公平比较,基线模型在章节3.2中提到的相同智能体循环下进行评估。
评估指标:我们使用三个关键指标评估性能:
通过率(Pass Rate):智能体生成的内核成功编译并通过功能正确性检查的任务百分比。
更快率(Faster Rate):生成的内核正确且执行时间快于基线(Eager和Compile模式)的任务百分比。
加速比(Speed-up):相对于基线执行时间的几何平均加速比,仅针对正确解决方案计算。
对于每个任务,我们从智能体的交互轨迹中提取表现最佳的解决方案,具体来说是相对于torch.compile实现最大加速比的那个。
4.2 主要结果
表1总结了CUDA Agent在KernelBench上与强大的专有模型基线相比的性能。我们的分析得出三个主要见解:
与强大的专有模型相比,CUDA Agent提供了显著更强的CUDA内核编码性能。尽管Claude Opus 4.5和Gemini 3 Pro取得了可观的通过率(91.2%-95.2%),但它们的更快率仍然较低(66%-69%),这表明通用LLM通常生成朴素的内核,无法超越
torch.compile。相比之下,CUDA Agent达到了98.8%的通过率和96.8%的更快率,表明专门的RL训练能够实现持续正确且高度优化的CUDA实现。与静态的
torch.compile相比,CUDA Agent证明了学习到的优化策略可以持续超越静态编译器启发式方法,特别是在算子融合等复杂场景中。这在Level 2任务(算子序列)中最为明显,CUDA Agent实现了完美的100%更快率和相对于torch.compile的巨大2.80倍加速。传统编译器依赖预定义的、基于规则的融合模式,这些模式通常难以处理非平凡的算子组合。相比之下,CUDA Agent通过其迭代的智能体循环探索了更大的设计空间,发现了静态后端无法触及的硬件特定的内存访问模式和分块策略。
4.3 消融研究
4.3.1 技能集成智能体循环的影响
为了评估交互式环境在策略学习中的关键作用,我们在两种不同的协议下训练了两个独立的模型:
单轮模型:使用标准代码生成目标进行训练,模型在单轮中预测最终内核和绑定,没有执行反馈,作为我们训练管道中的热身阶段。
CUDA Agent:遵循完整的训练管道进行训练,允许策略在多轮交互中看到编译错误和分析器反馈。
表2突显了单轮代码生成的局限性,并强调了我们的技能集成智能体循环的必要性。移除智能体循环导致正确性和优化质量都大幅下降。更重要的是,生成的内核不仅优化不足,而且性能常常退步。通过暴露于编译错误、运行时失败和分析器反馈,智能体可以迭代地诊断错误并在多轮中精化转换。
4.3.2 奖励设计的影响
如表2所示,奖励设计对优化结果有显著影响。用原始的加速比奖励替换我们的鲁棒奖励计划(w/o Robust Reward)会产生相当的功能正确性,但优化性能显著减弱。
这些结果表明,归一化的、基于里程碑的奖励与产生持续更快的内核的目标更一致。通过将奖励分配给明确的性能目标,而不是直接对噪声较大的运行时间比率进行回归,策略能更可靠地发现那些能转化为相对于Eager执行和编译器基线的真实加速的转换。
4.3.3 多阶段训练的影响
如表2所示,移除RFT或价值预训练都会导致优化性能大幅下降,尽管通过率基本保持不变。更重要的是,两种消融实验都表现出训练不稳定性和最终的崩溃,这促使我们在下面更仔细地分析这两个阶段。
拒绝采样微调(RFT)提供了一个关键的先验,可以防止策略崩溃。如图4a所示,移除RFT阶段会导致训练奖励迅速且灾难性地崩溃。为了诊断根本原因,我们检查了图4b中的策略熵,发现一个与奖励崩溃同时发生的急剧熵增。这种熵增表明策略分布变得越来越分散,产生不连贯和结构不良的输出。通过用强大的行为先验初始化策略,RFT在强化学习期间约束了熵的增长,并将优化轨迹保持在一个结构良好的输出分布内。
价值预训练对于提供可靠的优势估计和防止病态的搜索行为是必不可少的。没有初始化的critic,模型无法捕捉多轮交互状态的价值景观(图5a)。这种糟糕的估计导致轨迹长度爆炸(图5b),因为未初始化的critic无法惩罚徒劳或冗余的搜索路径。价值预训练确保critic能够立即提供准确的反馈,引导智能体走向高效的优化路径,并避免近乎无限的交互循环带来的计算不稳定性。
5. 结论
我们提出了CUDA Agent,一个大规模智能体强化学习系统,它赋予大型语言模型在现实的、执行驱动的开发工作流中生成和优化CUDA内核的能力。通过联合扩展数据合成、智能体环境和面向稳定性的RL训练,CUDA Agent将LLM从语法代码生成推向硬件感知的性能优化,在KernelBench上相对于torch.compile和强大的专有模型取得了持续的提升。这些结果表明了一个更广泛的后果:为基础模型配备结构化的环境和可靠的基于执行的奖励,可以将它们从被动的代码生成器转变为主动的系统优化器,为在GPU计算中自动化性能关键的软件开发开辟了一条道路。
