跳到正文
Hacker News· frostbyte7·· 6 天前AI 评分74

NVIDIA 与 CMU 的 CAKE:让智能体写 GPU kernel,编译器同步进化

Cake from Scratch: Agents Write GPU Kernels

AI 导读

NVIDIA 与卡内基梅隆大学的论文 CAKE(Compiler–Agent co-design for frontier Kernel Evolution,arXiv 2608.12629)提出让智能体编写可静态检查的 CAKE IR 调度语言,由编译器返回定位到资源、角色或流水线阶段的诊断,并在反复失败时反过来修改编译器本身。

正文 · AI 翻译

译文尚不完整,完整内容请切换到原文。

2022 年,FlashAttention 让 transformer 的注意力层在同一块 GPU 上快得多,而且没有做任何近似:它计算的注意力公式完全相同,只是换了计算顺序,因此其结果与常规方法仅相差浮点舍入误差。它没有发现新公式,而是找到了更好的调度:哪块数据放在哪种内存里、芯片的哪部分处理它、以什么顺序处理。此后又推出了三个版本,每个版本都将更好的调度与适合当时硬件的算法改动相结合;最后两个版本是为新一代 GPU 编写的。

这类程序被称为GPU 内核,AI 领域的大量资金都花在这里:大语言模型(LLM)读取或写入的每个 token 都要经过数十个这样的内核。近来,AI 编程智能体也开始编写内核。它们提出代码、编译、测试、计时,然后重试。在许多基准任务上,这个循环如今能产出正确且快速的代码。但在最难的内核上——即专家为每一代新 GPU 手写的那些——它仍然力有不逮。

CAKE(面向前沿内核演化的编译器–智能体协同设计)是 NVIDIA 与卡内基梅隆大学的一篇论文,它主张仅改进智能体是不够的。内核智能体方面的工作与 GPU 编程语言方面的工作一直各自推进,用论文的话说,“二者之间的鸿沟正是专家级内核迷失之处。”这条鸿沟就是智能体所处的世界:它们书写的语言以及它们获得的反馈。如今的智能体编写原始 CUDA 或某种高级内核语言,而环境给出的回应不过是“崩溃了”“输出错了”或“耗时 0.94 毫秒”。CAKE 改变了这两端。智能体用一种新的、可检查的语言来编写 GPU 调度,编译器则给出诊断信息,指向受影响的资源、角色或流水线阶段。(论文未公布其消息格式;以下是我们自己按此精神举的一个例子:这个屏障,在这个流水线阶段,被等待但从未被发出信号。)接着 CAKE 更进一步。当智能体反复以同样方式失败时,编译器本身会被修改:该失败会变成一项新检查、一个新语言特性,或一个修正后的性能模型。

论文给出的头条结果如下:

1.144× 对 0.928×从零开始的 Flash-KMeans8000 万 token 后的最佳内核,CAKE IR 对直接 CUDA/PTX,相对于调优基线(3 次运行的中位数,B200)

2.05×Kimi Delta Attention 预填充提示处理内核对官方 FlashKDA 内核,在 6 种输入规模上取平均;已在端到端服务中验证

1.42×–2.12×KNN 与 KMeans 库三个内核系列,按输入规模选择内核;在 112、198 和 124 种规模上平均分别为 1.418×、2.116× 和 1.803×(GB200);单一规模波动更大

4 个 PR上游至 FlashInferKDA 预填充与解码、TinyGEMM,以及一个 Alpha-MoE 巨型内核;四个均已于 2026 年 9 月 13 日合并(据 GitHub)

如果这些标签目前还看不出什么意义,那也在意料之中;每张卡片都有自己的章节(第一张见§12,第二张和第四张见§13,第三张见§14)。有两个词到处出现。一个形状(shape)是指一组输入尺寸的选择,例如32个批次、每批65,536个点。当一个结果覆盖多个形状时,加速比用几何平均来平均,这是比值的标准平均方式:2×的收益和0.5×的损失平均后是1×,而普通平均会声称是1.25×。

本文从头讲解CAKE。它基于论文(“CAKE: Compiler–Agent Co-Design for Frontier Kernel Evolution”,arXiv 2608.12629,2026年8月12日提交)、其LaTeX源码和图表,以及论文引用的公开来源;硬件背景来自NVIDIA的文档和独立测量,上游细节来自FlashInfer和FlashLib的拉取请求以及Kimi-K3模型卡,均在用到的地方给出链接。除非另有说明,数字均来自论文。我们从图表中读出的值标记为“近似”;我们自己计算出的数字标记为“推导”;我们对证据的解读标记为解释。论文未提及之处,我们会说明,§16汇总了这些空白。

本文面向谁

大多数关于GPU内核的论文都假设你已经会写内核。本文不是这样。如果你懂一些Python,并且大致知道神经网络做什么,你就能跟上。我们借用了深受喜爱的一本操作系统教材——Remzi和Andrea Arpaci-Dusseau所著的Operating Systems: Three Easy Pieces——的形式,并采用了它的三个习惯:

  • 三个部分。本文分三部分,每部分都以一段教授与学生之间的简短对话开头。对话提出问题并勾勒直觉;每个事实都在正文中展开。
  • 关键问题。处理真实问题的章节会用一个标有关键问题的方框来点明它,这样你就知道我们要解决什么。
  • 旁注与提示。旁注框里是可以跳过的背景知识。提示框里是超越本文的一般性经验。

如果你已经很懂GPU,可以跳过第一部分,从§5开始。

路线图

  • 第一部分 · 机器(§1–4):什么是GPU内核,为什么速度主要取决于数据移动,现代GPU如何变成异步流水线,以及专家的内核与仅仅正确的内核之间有什么区别。
  • 第二部分 · 语言(§5–8):如今智能体如何编写内核、它们得到什么反馈;GPU语言的选择菜单,以及为什么没有一种适合智能体;逐行解读CAKE IR;以及它的来源。
  • 第三部分 · 循环(§9–16):CAKE在内核接触GPU之前运行的检查;一轮演化;编译器本身如何演化;受控实验;生产结果;把调优好的内核变成库;相关工作;以及局限。
  • 尾声:一段收尾对话、总结、思考题,以及带注释的来源。

第一部分 · 机器——一段对话

教授欢迎回来。今天讲:一篇关于编写GPU内核的AI智能体的论文。

学生我每天都用PyTorch,从没写过内核。我需要写吗?

教授你一直在用它们。每个torch.matmul最终都会变成某人写的一个内核。问题是谁来为下一代GPU写下一个内核。

学生编译器不能直接生成吗?编译器不就是干这个的。

教授 对于简单操作,是的。但对于重要的操作,最快的代码仍然来自专家。他们了解芯片的一些特性,而编译器目前还无法很好地处理这些。

学生 比如呢?

教授 比如哪组线程加载数据,而另一组线程进行乘法运算,以及它们如何互相通知“你的数据准备好了”。如果搞错了,程序不会礼貌地崩溃。它会挂起,或者悄悄产生垃圾结果。

学生 那调试起来听起来很痛苦。而AI应该来做这个?

教授 这就是论文的问题。但首先你需要了解机器。让我们从内核到底是什么开始。

1. 什么是GPU内核

GPU是一种芯片,旨在同时对大量数据执行相同的简单操作。本节介绍后续文章所需的词汇:线程、线程束、线程块、流式多处理器,以及内核这个词本身。

1.1 许多小工人,而不是几个快工人

CPU有少量强大的核心,每个核心都旨在尽快完成一个指令序列。GPU则做出了相反的赌注。它有数千个更简单的通道,通过让所有通道同时忙碌来取胜。没人在意一次乘法需要多长时间;重要的是每秒完成多少次。这就是为延迟优化和为吞吐量优化之间的区别。

内核是在GPU上运行的函数。你编写一次,GPU为每个线程运行一个副本,通常有数百万个线程,每个线程处理数据的不同部分。将两个一百万个数字的向量相加是经典的首个内核:线程i读取a[i]和b[i],将它们相加,然后写入c[i]。

线程按层次组织,层次结构与硬件匹配:

  • 线程束。 NVIDIA GPU以32个线程为一组运行,称为线程束。一个线程束的32个线程同时执行相同的指令,每个线程处理自己的数据。
  • 线程块。 线程束被分组为线程块(也称为CTA,协作线程数组)。一个块的线程在芯片的同一部分运行,可以共享快速暂存内存。
  • 流式多处理器(SM)。 芯片由许多相同的SM组成。每个SM一次运行一个或多个线程块,拥有自己的寄存器、暂存内存和数学单元。NVIDIA B200有148个SM。
  • 网格。 一次内核启动会创建一个线程块网格,硬件将它们分布到各个SM上。
Four panels about an NVIDIA B200. A warp is 32 threads that execute the same instruction, each on its own data. Warps form thread blocks, and the hardware spreads the blocks of a kernel launch across the streaming multiprocessors (SMs); a B200 has 148 SMs. Inside each SM are registers (256 KB), tensor memory (256 KB), shared memory (up to 228 KB) and tensor cores, the matrix-multiply units. The memory ladder runs from registers and tensor memory (256 KB per SM each) through shared memory (up to 228 KB per SM) and the 126 MB L2 cache shared by all SMs down to 180 GB of HBM at up to 8 TB/s: fast on-chip storage of a few hundred kilobytes per SM, then larger and slower memory off the SMs.
第一部分词汇在一张图中:32线程的线程束形成线程块,硬件将块分布到B200的148个SM上,内存形成从快速片上存储(每个SM几百KB)到更大、更慢的L2和HBM的阶梯。打开全尺寸SVG。

1.2 为什么矩阵乘法是重要的内核

向量加法很简单,因为每个输出只需要两个输入。矩阵乘法不同。要计算输出矩阵C = A × B,每个输出元素需要A的整行和B的整列,而每个输入元素被多个输出使用。这种重用既是机会也是难点:好的内核加载每个输入一次并多次重用,而差的内核则不断从远处获取相同的数字。

Transformer 所做的几乎所有事情都建立在矩阵乘法之上:注意力机制中的投影、前馈层、混合专家(MoE)模型中的专家。现代 GPU 有专门用于这些运算的单元,称为张量核心,它们能在一条指令中完成小矩阵块的乘法。内核开发者将矩阵乘法称为 GEMM(通用矩阵乘法),将张量核心操作称为 MMA(矩阵乘法累加,C += A × B)。CAKE 论文中的大多数内核,本质上都是精心设计的、让张量核心持续运转的方法。

2. 快速内核的关键在于数据搬运

对新手来说,第一个意外是:在现代 GPU 上,算术运算很少是困难的部分。及时把数据送到算术单元才是。

2.1 内存阶梯

GPU 有几种内存,像梯子的横档一样排列。横档离数学运算单元越近,速度越快,而且大体上容量越小:片上横档每个 SM 只有几百 KB,片外横档则有数 MB 到数 GB。

