news 2026/10/15 8:19:59

CAKE:Agent与编译器闭环协同的GPU Kernel自动生成新范式

作者头像

张小明

前端开发工程师

1.2k 24
文章封面图
CAKE:Agent与编译器闭环协同的GPU Kernel自动生成新范式

GPU Kernel 的自动生成这两年一直是系统与编译交叉领域的热点,但大多数方案要么把 Agent 当成一个"黑盒代码生成器",要么把编译器当成一个"事后验证工具",两者各干各的,最后生成的 Kernel 性能上限往往卡在某个瓶颈上。CAKE 这篇工作有意思的地方在于,它把 Agent 和编译器放在同一个闭环里,让两边互相"教"对方,最终在部分 Kernel 上跑出了超过专家手写的性能。我读完这篇论文之后,结合自己之前折腾 Agent 生成代码和编译器优化的一些经验,把它的核心思路、技术细节和可复现的实操路径整理出来,适合对 Agent 开发、编译器优化、GPU 高性能计算感兴趣的读者参考。

1. 为什么"Agent 写 Kernel"这件事一直没做透

1.1 纯 Agent 生成 Kernel 的天花板在哪

先说一个我自己的观察。用大模型直接生成 CUDA Kernel,早期大家玩的都是"给一段 PyTorch 代码,让模型翻译成 CUDA"。这种做法在简单算子(element-wise、reduction)上效果还行,但一旦碰到需要精细控制 shared memory、warp shuffle、bank conflict 规避的算子,生成结果就很容易"看起来对、跑起来慢"。

根本原因在于:语言模型对 GPU 硬件的理解是"统计性"的,它知道常见的优化模式长什么样,但它不知道当前这块卡、这个 shape、这个 dtype 组合下,哪种 tile size 才是最优的。它没有反馈信号,只能靠"猜"。

我试过给模型喂大量 CUDA 优化手册和开源 Kernel 代码,让它模仿专家写法。结果确实比裸生成好一些,但提升有限——因为专家写法本身是"针对特定场景调出来的",换个 shape 就未必适用。模型学到的是"形",不是"神"。

1.2 编译器优化为什么也卡住了

另一条路是走编译器自动优化。TVM、Triton、Ansor 这类工作本质上是在一个搜索空间里找最优 schedule。它们的优势是有明确的代价模型和搜索策略,能保证找到的解至少在搜索空间内是较优的。

但编译器的问题在于:搜索空间是人设计的。如果某个优化手法(比如某种特殊的 swizzle 模式)没有被写进搜索空间,编译器就永远找不到它。而专家手写 Kernel 时,很多"神来之笔"恰恰来自搜索空间之外的直觉。

所以两条路各有各的墙:Agent 缺反馈,编译器缺创造力。

1.3 CAKE 的核心洞察:让两者互补

CAKE 的思路很直接——既然 Agent 有创造力但缺反馈,编译器有反馈但缺创造力,那就让它们组成一个闭环:

  • Agent 负责提出"候选优化策略",包括一些编译器搜索空间里没有的写法;
  • 编译器负责把这些策略落地成可编译的 IR,并给出性能反馈;
  • 反馈再喂回 Agent,让它下一轮提出更好的策略。

这个闭环的关键在于中间那层"IR"——论文里叫 Cake IR。它不是普通的 CUDA 源码,也不是 TVM 的 TIR,而是一种专门为"Agent 可读、编译器可执行"设计的中间表示。这一点后面会详细拆。

2. Cake IR:Agent 和编译器之间的"共同语言"

2.1 为什么不能直接用 CUDA 源码做交互

你可能会问:Agent 直接生成 CUDA 源码,编译器直接编译不就行了,为什么要搞一层 IR?

我一开始也有这个疑问,后来想明白了。直接用 CUDA 源码做交互有三个致命问题:

第一,反馈粒度太粗。编译器只能告诉你"这个 Kernel 跑出来是 2.3ms",但没法告诉你"是因为你的 tile size 选错了导致 shared memory 利用率低"。Agent 拿到一个标量性能数字,根本不知道该改哪里。

第二,搜索效率太低。CUDA 源码是一个巨大的文本空间,Agent 每次改动可能引入语法错误、逻辑错误,编译失败率高,大量迭代浪费在"修 bug"上而不是"优化"上。

第三,无法复用编译器已有的优化。编译器内部有大量成熟的 pass(循环展开、向量化、内存提升),如果 Agent 直接写最终 CUDA,这些 pass 就用不上了。

2.2 Cake IR 的设计取舍

Cake IR 的设计目标很明确:让 Agent 只需要表达"优化意图",把"优化实现"留给编译器。

