AI Agent开始自己优化CUDA内核

Agentic CUDA Kernel Optimizer 将代码生成、编译调试、性能分析和迭代搜索串成闭环,让 AI Agent 从“写 CUDA 代码”进一步变成“为特定 GPU 找更快实现”。它的价值不在于替代编译器,而在于自动探索编译器通常不会主动尝试的优化路径。
AI Agent开始自己优化CUDA内核
Agentic CUDA Kernel Optimizer 最近在 Hacker News 上引发关注,它试图让 AI Agent 自动完成 CUDA 内核的生成、验证、性能分析和迭代优化。 这类项目的意义不在于又出现了一个“能写 CUDA 的大模型”,而在于它把 CUDA 优化从一次性的代码生成任务,变成了一个可以反复试错、测量和回滚的工程闭环。

CUDA 内核是运行在 NVIDIA GPU 上、负责执行具体并行计算任务的函数。深度学习框架中的矩阵乘法、归约、注意力、卷积和张量变换,最终都会被拆解成一个个内核;模型训练和推理的速度,很大程度上取决于这些内核是否能充分利用 GPU 的计算单元、显存带宽和片上缓存。
这也是 CUDA 优化长期依赖专家的原因。表面上看,任务只是把 PyTorch 算子改写成 CUDA C++,但真正影响速度的往往是线程块大小、内存访问是否合并、共享内存占用、寄存器压力、Occupancy、指令级并行,以及不同 GPU 架构上的 Tensor Core 使用方式。代码能编译、结果也正确,并不意味着它跑得快。
发生了什么
Agentic CUDA Kernel Optimizer 的核心做法,是让模型在真实的编译和运行环境中持续修改内核,并用正确性检查与运行时间作为反馈。 项目公开在 GitHub 上,定位更接近一个可运行的实验性 Agent 框架,而不是已经包装好的商业产品。
它将一个传统的 CUDA 优化任务拆成了几步:
- 读取待优化的 PyTorch 算子或基准实现;
- 生成或修改 CUDA 内核代码;
- 编译扩展并处理编译错误;
- 用测试样例检查数值正确性;
- 在目标 GPU 上运行基准测试;
- 根据延迟、吞吐量和 Profiler 信息继续修改;
- 保存更快的候选实现,并淘汰错误或变慢的版本。
这个循环看起来并不神秘,但它改变了模型的工作方式。普通代码生成模型通常只负责“给出一段看起来合理的代码”,真正的编译、调试和测速由人完成。Agentic 优化器则把这些步骤都纳入任务本身,模型不必在第一次输出时猜中最终答案,而是可以像工程师一样通过实验逐步逼近更优实现。
这类 Agent 的关键不是生成一次,而是能够在几十轮甚至更长的反馈循环中保持目标不变。 目标同时包括“结果正确”和“运行更快”,任何一个条件不满足,候选内核都不能进入下一轮。
为什么编译器还不够
torch.compile 是 PyTorch 的编译优化系统,它会捕获模型计算图并尝试自动生成更高效的执行代码,但它并不等于针对每个算子进行开放式的 CUDA 研发。 编译器擅长使用既有规则进行融合、布局变换和代码生成,却通常不会像人类专家那样主动提出大量结构差异很大的实现方案,再逐个在硬件上验证。
可以把两者理解成不同层次的工具:
- 编译器像一条高度自动化的生产线,能稳定地把常见输入转换成质量不错的机器代码;
- CUDA Agent 像一个会写代码、会跑实验、会看性能数据的研发工程师,可能花更多时间尝试非常规方案;
- 传统 CUDA 专家则会结合算子数学结构、GPU 微架构和历史经验,直接缩小搜索空间。
在规则成熟的算子上,编译器通常已经足够好,Agent 反复尝试未必能带来收益。在新模型结构、非标准融合模式、特殊张量形状或特定硬件上,Agent 的搜索空间反而可能更有价值。
| 方案 | 主要工作方式 | 优势 | 局限 | 更适合的场景 |
|---|---|---|---|---|
| torch.compile | 基于计算图和编译规则自动优化 | 集成成本低,稳定性较好 | 对非常规内核和长尾形状的探索有限 | 通用模型部署、常见算子 |
| 传统 CUDA 专家 | 手工设计线程布局、内存访问和指令路径 | 能结合硬件细节做深度优化 | 人力成本高,难以覆盖大量算子 | 核心算子、极致性能场景 |
| Agentic CUDA Optimizer | 生成代码并通过编译、测试、Profiling 迭代 | 可自动探索多个实现,适合长尾问题 | 计算成本高,结果依赖验证环境 | 新算子、定制模型、特定 GPU |
| 固定循环式代码生成 | 生成后按固定规则修改若干次 | 实现简单,容易搭建 | 缺乏长期策略,容易重复无效尝试 | 原型验证和小规模实验 |
Agent 的优势不在于普遍击败 torch.compile,而在于为编译器覆盖不足的区域增加一个搜索层。 如果某个算子已经被 CUDA 库、CUTLASS 或框架编译器优化得很好,Agent 很可能只能带来边际收益;如果算子组合复杂、输入形状特殊,或者模型刚刚引入新的计算模式,它就有机会找到人工规则没有覆盖的路径。
性能反馈比代码生成更重要
CUDA 优化中的性能反馈,是指通过真实硬件运行时间、吞吐量、内存带宽和硬件计数器判断一个候选内核是否真的更好。 仅凭代码表面结构判断优劣非常危险,因为减少几行代码、增加并行线程,甚至使用更多共享内存,都可能让实际运行速度变慢。
一个典型例子是矩阵乘法。增加线程数量可能提高并行度,但也会增加寄存器使用和线程块调度压力;使用共享内存可以减少重复读取显存,却可能因为同步操作和容量限制抵消收益;让线程访问连续地址通常有利于显存合并访问,但某些张量布局下,重新排列数据的成本可能高于收益。
因此,Agent 需要的不只是编译器,还需要一套可靠的实验环境:
- 正确性验证:将候选内核输出与 PyTorch 或参考实现逐元素比较,并覆盖不同输入形状、数据类型和边界条件。
- 性能基准:重复运行并进行预热,避免首次编译、显存分配或 GPU 频率变化污染数据。
- Profiler 反馈:观察内存吞吐、计算利用率、Occupancy、寄存器和共享内存使用情况。
- 失败处理:区分语法错误、编译错误、运行时错误、数值误差和单纯的性能退化。
- 候选管理:保存历史上最快且正确的版本,避免后续探索把已经获得的收益覆盖掉。
这个过程更像自动化实验室,而不是聊天机器人。模型输出只是实验假设,编译器和 GPU 才是裁判。
从固定循环到 Agentic RL
Agentic Reinforcement Learning 是让模型在多步行动中根据环境奖励学习策略的训练方式,在 CUDA 优化中,奖励通常同时受正确性和运行性能约束。 近期公开的 CUDA Agent 研究进一步把这一思路扩展到大规模训练:通过数据合成、带工具的 CUDA 开发环境和长上下文强化学习,让模型获得更稳定的内核优化能力。
据公开资料,相关训练流程构造了约 6000 个训练样本。部分任务由最多 5 个 PyTorch 算子组合而成,并筛选原始运行时间在 1 毫秒到 100 毫秒之间的任务。这个范围有现实考虑:运行时间太短,测量噪声会掩盖优化收益;运行时间太长,则会显著拖高训练和评估成本。
任务还加入了反作弊检查,避免模型通过改变输入、跳过计算、返回缓存结果或破坏测试逻辑来获得虚假的加速。候选内核需要通过正确性检查,并达到相对于 torch.compile 至少 5% 的性能提升,才可能被视为有效优化目标。
这里的 5% 不是“所有任务平均提升 5%”的结论,而更像任务筛选和奖励设计中的门槛。公开材料没有给出一个足以代表所有 GPU、所有算子和所有输入形状的统一提升数字,因此不能简单宣称 CUDA Agent 已经全面超过编译器或人工专家。
强化学习的价值,在于让模型逐渐学会哪些修改值得尝试,而不是只记住某一段 CUDA 模板。 例如,模型可能学会在归约算子中考虑 warp-level primitive,在矩阵计算中优先检查 Tensor Core 路径,在内存受限任务中减少不必要的数据搬运,或者根据寄存器压力调整线程块配置。
KernelBench 是怎样的考场
KernelBench 是用于评估 GPU 内核生成与优化能力的基准测试集合,它把 PyTorch 算子作为问题输入,再比较生成内核的正确性与运行性能。 相关系统会在 KernelBench 上与 torch.compile 以及高能力闭源模型进行比较。
这类基准比“代码看起来像不像 CUDA”严格得多。一个候选实现必须首先在数值上正确,然后才能进入性能比较。不同输入尺寸、不同数据类型和不同 GPU 也可能改变排名,所以测试报告至少需要说明硬件、CUDA 版本、编译选项、测量方法和统计口径。
对于 Agentic CUDA Kernel Optimizer,真正值得观察的不是某一个演示任务的最低延迟,而是以下几项:
- 能否在不同算子类型之间迁移优化经验;
- 能否处理长尾输入形状,而不是只针对一个固定尺寸调参;
- 能否在多轮修改后避免陷入重复尝试;
- 能否识别性能噪声,不把偶然的低延迟当成稳定收益;
- 能否在换一张 GPU 或升级 CUDA 后保持正确性和收益;
- 总搜索成本是否低于人工优化节省的部署时间。
这最后一点经常被忽略。一个内核快了 20%,但 Agent 为此消耗数千次编译和运行,未必适合线上服务。只有当内核会被高频调用、优化结果可以复用,或者它能缩短关键模型的上线周期,自动搜索才具有明确的经济价值。
它最可能先改变什么
Agentic CUDA 优化最先适合解决的,是数量多、价值分散但人工不值得逐个处理的定制内核问题。 大型模型团队经常会遇到这种情况:核心注意力算子已经有成熟实现,但模型外围还有大量形状特殊的张量变换、掩码、归一化和融合算子。每个算子单独看只占几毫秒,却可能在高并发推理中累积成明显延迟。
它也适合以下场景:
- 新模型架构尚未被主流编译器充分支持;
- 同一个算子在不同批大小和序列长度下需要不同实现;
- 企业使用定制 GPU 或固定硬件,允许针对单一环境调优;
- 推理服务需要把多个小算子融合,减少 kernel launch 和中间张量写回;
- 研究人员想快速判断一个手工优化方向是否值得继续投入。
但它并不适合直接替代成熟的基础库。NVIDIA 官方 CUDA 库、cuBLAS、cuDNN、CUTLASS 和框架内置实现经过长期优化,覆盖了大量常见工作负载。Agent 生成的代码还要面对可维护性、跨版本兼容性、编译时间、部署安全和故障定位等问题。
还有三道硬门槛
第一道门槛是硬件依赖。 CUDA 内核的最佳方案通常不是跨 GPU 通用的。针对 H100 找到的线程布局和共享内存策略,换到 A100、消费级 RTX 或更新架构后可能失效。Agent 如果只在单一 GPU 上训练和评估,得到的更像“硬件特化程序”,而不是通用 CUDA 优化器。
第二道门槛是评测可靠性。 GPU 运行时间会受到频率、温度、并发进程和显存状态影响。短内核尤其容易出现测量抖动。没有多次采样、置信区间和稳定性检查,所谓 3% 或 5% 的提升可能只是噪声。
第三道门槛是搜索成本。 每次尝试都可能涉及代码生成、编译、加载和多轮执行。相比普通文本生成,这类任务的成本主要不在模型输出,而在真实 GPU 实验。未来系统需要学会提前判断“这个方向大概率没有收益”,把算力用在更有希望的候选上。
此外,内核正确性不能只靠少量随机样例。边界尺寸、空张量、非连续张量、半精度溢出、不同布局和极端数值都可能触发隐藏错误。一个在基准集上跑得很快但偶尔产生错误结果的内核,不能进入生产环境。
OpenAI Hub 判断
Agentic CUDA Kernel Optimizer 的重要性,在于它把 GPU 性能优化从“专家手艺”推进成了“可自动运行的搜索任务”,但它距离可靠的通用 CUDA 工程师仍有明显距离。
短期看,最现实的产品形态不是让 Agent 从零接管整个 GPU 软件栈,而是作为性能工程师的自动化助手:先分析热点,再提出若干候选实现,自动完成编译和回归测试,最后把经过验证的代码与性能报告交给人审查。这样既能利用 Agent 的搜索能力,也能保留工程师对边界条件和部署风险的控制。
中期看,真正有竞争力的系统需要同时具备三种能力:理解 PyTorch 或其他框架的计算图,理解 GPU 微架构和 Profiling 指标,以及根据历史实验建立可迁移的优化策略。只会调用编译器的 Agent 不够,只会写 CUDA 的模型也不够,关键是能把“代码修改”与“硬件反馈”长期连接起来。
这场竞争的终点不是让 AI 写出更多 CUDA,而是让它知道什么时候该融合、什么时候该换布局、什么时候该牺牲通用性换取特定硬件上的收益。 如果这些判断可以通过大规模实验和强化学习稳定获得,AI 基础设施的优化周期就可能从数周的人工调优压缩到数小时或数天。
不过,现阶段读者应该把它看作值得跟踪的工程方向和研究原型,而不是可以无条件替代 torch.compile、CUTLASS 或 CUDA 专家的成熟方案。判断它有没有实际价值,最终要回到三个数字:在指定 GPU 上的稳定加速比例、找到结果所消耗的 GPU 时间,以及生成内核在真实模型和真实输入分布中的回归表现。
参考来源
- Agentic CUDA Kernel Optimizer GitHub 项目:项目源码与基本介绍,展示了让 Agent 自动生成、编译、测试和优化 CUDA 内核的实践方向。
- KernelBench GitHub 项目:用于评估 PyTorch 算子对应 GPU 内核正确性与性能的基准测试项目。
- CUTLASS GitHub 项目:NVIDIA 开源的 CUDA 模板库,覆盖高性能矩阵乘法和相关线性代数内核,是理解手工 CUDA 优化的重要参考。