横档 位置 容量(B200) 用途
寄存器 每个 SM 内部 每个 SM 256 KB,在其线程间分配 每个线程的工作值
张量内存(TMEM) 每个 SM 内部,数据中心 Blackwell(B200、GB200)新增 每个 SM 256 KB 张量核心矩阵乘法的累加器
共享内存(SMEM) 每个 SM 内部 每个 SM 最多 228 KB 一个由块内线程共享、手工编程的暂存区
L2 缓存 所有 SM 共享 126 MB 硬件管理的缓存
HBM(高带宽内存,也称全局内存) 芯片旁边的 DRAM 堆栈 180 GB,约 8 TB/s 张量所在之处

这些硬件数字的来源

  • 寄存器(每个 SM 64K 个 32 位寄存器)、共享内存(每个 SM 228 KB,每个块最多 227 KB)和 126 MB 的 L2:来自 NVIDIA 的 Blackwell 调优指南以及 CUDA 编程指南的计算能力表。调优指南给出了 GB200 的 L2 大小;Chips and Cheese 在 B200 上测得了同样的 126 MB。
  • 张量内存:PTX ISA 描述每个 CTA 有 128 行 × 512 列 32 位单元,即 256 KB。
  • HBM:每个 B200 为 180 GB,最高 8 TB/s(NVIDIA 参考架构)。稠密 BF16:一个 8 卡 HGX B200 在稀疏模式下额定为 36 PFLOPS,稠密是其一半,因此每 GPU 为 2.25 PFLOPS。
  • SM 数量:NVIDIA 的数据表没有列出。148 这个数字来自 FlashAttention-4 论文和独立测量。GB200 中的 Blackwell GPU 是稍大的配置(186 GB HBM,2.5 PFLOPS 稠密 BF16)。

每秒八万亿字节听起来很庞大,直到你把它和数学运算比较。B200 的张量核心每秒可以完成约 2.25 千万亿(2.25 × 1015)次稠密 BF16 运算。(BF16,即 bfloat16,是一种 16 位浮点格式。深度学习使用它,是因为与普通的 32 位浮点数 FP32 相比,它能把内存流量减半,而网络能容忍其精度损失;FP8 和 FP4 则降到 8 位和 4 位。“稠密”意味着不使用 NVIDIA 用来将宣传数字翻倍的稀疏技巧。)用前者除以后者:要让张量核心保持忙碌,内核从 HBM 每读取一个字节,就必须完成大约 280 次运算(推导得出)。每字节工作量更少的内核,无论其数学多么巧妙,时间都花在等待内存上。

这个运算量与字节数的比值被称为算术强度,而一个内核要么受限于内存带宽、要么受限于计算能力(取决于哪个先耗尽)这一理念被称为roofline 模型。两个向量相加,每移动 12 字节(两次 4 字节读取、一次写入)才做一次加法;它永远是内存受限的。大型矩阵乘法可以达到每字节数千次运算,但前提是内核能很好地复用数据。

2.2 分块:每个快速内核背后的诀窍

提高复用率的标准方法是分块。线程块不是逐个元素地取输入,而是把 A 的一块和 B 的一块从 HBM 复制到共享内存,用张量核心做乘法,累加部分结果,然后继续处理下一对块。从 HBM 取出的每个数字都会被输出块的一整行或一整列复用:在 128 × 128 的块下,复用 128 次而不是一次。

分块也是 FlashAttention 的核心。它计算出的注意力与其他人相同,但它是逐块处理输入并融合各个步骤,因此那张巨大的注意力分数中间表(每对 token 对应一个数字)根本无需写入 HBM。

分块带来了一个新问题:当线程块等待下一块数据到达时,张量核心处于空闲状态。解决办法是在计算当前块的同时取下一块,使用两个或更多块缓冲区。这被称为双缓冲,更一般地称为软件流水线。它隐藏的是内存延迟而非减少流量,后来的注意力内核严重依赖它:为 Hopper 编写的 FlashAttention-3 就是围绕复制与乘法的重叠来构建的。

3. GPU 变成流水线

近期的 NVIDIA GPU 增加了硬件,使分块和流水线变得快得多,也难编程得多。本节将梳理这一变化,因为 CAKE 的设计与之紧密相关。

3.1 从“每个线程做所有事”到专门化角色

在较老的 GPU 内核中,每个 warp 都做所有工作:加载块的一部分,等待块中其余部分,计算,重复。此后每一代 GPU 都增加了硬件,能异步地自行完成其中一项工作,而线程则去做别的事。你不需要记住下面的名称。重要的是趋势:每一代都把又一项工作交给一个在后台工作的专用单元,于是线程不再做这项工作,而是开始协调它。后面会再次出现的名称有 TMA(复制引擎)、TMEM(累加器内存),以及宣告复制完成的屏障。

  • Ampere(A100,2020)增加了 cp.async,它在后台把数据从全局内存复制到共享内存。其张量核心由 warp 级的 mma.sync 指令驱动,数据由 ldmatrix 从共享内存送入。
  • Hopper(H100,2022)增加了张量内存加速器(TMA),这是一个复制引擎,用一条指令就能移动整个多维块;WGMMA,一条由四个 warp 组成的 warpgroup 发出的张量核心指令;异步事务屏障(mbarrier,在复制数据到达时按字节计数)用于发出块已到达的信号;以及线程块集群,其块可以读取彼此的共享内存。
  • Blackwell(B200,2024–25)加入了 tcgen05.mma,这是一条张量核心指令,由单个线程代表整个块发出,还加入了张量内存(TMEM),每个 SM 拥有专用的 256 KB,用于存放累加器,使其不再堵塞寄存器。两个相邻的 SM 甚至可以协作完成一次更大的乘法(即“2-CTA”模式)。

由于复制引擎和张量核心各自独立运行,自然的设计就是流水线。一些 warp 只负责发出复制指令(生产者);一个 warp 只负责发出矩阵乘法指令(MMA warp);其他 warp 则取走完成的结果,应用任何最终的逐元素步骤,例如缩放或激活函数,并将其写出。最后一个阶段称为尾声(epilogue)。这种设计称为warp 特化,每项工作都是一个warp 角色。

Animated diagram of a warp-specialized matrix-multiply main loop on one SM. Tiles k, k+1 and k+2 in HBM are copied by TMA into a three-stage shared-memory ring buffer; a producer warp issues the copies. Each stage has a full barrier (data landed, safe to read) and an empty barrier (data used, safe to overwrite). Beat 1: the producer copies tile k into stage 0 and arrives on its full barrier. Beat 2: the MMA warp waits on full, multiplies stage 0 into the tensor-memory accumulator and then arrives on empty, while tile k+1 is being copied into stage 1. Beat 3: steady state, with tile k+2 being copied into stage 2, tile k+1 being multiplied and stage 0 free again; after the last tile, the epilogue warps store the finished result back to HBM. Beat 4: stage 0 now holds tile k+3; if the producer skips its wait on the empty barrier, tile k+4 overwrites stage 1 while the MMA warp is still reading k+1, a race that gives silently wrong results or a hang.
一个作为流水线的 warp 特化主循环:生产者填充环形缓冲区各阶段,full 屏障表示一个 tile 已到达,MMA warp 消费它,empty 屏障释放该阶段,而一次被跳过的等待会让 tile k+4 覆盖一个仍在被读取的阶段。打开全尺寸 SVG。

3.2 交接:环形缓冲区与屏障

流水线需要一种让各工位相互交接工作的方式。在 warp 特化内核中,共享内存 tile 缓冲区构成一个由阶段组成的环形缓冲区:有三个阶段时,生产者可以填充阶段 2,同时张量核心处理阶段 1,而刚刚释放的阶段 0 则等待它的下一个 tile。

交接使用屏障。屏障是一个会合点,被设置为预期一定数量的到达,即其到达计数。一个在其上等待的 warp 会被挂起,直到发生了那么多次到达(对于 TMA 复制,直到所有复制的字节都已到达);然后屏障完成当前阶段并为下一阶段重置。NVIDIA 将向屏障发信号称为在其上到达,图中也使用该词。

每个阶段有两个屏障。生产者在其复制到达时到达该阶段的 full 屏障;消费者在读取前等待 full。消费者在完成该阶段时到达该阶段的 empty 屏障;生产者在覆盖它之前等待 empty。由于环形会绕回,每个屏障每圈被复用一次,因此等待者必须说明它在等待哪一圈。硬件用一个每圈翻转的 1 位相位来跟踪这一点;在错误的相位上等待要么过早返回,要么永不返回。

还有两个细节稍后很重要。GPU 通过两条独立的路径访问内存:线程执行的普通加载和存储(通用代理),以及复制引擎和张量核心(异步代理)。想象两个快递员向同一张桌子送货:没有任何机制保证一个快递员能看到另一个快递员按顺序送达的货物。因此,当线程用普通存储写入数据,而后续的异步操作要读取这些数据时(例如共享内存中的值被 TMA 存储或张量核心乘法消费),内核需要显式的代理栅栏,否则硬件可能会重排这些访问。反方向上,异步拷贝的完成通过其屏障被观察到时,数据已经可见。而每个 TMA 拷贝都需要一个描述符,它编码了瓦片的形状、步幅(内存中连续行和列之间的距离)以及布局。

3.3 为什么这很容易出错

这些细节中的每一个都是一种失败方式,而失败是残酷的:

  • 缺少等待会让消费者在数据到达之前读取阶段。不会崩溃;结果只是错误的,有时只对某些输入出错。
  • 屏障初始化的到达计数过高,或者等待一个永远不会完成的阶段,会让一个 warp 永远等待。内核挂起,唯一的症状是 GPU 停止响应。(计数过低,或阶段过期,则相反:等待提前返回,bug 变成静默的竞态。)
  • 缺少代理栅栏只在特定时序下产生陈旧数据,因此 bug 在一次运行中出现,在下一次运行中消失。
  • TMA 描述符的瓦片形状与其填充的共享内存缓冲区不一致,会破坏属于别人的内存。

有工具可以提供帮助。NVIDIA 的 compute-sanitizer 可以在运行时检测一些竞态和越界访问,但仅限于你恰好运行的输入,而且它不会解释哪个设计决策导致了问题。记住这一点:崩溃或挂起说明出了问题,而不是为什么。这个差距是第二部分主题。

4. 专家内核与正确内核的区别

把这些拼在一起,你就能看到 CAKE 论文所说的专家内核是什么意思。其摘要列出了三个决策,它们“将专家内核与仅仅正确的内核区分开来”:

  1. Warp 特化。哪些 warp 承担哪些角色:多少生产者,多少消费者,是否两组消费者交替(“乒乓”),以便一组计算而另一组写入结果。
  2. 屏障编排。多少个流水线阶段,哪个屏障控制哪个交接,每个等待期望哪个阶段,栅栏放在哪里。
  3. 内存层级放置。哪些数据放在寄存器中,哪些放在共享内存中,哪些放在张量内存中,以及以什么布局,使得复制引擎、张量核心和输出路径都同意每个字节的位置。