具体来说,Cake IR 把 Kernel 拆成几个层次:

层次内容谁负责
计算图算子的数学定义、数据依赖从原始框架自动提取
调度原语tile size、并行维度、内存层级Agent 提出候选
循环结构循环嵌套、展开因子、向量化编译器根据调度原语生成
硬件映射thread/block 映射、shared memory 分配编译器负责

Agent 主要操作的是"调度原语"这一层。它不需要写具体的循环代码,只需要说"我要在 M 维度上 tile 成 64,N 维度上 tile 成 128,K 维度上做 8 路展开",编译器就能把这些意图翻译成完整的 Kernel。

这样做的好处是:Agent 的输出空间被大幅压缩,每次迭代的有效信息密度更高,编译失败率也低得多。

2.3 Cake IR 的可读性设计

论文里特别强调了一点:Cake IR 是"Agent 可读"的。这意味着它的语法要足够简洁、语义要足够清晰,让语言模型能准确理解当前 Kernel 的状态。

我看了论文里给的 IR 示例,整体风格接近一种"带注释的伪代码",比 TIR 友好很多。比如一个 matmul 的调度原语大概长这样:

schedule matmul { tile: [M=64, N=128, K=32] parallel: [block: (M, N), thread: (M_inner, N_inner)] memory: { A: shared, B: shared, C: register } unroll: K by 8 vectorize: load by 4 }

这种表示对模型来说很好理解,因为它接近自然语言的描述方式,但又足够形式化,编译器能直接解析。

提示:如果你自己要做类似的 Agent+编译器闭环,IR 的可读性设计比性能优化更优先。因为 Agent 的理解准确率直接决定了整个闭环的收敛速度。

3. 闭环是怎么跑起来的:Agent 与编译器的协作机制

3.1 一轮迭代的完整流程

CAKE 的一轮迭代大致分四步:

  1. Agent 提出候选调度:基于当前 Kernel 状态和历史反馈,Agent 生成若干组调度原语候选。
  2. 编译器编译并执行:每组候选被翻译成完整 Kernel,编译、运行、测性能。
  3. 反馈收集:不只是收集最终耗时,还收集编译器内部的中间指标(如 shared memory 使用率、寄存器压力、occupancy)。
  4. Agent 更新策略:Agent 根据反馈调整下一轮候选。

这个流程听起来简单,但每一步都有坑。我下面拆开讲。

3.2 Agent 侧:怎么让模型"看懂"反馈

这是整个系统里我觉得最难的部分。性能反馈是一个多维向量,直接丢给模型,模型很难从中提取有效信息。

CAKE 的做法是把反馈做"结构化归因"。比如性能不达标时,系统会分析是哪个环节出了问题:

  • 如果是 occupancy 低,反馈里会明确说"寄存器压力过高,建议减少 unroll 因子";
  • 如果是内存带宽瓶颈,反馈里会提示"考虑增加 shared memory 复用";
  • 如果是计算瓶颈,反馈里会建议"尝试更大的 tile size 提升计算密度"。

这种归因不是简单的规则匹配,而是结合了编译器的静态分析和运行时的 profiling 数据。Agent 拿到这种"带诊断的反馈",下一轮调整的针对性就强很多。

我自己做类似系统时的经验是:反馈的信息密度比反馈的频率更重要。与其让 Agent 快速迭代 100 轮但每轮反馈都是"慢了",不如让它迭代 20 轮但每轮反馈都告诉它"慢在哪、为什么慢、往哪个方向调"。

3.3 编译器侧:怎么保证"意图"被正确落地

编译器这边要解决的是"意图到实现"的翻译问题。Agent 说"tile 成 64x128",编译器要决定:

  • 这 64x128 的 tile 怎么映射到 thread block?
  • shared memory 怎么分配?要不要 padding 避免 bank conflict?
  • 寄存器怎么分配?会不会 spill?

这些决策编译器不能随便做,因为不同的决策会导致完全不同的性能。CAKE 的方案是让编译器生成多个"实现变体",然后快速测一遍选最优的。这本质上是在编译器内部做了一层小规模的 auto-tuning。

这里有个工程上的细节值得注意:编译变体的测试要足够快。如果每个变体都要完整编译+运行,迭代速度会慢到无法接受。CAKE 用了缓存和增量编译来加速,具体做法是把不变的部分(如计算图)缓存起来,只重新编译变化的部分(调度原语)。

3.4 收敛判据:什么时候停

闭环系统必须有个停止条件,否则会一直跑下去。CAKE 的停止条件我理解是复合的:

  • 性能提升连续 N 轮低于阈值;
  • 或者达到预设的迭代上限;
  • 或者 Agent 提出的候选已经无法带来新的编译变体。

