Workflow
GPU内核优化
icon
搜索文档
字节清华智能体自动写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]