这些决策共同构成了内核的调度:驱动机器的计划,而不是计算什么。矩阵乘法的数学可以写在一行里。快速 Blackwell 矩阵乘法的调度有数百行,改变一个决策通常会迫使其他决策也改变。

第一部分总结

内核是数千个 GPU 线程并行执行的函数。它的速度主要取决于数据移动:通过一系列内存层级,以分块方式传输,在当前分块处理时预加载下一个分块。现代 NVIDIA GPU 将其转变为一条由专用 warp 组成的流水线,通过屏障将分块互相传递。专家的艺术在于调度(角色、屏障、内存放置),而其中的错误往往导致静默的错误答案或挂起,而非有用的错误信息。


第二部分 · 语言——对话

学生 所以智能体编写内核、运行它,然后查看发生了什么。这有什么问题?我就是这样调试的。

教授 当你的代码出错时,你看什么?

学生 堆栈跟踪。行号。那个为 None 的变量。

教授 现在想象一下,对于最严重的错误,你唯一的反馈是“它卡住了”、“答案错误”或“它花了 0.94 毫秒”。没有行号,没有变量。

学生 我会随机改东西然后祈祷。

教授 这大致就是内核智能体在最难的内核上的处境。所以,两个问题。智能体应该用什么语言编写?当出错时,编译器应该告诉它什么?

学生 它不应该像专家一样直接写 CUDA 吗?

教授 CUDA 让你能表达一切,包括每一个错误。更高级的语言保护你,但不会让你说出专家所说的话。论文的答案试图两者兼得:一种能说出专家所说的话的语言,其形式是编译器可以在任何运行之前检查的。让我们看看如何做到。

5. 智能体如今如何编写内核

编写内核的智能体很快出现。本节描述它们共享的循环、该循环擅长什么,以及 CAKE 所起始的反馈问题。

5.1 提议–测试–测量循环

这一设定由 KernelBench(2025)推广:给模型一个 PyTorch 算子,要求它给出一个计算相同结果的更快的 GPU 内核。此后,许多系统为 GPU 和其他 AI 加速器构建了类似的循环。

其中许多从生物学中借用了进化的思想,CAKE 的标题也以同样的方式使用这个词:保留一组候选内核,让 LLM 修改(“变异”)有前景的候选,测试并计时每个新候选,保留最快且正确的,然后重复。这些系统的区别在于它们改进循环的哪一部分:

  • 循环本身。 LLM 根据编译器错误、测试结果或分析器输出修订内核(AccelOpt,用于 AWS Trainium;Autocomp,用于张量加速器;KernelAgent,用于 GPU)。
  • 记忆。 KernelBlaster 添加了一个持久的知识库,存储过去的优化。
  • 搜索。 KernelEvolve 和 EvoEngineer 支持对候选内核进行进化搜索;AVO(来自 CAKE 的几位作者)用编码智能体取代固定的变异规则;K-Search 将规划与实现分离。
  • 模型。 AutoTriton 和 CUDA Agent 通过强化学习(从“正确且快速”等奖励中学习)训练模型本身。

在多样性之下,环境是相同的。智能体提出代码,代码被编译,数值测试将其输出与参考进行比较,基准测试测量其延迟,然后智能体选择下一次编辑。CAKE 论文的总结是:这些方法“演化搜索过程、累积记忆或模型权重,同时保留选定的 DSL [领域特定语言] 和评估环境。”

5.2 环境反馈了什么

这个循环对于局部调优很有效:尝试不同的 tile 大小、展开循环、融合两个操作。但它在 §4 中的决策上却举步维艰,原因论文用一句话概括:编译器错误、正确性结果和端到端计时“从不说出是哪个程序决策导致了同步失败、硬件契约违规或流水线停顿。”

想想每个信号告诉一个刚刚重写了流水线内核的智能体什么:

  • 挂起。 某个 warp 正在等待一个永远不会完成的屏障。哪个屏障?哪个阶段?哪个相位?进程就这样停止了。
  • 输出错误。 竞态、错误的布局、缺失的 fence,或者单纯的索引错误。测试可以说出哪些元素不同,但不能说是哪个决策导致的。
  • 一个时间。 0.94 毫秒。内核是受限于内存带宽、张量核心发射、流水线太浅,还是序列化两个角色的屏障?一个数字无法说明。
  • 分析器计数器。 有些循环还会传递分析器输出。它可以说内核在等待内存或张量核心,但不能说是哪个程序决策(阶段数、角色划分、缓冲区放置)造成的。

还有第二个更安静的问题。这些信号“无法在前沿工作负载暴露出缺失能力时增长。”如果语言无法表达新工作负载所需的调度,每次尝试都会因相同的结构性原因失败,而环境永远不会从中学习。

5.3 CAKE 的三项承诺

论文的观察是,编译器已经包含了专家头脑中保留的大部分机制:“结构化操作词汇表、资源模型、合法性检查、静态分析、成本模型、降级规则。”(降级是编译器将程序翻译成 GPU 自身代码的过程;§7.1。)CAKE 提出的问题是“如何让这些机制面向智能体,以及当前沿工作负载暴露差距时如何改进它。”它的答案有三个部分:

  1. 智能体编辑类型化 IR,而不是原始 CUDA。 IR(中间表示)是一种为编译器分析而设计的程序格式。CAKE 的是一种调度语言,其中硬件决策是显式的,因此可以在生成任何代码之前进行检查。
  2. 编译器返回局部化诊断,而不是通过/失败位。 正确性和安全性发现指向受影响的资源、角色或阶段,并指出被违反的规则类型;性能报告估计成本并指出广泛的瓶颈类别。廉价分析在候选者消耗 GPU 时间之前对其进行过滤。
  3. 测试框架本身也在进化。测试框架是围绕智能体、评判其工作的一切:检查、测试和基准(§9)。其中两部分在此至关重要。验证器读取调度方案,并在任何编译发生之前拒绝非法方案,很像一个 linter。成本模型预测候选方案运行有多快,而无需实际运行,就像天气预报:成本低,通常接近,但你仍要检查真实 GPU。当智能体反复以同样方式失败时,测试框架就会改变:新的验证器规则、新的原语(语言的内置操作)、重新校准的成本模型(根据新测量重新拟合),或可复用的优化策略,每项变更都由针对内核语料库的测试把关。
Two panels. The first, kernel-only evolution as used today: an agent writes a CUDA candidate, which is compiled, tested and timed; what comes back is a compile error, a hang or wrong output, or a single time such as 0.94 ms, that is, pass or fail and one number, which says that something is wrong but not why. The second, CAKE: the agent writes a CAKE IR candidate and the harness (static checks, a cost model, code generation and a GPU run) returns localized findings, for example that barrier q_full is waited on at stage 2 but never arrived on, or that the estimated bottleneck is the tensor-core issue rate. An outer loop turns capability gaps into compiler evolution and a harness update to the IR, analyses and lowering. The inner loop evolves kernels; the outer loop evolves the compiler. The finding texts are illustrative.
如今的核智能体只能听到通过/失败和一个时间;CAKE 的测试框架返回定位后的发现,而一个外层循环将能力缺口转化为编译器变更。发现文本仅为示意。打开全尺寸 SVG。

论文将这一结果称为编译器–智能体协同设计。智能体的环境(语言、分析、成本模型)围绕智能体的需求来设计,主要由智能体维护,并随着它们遇到的情况不断变化。智能体本身,即其模型和脚手架,保持固定(§10.2)。接下来三节涵盖第一项承诺:语言。

6. 智能体应该写哪种语言?

编写 GPU 内核的方式有很多。它们在一个问题上差异最大:谁来决定调度,程序员还是编译器?本节梳理这些选项,采用 CAKE 论文自己的分类,并带着一个问题:

6.1 高层:tile 语言

Triton(2019)通过让程序员以tile为单位思考,使内核编写变得容易:你编写对矩阵块的操作,编译器决定线程、内存和指令如何执行。更新的语言如 Helion、TileLang 和 NVIDIA 的 cuTile 在不同抽象层次上遵循类似理念。

这对人类非常友好,并能给出正确、通常快速的代码。CAKE 论文的反对意见很具体:tile 语言“隐藏了将专家内核与仅仅正确的内核区分开来的 warp 特化、屏障编排和内存层级放置。”编写 Triton 的智能体可以要求流水线深度(num_stages=3),并在 Blackwell 上开启自动 warp 特化。它无法说明哪些 warp 承担哪个角色、哪个屏障控制哪次交接,或累加器位于 TMEM 中。编译器做出这些选择,智能体只能希望它们与专家的选择一致。

6.2 低层:CUTLASS、CuTe DSL 和原生 CUDA

另一端是 NVIDIA 的 CUTLASS 库及其 Python 前端 CuTe DSL。这些暴露了硬件提供的一切。代价是布局代数:一个用于描述张量逻辑坐标如何映射到内存和线程的数学框架。再往下是带有内联 PTX(NVIDIA 类似汇编的虚拟指令集)的 CUDA C++,你需要手动管理地址、屏障相位位和描述符编码。

根据论文,像 CuTe DSL 这样的低级 DSL“暴露了这种控制,但要求一种布局演算,使得智能体的错误既容易发生又难以定位。”许多布局错误都能通过编译。CuTe 会拒绝不匹配的秩和大小,但一个形状正确而步幅或 swizzle 错误的布局会把正确的数字放在错误的位置,这会在之后表现为输出错误,或表现为内核运行缓慢。

6.3 中间地带:Gluon

Gluon 介于两个阵营之间。它复用 Triton 的编译器栈,但暴露了对布局、内存移动和异步性的更低层控制。它把旋钮转向控制,也随之转向控制的负担。

6.4 CAKE 瞄准的缺口

把各种选项并排放在一起,缺口就出现了。那些让你表达专家级调度的语言,会让你自己管理布局和原始细节,而许多最严重的错误(等待了错误的阶段、缺少 fence、形状正确但映射错误的布局)只在运行时才暴露出来。那些保护你的语言则把大部分调度留在编译器手中,最多只提供诸如流水线深度之类的提示。

A simplified table, following the CAKE paper's classification, asks three questions of four kinds of GPU kernel language, ordered from high level to low level, and of CAKE IR, set apart below them in blue: do you write the warp roles, barriers and memory tiers yourself; is a layout algebra needed; and is the program checked before it runs? Tile DSLs (Triton, Helion, TileLang, cuTile): no, the compiler decides; no; partly, type checks. Gluon, which reuses Triton's compiler stack: yes; partly, explicit layouts; partly, compiler checks. CuTe DSL / CUTLASS (a Python DSL and C++ templates): yes; yes, CuTe layouts; partly, shapes are checked early but synchronization only at run time. CUDA C++ with inline PTX, the GPU's assembly: yes; no, raw addresses; no, run it and see. CAKE IR, a typed schedule language: yes; no, concrete offsets; yes, a verifier and a cost model reason about it before it is compiled. High-level languages are correct but keep the expert's schedule out of reach; low-level ones make everything reachable, and every mistake too; Gluon sits between the two; CAKE IR gives an explicit schedule with no layout algebra, checked before it runs. Columns 1 and 2 follow the CAKE paper's classification; column 3 is our simplification, and each tool has more nuance than the table shows.
谁来决定调度?各 GPU 语言家族在控制、布局代数和检查方面的立场(第 1–2 列遵循 CAKE 论文的分类;第 3 列是我们的简化)。打开全尺寸 SVG。