实际用的时候,我建议再加一条:性能已经超过参考实现(如 cuBLAS)时可以考虑提前停。因为继续优化的边际收益可能很低,而迭代成本是实打实的。

4. 实测效果:CAKE 到底比专家强在哪

4.1 论文里的对比数据说明了什么

论文里给的对比主要是三类基线:专家手写 Kernel、TVM/Ansor 自动调度、以及纯 Agent 生成。CAKE 在部分算子上超过了专家手写,这个结论乍看很惊人,但拆开看其实合理。

超过专家的场景通常有这些特征:

  • shape 比较特殊:专家手写 Kernel 往往是针对常见 shape 调的,碰到非常规 shape 就不一定最优;
  • 硬件比较新:专家的经验可能来自上一代卡,新卡的特性(如更大的 shared memory、新的指令)没被充分利用;
  • 优化空间大:算子的计算模式允许很多种调度方式,专家只试了其中几种。

而在那些"专家已经调了很多年"的经典算子上,CAKE 基本是打平或略逊。这说明 CAKE 的优势不是"全面碾压",而是"在专家覆盖不到的地方补位"。

4.2 一个具体的优化案例拆解

论文里提到一个 matmul 的案例,我印象比较深。专家手写版本用的是经典的 128x128 tile + 双缓冲,性能已经很好。CAKE 最终找到的方案是 64x256 tile + 三缓冲 + 非对称 unroll。

这个方案为什么更好?我分析下来有几点:

  • 64x256 的 tile 在这个特定 shape 下,shared memory 利用率更高,因为 N 维度更大,B 矩阵的复用次数更多;
  • 三缓冲让流水线更满,掩盖了更多的内存延迟;
  • 非对称 unroll(M 方向 unroll 4,N 方向 unroll 8)匹配了实际的指令吞吐特性。

这些决策单看每一条都不算"黑科技",但组合起来就超过了专家版本。而专家之所以没这么写,很可能是因为他的经验里 128x128 是"默认最优",没有针对这个 shape 重新搜索。

4.3 我自己的复现尝试和踩坑

我按论文思路搭了一个简化版的闭环,用一个小模型做 Agent,编译器用 Triton 做后端。踩了几个坑:

坑一:Agent 输出格式不稳定。模型有时候会输出不符合 IR 语法的内容,导致解析失败。解决办法是在 prompt 里给足 few-shot 示例,并且在解析层做容错——解析失败时不是直接报错,而是让 Agent 重新生成。

坑二:反馈延迟太高。一开始每轮迭代要跑完整的编译+profiling,一轮要几分钟,跑 50 轮就是几个小时。后来改成先用编译器的静态代价模型做粗筛,只对 top-3 候选做实际运行,速度提升了一个数量级。

坑三:Agent 会"过拟合"反馈。有几次 Agent 发现某个调度原语能小幅提升性能,就反复在这个方向微调,陷入了局部最优。解决办法是引入一定的随机性,强制 Agent 每几轮探索一个"看起来不太相关"的方向。

5. 如果你想自己搭一套类似的系统

5.1 最小可行系统的组件清单

不需要一上来就复现完整 CAKE,可以先搭一个最小闭环:

组件作用推荐方案
Agent提出调度候选任意支持 function calling 的模型
IR 层结构化表达调度意图自定义 JSON schema 或类 DSL
编译器后端把 IR 翻译成可执行 KernelTriton 或 TVM
反馈收集性能数据 + 诊断信息profiling + 静态分析
迭代控制管理闭环流程简单的 Python 脚本即可

关键是 IR 层要设计好,它是整个系统的"关节"。

5.2 IR 设计的几条实操建议

  • 粒度要适中:太细(比如让 Agent 指定每条指令)Agent 学不会,太粗(比如只说"优化这个 Kernel")编译器没法执行。建议粒度定在"调度原语"这一层。
  • 要有默认值:Agent 没指定的参数,编译器要有合理的默认值,避免因为缺参数导致编译失败。
  • 要可验证:IR 解析后要做合法性检查,比如 tile size 不能超过 shared memory 容量,并行维度不能超过硬件限制。
  • 要可扩展:预留自定义原语的接口,方便后续加入新的优化手法。

5.3 反馈设计的核心原则

反馈设计我总结成一句话:让 Agent 知道"改哪里、往哪改、为什么"。

具体来说,反馈里应该包含:

  • 当前性能指标(耗时、吞吐、occupancy 等);
  • 与上一轮的对比(哪些指标变好了,哪些变差了);
  • 瓶颈诊断(当前瓶颈在计算、内存还是同步);
  • 调整建议(基于诊断给出的方向性提示)。

