CUDA Agent
搜索文档
字节清华智能体自动写CUDA内核,比torch.compile加速2.11倍
量子位· 2026-03-03 15:02
核心观点 - 字节Seed与清华AIR团队合作,开发出名为CUDA Agent的强化学习智能体系统,能够自动生成和优化高性能GPU代码,在标准测试集上性能全面超越现有商业大模型方案 [1][2][3] 技术方案与系统设计 - 该系统是一个完整的大规模智能体强化学习系统,旨在通过强化学习让模型真正学会生成和优化高性能CUDA内核 [9] - 系统包含三个核心组件:一套可扩展的数据合成流水线、一个配备技能规范文档的CUDA开发环境、一个为长上下文场景设计的RL训练算法 [18] - 智能体工作环境采用ReAct风格的交互循环,配合编码工具和一份名为SKILL.md的CUDA技能规范文档 [16] - 标准流程包括:对原生PyTorch代码做性能分析,实现CUDA内核和绑定代码,在GPU沙箱中编译运行,再根据结果迭代优化 [17] - 通过验证的标准是:正确性检查全部通过,同时速度超过torch.compile至少5% [17] 训练数据构建 - 训练数据构建经历了三个阶段:种子算子挖掘、组合合成、执行驱动过滤与反作弊检查 [10][11][12] - 种子算子来源是torch和transformers,每个算子被表示为一个包含初始化方法和前向方法的Python类 [11] - 通过随机采样最多5个torch算子并按顺序拼接成融合任务来组合合成 [11] - 执行驱动过滤只保留在eager模式和compile模式下都能正常运行的任务,同时去除含有随机算子的样本 [11] - 反作弊检查排除那些对不同输入产生常数输出或无法区分输出的任务 [12] - eager运行时间被控制在1毫秒到100毫秒的范围内,与KernelBench高度相似的样本也被移除以降低数据污染风险 [13] - 经过筛选,最终留下6000个训练样本,命名为CUDA-Agent-Ops-6K,并已对外开源发布 [14][32] 训练过程与约束 - 为防止模型走捷径,团队设置了一整套约束:验证脚本和性能分析脚本被设为受保护状态;回退函数调用被禁止;正确性检查使用5组不同输入;性能分析引入同步预热步骤以消除测量噪声;网络检索功能完全关闭 [21] - 奖励机制采用基于里程碑的离散奖励,分别对应正确性达标和速度提升两个阶段,确保奖励信号反映真实的内核质量 [22] - 训练流水线被拆成多个阶段以应对长上下文RL不稳定的问题 [23] - 第一阶段是单轮PPO热身,模型在非交互式设置下学习基础的CUDA代码生成能力 [24] - 进入多轮智能体RL前,actor和critic需要分别初始化:Actor使用拒绝微调(RFT)从大量采样轨迹中筛选出取得正向结果的样本进行训练;Critic通过价值预训练完成初始化 [25][26] - 该分阶段设计支撑了极大的上下文规模:训练时上下文窗口达到128K token,单条轨迹最多150轮交互;评估时交互轮数放宽到200轮 [27] 性能表现 - 在KernelBench的整体评测上,CUDA Agent的通过率为98.8% [28] - 相对torch.compile,faster rate达到96.8%,几何平均加速比为2.11倍;相对eager模式,faster rate为98.4%,加速比为2.60倍 [28] - 分级来看,Level-1和Level-2的faster rate(相对compile)均为100% [29] - 在Level-2的算子序列任务上,几何平均加速比达到了2.80倍 [29] - 最难的Level-3,通过率94%,faster rate(相对compile)为90%,几何平均加速比1.52倍 [29] - 与强商业模型的对比拉开明显差距:Claude Opus 4.5和Gemini 3 Pro在KernelBench整体上的faster rate(相对compile)约为66.4%到69.6%,几何平均加速比在1.42倍到1.46倍之间,而CUDA Agent的整体faster rate为96.8%,加速比2.11倍 [30] - 在Level-3难题上,CUDA Agent的90% faster rate领先最强商业基线约40个百分点 [30]
英伟达护城河被AI攻破,字节清华CUDA Agent,让人人能搓CUDA内核
机器之心· 2026-03-03 10:55
研究背景与核心突破 - 字节跳动与清华大学的研究团队开发了名为CUDA Agent的AI模型,该模型能够编写并优化CUDA内核,其核心突破在于将优化目标从代码正确性转向硬件性能,通过强化学习直接基于GPU的实际运行速度进行训练[1][6][7] - 此前,GPT、Claude等大模型生成的CUDA代码仅能保证“正确”,但无法实现深度优化,CUDA Agent解决了这一痛点,其性能在简单/中等内核上比torch.compile高出2倍,在复杂内核上高出约92%,即使在最难的设置下也比Claude Opus 4.5和Gemini 3 Pro高出约40%[2][5] 技术方法与系统设计 - 该研究构建了一个大规模的智能体强化学习系统,包含三个核心部分:可扩展的数据合成机制、一个集成技能增强且具备可靠验证与性能分析能力的CUDA开发环境,以及用于稳定长上下文训练的强化学习算法技术[12] - 数据合成通过三阶段管线完成:从torch和transformers中挖掘种子算子,基于LLM进行组合式合成(最多采样5个算子),并通过执行结果筛选,最终构建了包含6000条训练样本的CUDA-Agent-Ops-6K数据集[14][16] - 智能体环境采用ReAct风格的工作流,支持迭代式的编码-编译-调试循环以及基于性能分析器的优化过程,训练流程采用分阶段设计,包括单轮PPO预热、actor/critic初始化和完整的多轮强化学习,以稳定长时序任务[17][19][22] 性能评估结果 - 在KernelBench基准测试中,CUDA Agent取得了SOTA成绩:在Level-1、Level-2和Level-3三个划分上,相比torch.compile的加速比例(faster rate)分别达到100%、100%和92%[10] - 整体来看,CUDA Agent的通过率为98.8%,相对于torch.compile的加速达成率为96.8%,几何平均加速比为2.11倍;相对于原生Eager模式的几何平均加速比为2.60倍[25] - 在最高难度的Level-3任务上,CUDA Agent相对于torch.compile的加速达成率达到90%,相比最强的专有基线(Claude Opus 4.5的50%)高出约40个百分点;在Level-2的算子序列任务上,其加速达成率达到100%,几何平均加速比达到2.80倍[26] 行业影响与未来展望 - 该技术证明了大型语言模型可以通过基于硬件反馈的强化学习,掌握极高门槛的“硬件直觉”,有望打破传统编译器(如torch.compile或Triton)的优化瓶颈[29] - 研究团队认为,一条通向全自动、高度性能优化计算基础设施的道路正在出现[30]