7. CAKE IR,逐行解读

CAKE IR 是智能体编写的语言。理解它的最好方式就是读一些。

7.1 调度声明什么,lowering 推导什么

该设计建立在分工之上。用论文的话说,“调度声明什么应当发生,lowering 推导如何发生。”调度记录哪些 warp 承担哪些角色、哪些缓冲区以多深的方式暂存、哪个屏障门控哪次交接,以及哪种指令形式消费哪个操作数。Lowering,即编译器向下翻译到 CUDA 和 PTX 的过程,计算机械性的后果:“屏障地址、阶段位、TMEM 偏移、描述符编码和 warp 身份都从声明中计算得出,而不是由智能体写出来。”

以下是论文用来展示这一思想的片段。论文简单地称之为“一个 CAKE IR 调度片段”;它的名称 fmha_fwd(融合多头注意力,前向传播)以及对张量内存的使用表明,它是以 Blackwell 注意力内核的开头为蓝本建模的(我们的解读):

@cake.schedule()
def fmha_fwd(lm, Q: LM.tma3d, O: LM.tma2d, seqlen_q: LM.i32):
    # declare resources: named resources, not raw addresses
    pool     = lm.smem(98304)
    smem_q   = pool.view(offset=0, shape=(128,128), dtype=lm.bf16, stage=3)
    tmem_acc = lm.tmem(cols=0, width=128, shape=(128,128), dtype=lm.f32)

    # declare roles: warp groups with assigned work
    load = lm.role(warps=[0])
    mma  = lm.role(warps=[1])
    pipe = lm.pipeline(stages=3)
    q_full = lm.barrier(count=3, prod=[load], cons=[mma],
                        init_count=1, pipeline=pipe)

    with load:
        for stage in lm.range(0, 3):
            smem_q.tma_load(Q, coords=(0,0,stage), stage=stage, barrier=q_full)

    with mma:
        for stage in lm.range(0, 3):
            lm.wait(q_full, stage=stage)
            lm.fence_proxy()
            lm.mma(tmem_acc, smem_q[stage], smem_q[stage], init=(stage == 0))

分四步来读:

  • 签名。@cake.schedule() 将该函数标记为 CAKE 调度。lm 是构建器,每个资源、角色和操作都通过它来声明,就像你使用 np. 或 torch. 那样。从它们的名称来看,Q: LM.tma3d 是 TMA 复制引擎读取的 3-D 张量,O: LM.tma2d 是 2-D 输出,seqlen_q: LM.i32 是 32 位整数;O 和 seqlen_q 在此片段中未使用。
  • 资源。lm.smem(98304) 预留了一个 96 KB 的共享内存池。pool.view(...) 划出一个命名的 128 × 128 BF16 值缓冲区,具有三个阶段:三个阶段 × 128 × 128 个值 × 2 字节 = 96 KB,因此这些阶段恰好填满该池。lm.tmem(...) 在张量内存中预留一个 128 × 128 FP32 累加器。这里没有任何东西是地址;每个缓冲区都有名称、形状、类型以及编译器可见的生命周期。
  • 角色与同步。lm.role(warps=[0]) 和 lm.role(warps=[1]) 命名了两个 warp 角色,一个加载器和一个张量核心发射器。lm.pipeline(stages=3) 声明了流水线,lm.barrier(...) 声明了它们之间的交接:一个屏障,其生产者是 load,消费者是 mma。我们对参数的理解是:count=3 为每个阶段提供一个屏障,init_count=1 是到达计数,因此一次到达(拷贝落地)就完成了每个阶段。
  • 每个角色的工作。with load: 不会打开任何东西,不像 Python 中的 with open(...):它的意思是“里面的代码只在 load 角色的 warp 上运行”。lm.range 是调度器自己的循环,不是 Python 的 range。在 with load: 内部,加载器为每个阶段发出一次 TMA 拷贝,并在每次拷贝落地时让硬件到达 q_full。在 with mma: 内部,张量核心 warp 在每个阶段等待 q_full,调用 lm.fence_proxy(),并将矩阵乘法发射到 TMEM 累加器中,在第一个阶段初始化它。从名字来看,fence_proxy 是一个代理栅栏(§3.2);论文没有说明它降低为什么,而且由于 TMA 写入和张量核心读取都通过异步代理,仅凭片段无法说明为什么这里需要它。

注意缺失了什么。没有屏障地址,没有相位位,没有描述符编码,也没有计算哪个线程在哪个 warp 中;编译器从声明中推导出所有这些。缓冲区确实带有一些简单的承诺,例如视图在池中的偏移量和累加器在张量内存中的起始列。§7.3 解释了为什么 CAKE 直接要求这些,编译器从中推导出其余寻址。该片段是一个教学示例,不是完整的内核:它将 Q 瓦片与自身相乘,而真正的注意力内核还会加载 K 和 V 并运行 softmax。

你可能还会注意到一个缺失的屏障。§3.2 说每个阶段需要一个 full 和一个 empty 屏障,但这个片段只声明了 q_full。它恰好填充其三个阶段各一次,并且从不环绕环形缓冲区,因此没有槽位会被覆盖。一个真正的流式主循环会通过环形缓冲区传输许多瓦片,它还会声明一个空屏障,由 mma 到达并由 load 等待:正是这种等待的缺失导致了 §3 图中的竞态。

The CAKE IR schedule fragment from the paper's Figure 3, split into three bracketed blocks. Resources: a 96 KB shared-memory pool, a three-stage 128×128 BF16 buffer smem_q, and a 128×128 FP32 accumulator in tensor memory. Roles and sync: load on warp 0, mma on warp 1, and a three-stage barrier q_full handing tiles from load to mma. Work per role: load issues three TMA copies; mma waits on q_full, issues a proxy fence, and multiplies. A side panel lists what lowering derives: barrier addresses, phase bits, TMEM offsets, descriptor encodings and warp identity, emitted as CUDA/PTX.
论文的示例调度分为三个块:资源、角色及其之间的屏障,以及每个角色的工作;降低过程推导出机械细节(简化自论文的图 3,省略了类型注释和注释)。打开全尺寸 SVG。

7.2 使其可检查的四个属性

论文列出了四个“完成工作”的属性:

  • 类型检查的词汇表。计算、内存移动、同步、数学和 warp 控制使用一组固定的 IR 操作,而不是嵌入的 C 或 PTX 字符串。类型错误的程序在构建时就会被拒绝。
  • 声明的资源。内存区域、同步对象和流水线只声明一次,因此“IR 知道每个缓冲区的形状、数据类型和生命周期。”
  • 显式角色。Warp 组被命名,“每个跨角色交接都是可见的,而不是隐式约定。”
  • 自动推导的元数据。声明的机械后果是“降低的,而不是编写的。”

好处是分析可以在代码生成之前对调度决策进行推理,并且测试框架可以将发现与受影响的资源、角色或阶段联系起来,“而不是只返回后端错误或挂起。”

7.3 有意不提供布局代数

CAKE“刻意不把布局作为一等抽象”。智能体不是把布局当作数学对象来操作,而是写下具体的承诺:“一个 SMEM 视图偏移、一个操作数字节偏移、一个 TMEM 列范围、一个 swizzle 标签、一个 TMA 描述符坐标。”然后编译器“承担起判断这些承诺是否合法的重担”:它检查这些承诺在程序的数据流上是否保持一致(生产者写入的就是消费者期望读取的),以及它们是否满足目标指令的规则。诊断信息会指出有问题的 IR 决策以及不匹配的大类。

§7.1 片段已经包含两个这样的承诺:pool.view(offset=0, ...) 表示环形缓冲区从池的字节 0 开始,lm.tmem(cols=0, width=128, ...) 表示累加器占用张量内存的第 0 到 127 列。swizzle 是对元素在共享内存中位置的固定重排,使得一起读取的线程不会在同一内存 bank 上冲突。例如,如果 TMA 拷贝以某种 swizzle 写入一个 tile,而 MMA 读取时假设的是另一种,编译器会在任何东西运行之前就指出不一致的声明。(这个例子是我们举的;论文只按类别描述这些检查。)

这与 CuTe 的约定正好相反。CuTe 要求作者是布局专家,并以表达力作为回报。CAKE 只要求作者提供简单事实,并让编译器负责发现矛盾。论文只通过其契约和覆盖范围来描述其布局验证器,而不涉及其内部实现。

7.4 一种语言,多种 GPU

同一套调度语言面向从 Ampere 到 Blackwell 的 NVIDIA GPU。角色–屏障–流水线调度的结构保持不变;允许哪些指令以及如何降低,仍然因目标而异。(sm_ 名称是 NVIDIA 对 GPU 世代的代号:sm_80 是 Ampere,sm_90a 是 Hopper,sm_100a 是 Blackwell;末尾的 a 表示该芯片特有的功能。)你不需要了解下面的细节;右列只是展示在相同调度结构下指令集的变化。

CAKE 面向哪些 GPU,以及它们之间有什么变化

目标 示例 GPU 张量核心路径与显著特性
sm_80 A100 mma.sync、ldmatrix、cp.async;无 TMA、集群或 TMEM
sm_89 L40S、RTX 6000 Ada 与 sm_80 相同,外加 FP8 张量核心;无 TMA 或 TMEM
sm_90a H100、H200 WGMMA + mma.sync;TMA、集群、异步屏障、分布式共享内存
sm_100a B200 tcgen05.mma + TMEM;2-CTA MMA;tcgen05.{ld,cp,shift}
sm_103a B300 增加 tcgen05.ld.red 和 K = 96 块缩放 MMA
sm_120a RTX 5090、RTX PRO 6000 mma.sync + ldmatrix、TMA、集群、分布式共享内存;无 tcgen05 或 TMEM
sm_121a DGX Spark (GB10) 与 sm_120a 相同的指令,单独的二进制文件,不同的特殊功能单元速率

有两条策略很突出。编译器“要求精确的目标匹配”:它会报告缺少设备或工具链,而不是悄悄地把调度降低到较旧的架构。性能估算只出现在经过校准的地方:B200 是测量基线,H100 单独校准,其他目标则报告覆盖范围限制,而不是借用数字。在静态检查之后,CAKE IR 确定性地降低为可读的 CUDA/PTX,然后通过 NVIDIA 的标准工具链。生成的源代码仍然可作为逃生通道使用;默认情况下,决策保留在 IR 中,以便可以对其进行分析。

8. CAKE IR 从何而来

你可能会以为这样的语言是在白板上设计出来然后再实现的。CAKE IR 则是逐步生长出来的。

8.1 从专家内核中生长出来