这四类信息组合起来,Agent 的迭代效率会高很多。

6. 这套思路还能往哪些方向延展

6.1 从单 Kernel 到整个算子图

CAKE 目前主要针对单个 Kernel 做优化。但实际模型里,算子之间是有融合机会的。如果把这个闭环扩展到算子图层面,Agent 不仅能优化单个 Kernel 的调度,还能决定哪些算子该融合、怎么融合,优化空间会大很多。

6.2 跨硬件平台的迁移

不同 GPU 架构的特性差异很大(shared memory 大小、warp 调度策略、指令集)。CAKE 的闭环机制天然适合做跨平台迁移:把在一个平台上搜到的最优调度作为起点,在新平台上继续迭代。这比从零开始搜要快得多。

6.3 Agent 能力的持续进化

论文标题里"共同进化"这个词我觉得很关键。随着闭环跑的次数增多,Agent 其实在积累"哪种调度在哪种场景下有效"的经验。如果把这些经验沉淀下来(比如做成一个检索库或微调数据),Agent 的能力会随时间提升,而不是每次都从零开始。

我自己在实际操作中的体会是:这套系统的价值不在于"一次性能超过专家",而在于它能持续地从每次迭代中学习,把专家的隐性经验显性化、结构化。专家调 Kernel 靠的是多年积累的直觉,而 CAKE 这类系统做的事情,是把这种直觉拆解成可搜索、可复用、可迁移的知识。这个过程本身,可能比最终的性能数字更有意义。

最后分享一个小技巧:如果你打算尝试类似的方向,不要一上来就追求"超过专家"。先把闭环跑通,哪怕性能只有专家的 70%,只要迭代速度够快、反馈够清晰,后面追上来的速度会超出你的预期。真正难的不是让 Agent 写出好 Kernel,而是让整个系统"知道自己在往哪个方向变好"。

版权声明: 本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若内容造成侵权/违法违规/事实不符,请联系邮箱:809451989@qq.com进行投诉反馈,一经查实,立即删除!
网站建设 2026/10/15 8:19:50

喷雾燃烧仿真技术:模型框架与工程实践

1. 喷雾燃烧仿真技术概述在工业燃烧器和航空发动机领域,喷雾燃烧是最常见的燃烧形式之一。当液体燃料通过喷嘴雾化成微小液滴并与氧化剂混合燃烧时,会形成复杂的多相湍流燃烧过程。这种燃烧方式在燃气轮机、柴油机和火箭发动机中广泛应用,其燃…

作者头像 李华
网站建设 2026/10/15 8:19:37

Matlab在多能流计算中的应用与优化

1. 项目背景与核心价值区域综合能源系统(RIES)是当前能源领域的前沿研究方向,它打破了传统电、热、气等能源系统单独规划运行的局限。我在参与某工业园区能源系统改造项目时,深刻体会到多能耦合计算的必要性——当光伏发电过剩时&…

作者头像 李华
网站建设 2026/10/15 8:19:35

电热综合能源系统调度优化与Matlab实现

1. 项目背景与核心价值电热综合能源系统作为区域能源供应的重要形式,正在经历从传统化石能源向可再生能源转型的关键阶段。这个转型过程中最大的技术挑战在于:风电、光伏等可再生能源具有显著的波动性和不确定性,而供热系统又存在较强的热惯性…

作者头像 李华
网站建设 2026/10/15 8:18:46

攻防世界Misc高手进阶区:隐性攻防能力训练体系解析

1. 这不是CTF题库,而是一套面向实战的隐性能力训练体系“攻防世界 misc高手进阶区”——光看标题,很多人第一反应是:又一个CTF在线练习平台的Misc(杂项)分类?点进去发现题目零散、描述简略、没有标准解法提…

作者头像 李华
网站建设 2026/10/15 8:17:48

OpenCV源码调试实战:从下载编译到IDE单步追踪

1. 项目概述:为什么一个“下载安装”要专门调试源码?OpenCV下载安装与源码调试——这八个字看起来像入门第一课,但实际踩过坑的人才知道,它根本不是点几下Next就能完事的流程。我带过三届某高校实验室的图像处理方向本科生&#x…

作者头像 李华
网站建设 2026/10/15 8:17:22

Unity Spark实测:用自然语言让AI生成可玩3D游戏原型

几年前我们聊"用AI做游戏",大家想的还是让AI生成几张角色原画、写一段对话文本,或者是让某个智能体在2D网格里自己走两步。但今天打开浏览器,用自然语言描述一个场景,AI能直接给你搭出一个可运行、可游玩的3D原型&#…

作者头像 李华