跳到主内容
精选85MarkTechPost(RSS)论文研究

字节Seed与清华AIR发布CUDA Agent,用强化学习生成GPU内核

ByteDance Seed and Tsinghua AIR Introduces CUDA Agent: A Large-Scale Agentic RL System for CUDA Kernel Generation

原文
推荐理由

做AI编译器或算子优化的同学必看,这套把强化学习搬进真实CUDA开发环境的做法,直接把内核性能干到编译器的2倍以上,值得照着复现一遍。

ByteDance Seed and Tsinghua AIR have released CUDA Agent, an agentic reinforcement learning system that trains a large language model to write GPU kernels that beat a compiler. The gap it targets is narrow but stubborn: frontier models already produce correct CUDA, they just produce slow CUDA. On KernelBench, the base model Seed1.6 passes 74.0% of tasks yet outruns torch.compile on only 27.2% of them, at a 0.69× geometric-mean speedup which means its kernels are, on average, slower than what the compiler generates on its own. CUDA Agent closes that gap by putting the model inside a real CUDA development environment with profiling, correctness checks and a permission-locked sandbox, then training it with PPO for 150 steps at a 131,072-token context. The result is a 98.8% pass rate and a 96.8% faster-than-torch.compile rate across the 250-task benchmark, at 2.11× geomean over compile — roughly 40 points ahead of Claude Opus 4.5 and Gemini 3 Pro on the hardest Level-3 split.

字节跳动Seed与清华大学AIR联合发布了CUDA Agent,这是一个代理强化学习系统,训练大型语言模型编写GPU内核,其性能超越编译器。它针对的差距虽窄但顽固:前沿模型已能生成正确的CUDA代码,只是生成的CUDA运行缓慢。在KernelBench上,基础模型Seed1.6通过74.0%的任务,但仅在27.2%的任务上超越torch.compile,几何平均加速比为0.69倍,意味着其内核平均比编译器自行生成的还要慢。CUDA Agent通过将模型置于真实的CUDA开发环境中,配备性能剖析、正确性检查和权限锁定沙箱,并用PPO在131,072令牌上下文中训练150步,从而弥合了这一差距。结果是在250项任务基准上,通过率达到98.8%,超越torch.compile的比例为96.8%,几何平均加速比达2.11倍——在最难的Level-3子集上,比Claude Opus 4.5和Gemini 3 Pro领先约40个百分点。

Is it deployable?

它是否可部署?

Partly, but the trained agent is not released. It is built on Seed1.6, a proprietary MoE model with 23B active and 230B total parameters, and the paper ships no weights. Public: the CUDA-Agent-Ops-6K dataset, the SKILL.md spec and the reward and warm-up recipes.

部分可部署,但训练好的代理未发布。它基于Seed1.6,这是一个专有的MoE模型,拥有230亿激活参数和2300亿总参数,论文未提供权重。公开内容:CUDA-Agent-Ops-6K数据集、SKILL.md规范以及奖励和预热配方。

Which companies: The profiling sandbox alone used 128 NVIDIA H20 GPUs, which puts full replication inside frontier labs, GPU clouds and large infrastructure teams. Mid-size teams can still adopt the parts — dataset, milestone reward, anti-reward-hacking constraints, skill spec — on top of an open base model.

哪些公司:仅性能剖析沙箱就使用了128块NVIDIA H20 GPU,这使得全面复现仅限于前沿实验室、GPU云和大型基础设施团队。中型团队仍可采用部分组件——数据集、里程碑奖励、防奖励黑客约束、技能规范——基于开源基础模型之上。

Industries and applications: AI infrastructure and inference serving, GPU cloud, autonomous driving, quantitative trading, medical imaging and recommendation systems — anywhere fused kernels sit on a latency-critical path. Uses include fusing operator sequences torch.compile handles poorly, cutting cost per token, and re-tuning kernels across GPU generations.

行业和应用:AI基础设施和推理服务、GPU云、自动驾驶、量化交易、医学影像和推荐系统——凡是融合内核处于延迟关键路径上的场景。用途包括融合torch.compile处理不佳的算子序列、降低每令牌成本,以及跨GPU代际重新调优内核。

Data synthesis

数据合成

The research team crawls reference operators from the torch and transformers libraries. An LLM then samples up to five torch operator classes and stacks them into one fused layer. A filter keeps only operators that execute in both eager and compile modes, are deterministic, produce non-constant outputs, and run between 1 ms and 100 ms in eager mode. Samples with AST similarity above 0.9 to any KernelBench task are removed. The result is CUDA-Agent-Ops-6K: 6,000 samples, 83.77% of them two-operator compositions.