“CAKE 并非始于预定义的 CAKE IR 词汇表。”其起始材料是一个生产级内核语料库和一组硬件设计原则。附录描述了五个步骤:

  1. 语料收集。来自生产库的高质量 CUDA 内核:Alpha-MoE、CUTLASS、cuTile、DeepGEMM、FlashAttention-4、Flash-KMeans、FlashInfer、SonicMoE 和 TileLang。用其他语言编写的内核首先由编码代理通过内联 PTX 翻译为 CUDA。
  2. 抽象提取。代理分析语料库,并将反复出现的模式(屏障编排、流水线分级、warp 角色划分、TMA 描述符设置、TMEM 累加器生命周期)总结为候选抽象。
  3. 硬件驱动的设计。人类专业知识引导这些抽象走向 Blackwell 编程模型:将 TMEM 作为一等资源,将 warp 特化作为主要的并行形式,将异步屏障作为同步原语,以及用于协调多个 SM 的集群作用域操作。
  4. 原则驱动的迭代。每个候选抽象都对照八项设计原则(见下文)进行检查;违反者被改进或拒绝。
  5. 移植驱动的扩展。新内核被持续移植。每次移植要么成功,从而验证抽象,要么暴露缺口,从而触发扩展 IR 或其 lowering 的提案。

这个循环永不闭合:“每个新内核家族都会对 IR 进行压力测试,并推动进一步演进。”整个过程中的设计目标都很严苛:IR 必须“复现专家编写内核的物理调度和性能”。根据论文,该框架同样“主要由代理在人工合并门禁下维护”。

8.2 八项原则

每个候选抽象都对照这八项原则进行检查,这些原则摘自论文附录:

原则 它要求什么
P1 符合人体工程学 让 NumPy 和 PyTorch 用户感到熟悉;避免不必要的簿记。
P2 性能透明 让与性能相关的硬件决策保持可见,并让 lowering 可检查。
P3 规范化 每个操作只有一种规范形式,而不是多种等价写法。
P4 静态类型检查 类型规则约束 lowering,并在构造期间拒绝类型错误的程序。
P5 分析友好 暴露静态分析所需的信息。
P6 测试门禁 针对内核矩阵测试评估每一项 IR 变更,以进行分析和编译。
P7 分析一致 每当 IR 的数据模型发生变化时,就更改分析。
P8 硬件接地 记录每个操作的预期硬件行为。

P5 和 P7 使新特性保持可分析,P6 防止回归,P2 和 P8 使到硬件的映射对“人类和代理同样”清晰可读。论文将此与数学家陶哲轩关于 AI 生成证明的观察联系起来:当生成产物变得廉价时,瓶颈就转移到验证和理解它们上。CAKE 的版本是:当代理能够批量生产内核时,重要的是每个内核都能被检查和阅读。

第二部分总结

如今的 kernel agent 生活在一个黑盒中,它只给出通过/失败和一个时间。CAKE 为它们提供了一种语言,专家的决策(角色、缓冲区、屏障、流水线)在其中是显式且有类型的,而容易出错的簿记工作(屏障地址、相位位、描述符编码、warp 身份)则由编译器推导。它刻意避免布局代数,而是让编译器检查具体的布局承诺。这门语言是从生产 kernel 中生长出来的,遵循八项原则,使其在成长过程中保持可分析性。


第三部分 · 循环——一场对话

学生 好,所以 agent 编写这种调度语言。然后呢?它最终还是要跑在 GPU 上。

教授 最终是的。但为了找一个缺失的屏障而跑一次 GPU,代价太高了。什么更便宜?

学生 先读代码?像 linter 那样?

教授 正是:在编译之前运行检查,并指出哪一个屏障、哪一个阶段。然后一个成本模型——对每个幸存者运行速度的廉价估计——挑出值得计时的那些。只有到那时才上 GPU。

学生 那当检查本身漏掉了一整类 bug 时呢?

教授 那你就发现了这篇论文最有趣的想法。这个 bug 变成一项新检查。编译器学会了。

学生 这一切真的能带来更快的 kernel 吗?还是只是更友好的错误信息?

教授 好。保持这种怀疑。我们会看证据,也会看它没有展示的东西。

9. 测试框架:GPU 之前的检查

测试框架是 agent 周围的一切:分析、测试、基准,以及什么算作结果的规则。本节描述 CAKE 的测试框架检查什么,以及它反馈什么。

9.1 七类检查

论文按其对 agent 的作用来描述其分析套件,而不是按内部编译器 pass。共有七个类别,分三种处置方式:拒绝候选的门禁、描述候选的报告,以及建议改进的提示。

类别 处置 目的 示例(我们的,仅作说明)
程序安全 编译前门禁 发现同步、顺序和内存使用方面的隐患 消费者等待一个生产者从未到达的屏障(§3.3)
硬件符合性 编译前门禁 强制执行受支持的资源、指令和架构契约 需要的共享内存超过 SM 容量的缓冲区,或在 Hopper 目标上使用仅限 Blackwell 的指令
数据一致性 编译前门禁 检查数据流,以及生产者和消费者对表示形式达成一致 TMA 以一种 swizzle 写入 tile,而 MMA 以另一种假设读取它
调度语义 编译前门禁 检查所声明调度的结构不变量 一个三阶段流水线,其循环却触及第四个阶段
数值验证 执行门禁 将编译输出与权威外部参考进行比较 输出与参考的差异超过容差
性能分析 报告 估计成本并识别大致的瓶颈类别 “可能受内存带宽限制”
优化指导 提示 在不阻塞编译的情况下建议有前景的修改 “更深的流水线可能隐藏拷贝延迟”

四个预编译门会拒绝“许多在数学上看似合理、但与目标执行模型不兼容的候选方案”。一项发现“会指出受影响的程序区域以及所违反的契约类别”。将其与 §3.3 中的挂起情形对比:不是 GPU 冻结,而是智能体听到类似消费者在等待此屏障的第 2 阶段,但没有任何生产者会到达那里这样的信息。这个例子是我们的;论文并未公布其消息格式。

9.2 正确性与性能

数值正确性会针对参考实现,在不同形状和输入分布下进行检查,最终验收要求在目标框架中进行端到端评估(例如,由 SGLang 这一流行的开源 LLM 服务系统所服务的完整模型)。

性能分为两层。经过校准的成本模型会估计候选方案的运行速度,并指出其可能的瓶颈,因此测试框架可以在不运行所有候选方案的情况下对其进行排序和筛选。但该模型只是一个过滤器:“设备上的测量和分析仍然是最终的基准真相。”在评估中,内核计时使用 CUPTI(NVIDIA 的分析接口)在 GPU 上测量,并在每个计时样本前刷新 L2 缓存,以免一次运行预热缓存而影响下一次。(§13.2 中有少数数字采用不同的测量方式:API 级别的 Alpha-MoE 加速比和 SGLang 服务吞吐量。)

论文对其静态分析所承诺的内容十分谨慎。它们只是一个“仅在其建模领域内”有效的门;“不能证明全局 GPU 正确性,也无法捕捉所有微架构行为”;“假阳性和假阴性都可能发生”。缺失的信息会“被明确报告,而不是被视为检查成功”。

10. 一轮进化

有了语言和测试框架,一次内核进化运行就具有简单的形态。本节跟踪其中一轮。

10.1 四个阶段

论文将一次运行描述为四个阶段:

  1. 生成结构上不同的 CAKE IR 候选方案:不同的角色划分、流水线深度、内存放置,而不仅仅是不同的常量。
  2. 过滤在花费任何 GPU 时间之前,通过 IR 构建检查、验证器的硬性门以及成本模型排序对它们进行筛选。
  3. 评估让幸存者接受外部预言机(一个可信的参考实现,其输出定义了正确答案)的检验,并附上基准测试和分析器证据。
  4. 将证据路由。根据诊断结果,一项发现会被送回候选方案(修复它)、验证器(添加规则)、成本模型(重新校准它)或 IR 词汇表(添加原语)。
A funnel with illustrative counts. Twelve structurally distinct CAKE IR candidates are generated. Build and type checks reject ill-typed ones, leaving nine; the verifier's hard gates (program safety, hardware conformance, data consistency, schedule semantics) reject schedules with a reason such as a missing wait or a tensor-memory overflow, leaving five; the calibrated cost model ranks them and keeps three, all without spending GPU time. The survivors run on the GPU against the correctness oracle with CUPTI timing and profiling, which picks one best kernel. Findings from every stage are then routed to whoever should learn from them: the candidate (repair it), the verifier (add a rule), the cost model (recalibrate it) or the IR (add a primitive).
一轮进化:廉价的静态检查以给出理由的方式拒绝非法调度,成本模型对幸存者排序,只有最好的少数几个在 GPU 上运行,而每个阶段的发现都会被路由到系统中应当从中学习的部分(数量仅为示意)。打开全尺寸 SVG。

第四阶段正是 CAKE 与带有更好错误消息的搜索循环的不同之处。一项发现不只是被使用后丢弃;它会被发送到系统中应当从中学习的相应部分。

10.2 保持不变的部分

每一次运行都锚定于两件事。第一是工作负载契约,即“稳定权威”:它固定了形状、正确性判定标准及其容差(输出必须与参考值有多接近,如torch.allclose所示)、硬件,以及智能体可以查看哪些参考。结果会被保留,因此决策可被审计,反复出现的发现可被复用。

第二是模型。论文中所有智能体任务都使用同一个模型 GPT-5.6-sol,推理强度为xhigh。“保持模型和智能体脚手架固定,使得比较……可归因于环境而非模型能力。”(脚手架是围绕 LLM 的程序,为它提供编辑文件和运行代码的工具,并驱动循环。)无论实验的两组之间出现什么差异,都不是因为其中一组拥有更聪明的模型。

11. 编译器也在进化

到目前为止,智能体是在一个恰好设计良好的固定环境内进化内核。CAKE 的最后一项承诺是:环境并非固定不变。

11.1 通往编译器的两条路径

内核候选、验证结果、基准测试和失败报告都是改变编译器的证据。论文描述了两条路径:

  • 路径 1:挖掘专家内核。智能体检查生产内核和硬件文档,找出 IR 缺失的 Blackwell 模式(“新的指令形式、资源类型、描述符变体、同步惯用法”),并撰写编译器变更提案。每个提案在实施前都要对照 §8.2 的设计原则进行检查,包括性能透明度和对验证的友好性。
  • 路径 2:提炼失败。智能体利用失败候选的反馈(“sanitizer 报告、失败案例、正确性不匹配、调试日志”),将反复出现或代价高昂的失败模式转化为新的分析:“不透明的运行时崩溃变成验证器规则,重复出现的非法 lowering 模式变成静态检查,系统性的错误预测变成校准目标。”
Flow diagram of how CAKE's compiler evolves. Path 1 mines production kernels and hardware documentation for patterns the IR is missing: new instruction forms, resource types, descriptor variants, synchronization idioms. Path 2 distills failed candidates: an opaque crash becomes a verifier rule, an illegal lowering a static check, a misprediction a calibration target. Both feed a compiler change proposal, which is checked against the IR design principles, implemented together with its analyses, and must pass corpus tests (400+ static and compile cases, 399 GPU correctness cases) at a human merge gate before the updated harness meets the next kernel family.
外层循环:智能体挖掘专家内核中 IR 尚无法表达的模式,并将反复出现的失败提炼为检查;每一次变更都必须在人工合并门禁后通过整个内核语料库。打开完整尺寸的 SVG。

11.2 为什么原语和分析必须一起改变

这两条路径相互促进。新的原语告诉编译器更多关于硬件的信息,从而使更强的分析成为可能;新的分析约束了未来原语可能的样子。论文坚持认为,原语及其分析“必须一起进化”:“没有效果和合法性规则的语法会使 IR 更难分析,而没有语料库验证的新验证器规则可能会拒绝有效内核。”通俗地说:向语言添加一个新操作却不告诉检查器它读取、写入和需要什么,检查器就会对使用它的每个程序视而不见;添加一个新检查却不在已知良好的内核上运行它,它可能开始拒绝正确的程序。

这就是为什么每一次编译器变更都要在整个内核语料库上进行测试门禁。语料库很大:超过 400 个静态和编译用例,以及 399 个 GPU 正确性用例,涵盖约 28 个内核家族(§13.4)。拒绝已知良好内核的新规则不会被合并。

谁来做这些改动?Agent 根据人类编写的高层描述来实现和维护分析,而人类负责批准合并:论文称编译器演化“在合并关卡处仍由人类引导”。

12. 受控实验:从零开始的 Flash-KMeans

对系统的描述是廉价的;受控比较则不然。本节介绍论文中最干净的实验:同一个 agent、同一个模型、同一个任务,分别用 CAKE IR 或直接用 CUDA/PTX 编写。

12.1 工作负载:k-means 的一步

k-means 是机器学习中最古老的算法之一。你有 N 个点,想要分成 K 组。先对组中心(质心)做 K 个猜测,然后重复两个步骤:分配每个点到最近的质心,更新每个质心为分配给它的点的平均值。每次重复称为一次 Lloyd 迭代。

Flash-KMeans 是一种快速、精确的 GPU 实现,它的重要性源于一个你可能意想不到的原因:视频生成。Sparse VideoGen2 用 k-means 对相似 token 进行分组,以便稀疏注意力能在相关 token 的连续块上工作(“语义感知置换”)。在那里,k-means 在模型的前向传播中运行,因此其耗时直接加到生成时间上。(稀疏注意力只对最重要的 token 对计算注意力。)

According to the paper, a Lloyd iteration is dominated by two BF16 kernels that account for more than 95% of the end-to-end time:

  • assign computes the squared distance from every point to all K centroids and returns the index of the nearest one. It is compute-bound: a matrix multiplication followed by a reduction.
  • centroid_update sums the points of each cluster and counts them. It is limited by memory bandwidth and by contention on atomic updates (many threads adding into the same cluster’s running sum at once, which the hardware must apply one at a time).

The experiment focuses on assign, which “exercises the tensor-core pipeline and scalar epilogue” (here the epilogue picks the nearest centroid), at one fixed shape: B = 32 batches, N = 65,536 points, K = 1,024 centroids, D = 128 dimensions, BF16 inputs and FP32 accumulation. The bar is high: the baseline is a tuned Triton implementation from the FlashML project (whose library, FlashLib, collects fast GPU versions of classical machine-learning operators), measured at 0.938 ms.

12.2 The rules: what the agent may and may not see

The paper is careful about a subtle form of cheating. If the agent could read an expert’s CUDA for this kernel, it could copy the schedule instead of discovering it. So in these clean-start runs the agent may see “the mathematical specification, evaluation contract, correctness oracle, and high-level code, but not low-level target implementations such as CUDA, PTX, SASS, or equivalent generated source.” (SASS is the GPU’s actual machine code, one level below PTX.) An existing implementation may be run as a black-box timing baseline, but its internals stay hidden. The restriction was enforced in isolated environments and audited afterward.

The two arms differ in what the agent writes, and therefore in the compiler feedback it gets back:

  • Treatment: the agent writes typed CAKE IR, with the harness of Part III.
  • Control: the agent writes CUDA C++ and inline PTX directly, under the same reference rules.

Everything else is fixed: the coding agent and its scaffold, the model and reasoning effort, the task statement, the correctness oracle, the benchmark harness, and the target shape. Each arm runs three times, with a budget of 80 million tokens.

12.3 The result

Table 2 of the paper gives the outcome at the full 80M-token budget, three runs per arm:

Representation Reached a plateau by 80M tokens Active evolve time (hours) Best speed at 80M tokens (× baseline)
CAKE IR 3 of 3 runs 1.89 [1.02, 2.33] 1.144 [1.041, 1.205]
Direct CUDA/PTX 0 of 3 runs 3.73 [3.59, 4.34] 0.928 [0.852, 1.151]

Entries are medians over three runs, with [minimum, maximum] in brackets. With CAKE IR, the median best kernel runs 14.4% faster than the tuned baseline (1.144×); with direct CUDA/PTX, it runs at 92.8% of the baseline’s speed. CAKE IR runs met the paper’s “prespecified plateau criterion” in all three runs; the CUDA/PTX runs met it in none. A run reaches a plateau when its best speed has stopped improving, by a rule the paper fixed in advance but does not publish. So all three CAKE runs had leveled off by 80M tokens, while all three CUDA/PTX runs were still improving when the budget ran out, and we cannot say where they would have ended. The CAKE runs also used about half the active time; since the paper does not break that time down (its stopping and timing accounting is “retained in the artifact”), we cannot tell whether they finished sooner because they stopped at a plateau or because each step was quicker.

A chart of the best validated speed so far, relative to the tuned FlashML KMeans baseline (0.938 ms), against cumulative tokens from 10 to 80 million, for the assign kernel at one shape on B200, three runs per arm. The curves are three-run means read approximately from the paper's Figure 5. CAKE IR rises from about 0.34× at 10M tokens to about 0.8× at 20M and crosses the baseline by about 55M tokens; direct CUDA/PTX rises from about 0.14× to about 0.97× at 80M and stays below the baseline. At 80M tokens the median best speed is 1.144× [1.041, 1.205] for CAKE IR, with 3 of 3 runs at a plateau and 1.89 hours of active evolution, versus 0.928× [0.852, 1.151] for direct CUDA/PTX, with 0 of 3 runs at a plateau and 3.73 hours.
The controlled experiment: with the same model, task and oracle, the CAKE IR mean passes the tuned baseline by about 55M tokens while the CUDA/PTX mean stays below it at 80M (its best single run reaches 1.151×). Curves are approximate three-run means read from the paper's Figure 5; the card shows the paper's Table 2. Open the full-size SVG.

The trajectory tells the same story over time. At each 5-million-token checkpoint, every run contributes its best validated speed so far. The mean of the CAKE IR runs “crosses the tuned FlashML baseline by 55 million tokens and continues to improve, while the direct CUDA/PTX mean remains below baseline at the 80-million-token cutoff.” Reading the paper’s plot (approximate values), the CAKE IR mean is already near 0.8× by 20M tokens, when the CUDA/PTX mean is still around 0.2×.

12.4 How to read it

Our interpretation, with the caveats the numbers themselves carry:

  • Within the same budget, the environment changed the outcome. With CAKE IR and its harness, all three runs beat a tuned baseline; the CUDA runs, on median, did not, and they were still climbing. Which part of that environment did the work (the language, the checks, the cost model, or compiler evolution) is still open (§16.2).
  • But direct CUDA is not hopeless. The best CUDA/PTX run reached 1.151×, above the CAKE median. The difference is consistency and pace: CUDA can get there, but less reliably and more slowly. With three runs per arm, the spread matters as much as the median.
  • It is one kernel, at one shape, with one model. The experiment isolates the effect of the environment well, but it does not tell us how large the effect is for other kernels or other models.

13. Beyond the Benchmark

The controlled experiment shows the effect cleanly; the rest of the evaluation shows the system at work on real kernels. It answers two more questions: can CAKE discover a kernel’s schedule when it may not look at any existing low-level implementation (fast ones may exist, but only as black-box timing baselines), and can it reproduce, or beat, kernels that experts already tuned?

13.1 Frontier kernels: new architectures, no reference to copy

The paper calls a kernel frontier “in the operational sense that the agent must discover its physical schedule without inspecting a low-level target implementation.” This is where a co-evolving IR should help most, and also where it is most exposed: a missing capability shows up as a schedule “the agent cannot express at all.”