研究团队从torch和transformers库中爬取参考算子。随后,一个LLM采样最多五个torch算子类,并将它们堆叠成一个融合层。过滤器仅保留在即时模式和编译模式下都能执行、确定性、产生非恒定输出,且在即时模式下运行时间在1毫秒至100毫秒之间的算子。与任何KernelBench任务AST相似度高于0.9的样本被移除。结果是CUDA-Agent-Ops-6K:6000个样本,其中83.77%为双算子组合。

Environment and reward

环境与奖励

The agent loop mirrors OpenHands tooling — Bash, Read/Write, Edit/MultiEdit, Glob, Grep, NotebookEdit, BashOutput, KillBash — under a ReAct pattern. CUDA instructions ship in the Agent Skills format. SKILL.md tells the model to profile the PyTorch model, rewrite model_new.py with custom kernels, compile in a GPU sandbox, and iterate until the kernel is at least 5% faster than torch.compile at atol=1e-2, rtol=1e-2.

代理循环镜像了OpenHands工具——Bash、读/写、编辑/多编辑、Glob、Grep、NotebookEdit、BashOutput、KillBash——在ReAct模式下运行。CUDA指令以Agent Skills格式提供。SKILL.md告诉模型对PyTorch模型进行性能分析,用自定义内核重写model_new.py,在GPU沙箱中编译,并迭代直到内核在atol=1e-2、rtol=1e-2下比torch.compile快至少5%。

Reward hacking gets five countermeasures: permission-locked verification and profiling scripts, context managers that forbid torch.nn.functional fallbacks, checks against five random inputs, profiling with device synchronization and warm-up, and no web search tool.

奖励黑客行为有五种对策:权限锁定的验证和性能分析脚本、禁止torch.nn.functional回退的上下文管理器、对五个随机输入的检查、带设备同步和预热的性能分析,以及无网络搜索工具。

The reward is discrete rather than a raw speedup ratio. r ∈ {−1, 1, 2, 3}: −1 on correctness failure, 3 if the kernel clears both eager and torch.compile by more than 5%, 2 if it clears eager only, 1 otherwise.

奖励是离散的,而不是原始加速比。r ∈ {−1, 1, 2, 3}:正确性失败时为−1,如果内核比eager和torch.compile都快超过5%则为3,仅比eager快则为2,否则为1。

https://arxiv.org/pdf/2602.24286v1

https://arxiv.org/pdf/2602.24286v1

Results

结果

Table 1, overall: 98.8% pass rate, 98.4% faster than eager, 96.8% faster than torch.compile, at 2.60× and 2.11× geomean respectively. Level 2 (operator sequences) is the strongest split: 100% pass, 100% faster rate, 2.80× over torch.compile. Level 3 lands at 94.0% pass, 90.0% faster rate and 1.52×, roughly 40 points above Claude Opus 4.5 (50.0%) and Gemini 3 Pro (52.0%) on faster rate versus compile.

表1,总体:通过率98.8%,比eager快98.4%,比torch.compile快96.8%,几何平均分别为2.60×和2.11×。第2级(算子序列)是最强的分割:通过率100%,更快率100%,比torch.compile快2.80×。第3级通过率94.0%,更快率90.0%,1.52×,在更快率对比编译方面比Claude Opus 4.5(50.0%)和Gemini 3 Pro(52.0%)高出约40个百分点。

One inconsistency: the abstract and introduction state 100% / 100% / 92% faster rates for Levels 1–3, while Table 1 reports 97.0% / 100.0% / 90.0%. Table 1 is the main results table.

一处不一致:摘要和引言声称第1-3级的更快率为100% / 100% / 92%,而表1报告为97.0% / 100.0% / 90.0%。表1是主要结果表。

Ablations are blunt. Removing the agent loop drops faster rate versus compile from 96.8% to 14.1%. A raw speedup reward gives 60.4%, no RFT gives 49.8% plus reward collapse, no value pretraining gives 50.9% plus runaway trajectories.

消融实验很直接。移除代理循环将对比编译的更快率从96.8%降至14.1%。原始加速奖励给出60.4%,无RFT给出49.8%加上奖励崩溃,无价值预训练给出50.9%加上失控轨迹。

Case studies show what the policy learns. A diagonal matmul rewritten as row-wise scaling: 73.31× over torch.compile. A matmul-divide-sum-scale chain reordered and fused: 24.04×. A ResNet BasicBlock with BatchNorm folded into convolution andcudnnConvolutionBiasActivationForward: 3.59×.

案例研究展示了策略学到的东西。一个对角矩阵乘法重写为逐行缩放:比torch.compile快73.31×。一个矩阵乘法-除法-求和-缩放链被重排并融合:24.04×。一个ResNet BasicBlock,BatchNorm折叠进卷积和cudnnConvolutionBiasActivationForward:3.59×。

更进一步:量化金融体系

看懂新闻只是起点——沿量化金融路径,把它变成能交付的工程能力

进入量化体系 →

相似阅读

另一事件,读法相近