Kimi Delta Attention. The official FlashKDA kernels served only as a black-box timing baseline; their source and generated code were not given to the agent. The agent-generated prefill kernel, compatible with FlashKDA and covering fixed-length, packed variable-length, and tail inputs (batches whose sequences share one length, batches of different-length sequences packed end to end, and the leftover tokens at the end of a sequence), reaches a 2.05× geometric-mean speedup over that baseline across six B200 BF16 shapes. The paper calls it “bitwise correct on its validation contract” (the pull request itself reports outputs that match FlashKDA within a tolerance of 10−2), and it was verified in end-to-end serving of Kimi-K3 under SGLang. Separate decode kernels reach a 1.14× geometric-mean speedup over upstream FlashInfer (the library’s current official kernels) across 30 public-API shapes. Both went into FlashInfer, an open-source kernel library used inside LLM serving engines such as SGLang and vLLM, as generated CUDA (pull requests #4262 and #4279, merged in August 2026), “so downstream users take on no dependency on CAKE.”

Gated DeltaNet and sparse attention. Against FlashInfer, CAKE’s Gated DeltaNet prefill and speculative-decoding paths “improve performance while preserving the model’s recurrent state,” and MiniMax sparse attention shows that the representation supports sparse-attention families across prefill and decode. The paper gives no speedup numbers for these two. They are “dispatch families rather than single kernels”: several physical schedules, each its own CAKE IR program, behind one entry point (§14).

13.2 Production kernels, evolved from a reference

TinyGEMM. Here the agent starts from an existing kernel: FlashInfer’s BF16 matrix multiply for tiny batches (the paper calls it small-M), derived from TensorRT-LLM, NVIDIA’s own LLM-inference library. Tiny batches are typical when a server decodes a handful of requests at once. The agent produced an adaptive family of shallow and deep pipelines, including variants that use programmatic dependent launch (PDL, which lets the next kernel start its preamble before the previous one finishes) and paths for batch sizes below eight. FlashInfer PR #4274 reports an 18–23% geometric-mean reduction in kernel time across 35 canonical shapes and a broader regression suite. Greedy decoding (always taking the single most likely next token) produced bit-for-bit identical outputs for GPT-OSS-20B and GPT-OSS-120B on B200 and on GB300, a newer Blackwell system. An SGLang experiment with GPT-OSS-120B measured up to 7.6% more output tokens per second with 128 requests in flight on one GPU, and no difference beyond noise when the model was split across four GPUs (four-way tensor parallelism).

Alpha-MoE. In a mixture-of-experts layer, a small router sends each token to a few of many “expert” feed-forward networks and weights their answers. The original Alpha-MoE does the whole layer in one GPU program, a megakernel, written for Hopper with 8-bit weights and activations (W8A8). CAKE agents rewrote it for Blackwell. The megakernel gathers the tokens routed to each expert, runs the expert’s two matrix multiplies with the activation function between them, converts the intermediate result back to 8 bits (requantization), and adds each expert’s output into its token’s result, weighted by the router’s score. Against FlashInfer’s TensorRT-LLM-derived pre-routed API, the end-to-end API-level speedups are 6.204× at N = 256 and 4.025× at N = 512; measured as GPU span, they are 1.215× and 1.170×. (The paper does not say what N counts here; it is not the N of the k-means experiment.) The kernel was merged into FlashInfer in September 2026 (PR #4287), but the pull request reports a different comparison: against stock SGLang’s five-kernel Triton MoE path on GB300, its fused kernel is 1.37×–1.75× faster in GPU time. The paper’s 6.204× and 4.025× do not appear there.

13.3 Reproducing expert kernels

The last question is whether, when expert structure already exists and the agent may look at it, the harness helps agents preserve correctness while matching highly optimized kernels. The references come from FlashAttention-4, TensorRT-LLM, DeepGEMM, CUTLASS, and FlashInfer. All runs are on B200 at a fixed shape, use median CUPTI GPU span, and pass kernel-specific correctness gates. You do not need the variant names: read “Relative performance” (above 1.0× means CAKE’s kernel is faster than the expert’s) and compare the line counts (“device” lines are the code that runs on the GPU).

Family Variant Relative performance CAKE IR lines Reference device lines
FlashAttention-4 forward, BF16, non-causal 1.0045× 430 2,369
FlashAttention-4 backward, BF16, non-causal 1.0470× 514 2,552
TensorRT-LLM GQA decode, FP16 1.043× 783 8,515
DeepGEMM 1D1D GEMM, FP8 1.0370× 221 516
DeepGEMM grouped GEMM, BF16, masked 1.0174× 401 442
DeepGEMM MQA indexer, FP8 1.2700× 480 704
DeepGEMM MQA indexer, FP4 1.2730× 392 704
DeepGEMM paged MQA indexer, FP4 1.0036× 395 779
CUTLASS MLA decode, BF16, TMA 1.2174× 845 1,860
DeepSeek-V4 sparse MLA decode, BF16 1.1297× 1,299 13,609
DeepSeek-V4 sparse MLA decode, FP8 0.9649× 1,393 8,942

Ten of the eleven comparisons meet or beat the reference; the last reaches 96.5%. The paper adds two honest qualifications. Entries below the reference “generally reflect compiler-integration maturity”: a feature the kernel wants is still being integrated into the compiler, so the submitted kernel uses the closest supported strategy. And the strongest wins, the two MQA indexers at about 1.27×, “are not faithful transcriptions of the reference kernels”: the agent explored optimizations the original lacked, so above-parity results “reflect search rather than transcription fidelity alone.” The line counts are descriptive only; the languages and counting scopes differ, and they show just that these schedules “can be represented compactly.”

Shapes and counting details

  • S1 (FA4): B = 4, H = 32, S = 8,192, D = 128. S2 (TRT-LLM GQA): B = 128, 64 query heads, 8 KV heads, KV length 4,096, D = 128, page size 16. S3 (FP8 GEMM): M = 4,096, N = 7,168, K = 4,096. S4 (grouped GEMM): 256 groups × M = 128, N = 4,096, K = 7,168. S5 (MQA indexers): H = 32, D = 128, query length 1,024, KV length 2,048. S6 (paged indexer): B = 256, H = 64, D = 128, average KV 4,096, block 64. S7 (CUTLASS MLA): DeepSeek-V3, B = 128, KV length 4,096, page 128. S8 (DeepSeek-V4 sparse MLA): B = 3, ragged query lengths [3, 4, 5], H = 128, Dqk = Dv = 512.
  • The paper’s table also lists “reference + support” line counts (for example 3,039 and 3,690 for the two FA4 kernels). For DeepSeek-V4, reference lines are restricted to source reachable under the fixed S8 route and runtime constants.
  • GQA (grouped-query attention) lets several query heads share one key/value head; “causal” attention lets a token look only at earlier tokens. MQA (multi-query attention) indexers score which past tokens a query should attend to, as in DeepSeek-V3.2’s lightning indexer; MLA is DeepSeek’s multi-head latent attention. “Paged” kernels read a KV cache stored in fixed-size blocks. Our GLM-5.3 and DeepSeek-V4 deep dives explain both.

13.4 The breadth of the corpus

The headline experiments hide how much the harness now sustains. The validated corpus contains more than 400 static and compile cases and 399 GPU correctness cases across roughly 28 families, including attention, dense and sparse GEMM, MoE, quantization, normalization, state-space models, KNN (k-nearest-neighbour search), and KMeans, with architecture-specific paths from Ampere through Blackwell. It includes more than 100 ports of TensorRT-LLM kernels for attention, MLA decode, and MoE.

One property stands out: composition. Because roles, barriers, and buffers are declared rather than implied, schedules that would normally be separate kernels can be written as one device program. BatchAttention combines decode and prefill work in one kernel; the Alpha-MoE and mega-MoE families fuse routing, expert computation, and output accumulation “without materializing intermediate results” (that is, without writing them out to memory). The Alpha-MoE rewrite also shows that “schedule structure can survive even when target instructions change,” from Hopper’s to Blackwell’s.

14. From a Tuned Shape to a Library

Everything so far optimizes a kernel for a particular shape. A library has to accept whatever shape the caller passes. This section covers the step that benchmark papers usually skip, and that CAKE treats as a stage of its own.

14.1 A different objective, so a separate stage

“Closing that gap is not a matter of running the inner loop on more shapes; it is a separate stage with a different objective, a different ranking signal, and a different failure mode.” An exact shape gives evolution a clean target and permits aggressive specialization; scoring the inner loop on broad coverage would blur that signal. So generalization begins only after strong per-shape seeds exist, and it is scored on performance including the dispatcher over a fixed workload. Incorrect or slow seeds go back to the inner loop “rather than being hidden behind routing.”

14.2 Building a dispatcher

A dispatcher is a small piece of code in front of several kernels that looks at the input shape and calls the right one. In Python it would be an if/elif/else chain (an illustrative sketch, not CAKE’s code):

def kmeans_assign(x, centroids):
    D = x.shape[-1]
    if D == 352:
        return split_k_d352(x, centroids)
    elif 144 <= D <= 176:
        return padded_tail(x, centroids)
    # ... more routes ...
    else:
        return fallback(x, centroids)

Each condition is a guard, each kernel behind one is a route, and the final else is the fallback. (“Split-K” splits the long reduction dimension across several thread blocks and adds up their partial results, which helps when the other dimensions are too small to fill the GPU; a “padded tail” route rounds an awkward size up to a convenient one.)

The generalization stage groups the measured seeds into shape buckets, produces specialized or shared variants, and orders their guards (conditions on the shape, such as “D = 352”) behind an explicit fallback that handles anything the guards do not match. Tuning may change implementation parameters, but never the input shape. Before any aggregate number is reported, validation covers representative and held-out inputs, boundary and tail cases, guards that overlap or leave gaps, and the fallback path itself.

How CAKE turns kernels tuned for single shapes into one library call. Evolution tunes a seed, a CAKE IR program, for one exact shape, but a library call such as kmeans(B, N, K, D) can arrive with any shape. The generalization stage groups seeds into shape buckets and orders guards on the shape in front of an explicit fallback; guards are tried top to bottom and the first match wins. Using routes from the paper's KMeans portfolio (the guard order shown is illustrative): D = 288 and D = 352 go to split-K routes (14 and 6 shapes), D from 144 to 176 to a padded-tail route (8 shapes), D = 112 to a specialized route (6 shapes), six more routes cover 14 shapes, a general route covers 64 and the fallback 12, for 12 routes and 124 shapes. Each route is a separate CAKE IR program, so it can be analyzed and benchmarked on its own. Before any aggregate is reported, validation covers representative and held-out inputs, boundary and tail cases, overlapping or missing guards, and the fallback path; the shape set is fixed before tuning, and guards may split it but not add evaluation shapes. On GB200 the GPU-span geometric-mean speedup over the reference is 1.418× for KNN build (112 shapes, 8 outer families, 80 routes), 2.116× for KNN search (198 shapes) and 1.803× for KMeans (124 shapes, 12 routes), with no incorrect outputs and KNN recall 1.0.
From one tuned shape to a library call: ordered guards send each call to one of several separately checkable CAKE IR routes or an explicit fallback. The KMeans routes and shape counts are from the paper; the guard conditions are inferred from the route names. Open the full-size SVG.

One policy keeps routing from turning into a zoo: the stage “reuses a single physical schedule across as much of the shape domain as it can, and introduces another only when the domain requires a material schedule change.” Each route is its own CAKE IR program, so every alternative can still be analyzed and benchmarked on its own.

14.3 Not grading your own homework

The subtle danger in this stage is evaluation leakage: if the same shapes are used to tune the dispatcher and to report its speed, the dispatcher can simply memorize them. CAKE declares the valid shape domain before tuning. Dispatcher predicates “may partition that domain, but they may not introduce convenient new evaluation rows,” and coverage expands only “through deterministic unseen shards of the same source”: new test shapes are drawn by a fixed rule from the same source and kept unseen during tuning.

14.4 Results on real libraries

The classic machine-learning workloads give the cleanest reading, because their portfolios are too large for one shape to carry the result. KNN (k-nearest neighbours) finds, for each query point, the closest points in a data set: “KNN build” usually means computing it among all points of one set (building a neighbour graph), “KNN search” for separate query points. On GB200, the generalized KNN build, KNN search, and KMeans kernels that CAKE contributed to FlashLib (merged in August 2026 as pull requests #15, #16 and #18; the paper itself cites no PR numbers) reach these results in the paper’s measurements, where Gspan is the unweighted geometric mean, over shapes, of the reference’s median CUPTI GPU span divided by CAKE’s:

Library function Shapes Gspan
KNN build 112 1.418×
KNN search 198 2.116×
KMeans 124 1.803×

There were no incorrect outputs, and KNN recall was 1.0 (every true nearest neighbour was found). The paper warns against comparing these numbers with the clean-start experiment of §12: the measurements “use different hosts, shape distributions, baselines, and protocols,” so the difference “is not, by itself, a measured cost of generalization.”

Route-level breakdown (the paper's appendix figure)

KNN build (baseline FlashLib 0.2.0): 8 outer families covering 112 shapes, with 80 distinct final routes observed.

Outer family Shapes / routes Gspan
D64 7 / 5 1.130×
D128 low-K BF16 38 / 17 1.258×
D128 low-K FP16 1 / 1 1.353×
D128 mid-K 32 / 30 1.425×
D128 large-K 10 / 5 1.909×
D192 1 / 1 1.432×
D256 8 / 8 1.386×
high-D 15 / 13 1.758×

Flash-KMeans (baseline: the authors’ tuned FlashLib implementation): 12 final routes covering 124 shapes.

Final route Shapes Gspan
gap-fused general 64 1.753×
D288 split-K 14 2.227×
aligned-shape fallback 12 1.439×
D144–176 padded tail 8 1.786×
D112 specialized 6 1.414×
D352 split-K 6 3.757×
D480 split-K 6 2.597×
micro-D hybrid 2 1.309×
micro-D pipelined 2 1.095×
large paired 2 1.217×
D128 specialized 1 1.036×
D224 TMEM 1 1.037×

Values are per-row geometric means from same-session CUPTI measurements.

CAKE touches four research areas: GPU languages (§6), compilers, kernel agents, and self-improving systems. The cleanest way to place it is to ask of each line of work: what gets evolved, and what stays fixed?

Line of work Examples What changes What stays fixed
Kernel agents with search or memory KernelEvolve, EvoEngineer, AVO, K-Search, KernelBlaster the search process, the accumulated memory the language and the evaluation environment
Kernel models trained with RL AutoTriton, CUDA Agent the model’s weights the language and the evaluation environment
LLM-driven program search FunSearch, AlphaEvolve the programs, through LLM-written mutations the evaluator
Self-improving agents ADAS (Automated Design of Agentic Systems), Darwin Gödel Machine, Meta-Harness the agent’s own code, or the code around a fixed model the foundation model
CAKE   the compiler harness: IR vocabulary, analyses, cost calibration the coding agent and the model

In the paper’s words, CAKE “targets that complementary layer: it changes the representation being searched and the structured compiler evidence returned to the agent.” Nothing stops the two from combining: a stronger search procedure or a better-trained model could run inside CAKE’s environment.

On the compiler side, TVM, XLA, MLIR, TensorIR, Ansor, and MetaSchedule structure programs so they can be analyzed and optimized automatically, and Graphene, Twill, and Tawa model asynchronous GPU execution, pipelining, or warp specialization. CAKE “shares the principle that structured programs enable useful analysis, but places compiler findings inside an agent evolution loop and lets recurring kernel evidence evolve the harness.” More broadly, the paper cites the “self-defining systems” agenda, which considers AI-operated systems that can change their own mechanisms and abstractions, and then separates CAKE from all of these evolving systems by what is evolved: the agent and model stay fixed, and only a domain-specific compiler harness changes, under corpus tests and human merge gates.

16. Limits and Open Questions

A good paper says where it stops. CAKE does, and a careful reader can add a few more questions.

16.1 What the paper says about its own limits

  • NVIDIA only. CAKE targets Ampere through Blackwell. The schedule language, role model, and analyses carry across those generations, but instruction forms, legality rules, and cost anchors are architecture-specific, and “that cost is the honest measure of transfer.” For non-NVIDIA hardware it “remains unmeasured,” and the backend would have to be rebuilt.
  • Uneven evidence. Most performance results are on B200. The timing model is calibrated only for B200 and H100 and “declines to predict elsewhere.”
  • Incomplete analyses. The static analyses and performance models are intentionally incomplete; they rank and filter, and GPU execution stays the ground truth.
  • Humans in the loop. Compiler evolution is “still human-guided at merge gates.”

16.2 Questions the paper leaves open

These are ours, not the paper’s:

  1. Can others use it? The paper does not say that CAKE’s compiler, IR, or harness is publicly released, and we found no release as of October 2026 (only unofficial community re-implementations). The generated code is public: the four FlashInfer pull requests and three FlashLib pull requests are merged, and a FlashInfer tracker issue indexes many more CAKE-generated kernels.
  2. How general is the controlled result? The clean-start comparison is three runs per arm, one kernel, one shape, and one model (GPT-5.6-sol). We do not know how the gap changes with other kernels, other models, or a larger budget. (The FlashInfer tracker describes CAKE-generated kernels produced with Claude and Codex agents, which suggests the harness is not tied to one model, but no controlled comparison is reported.)
  3. Which part does the work? The language, the pre-compile checks, the cost model, and compiler evolution all change at once between the two arms. There is no ablation that, for example, keeps CAKE IR but removes the cost model, or freezes the compiler.
  4. Why not a tile-language arm? The control arm writes CUDA/PTX. A third arm writing Triton or CuTe DSL under the same rules would show where tile languages and layout algebras land.
  5. What does compiler evolution cost? The paper does not report how many compiler changes were made, how many came from each path, how often the corpus gate rejected a change, or how much human review the merge gates took.
  6. How accurate are the checks? False-positive and false-negative rates for the verifier, and the accuracy of the cost model’s predictions, are not reported.
  7. Missing numbers, and two mismatches. The Gated DeltaNet and MiniMax sparse-attention results are described without speedups, and the plateau criterion is not defined in the text. Two upstream records also differ from the paper: the merged Alpha-MoE pull request reports a different baseline and different speedups (§13.2), and the KDA prefill pull request reports agreement within a tolerance where the paper says “bitwise correct.”

None of these undercuts the main idea. They mark where the next paper, or an independent reproduction, could add the most.

Summary of Part III

CAKE’s harness checks a schedule before it reaches the GPU (program safety, hardware conformance, data consistency, schedule semantics), ranks the survivors with a calibrated cost model, and treats on-device measurement as ground truth. Each round sends its findings to the candidate, the verifier, the cost model, or the IR, and recurring failures become new checks or primitives that must pass the whole kernel corpus behind a human merge gate. With the model fixed, CAKE IR and its harness beat direct CUDA/PTX in a small controlled study, produced new kernels for architectures such as KDA without seeing low-level implementations, matched or beat ten of eleven expert kernels, and turned single-shape seeds into dispatcher-backed library families. The open questions are about generality, attribution, and cost, not about whether the idea works.


Epilogue — a dialogue

Professor So. What did you learn?

Student That a fast kernel is mostly a good schedule: who loads, who multiplies, how they hand tiles to each other, and where every byte lives.

Professor And why agents struggled to write one?

Student They were writing in a language that either hid the schedule or made them manage every address by hand. And all they heard back was "crashed" or a number.

Professor And CAKE's fix?

Student A language where the decisions are explicit and the bookkeeping is derived, checks that name the broken barrier before the GPU ever runs, and a compiler that turns repeated failures into new checks. The agent gets better because its world gets better.

Professor And what would you still want to know?

Student Whether it holds with other models and other kernels, which piece matters most, and whether I can try it myself.

Professor Good. Now you're reading papers properly.

Summary

  • The problem. Expert GPU kernels are distinguished by their schedule (warp roles, barrier choreography, memory placement), and mistakes in a schedule cause silent wrong answers or hangs. Kernel agents have improved their search, memory, and models, but they still write in languages that either hide the schedule or expose it without guardrails, and they get back only errors, pass/fail results, and timings.
  • The idea. Co-design the agent’s environment with the agent. Agents write CAKE IR, a typed, hardware-explicit schedule language with no layout algebra, in which resources, roles, and handoffs are declared and the mechanical details are derived. A harness checks safety, hardware conformance, data consistency, and schedule semantics before compilation, ranks candidates with a calibrated cost model, and treats GPU measurement as ground truth. The compiler evolves: recurring failures and newly found expert patterns become verifier rules, primitives, and calibrations, gated by a corpus of more than 400 static and compile cases and 399 GPU correctness cases.
  • The evidence. With the model fixed, CAKE IR with its harness beats direct CUDA/PTX on a clean-start Flash-KMeans task (median 1.144× vs 0.928× of a tuned baseline at 80M tokens; 3/3 vs 0/3 runs plateaued; about half the active time). Without seeing any low-level implementation (CUDA, PTX, or SASS), it produced a KDA prefill kernel 2.05× faster than the official kernels (geometric mean over six shapes), validated in end-to-end serving. It matches or beats 10 of 11 expert kernels, and its dispatcher-backed KNN and KMeans families are faster by 1.42× to 2.12× (geometric means per family) across more than 400 shapes. Four kernels are merged into FlashInfer.
  • The limits. NVIDIA GPUs only, mostly B200 evidence, a small controlled study, incomplete analyses by design, human merge gates, and no public release of CAKE itself (its generated kernels are public).

Questions to Ponder

OSTEP ends its chapters with homework. Ours has no simulator, but these are worth a few minutes each:

  1. In §3.2, a ring buffer with three stages uses a one-bit phase per barrier. Why is one bit enough? What would go wrong if the producer could run two full laps ahead of the consumer?
  2. CAKE refuses to step a schedule down to an older GPU (§7.4). What could go wrong if a Blackwell schedule were silently lowered to Hopper instructions?
  3. The pre-compile checks are allowed to have false positives (§9). In an evolution loop, which is worse: a check that wrongly rejects 5% of valid kernels, or one that misses 5% of broken ones? Does your answer change for a production compiler?
  4. Design the missing ablation from §16.2: which arms would you run to separate the effect of the language from the effect of the harness and of compiler evolution, and what would each arm keep fixed?
  5. §14.3 forbids the dispatcher from adding its own evaluation shapes. Give a concrete example of how a dispatcher could look 2× faster than it is if this rule were broken.
  6. The paper borrows Terence Tao’s point that cheap generation shifts the bottleneck to verification and understanding (§8.2). Where else in machine-learning systems do you see the same shift happening?

Sources and Further Reading

The paper

  • CAKE: Compiler–Agent Co-Design for Frontier Kernel Evolution (arXiv:2608.12629), Zihao Ye, Yingyi Huang, Hongyi Jin, Bohan Hou, Junru Shao, Zhongming Yu, Jinqi Chen, Meghan Cowan, Shiyi Cao, Shanli Xing, Hanfeng Chen, Vinod Grover, Tianqi Chen, and Luis Ceze (NVIDIA and Carnegie Mellon University). Every number in this post comes from here unless marked otherwise. Read the paper’s §2–3 (program representation and compiler harness) for the design, and its appendix for the eight principles and the full kernel tables.

The upstream kernels

  • FlashInfer pull requests #4262 (KDA prefill), #4279 (KDA decode), #4274 (TinyGEMM), and #4287 (Alpha-MoE), all merged: generated CUDA you can read without CAKE. The CAKE kernel tracker lists many more.
  • FlashLib pull requests #15, #16 and #18: the dispatcher-backed KNN and KMeans kernels of §14.

Background on GPUs and kernels

Kernel agents and evolutionary search

The workloads

The writing style

  • Operating Systems: Three Easy Pieces by Remzi H. Arpaci-Dusseau and Andrea C. Arpaci-Dusseau, free online. The dialogues, crux boxes, asides, and tips in this post are borrowed, with gratitude, from its format.

On this blog

来源:Hacker News · gangliao.me