2022 年,FlashAttention 让 Transformer 的注意力层在同一块 GPU 上快了许多,而且计算结果一个数都没变。它没有找到新公式,找到的是更好的调度(schedule):哪块数据放在哪一级存储里,由芯片的哪个部分来处理,按什么顺序进行。之后又出了三个版本,每一版主要都是更好的调度,后两版则分别为新一代 GPU 而写。
这类程序叫作 GPU kernel,AI 领域很大一部分钱就花在它们身上:大语言模型(LLM)读入或生成的每一个 token,都要经过几十个 kernel。近来,AI 编程智能体也开始写 kernel 了。它们给出一版代码,编译、测试、计时,然后再来一轮。在许多基准任务上,这个循环如今已能产出既正确又快的代码;但碰上最难的那类 kernel,也就是专家为每一代新 GPU 亲手编写的那些,它仍然差一截。
CAKE(Compiler–Agent co-design for frontier Kernel Evolution,面向前沿 kernel 进化的编译器–智能体协同设计)是 NVIDIA 与卡内基梅隆大学的一篇论文。它认为,主要瓶颈已经不在智能体身上,而在智能体所处的世界:它们写代码用的语言,以及它们收到的反馈。今天的智能体写的要么是原生 CUDA,要么是某种高层 kernel 语言,而环境的回应,不外乎“崩溃了”或者“耗时 0.94 ms”。CAKE 把这两头都改了。智能体改用一种新的、可以检查的语言来写 GPU 调度;编译器则给出具体的诊断:这个流水线级里的这个屏障,有人在等,却始终没人发信号。CAKE 还往前多走了一步:如果智能体总在同一类问题上栽跟头,就去改编译器本身,把这种失败变成一项新的检查、一个新的语言特性,或是一个修正过的性能模型。
用论文自己的数字来说,核心结果如下:
本文从最基础的概念讲起,一步步解释 CAKE。我们依据的是论文《CAKE: Compiler–Agent Co-Design for Frontier Kernel Evolution》(arXiv 2608.12629,提交于 2026 年 8 月 12 日)、它的 LaTeX 源码和插图,还有它引用的公开资料。除非另有说明,数字都来自论文。从图上读出的数值标为“近似值”;我们自己对证据的解读标为“我们的理解”。论文没有说的,我们也会明说,并在 §16 汇总这些空白。
这篇文章写给谁
大多数讲 GPU kernel 的论文,都默认你本来就会写 kernel。本文不做这个假设。只要你会一点 Python,大致知道神经网络是做什么的,就能跟得上。我们借鉴了一本深受读者喜爱的操作系统教材的体例:Remzi 和 Andrea Arpaci-Dusseau 的 Operating Systems: Three Easy Pieces(中文版书名《操作系统导论》),并沿用它的三个做法:
- 三个部分。全文分为三部分,每部分都以教授和学生的一段简短对话开场。对话只负责提问,回答交给正文。
- 关键问题。每一章都会把本章的核心难题写进一个标着“关键问题”的方框,让你随时清楚我们要解决的是什么。
- 补充与提示。“补充”方框里是可以跳过的背景知识;“提示”方框总结一条通用经验,它不会随这篇论文一起过时。
如果你已经很熟悉 GPU,可以跳过第一部分,直接从 §5 读起。
路线图
- 第一部分 · 机器(§1–4):什么是 GPU kernel;为什么速度主要取决于数据搬运;现代 GPU 如何变成了异步的装配线;以及专家写的 kernel 和仅仅正确的 kernel 差在哪里。
- 第二部分 · 语言(§5–8):今天的智能体如何写 kernel、能得到什么反馈;可选的 GPU 语言有哪些,为什么没有一种真正适合智能体;逐行解读 CAKE IR;以及它从何而来。
- 第三部分 · 循环(§9–16):在 kernel 碰到 GPU 之前,CAKE 先要做哪些检查;一轮进化怎样进行;编译器自身如何进化;对照实验;生产环境中的成果;如何把一个调优好的 kernel 做成一个库;相关工作;以及局限所在。
- 尾声:一段收尾的对话、全文小结、几个值得思考的问题,以及附有说明的参考资料。
第一部分 · 机器 —— 一段对话
教授 欢迎回来。今天讲一篇论文,主角是会写 GPU kernel 的 AI 智能体。
学生 我天天用 PyTorch,可从来没写过 kernel。我用得着它吗?
教授 你一直都在用。每一次 torch.matmul,最后跑的都是某个人写好的 kernel。问题是,下一代 GPU 上的新 kernel 由谁来写。
学生 让编译器直接生成不就行了?编译器不就是干这个的吗?
教授 简单的运算可以。可对那些真正重要的运算,最快的代码至今仍出自专家之手。芯片上有些事该怎么安排,专家心里有数,编译器眼下还拿不好主意。
学生 比如说?
教授 比如,让哪一组线程搬数据,同时让另一组做乘法,它们之间又怎么互相通知“你的数据好了”。这些一旦弄错,程序可不会客客气气地崩溃。它要么直接挂死,要么悄无声息地算出一堆垃圾。
学生 这调试起来得多痛苦啊。还要让 AI 来干?
教授 这正是这篇论文要回答的问题。不过,你得先看清这台机器。我们就从 kernel 到底是什么说起。
1. GPU kernel 是什么
GPU 这种芯片,生来就是为了同时对海量数据执行同一种简单运算。本节先介绍后文要用到的词汇:线程、warp(线程束)、线程块、流式多处理器,以及 kernel 这个词本身。
1.1 不靠几个快工人,靠一大群小工人
CPU 只有寥寥几个强大的核心,每个核心都为尽快跑完一串指令而设计。GPU 则反其道而行之:它有成千上万条更简单的执行通道,靠让它们全部同时忙起来取胜。没人在乎一次乘法要花多长时间,重要的是每秒能完成多少次。这就是为延迟优化和为吞吐优化的区别。
kernel 是在 GPU 上运行的函数。你只写一份,GPU 会为每个线程各运行一份,线程数常常多达数百万,每个线程处理数据中不同的一部分。最经典的入门 kernel 是把两个各含一百万个数的向量相加:线程 i 读取 a[i] 和 b[i],把它们相加,再写入 c[i]。
线程按层级组织,而这套层级与硬件一一对应:
- Warp。NVIDIA GPU 把线程每 32 个编成一组来运行,这样的一组叫作 warp。同一个 warp 的 32 个线程在同一时刻执行同一条指令,各自处理自己的数据。
- 线程块。若干个 warp 组成一个线程块(也叫 CTA,全称 cooperative thread array,协作线程阵列)。同一个线程块的线程在芯片的同一区域上运行,可以共用一块高速暂存区。
- 流式多处理器(SM)。整颗芯片由许多个相同的 SM 组成。每个 SM 同一时间运行一个或多个线程块,拥有自己的寄存器、暂存区和运算单元。NVIDIA B200 有 148 个 SM。
- 网格(grid)。一次 kernel 启动会创建一个由线程块组成的网格,硬件再把这些线程块分散到各个 SM 上。
1.2 为什么矩阵乘法是最要紧的 kernel
向量加法很简单,因为每个输出恰好只需要两个输入。矩阵乘法则不同。要计算输出矩阵 C = A × B,每个输出元素都需要 A 的一整行和 B 的一整列,而每个输入元素又会被许多个输出用到。这种复用既是机会,也是难点:好的 kernel 把每块输入只加载一次,然后反复使用;差的 kernel 则一遍又一遍地从远处取回同样的数。
Transformer 做的几乎每件事都建立在矩阵乘法之上:注意力里的投影、前馈层,以及混合专家(MoE)模型里的专家。现代 GPU 为矩阵乘法配备了专门的硬件单元,叫作 Tensor Core(张量核心),一条指令就能完成一次小矩阵块的乘法。CAKE 论文里的大多数 kernel,归根结底都是一些精心设计的办法,好让 Tensor Core 始终“吃饱”。
2. 快的 kernel,拼的是数据搬运
先说一个让新手意外的事实:在现代 GPU 上,算术本身很少是难点,难的是及时把数送到运算单元手里。
2.1 内存阶梯
GPU 上有好几种内存,像梯子的横档一样一级一级排开。离运算单元越近的那一级,速度越快,容量也越小:
| 梯级 | 位置 | 容量(B200) | 用途 |
|---|---|---|---|
| 寄存器 | 每个 SM 内部 | 每个 SM 256 KB | 每个线程正在使用的数值 |
| 张量内存(TMEM) | 每个 SM 内部,Blackwell 新增 | 每个 SM 256 KB | Tensor Core 矩阵乘法的累加器 |
| 共享内存(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 编程指南的计算能力(compute capability)表。调优指南给出的是 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,稠密为其一半,折合每卡 2.25 PFLOPS。
- SM 数量:NVIDIA 的产品资料没有列出。148 这个数来自 FlashAttention-4 论文和独立实测。GB200 里的 Blackwell GPU 规格略高(186 GB HBM,稠密 BF16 2.5 PFLOPS)。
每秒 8 TB 听起来大得惊人,可一跟算力相比就相形见绌了。B200 的 Tensor Core 每秒大约能完成 2.25 千万亿次(2.25 × 1015)稠密 BF16 运算。两者一除就知道:要让 Tensor Core 一直忙着,kernel 每从 HBM 读一个字节,大约就得做 280 次运算。每字节做的运算比这少,kernel 的时间就会耗在等内存上,数学再巧妙也没用。
运算次数与字节数之比叫作算术强度(arithmetic intensity);而“kernel 要么受限于内存带宽,要么受限于算力,取决于哪个先耗尽”这一思想,叫作 roofline 模型。两个向量相加,每搬运 12 个字节(两次 4 字节读取、一次写入)才做一次加法,它永远是访存受限(memory-bound)的。大型矩阵乘法每字节可以做到数千次运算,但前提是 kernel 把数据复用得好。
2.2 分块:每个快 kernel 背后的诀窍
提高复用的标准做法是分块(tiling)。线程块不再逐个元素地读取输入,而是把 A 的一个 tile(块)和 B 的一个 tile 从 HBM 复制到共享内存,用 Tensor Core 把它们相乘,累加部分结果,然后接着处理下一对 tile。从 HBM 取来的每个数,都会被这个 tile 里的每一个输出复用。
分块又带来一个新问题:线程块等待下一个 tile 送达时,Tensor Core 只能闲着。解决办法是准备两个或更多的 tile 缓冲区,在计算当前 tile 的同时提前去取下一个 tile。这叫双缓冲(double buffering),更一般的叫法是软件流水线(software pipelining)。FlashAttention 和此前的注意力代码拉开差距,靠的正是这一类调度决策。FlashAttention 算的注意力和别人完全一样,它只是重新安排了 tile,让那个巨大的中间矩阵根本不必写入 HBM。
3. GPU 变成了一条装配线
近几代 NVIDIA GPU 加入的新硬件,让分块和流水线快了许多,也让编程难了许多。本节梳理这一变化,因为 CAKE 的设计紧紧跟随着它。
3.1 从“每个线程什么都干”到各司其职
在较早的 GPU kernel 里,每个 warp 什么活都干:加载 tile 的一部分,在屏障(barrier)处等待线程块里的其他线程,计算,如此反复。此后的每一代 GPU 都加入了新硬件,能在线程忙别的事情时,异步地独立完成其中某一项工作:
- Ampere(A100,2020 年)新增了
cp.async,它在后台把数据从全局显存复制到共享内存。这一代 GPU 的 Tensor Core 由 warp 级的mma.sync指令驱动,所需数据由ldmatrix从共享内存载入。 - Hopper(H100,2022 年)新增了 Tensor Memory Accelerator(TMA),这是一个复制引擎,一条指令就能搬运一整个多维 tile;WGMMA,一种由 4 个 warp 组成的 warpgroup 共同发出的 Tensor Core 指令;异步事务屏障(一种 mbarrier,会在复制数据陆续落地时清点到达的字节数),用来在 tile 到齐时发出信号;以及线程块集群(cluster),同一集群里的线程块可以读取彼此的共享内存。
- Blackwell(B200,2024–25 年)新增了
tcgen05.mma,一种由单个线程代表整个线程块发出的 Tensor Core 指令;以及张量内存(TMEM),即每个 SM 上专用的 256 KB 存储,用来存放累加器,累加器从此不再挤占寄存器。两个相邻的 SM 甚至可以协作完成一次更大的乘法(即“2-CTA”模式)。
既然复制引擎和 Tensor Core 都能自己运转,最自然的设计就是一条装配线:一些 warp 只负责发出复制,一个 warp 只负责发出矩阵乘法,其余的 warp 接过算好的结果,把它们写出去。这叫 warp 专用化(warp specialization),每一项工作就是一个 warp 角色。
3.2 交接:环形缓冲区与屏障
装配线上的各个工位之间,需要一种交接工作的方式。在 warp 专用化的 kernel 中,共享内存里的 tile 缓冲区组成一个分为若干级(stage)的环形缓冲区:以三级为例,生产者可以一边填充第 2 级,Tensor Core 一边处理第 1 级,而第 0 级正等着被释放。
每一级都有两个屏障。生产者的复制数据落地后,向这一级的 full 屏障发出信号;消费者读取之前,先在 full 上等待。消费者用完这一级后,向 empty 屏障发出信号;生产者覆盖这一级之前,先在 empty 上等待。由于环形缓冲区会绕回起点,每个屏障每转一圈就被复用一次,硬件用一个只有 1 比特的相位(phase)来跟踪圈数:等待方必须说明自己期待的是哪个相位,等错了相位,要么过早返回,要么永远不返回。
还有两个细节后面会用到。TMA 发起的复制和 Tensor Core 执行的乘法属于硬件的异步代理(async proxy),而普通的加载和存储属于通用代理(generic proxy);数据从一边跨到另一边时,kernel 需要显式插入一个代理栅栏(proxy fence),否则硬件可能会重排这些访问。另外,每次 TMA 复制都需要一个描述符(descriptor),它编码了 tile 的形状、步长(stride)以及在内存中的布局。
3.3 为什么这么容易出错
这些细节里的每一处都可能出错,而且出错的方式很残酷:
- 漏掉一次 wait,消费者就会在数据送达之前读取这一级。程序不会崩溃,结果却是错的,而且有时只在部分输入上出错。
- 屏障初始化时设错了到达计数(arrival count),或者在错误的相位上等待,都会让某个 warp 永远等下去。kernel 就此挂死,唯一的症状是 GPU 不再响应。
- 漏掉代理栅栏,只有在特定的时序下才会读到过期数据,于是 bug 这次运行出现,下次运行又消失了。
- TMA 描述符里的 tile 形状如果与它要填充的共享内存缓冲区对不上,就会破坏属于别人的内存。
也有工具能帮上忙。NVIDIA 的 compute-sanitizer 可以在运行时检测出部分数据竞争和越界访问,但只限于你碰巧运行的那些输入,而且它不会告诉你问题出自哪一个设计决策。请记住:崩溃或挂死只告诉你出了问题,不告诉你为什么。这个缺口,正是第二部分的主题。
4. 专家级 kernel 与正确的 kernel 差在哪里
把这些拼在一起,就能看懂 CAKE 论文所说的专家级 kernel 是什么意思。论文摘要点出了三项决策,正是它们“把专家级 kernel 与仅仅正确的 kernel 区分开来”:
- warp 专用化。哪些 warp 承担哪些角色:几个生产者、几个消费者,两组消费者是否轮流上阵(“乒乓”,ping-pong),让一组计算的同时另一组写出结果。
- 屏障编排。流水线分几级,哪个屏障为哪一次交接把关,每次等待期待的是哪个相位,栅栏插在哪里。
- 内存层级放置。哪些数据放在寄存器里,哪些放在共享内存里,哪些放在张量内存里,以什么布局存放,好让复制引擎、Tensor Core 和输出路径对每个字节放在哪里达成一致。
这些决策合起来,构成 kernel 的调度(schedule):它规定怎样驱动这台机器,而不是计算什么。矩阵乘法的数学一行就能写完;一个快速的 Blackwell 矩阵乘法,调度却要写上几百行,而且改动其中一项决策,通常会迫使其他决策跟着改。
第一部分小结
kernel 是由成千上万个 GPU 线程并行运行的函数。它的速度主要取决于数据搬运:沿着一级级内存,按 tile 搬运,并在处理当前 tile 的同时加载下一个。现代 NVIDIA GPU 把这一切变成了一条装配线,由各司其职的 warp 通过屏障相互传递 tile。专家的功夫在于调度(角色、屏障、内存放置),而调度上的错误往往导致悄无声息的错误结果或挂死,而不是一条有用的报错信息。
第二部分 · 语言 —— 一段对话
学生 也就是说,智能体写一个 kernel,跑一跑,看看结果。这有什么问题?我自己调试也是这么做的。
教授 你的代码出错时,都看些什么?
学生 看调用栈,看行号,看哪个变量成了 None。
教授 现在想象一下,你唯一能得到的反馈是“卡住了”,或者“用时 0.94 毫秒”。没有行号,也没有变量。
学生 那我只能乱改一通,碰碰运气。
教授 今天的 kernel 智能体,大致就是这个处境。所以有两个问题:智能体该用什么语言来写?出了错,编译器又该告诉它什么?
学生 像专家那样直接写 CUDA 不就行了?
教授 CUDA 什么都写得出来,包括每一种错误。更高层的语言会保护你,可专家写得出来的东西,它不让你写。论文的答案在两者之间。我们来看看具体在哪儿。
5. 今天的智能体如何写 kernel
写 kernel 的智能体出现得很快。本节介绍它们共用的那个循环、这个循环擅长什么,以及 CAKE 的出发点:反馈问题。
5.1 “提出–测试–测量”循环
这个任务设定因 KernelBench(2025)而流行开来:给模型一个 PyTorch 算子,让它写出一个计算结果相同、但更快的 GPU kernel。此后,许多系统围绕这个任务搭起了各自的循环。有的采用人工设计的循环,由 LLM 根据编译错误、正确性结果或性能分析器(profiler)的输出修改 kernel(AccelOpt、Autocomp、KernelAgent)。KernelBlaster 加入了一个持久化的知识库,收录以往的优化经验。KernelEvolve 和 EvoEngineer 在由候选 kernel 组成的种群上做进化搜索。AVO 出自 CAKE 的几位作者之手,它用编程智能体取代固定的变异规则。K-Search 把规划和实现分开。AutoTriton 和 CUDA Agent 则用强化学习训练模型本身。
花样虽多,底下的环境却是同一个:智能体提出代码,代码被编译,数值测试把输出和参考实现比对,基准测试测出延迟,智能体再决定下一处怎么改。CAKE 论文的概括是:这些方法“进化的是搜索过程、积累的记忆或模型权重,而选定的领域专用语言(DSL)和评测环境保持不变”。
5.2 环境回应了什么
这个循环很擅长局部调优:换一个 tile 大小,展开一个循环,融合两个操作。但碰上 §4 里的那些决策,它就力不从心了。原因论文一句话就说清了:编译错误、正确性结果和端到端计时“从不指出是哪个程序决策导致了同步失败、硬件契约违规或流水线停顿”。
设想一个智能体刚刚重写了一个带流水线的 kernel,看看每种信号能告诉它什么:
- 挂死。某个 warp 在等一个永远不会完成的屏障。哪个屏障?哪一级?哪个相位?进程就这么停住了。
- 输出错误。可能是数据竞争、布局错误、漏掉的栅栏,也可能只是普通的索引 bug。测试能告诉你哪些元素不对,却说不出是哪个决策造成的。
- 计时结果。0.94 ms。这个 kernel 是卡在显存带宽上,卡在 Tensor Core 指令发射上,是流水线太浅,还是某个屏障让两个角色串行了?一个数字回答不了。
还有第二个更隐蔽的问题:这些信号“无法在前沿工作负载暴露出能力缺失时随之成长”。如果语言表达不了新工作负载需要的调度,每次尝试都会因为同一个结构性原因失败,而环境始终无法从中吸取教训。
5.3 CAKE 的三个核心主张
论文的观察是:专家装在脑子里的那套机制,编译器大多已经具备,即“结构化的操作词汇表、资源模型、合法性检查、静态分析、代价模型、lowering(逐层翻译为底层代码)规则”。CAKE 要问的是:“如何让这套机制面向智能体,以及当前沿工作负载暴露出缺口时如何改进它。”它的回答分三部分:
- 智能体编辑的是带类型的中间表示(IR),而不是原始 CUDA。IR 是一种专为编译器分析而设计的程序格式。CAKE 的 IR 是一门调度语言,硬件决策在其中都是显式的,因此在生成任何代码之前就能检查它们。
- 编译器返回精确定位的诊断,而不是一个“通过/失败”比特。无论是正确性问题还是性能问题,诊断结果都会指向出错的角色、流水线级、缓冲区或指令;开销低的分析还会在候选 kernel 消耗 GPU 时间之前,先把它们筛一遍。
- harness(测评环境)本身也会进化。反复出现的失败会变成新的验证器规则、新的 IR 原语、重新校准的代价模型和可复用的策略;每一项改动都要经过 kernel 语料库上的测试把关。
论文把这种做法称为编译器–智能体协同设计(compiler–agent co-design):智能体和它所处的环境一起设计,环境还会随着智能体遇到的问题不断变化。接下来三节讲第一个主张:语言。
6. 智能体该用哪种语言?
写 GPU kernel 的方式有很多,它们最大的区别在于一个问题:调度由谁决定,程序员还是编译器?本节按 CAKE 论文自己的分类,把这份“菜单”从头看一遍。
6.1 高层:tile 级语言
Triton(2019)让程序员以 tile 为单位思考,从而降低了写 kernel 的门槛:你只写矩阵的一个块上要做什么,至于线程、各级内存和指令具体怎么执行,交给编译器决定。更新的语言,比如 Helion、TileLang 和 NVIDIA 的 cuTile,在不同的抽象层次上沿用了类似的思路。
这对人类来说非常好用,写出来的代码正确,而且往往很快。CAKE 论文的异议很具体:tile 级语言(tile DSL)“隐藏了 warp 专用化、屏障编排和内存层级放置,而正是这些把专家 kernel 和仅仅正确的 kernel 区分开来”。如果语言根本没法表达“一个生产者 warp、一个 MMA warp、一个三级环形缓冲区、累加器放在 TMEM 里”,用 Triton 写代码的智能体就提不出这样的要求,只能指望编译器碰巧这么选。
6.2 低层:CUTLASS、CuTe DSL 与原始 CUDA
另一端是 NVIDIA 的 CUTLASS 库及其 Python 前端 CuTe DSL。它们把硬件提供的一切都暴露出来,代价是一套布局代数(layout algebra):一个数学框架,用来描述张量的逻辑坐标如何映射到内存和线程上。再往下是带内联 PTX(NVIDIA 的一种类似汇编的虚拟指令集)的 CUDA C++,地址、屏障相位位和描述符编码都要你亲手管理。
论文认为,低层语言“暴露了这种控制,却要求掌握一套布局演算,使得智能体的错误既容易发生,又难以定位”。布局写错不会引发编译错误,而是让正确的数落到错误的位置上,到后面才表现为错误的输出。
6.3 两者之间:Gluon
Gluon 介于两大阵营之间。它复用 Triton 的编译器栈,但暴露出对布局、数据搬运和异步执行的更底层控制。它把刻度盘往“控制”一侧拨了拨,也就一并拨向了控制带来的负担。
6.4 CAKE 瞄准的空白
把这些选项并排一放,空白就显现出来了。能表达专家级调度的语言,要你自己管理布局和各种底层细节,而且要到运行时才能发现你的错误;能保护你的语言,又不把调度交到你手里。
7. 逐行读懂 CAKE IR
CAKE IR 就是智能体编写的语言。要理解它,最好的办法是读上一段。
7.1 调度声明什么,lowering 推导什么
这一设计建立在一种分工之上。用论文的话说:“调度声明要发生什么,lowering 推导如何发生。”调度记录的是:哪些 warp 担任哪些角色,哪些缓冲区分几级,哪个屏障把守哪次交接,哪种指令形式消费哪个操作数。lowering 是编译器把程序向下翻译成 CUDA 和 PTX 的过程,它负责算出这些决策的机械性后果:“屏障地址、相位位、TMEM 偏移、描述符编码和 warp 身份,全都根据声明计算得出,而不是由智能体逐一写出。”
下面是论文用来说明这一思路的片段,取自一个面向 Blackwell 的融合多头注意力(FMHA)前向 kernel:
@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))
分三块来读:
- 资源。
lm.smem(98304)预留一块 96 KB 的共享内存池。pool.view(...)从中划出一个具名缓冲区:128 × 128 个 BF16 值,分三级,也就是一个三槽的环形缓冲区。lm.tmem(...)在张量内存中预留一个 128 × 128 的 FP32 累加器。这里没有一处是地址;每个缓冲区都有名字、形状、类型和生命周期,编译器全都看得见。 - 角色与同步。
lm.role(warps=[0])和lm.role(warps=[1])命名了两个 warp 角色:一个负责加载,一个负责发射 Tensor Core 指令。lm.pipeline(stages=3)声明流水线,lm.barrier(...)声明两者之间的交接:一个三级屏障,生产者是load,消费者是mma。 - 各角色的工作。在
with load:中,加载 warp 每一级发起一次 TMA 拷贝,并让硬件在每次拷贝落地时通知q_full。在with mma:中,Tensor Core warp 逐级等待q_full,发出一个代理栅栏(在硬件两条内存访问路径之间保证顺序的指令,见 §3.2),然后发射矩阵乘法,结果写入 TMEM 累加器,并在第一级时将累加器初始化。
注意这里少了什么:没有屏障地址,没有相位位,没有描述符编码,也不用计算哪个线程属于哪个 warp;这些全由编译器根据声明推导出来。不过,缓冲区确实带着几项朴素的承诺,比如视图在内存池中的偏移,以及累加器在张量内存中的起始列。§7.3 会解释 CAKE 为什么直接要求写出这些值,其余的寻址信息则由编译器以它们为起点推导出来。这段代码只是教学示例,不是完整的 kernel:它把 Q tile 和它自己相乘,而真正的注意力 kernel 还要加载 K 和 V,并计算 softmax。
7.2 让它可以检查的四个特性
论文指出,真正“起作用”的是四个特性:
- 经过类型检查的词汇表。计算、数据搬运、同步、数学运算和 warp 控制,用的都是一组固定的 IR 操作,而不是嵌入的 C 或 PTX 字符串。类型不对的程序在构建过程中就会被拒绝。
- 预先声明的资源。内存区域、同步对象和流水线都只声明一次,因此“IR 知道每个缓冲区的形状、数据类型(dtype)和生命周期”。
- 显式的角色。warp 组都有名字,而且“每一次跨角色的交接都是可见的,而不是一种隐含的约定”。
- 自动推导的元数据。声明带来的机械性后果“由 lowering 生成,而不是由作者编写”。
这样做的回报是:分析可以在代码生成之前就对调度决策进行推理,harness 也能把一条诊断结果对应到受影响的资源、角色或流水线级,“而不是只返回一个后端错误或一次挂死”。
7.3 有意不用布局代数
CAKE 的做法是“刻意不把布局作为一等抽象”。智能体不用把布局当作数学对象来摆弄,只需写下具体的承诺:“一个 SMEM 视图偏移、一个操作数字节偏移、一段 TMEM 列范围、一个 swizzle 标签、一个 TMA 描述符坐标”。随后由编译器“承担起判断这些承诺是否合法的责任”:它检查这些承诺沿着程序的数据流是否始终一致(生产者写入的,正是消费者期望读到的),以及是否满足目标指令的规则。诊断信息会指出出错的 IR 决策,以及不匹配所属的大致类别。
这和 CuTe 的取舍正好相反。CuTe 要求作者精通布局,回报是强大的表达能力;CAKE 只要作者提供朴素的事实,由编译器负责揪出其中的矛盾。至于布局验证器,论文只描述了它的契约和覆盖范围,没有介绍内部实现。
7.4 一种语言,多款 GPU
同一种调度语言覆盖了从 Ampere 到 Blackwell 的 NVIDIA GPU。“角色–屏障–流水线”式调度的结构可以沿用;至于允许使用哪些指令、这些指令如何 lowering,仍然因目标而异:
| 目标 | 代表 GPU | Tensor Core 路径与主要特性 |
|---|---|---|
sm_80 |
A100 | mma.sync、ldmatrix、cp.async;没有 TMA、集群和 TMEM |
sm_89 |
L40S、RTX 6000 Ada | 同 sm_80,外加 FP8 Tensor Core;没有 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 的块缩放(block-scaled)MMA |
sm_120a |
RTX 5090、RTX PRO 6000 | mma.sync + ldmatrix、TMA、集群、分布式共享内存;没有 tcgen05 和 TMEM |
sm_121a |
DGX Spark(GB10) | 指令与 sm_120a 相同,但二进制单独一份,特殊函数单元速率不同 |
有两条策略值得一提。第一,编译器“要求目标精确匹配”:缺少对应的设备或工具链时,它会直接报告,而不是悄悄把调度 lowering 到更老的架构上。第二,性能估计只在经过校准的地方给出:B200 是实测基线,H100 单独校准,其他目标则报告覆盖范围受限,而不是借用别处的数字。通过静态检查之后,CAKE IR 会确定性地 lowering 为可读的 CUDA/PTX,再经过 NVIDIA 的标准工具链。生成的源码始终可以取用,算是一条逃生通道;但默认情况下,决策都留在 IR 里,以便分析。
8. CAKE IR 从何而来
你也许以为,这样一门语言是先在白板上设计好、再动手实现的。CAKE IR 却是“长”出来的。
8.1 从专家 kernel 里长出来
“CAKE 并不是从一套预先定义好的 CAKE IR 词汇表起步的。”它的原始素材是一个生产级 kernel 语料库,外加一组硬件设计原则。论文附录描述了五个步骤:
- 收集语料。从生产级库中收集高质量的 CUDA kernel:Alpha-MoE、CUTLASS、cuTile、DeepGEMM、FlashAttention-4、Flash-KMeans、FlashInfer、SonicMoE 和 TileLang。用其他语言写的 kernel,先由编程智能体翻译成带内联 PTX 的 CUDA。
- 提炼抽象。智能体分析语料,把反复出现的模式(屏障编排、流水线分级、warp 角色划分、TMA 描述符设置、TMEM 累加器的生命周期)归纳为候选抽象。
- 依据硬件做设计。人类专家的经验把这些抽象引向 Blackwell 的编程模型:TMEM 是一等资源,warp 专用化是主要的并行形式,异步屏障是同步原语,集群范围的操作用来协调多个 SM。
- 按原则迭代。每个候选抽象都要对照八条设计原则(见下文)检查;违反原则的,要么改进,要么淘汰。
- 靠移植扩展。新的 kernel 源源不断地移植进来。每次移植要么成功,从而验证了现有抽象;要么暴露出缺口,从而引出一份扩展 IR 或其 lowering 的提案。
这个循环永远不会结束:“每一个新的 kernel 家族都会对 IR 进行压力测试,并推动它进一步进化。”设计目标自始至终都定得很高:IR 必须能“复现专家手写 kernel 的物理调度和性能”。据论文介绍,harness 同样“主要由智能体维护,合并前由人类把关”。
8.2 八条原则
| 原则 | 具体要求 |
|---|---|
| P1 易用(Ergonomic) | 让 NumPy 和 PyTorch 用户用着眼熟;避免不必要的簿记。 |
| P2 性能透明(Performance-transparent) | 与性能相关的硬件决策要保持可见,lowering 过程要可以检查。 |
| P3 规范(Canonical) | 每个操作只有一种规范形式,而不是几种等价的写法。 |
| P4 静态类型检查(Statically type-checked) | 用类型规则约束 lowering,在构建阶段就拒绝类型错误的程序。 |
| P5 便于分析(Analysis-friendly) | 暴露静态分析所需的信息。 |
| P6 测试把关(Test-gated) | 每一处 IR 改动,都要用 kernel 矩阵测试检验分析和编译。 |
| P7 分析一致(Analysis-consistent) | IR 的数据模型一变,分析也要随之修改。 |
| P8 立足硬件(Hardware-grounded) | 为每个操作写明预期的硬件行为。 |
P5 和 P7 让新特性始终可以分析,P6 防止改动引入回归,P2 和 P8 则让 IR 到硬件的映射“对人类和智能体都清晰易读”。论文把这一点和数学家陶哲轩(Terence Tao)关于 AI 生成证明的一个观察联系起来:当生成产物变得廉价,瓶颈就转移到验证和理解这些产物上。放到 CAKE 这里就是:当智能体能批量产出 kernel,要紧的是每个 kernel 都能被检查、被读懂。
第二部分小结
今天的 kernel 智能体活在一个黑盒里,它只回答“通过/失败”和一个时间。CAKE 给了它们一门语言:专家的决策(角色、缓冲区、屏障、流水线)在其中是显式的、带类型的,而容易出错的簿记工作(地址、相位、偏移、编码)由编译器推导。它刻意避开布局代数,转而让编译器检查具体的布局承诺。这门语言是从生产级 kernel 中长出来的,八条原则让它在成长中始终保持可分析。
第三部分 · 循环 —— 一段对话
学生 好,智能体就用这门调度语言来写。然后呢?代码最后总还是得上 GPU 跑吧。
教授 最后是要跑。可是靠上 GPU 跑一趟来发现漏掉的屏障,代价太高了。有没有更便宜的办法?
学生 先读代码?像 linter 那样?
教授 没错。先做检查:在编译之前就运行,还能指出是哪个屏障、哪一级出了问题。接着用一个模型猜一猜,过了关的候选里哪些值得拿去计时。最后才轮到 GPU。
学生 要是检查本身漏掉了一整类 bug 呢?
教授 那你就找到这篇论文最有意思的想法了。这个 bug 会变成一条新的检查。编译器会长记性。
学生 这一套真能做出更快的 kernel 吗?还是只换来更好看的报错信息?
教授 很好,保持这份怀疑。我们会去看证据,也会看看这些证据说明不了什么。
9. Harness:上 GPU 之前的检查
所谓 harness(测评环境),就是围绕智能体的一整套东西:各项分析、测试、基准测试,以及判定什么结果才算数的规则。本节介绍 CAKE 的 harness 检查什么,又会向智能体反馈什么。
9.1 七类检查
论文是按功能来介绍这套分析的:看它能为智能体做什么,而不是看编译器内部有哪些 pass。检查一共七类,分属三种处置方式:拒绝候选的关卡(gate)、描述候选的报告(report),以及建议如何改进的提示(hint)。
| 类别 | 处置方式 | 作用 |
|---|---|---|
| 程序安全 | 编译前关卡 | 找出同步、顺序和内存使用方面的隐患 |
| 硬件合规 | 编译前关卡 | 强制遵守受支持的资源、指令与架构契约 |
| 数据一致性 | 编译前关卡 | 检查数据流,以及生产者与消费者对数据表示的约定是否一致 |
| 调度语义 | 编译前关卡 | 检查所声明调度的结构不变量 |
| 数值验证 | 执行关卡 | 把编译后的输出与权威的外部参考结果对比 |
| 性能分析 | 报告 | 估算开销,指出瓶颈大致属于哪一类 |
| 优化指导 | 提示 | 建议值得尝试的修改,但不阻止编译 |
四道编译前关卡会拒掉“许多在数学上说得通、却与目标执行模型不兼容的候选”。每条检查结果都会“指出受影响的程序区域,以及违反的是哪一类契约”。拿它和 §3.3 里的挂死比一比:智能体面对的不再是一块毫无反应的 GPU,而是类似这样的反馈:消费者在这个屏障的第 2 级上等待,但从来没有生产者到达那里。这个例子是我们自己举的,论文没有公开诊断消息的格式。
9.2 正确性与性能
数值正确性的检查方法,是在不同的形状和输入分布下与参考实现对比;最终验收还要求在目标框架中做端到端评测(例如用 SGLang 为完整模型提供推理服务)。
性能评估分为两层。经过校准的代价模型先估计候选能跑多快,并指出它可能的瓶颈,这样 harness 不必把每个候选都跑一遍,就能给它们排序、筛选。但这个模型只是一道筛子:“设备上的实测和性能剖析(profiling)仍是最终的评判依据。”论文评测中报告的每个数字,都是用 CUPTI(NVIDIA 的性能剖析接口)在 GPU 上实测的;每次计时采样之前都会清空 L2 缓存,免得上一次运行替下一次把缓存预热好。
10. 一轮进化
语言和 harness 都就位之后,一次 kernel 进化运行的流程就很简单了。本节跟着走完一轮。
10.1 四个阶段
论文把一次运行分为四个阶段:
- 生成结构各异的 CAKE IR 候选:角色划分、流水线深度、内存放置各不相同,而不只是换几个常数。
- 筛选这些候选:在花费任何 GPU 时间之前,先用 IR 构建检查、验证器的硬性关卡和代价模型排序把它们过一遍。
- 评估过关的候选:对照外部的正确性判据(oracle)检验,并收集基准测试和性能剖析证据。
- 分流证据。根据诊断结果,一条发现会流向候选本身(修复它)、验证器(新增一条规则)、代价模型(重新校准),或者 IR 词汇表(新增一个原语)。
正是第四个阶段,让 CAKE 不只是一个报错信息更好的搜索循环。一条发现不会用完就扔,而是被送到系统里最该从中学习的那个部分。
10.2 哪些保持不变
每次运行都有两个锚点。第一个是工作负载契约(workload contract),也就是论文所说的“稳定的权威”:它固定了形状、正确性判据及其容差、硬件,以及智能体可以查看哪些参考资料。运行结果都会保留下来,因此决策可供审计,反复出现的发现也能复用。
第二个是模型。论文中所有的智能体任务都使用同一个模型 GPT-5.6-sol,推理强度(reasoning effort)为 xhigh。“固定模型和智能体脚手架(scaffold),使这些比较……能够归因于环境,而不是模型能力。”一项实验的两组之间无论出现什么差别,都不会是因为其中一组用了更聪明的模型。
11. 编译器也在进化
到目前为止,智能体都是在一个固定的环境里进化 kernel,只不过这个环境恰好设计得很好。CAKE 的最后一项承诺是:环境本身并不固定。
11.1 进入编译器的两条路径
候选 kernel、验证结果、基准测试和失败报告,都是修改编译器的证据。论文描述了两条路径:
- 路径 1:挖掘专家 kernel。智能体研读生产级 kernel 和硬件文档,找出 IR 尚未覆盖的 Blackwell 模式(“新的指令形式、资源类型、描述符变体、同步惯用法”),并撰写编译器修改提案。每份提案在实现之前,都要对照 §8.2 的设计原则检查,其中包括性能透明和易于验证。
- 路径 2:从失败中提炼。智能体利用失败候选留下的反馈(“sanitizer 报告、失败用例、正确性不一致、调试日志”),把反复出现或代价高昂的失败模式提炼成新的分析:“一次原因不明的运行时崩溃变成一条验证器规则,一种反复出现的非法 lowering 模式变成一项静态检查,一种系统性的预测偏差变成一个校准目标。”
11.2 为什么原语和分析必须一起改变
两条路径相辅相成。新原语让编译器更了解硬件,从而能做更强的分析;新分析反过来又约束了未来原语可以是什么样子。论文坚持认为,一个原语和它的分析“必须共同进化”:“只有语法、没有副作用和合法性规则,会让 IR 更难分析;而未经语料库验证的新验证器规则,可能会拒掉有效的 kernel。”
所以,每一处编译器改动都要经过整个 kernel 语料库的测试把关。这个语料库规模可观:400 多个静态检查与编译用例,外加 399 个 GPU 正确性用例,覆盖约 28 个 kernel 家族(§13.4)。新规则只要拒掉一个已知正确的 kernel,就不会被合入。
这些改动由谁来做?人类写出高层描述,智能体据此实现并维护各项分析;合并则由人类批准。论文说,编译器的进化“在合并关卡上仍由人类把关”。
12. 对照实验:从零起步的 Flash-KMeans
描述一个系统很容易,做一次对照比较却不容易。本节讲的是论文里最干净的一个实验:同一个智能体、同一个模型、同一个任务,唯一的区别是用 CAKE IR 写,还是直接用 CUDA/PTX 写。
12.1 工作负载:k-means 的一步
k-means 是机器学习里最古老的算法之一。手上有 N 个点,想把它们分成 K 组。先猜 K 个组中心(称为质心,centroid)作为起点,然后反复执行两步:把每个点分配(assign)给离它最近的质心,再把每个质心更新(update)为分配给它的那些点的平均值。每重复一轮,称为一次 Lloyd 迭代。
Flash-KMeans 是一个快速、精确的 GPU 实现。它之所以重要,原因可能出乎你的意料:视频生成。Sparse VideoGen2 用 k-means 把相似的 token 聚到一起,让稀疏注意力能在由相关 token 组成的连续块上工作(“语义感知重排”,semantic-aware permutation)。在那里,k-means 运行在模型内部,所以它的速度就是模型的速度。
按论文的说法,一次 Lloyd 迭代的时间主要花在两个 BF16 kernel 上,两者合计占端到端时间的 95% 以上:
assign计算每个点到全部 K 个质心的平方距离,并返回最近那个质心的下标。它是计算受限(compute-bound)的:一次矩阵乘法,后接一次归约。centroid_update把每个簇里的点加起来,并统计点数。它受限于显存带宽,以及原子更新时的争用。
实验聚焦于 assign,它“考验 Tensor Core 流水线和标量 epilogue(收尾阶段)”。形状固定为:B = 32 个批次,N = 65,536 个点,K = 1,024 个质心,D = 128 维,BF16 输入,FP32 累加。标杆定得很高:基线是 FlashML 项目中一个经过调优的 Triton 实现(该项目的 FlashLib 库收集了各种经典机器学习算子的高速 GPU 版本),实测 0.938 ms。
12.2 规则:智能体能看什么,不能看什么
论文对一种隐蔽的作弊方式格外小心。如果智能体能读到专家为这个 kernel 写的 CUDA,它就可以照抄调度,而不必自己去发现。所以在这些从零起步(clean start)的运行中,智能体可以看到“数学规格说明、评测契约、正确性判据和高层代码,但看不到 CUDA、PTX、SASS 或等价生成代码这类底层目标实现”。现成的实现可以作为黑盒计时基线拿来运行,但其内部细节始终对智能体隐藏。这条限制在隔离环境中强制执行,事后还做了审计。
两组实验唯一的区别,是智能体写什么:
- 实验组:智能体写带类型的 CAKE IR,配合第三部分介绍的 harness。
- 对照组:智能体直接写 CUDA C++ 和内联 PTX,遵守同样的参考实现访问规则。
其余一切都固定不变:编程智能体及其脚手架、模型及推理强度、任务描述、正确性判据、基准测试 harness,以及目标形状。每组各跑三次,预算为 80M token。
12.3 结果
| 表示形式 | 80M token 内进入平台期 | 有效进化时间(小时) | 80M token 时的最佳速度(× 基线) |
|---|---|---|---|
| CAKE IR | 3 次中 3 次 | 1.89 [1.02, 2.33] | 1.144 [1.041, 1.205] |
| 直接写 CUDA/PTX | 3 次中 0 次 | 3.73 [3.59, 4.34] | 0.928 [0.852, 1.151] |
表中数值是三次运行的中位数,方括号里依次是最小值和最大值。用 CAKE IR 时,各次运行最佳 kernel 的中位数比调优基线快 14.4%(1.144×);直接写 CUDA/PTX 时,只有基线速度的 92.8%。CAKE IR 的三次运行全部满足论文“预先设定的平台期判据”,CUDA/PTX 一次也没有。CAKE IR 用掉的有效时间也只有约一半。
随时间展开来看,轨迹讲的是同一个故事。在每个 5M token 的检查点上,每次运行都贡献自己截至当时、经过验证的最佳速度。CAKE IR 各次运行的均值“到 55M token 时越过调优过的 FlashML 基线,并继续提升;而直接写 CUDA/PTX 的均值到 80M token 截止时仍低于基线”。从论文的图上读数(近似值):到 20M token 时,CAKE IR 的均值已接近 0.8×,而 CUDA/PTX 的均值还在 0.2× 左右。
12.4 怎么解读
下面是我们的理解,以及这些数字本身附带的限定条件:
- 语言改变的是结果,而不只是到达结果的速度。在同样的预算内,CAKE 的三次运行全部超过了调优基线;CUDA 的运行按中位数看没有做到。
- 但直接写 CUDA 并非毫无希望。CUDA/PTX 最好的一次运行达到 1.151×,比 CAKE 的中位数还高。差别在于稳定性和速度:CUDA 也能做到,只是不那么可靠,也更慢。每组只有三次运行,结果的分布范围和中位数同样重要。
- 这只是一个 kernel、一个形状、一个模型。这个实验很好地隔离出了表示形式的作用,但没有告诉我们,换成别的 kernel 或别的模型,这个作用有多大。论文正文也没有定义它的平台期判据,只说详细的停止与计时记录“保留在配套材料(artifact)中”。
13. 基准测试之外
对照实验干净地展示了效果;其余评测则展示这套系统在真实 kernel 上的工作情况。这部分评测还回答了另外两个问题:对于还没有人公开过快速底层实现的 kernel,CAKE 能不能找到它的调度?对于专家已经调优过的 kernel,它能不能复现,甚至超越?
13.1 前沿 kernel:新的模型架构,没有参考可抄
论文称一个 kernel 为前沿 kernel,是“在操作意义上而言:智能体必须在不查看底层目标实现的情况下,发现它的物理调度”。协同进化的 IR 在这里理应最能帮上忙,但这里也最容易暴露它的短板:缺了某种能力,就会表现为一个“智能体根本无法表达”的调度。
Kimi Delta Attention。官方的 FlashKDA kernel 只用作黑盒计时基线,它们的源码和生成代码都没有提供给智能体。智能体生成的 prefill kernel 与 FlashKDA 兼容,覆盖定长、打包变长(packed variable-length)和尾部输入,在六个 B200 BF16 形状上相对该基线取得 2.05× 的几何平均加速比。论文称它“在其验证契约上逐位正确”(pull request 本身报告的是:输出与 FlashKDA 在 10−2 的容差内一致),并且它在 SGLang 上端到端部署 Kimi-K3 时通过了验证。另有单独的 decode kernel,在 30 个公开 API 形状上相对上游 FlashInfer 达到 1.14×。两者都以生成的 CUDA 代码合入了 FlashInfer(pull request #4262 和 #4279,2026 年 8 月合入),“因此下游用户无需引入对 CAKE 的依赖”。
Gated DeltaNet 与稀疏注意力。与 FlashInfer 相比,CAKE 的 Gated DeltaNet prefill 路径和投机解码路径“在保持模型循环状态的同时提升了性能”;MiniMax 稀疏注意力则表明,这套表示形式在 prefill 和 decode 两侧都能支持稀疏注意力家族。论文没有给出这两项的加速数字。它们是“分发家族,而不是单个 kernel”:一个入口背后有多个物理调度,每个都是独立的 CAKE IR 程序(§14)。
13.2 生产级 kernel:从参考实现出发进化
TinyGEMM。这一次,智能体从一个现成的 kernel 出发:FlashInfer 中面向极小 batch 的 BF16 矩阵乘法(论文称之为小 M),源自 TensorRT-LLM。服务端同时只为少数几个请求做 decode 时,batch 通常就是这么小。智能体产出了一个自适应的 kernel 家族,既有浅流水线也有深流水线,其中包括使用 PDL(programmatic dependent launch,允许下一个 kernel 在上一个结束之前就开始执行自己的前导部分)的变体,以及专为 batch 小于 8 准备的路径。FlashInfer PR #4274 报告,在 35 个标准形状和一套覆盖面更广的回归测试上,kernel 耗时的几何平均降低了 18–23%。在 B200 和 GB300 上,GPT-OSS-20B 与 GPT-OSS-120B 的贪心解码输出保持逐位一致;一项用 GPT-OSS-120B 做的 SGLang 实验测得,单卡、并发 128 时输出吞吐最多提高 7.6%(四路张量并行时,差异在噪声范围内)。
Alpha-MoE。原版 Alpha-MoE 是为 Hopper 编写的“巨型 kernel”(megakernel),把混合专家(MoE)计算融合在一起,权重和激活都是 8 位(W8A8)。CAKE 的智能体把它改写到了 Blackwell 上。一个设备端程序就包办了按路由收集(routed gather)、两次投影、激活、重新量化,以及按路由权重累加输出。与 FlashInfer 中源自 TensorRT-LLM 的预路由(pre-routed)API 相比,端到端的 API 层加速比在 N = 256 时为 6.204×,在 N = 512 时为 4.025×;若按 GPU 跨度(GPU span)衡量,则分别为 1.215× 和 1.170×。这个 kernel 已于 2026 年 9 月合入 FlashInfer(PR #4287),但 pull request 里报告的是另一组对比:在 GB300 上,与 SGLang 原生的五个 Triton kernel 组成的 MoE 路径相比,融合后的 kernel 按 GPU 时间快 1.37×–1.75×。论文中的 6.204× 和 4.025× 并没有出现在其中。
13.3 复现专家 kernel
最后一个问题是:当专家设计的结构已经存在,而且智能体可以查看时,harness 能否帮助智能体在追平高度优化 kernel 的同时保持正确性。参考实现来自 FlashAttention-4、TensorRT-LLM、DeepGEMM、CUTLASS 和 FlashInfer。所有运行都在 B200 上以固定形状进行,用 CUPTI 测得的 GPU 跨度中位数计时,并且都通过了各 kernel 专属的正确性关卡。
| 家族 | 变体 | 相对性能 | CAKE IR 行数 | 参考实现设备端代码行数 |
|---|---|---|---|---|
| FlashAttention-4 | 前向,BF16,非因果 | 1.0045× | 430 | 2,369 |
| FlashAttention-4 | 反向,BF16,非因果 | 1.0470× | 514 | 2,552 |
| TensorRT-LLM GQA | decode,FP16 | 1.043× | 783 | 8,515 |
| DeepGEMM | 1D1D GEMM,FP8 | 1.0370× | 221 | 516 |
| DeepGEMM | 分组 GEMM,BF16,带掩码 | 1.0174× | 401 | 442 |
| DeepGEMM | MQA 索引器,FP8 | 1.2700× | 480 | 704 |
| DeepGEMM | MQA 索引器,FP4 | 1.2730× | 392 | 704 |
| DeepGEMM | 分页 MQA 索引器,FP4 | 1.0036× | 395 | 779 |
| CUTLASS MLA | decode,BF16,TMA | 1.2174× | 845 | 1,860 |
| DeepSeek-V4 稀疏 MLA | decode,BF16 | 1.1297× | 1,299 | 13,609 |
| DeepSeek-V4 稀疏 MLA | decode,FP8 | 0.9649× | 1,393 | 8,942 |
11 组对比中有 10 组持平或超过参考实现,剩下一组也达到了参考的 96.5%。论文还坦率地补充了两点限定。其一,低于参考的条目“通常反映的是编译器集成的成熟度”:kernel 想用的某项特性还处在集成到编译器的过程中,所以提交的 kernel 只能采用最接近的、已支持的策略。其二,领先最多的两项,也就是约 1.27× 的两个 MQA 索引器,“并不是对参考 kernel 的忠实转写”:智能体探索了原版没有的优化,所以超过持平线的结果“反映的是搜索,而不仅仅是转写的忠实度”。行数只起描述作用:两边语言不同,计数范围也不同,它们只能说明这些调度“可以被紧凑地表示”。
形状与行数统计细节
- S1(FA4):B = 4,H = 32,S = 8,192,D = 128。S2(TRT-LLM GQA):B = 128,64 个查询头,8 个 KV 头,KV 长度 4,096,D = 128,页大小 16。S3(FP8 GEMM):M = 4,096,N = 7,168,K = 4,096。S4(分组 GEMM):256 组 × M = 128,N = 4,096,K = 7,168。S5(MQA 索引器):H = 32,D = 128,查询长度 1,024,KV 长度 2,048。S6(分页索引器):B = 256,H = 64,D = 128,平均 KV 长度 4,096,块大小 64。S7(CUTLASS MLA):DeepSeek-V3,B = 128,KV 长度 4,096,页大小 128。S8(DeepSeek-V4 稀疏 MLA):B = 3,查询长度参差不齐([3, 4, 5]),H = 128,Dqk = Dv = 512。
- 论文的表格还列出了“参考实现 + 支撑代码”的行数(例如两个 FA4 kernel 分别为 3,039 和 3,690)。对 DeepSeek-V4,参考实现的行数只统计在固定的 S8 路由和运行时常量下可达的源码。
- MQA(多查询注意力)索引器为历史 token 打分,决定一个查询应该关注哪些 token,就像 DeepSeek-V3.2 的闪电索引器那样;MLA 是 DeepSeek 的多头潜在注意力。我们的 GLM-5.3 和 DeepSeek-V4 深度解读对两者都有讲解。
13.4 语料库的广度
前面几项重头实验,反而让人看不出 harness 如今撑起了多大的规模。经过验证的语料库包含 400 多个静态与编译用例,以及 399 个 GPU 正确性用例,覆盖约 28 个家族:注意力、稠密与稀疏 GEMM、MoE、量化、归一化、状态空间模型、KNN 和 KMeans,并带有从 Ampere 到 Blackwell 各代架构的专用路径。其中还有 100 多个移植自 TensorRT-LLM 的 kernel,涵盖注意力、MLA decode 和 MoE。
有一个特性格外突出:组合(composition)。由于角色、屏障和缓冲区都是显式声明的,而不是隐含的,原本要写成几个独立 kernel 的调度,可以写成一个设备端程序。BatchAttention 在一个 kernel 里同时处理 decode 和 prefill 的工作;Alpha-MoE 与 mega-MoE 家族把路由、专家计算和输出累加融合在一起,“无需把中间结果物化到内存中”。Alpha-MoE 的改写还表明,“即使目标指令变了,调度结构也能保留下来”:这里是从 Hopper 的指令换成了 Blackwell 的指令。
14. 从调好的单个形状到一个库
前面讲的一切,都是为某一个特定形状优化 kernel。库却必须来者不拒:调用者传进来什么形状,它都得接住。本节讲的这一步,基准测试类论文通常略过不提,CAKE 却把它单独当作一个阶段。
14.1 目标不同,所以自成一个阶段
“弥合这个差距,并不是让内层循环多跑一些形状就行;它是一个独立的阶段,目标不同,排序信号不同,失败模式也不同。”只盯一个确切的形状,进化就有了清晰的目标,也可以放手做激进的特化;如果改按覆盖面来给内层循环打分,这个信号就会被冲淡。所以,泛化要等每个形状都有了足够强的种子(seed)之后才开始,打分依据则是固定工作负载上连同分发器(dispatcher)在内的性能。不正确或太慢的种子会退回内层循环,“而不是藏到路由后面”。
14.2 构建分发器
泛化阶段先把测过的种子按形状分桶,生成特化的或共享的变体,再把它们的守卫条件(guard)依次排好,最后接一条显式的兜底路径(fallback)。守卫条件是对形状的判断,比如“D = 352”;兜底路径负责处理所有守卫条件都没匹配上的输入。调优可以改实现参数,但绝不能改输入形状。在报告任何汇总数字之前,验证要覆盖这几类情况:有代表性的输入和留出的(held-out)输入,边界情况和尾部情况,相互重叠或留有空隙的守卫条件,以及兜底路径本身。
有一条策略让路由不至于沦为大杂烩:这个阶段“在形状域中尽可能大的范围内复用同一个物理调度,只有当形状域要求调度发生实质性变化时,才引入另一个”。每条路由都是一个独立的 CAKE IR 程序,所以每个备选方案仍然可以单独分析、单独做基准测试。
14.3 不能自己出题自己考
这个阶段藏着一个不易察觉的危险:评测泄漏(evaluation leakage)。如果调优分发器和报告它的速度用的是同一批形状,分发器完全可以把这些形状直接背下来。CAKE 在调优之前就先声明有效的形状域。分发器的判定条件(predicate)“可以划分这个域,但不能引入讨巧的新评测条目”;覆盖范围也只能“通过同一来源中按确定性方式划出的、未见过的分片”来扩展。
14.4 在真实库上的结果
经典的机器学习工作负载给出的读数最干净,因为它们的 kernel 组合(portfolio)太大,单靠某一个形状撑不起整体结果。在 GB200 上,贡献给 FlashLib 的泛化版 KNN 构建、KNN 搜索和 KMeans kernel(pull request #15、#16 和 #18,2026 年 8 月合入)取得了下表的结果。Gspan 的算法是:在每个形状上,用参考实现的 CUPTI GPU 跨度中位数除以 CAKE 的,再对所有形状取不加权的几何平均。
| 库函数 | 形状数 | Gspan |
|---|---|---|
| KNN 构建 | 112 | 1.418× |
| KNN 搜索 | 198 | 2.116× |
| KMeans | 124 | 1.803× |
没有出现任何错误输出,KNN 的召回率为 1.0(真正的最近邻全部找到了)。论文特别提醒,不要把这些数字和 §12 的从零起步实验相比:两边的测量“使用了不同的主机、形状分布、基线和测量流程”,因此二者之间的差异“本身并不是实测得到的泛化代价”。
路由级明细(论文附录中的图)
KNN 构建(基线为 FlashLib 0.2.0):8 个外层族,共覆盖 112 个形状,实际观察到 80 条不同的最终路由。
| 外层族 | 形状数 / 路由数 | Gspan |
|---|---|---|
| D64 | 7 / 5 | 1.130× |
| D128 小 K BF16 | 38 / 17 | 1.258× |
| D128 小 K FP16 | 1 / 1 | 1.353× |
| D128 中等 K | 32 / 30 | 1.425× |
| D128 大 K | 10 / 5 | 1.909× |
| D192 | 1 / 1 | 1.432× |
| D256 | 8 / 8 | 1.386× |
| 高维 | 15 / 13 | 1.758× |
Flash-KMeans(基线为作者自己调优过的 FlashLib 实现):12 条最终路由,共覆盖 124 个形状。
| 最终路由 | 形状数 | Gspan |
|---|---|---|
| gap-fused 通用 | 64 | 1.753× |
| D288 split-K | 14 | 2.227× |
| 对齐形状兜底 | 12 | 1.439× |
| D144–176 尾部填充 | 8 | 1.786× |
| D112 特化 | 6 | 1.414× |
| D352 split-K | 6 | 3.757× |
| D480 split-K | 6 | 2.597× |
| 极小 D 混合 | 2 | 1.309× |
| 极小 D 流水线 | 2 | 1.095× |
| 大尺寸成对 | 2 | 1.217× |
| D128 特化 | 1 | 1.036× |
| D224 TMEM | 1 | 1.037× |
表中数值是每一行的几何平均,数据来自同一会话内的 CUPTI 测量。“Split-K”把归约维度切开分给多个线程块,再把它们的部分结果合并起来;当其他维度太小、填不满 GPU 时,这种做法很有帮助。
15. CAKE 在相关工作中的位置
CAKE 涉及三个研究领域:GPU 语言(§6)、编译器,以及自我改进的系统。要给它定位,最清楚的办法是对每一条研究路线都问一句:进化的是什么,固定不变的又是什么?
| 研究路线 | 代表工作 | 变的是什么 | 不变的是什么 |
|---|---|---|---|
| 带搜索或记忆的 kernel 智能体 | KernelEvolve、EvoEngineer、AVO、K-Search、KernelBlaster | 搜索过程、积累下来的记忆 | 语言和评测环境 |
| 用强化学习训练的 kernel 模型 | AutoTriton、CUDA Agent | 模型权重 | 语言和评测环境 |
| LLM 驱动的程序搜索 | FunSearch、AlphaEvolve | 程序本身(借助 LLM 写出的变异) | 评估器 |
| 自我改进的智能体 | ADAS、Darwin Gödel Machine、Meta-Harness | 智能体自身的代码,或固定模型外围的代码 | 基础模型 |
| CAKE | 编译器 harness:IR 词汇、分析、代价校准 | 编程智能体和模型 |
用论文的话说,CAKE“瞄准的是与之互补的那一层:它改变的是被搜索的表示形式,以及返回给智能体的结构化编译器证据”。两者完全可以结合:更强的搜索方法,或者训练得更好的模型,都可以放进 CAKE 的环境里运行。
编译器这一侧,TVM、XLA、MLIR、TensorIR、Ansor 和 MetaSchedule 把程序组织成结构化的形式,以便自动分析和优化;Graphene、Twill 和 Tawa 则对 GPU 的异步执行、流水线或 warp 专用化进行建模。CAKE“认同结构化程序能支撑有用的分析这一原则,但把编译器的发现放进了智能体的进化循环,并让反复出现的 kernel 证据推动 harness 进化”。更广义地看,它是“自我定义系统”(self-defining systems)这一研究议程的一个实例:由 AI 运作、能够改变自身机制和抽象的系统。
16. 局限与开放问题
好论文会交代自己止步于何处。CAKE 交代了,细心的读者还能再补上几个问题。
16.1 论文自述的局限
- 只支持 NVIDIA。CAKE 面向从 Ampere 到 Blackwell 的各代架构。调度语言、角色模型和各项分析可以在这几代之间通用,但指令形式、合法性规则和代价锚点都与具体架构绑定,而“这部分成本才是衡量迁移的诚实尺度”。对于非 NVIDIA 硬件,这一成本“仍未测量”,而且后端得重新构建。
- 证据分布不均。大部分性能结果来自 B200。耗时模型只针对 B200 和 H100 校准过,在其他硬件上则“拒绝预测”。
- 分析并不完备。静态分析和性能模型是有意做得不完备的:它们负责排序和过滤,GPU 上的实际执行始终是最终依据。
- 人仍在回路中。编译器进化“在合并门禁处仍由人来引导”。
16.2 论文留下的开放问题
以下问题是我们提出的,并非出自论文:
- 别人用得上吗?论文没有说 CAKE 的编译器、IR 或 harness 已经公开发布;截至 2026 年 10 月,我们也没有找到正式发布(只有一些非官方的社区复现)。公开的是生成出来的代码:四个 FlashInfer pull request 和三个 FlashLib pull request 都已合入,FlashInfer 还有一个跟踪 issue,索引了更多由 CAKE 生成的 kernel。
- 对照实验的结果有多普适?从零起步的对比,每组只跑了三次,而且只用了一个 kernel、一个形状和一个模型(GPT-5.6-sol)。换成别的 kernel、别的模型,或者给更大的预算,差距会怎样变化,我们不得而知。(FlashInfer 的跟踪 issue 提到,CAKE 生成的 kernel 也出自 Claude 和 Codex 智能体之手,这说明 harness 并不绑定某一个模型;但没有报告任何受控对比。)
- 到底是哪一部分在起作用?两组之间,语言、编译前检查、代价模型和编译器进化是同时改变的。论文没有做这样的消融实验:比如保留 CAKE IR 但去掉代价模型,或者把编译器冻结住。
- 为什么没有 tile 级语言这一组?对照组写的是 CUDA/PTX。如果在同样的规则下再加第三组,改用 Triton 或 CuTe DSL 来写,就能看出 tile 级语言和布局代数会落在什么位置。
- 编译器进化的代价有多大?论文没有报告编译器一共改了多少次、每条路径各贡献了多少、语料库门禁拒绝改动的频率,以及合并门禁耗费了多少人工审查。
- 这些检查有多准?论文没有报告验证器的误报率和漏报率,也没有报告代价模型预测的准确度。
- 缺失的数字,以及两处出入。Gated DeltaNet 和 MiniMax 稀疏注意力的结果只有文字描述,没有给出加速比;平台期判据在正文中也没有定义。另有两份上游记录与论文的说法不一致:已合入的 Alpha-MoE pull request 报告的是另一个基线和另一组加速比(§13.2);KDA prefill 的 pull request 报告的是容差范围内一致,而论文说的是“逐位正确”。
这些问题都不动摇论文的核心思想。它们标出的,正是下一篇论文或一次独立复现最能有所贡献的地方。
尾声 —— 一段对话
教授 好了。你都学到了什么?
学生 快的 kernel,很大程度上就是一份好的调度:谁来加载,谁来做乘法,彼此之间怎么交接 tile,每个字节放在哪里。
教授 那智能体为什么写得这么吃力?
学生 因为它们用的语言,要么把调度藏起来,要么逼它们亲手管理每一个地址。得到的反馈,也只有一句“崩了”,或者一个数字。
教授 那 CAKE 怎么解决?
学生 一门语言,决策都写在明面上,琐碎的记账交给编译器去推导;一套检查,GPU 还没开跑,就能点出是哪个屏障坏了;还有一个编译器,会把反复出现的失败变成新的检查。智能体变强,是因为它所处的世界变好了。
教授 那你还想知道些什么?
学生 换成别的模型、别的 kernel,结论还成不成立;哪一部分最关键;还有,我能不能自己动手试试。
教授 很好。现在你算是会读论文了。
全文小结
- 问题。专家级 GPU kernel 胜在调度(warp 角色、屏障编排、内存放置),而调度一旦出错,就会悄无声息地算出错误结果,或者干脆挂死。kernel 智能体在搜索、记忆和模型上都有了改进,但它们写代码用的语言,仍然要么把调度藏起来,要么把调度暴露出来却不设护栏;它们得到的反馈也只有报错、通过与否和耗时。
- 思路。智能体和它的工作环境要协同设计(co-design)。智能体编写 CAKE IR:一种带类型的调度语言,硬件决策都摆在明面上,也不用布局代数;资源、角色和交接都要显式声明,机械性的细节则自动推导出来。harness 在编译之前检查程序安全、硬件合规、数据一致性和调度语义,用校准过的代价模型给候选 kernel 排序,并以 GPU 实测为最终依据。编译器也在进化:反复出现的失败和新发现的专家模式,会变成验证器规则、原语和校准;每项改动都要通过语料库的把关,其中包括 400 多个静态与编译用例,以及 399 个 GPU 正确性用例。
- 证据。模型固定时,在从零起步的 Flash-KMeans 任务上,CAKE IR 胜过直接写 CUDA/PTX(80M token 时,速度中位数分别为调优基线的 1.144× 和 0.928×;进入平台期的运行分别为 3/3 和 0/3;有效进化时间约为后者的一半)。在看不到参考实现的情况下,它写出的 KDA prefill kernel 比官方 kernel 快 2.05×(六个形状上的几何平均),并在端到端推理服务中通过了验证。它追平或超越了 11 个专家 kernel 中的 10 个;有分发器支撑的 KNN 与 KMeans kernel 家族,在 400 多个形状上快 1.42×–2.12×。四个 kernel 已合入 FlashInfer。
- 局限。只支持 NVIDIA GPU;证据大多来自 B200;受控实验规模小;分析在设计上就不完备;合并环节仍由人工把关;CAKE 本身没有公开发布(它生成的 kernel 是公开的)。
思考题
《操作系统导论》(OSTEP)每一章都以作业收尾。我们的作业没有配套的模拟器,但下面每道题都值得花几分钟想一想:
- 在 §3.2 中,三级环形缓冲区给每个屏障配了一个 1 位的相位。为什么 1 位就够了?如果生产者能领先消费者整整两圈,会出什么问题?
- CAKE 拒绝把调度降级到更老的 GPU 上(§7.4)。如果编译器悄悄地把一份 Blackwell 调度 lowering 成 Hopper 指令,可能会出什么问题?
- 编译前检查允许出现误报(§9)。在进化循环里,一项检查错误地拒掉了 5% 的合法 kernel,另一项漏掉了 5% 有问题的 kernel,哪一项更糟?如果换成生产环境用的编译器,你的答案会变吗?
- 请把 §16.2 中缺失的那个消融实验设计出来:你会跑哪几组实验,才能把语言、harness 和编译器进化各自的作用区分开?每一组又要固定哪些条件?
- §14.3 禁止分发器自行添加评测形状。请举一个具体例子:如果打破这条规则,分发器怎么会看起来比实际快 2×?
- 论文借用了陶哲轩的观点:生成变得廉价之后,瓶颈会转移到验证和理解上(§8.2)。在机器学习系统里,你还在哪些地方看到了同样的转移?
参考资料与延伸阅读
论文
- 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 和 Luis Ceze(NVIDIA 与卡内基梅隆大学)。除非另有标注,本文的每个数字都出自这篇论文。设计部分请读论文 §2–3,八条原则和完整的 kernel 表格见附录。
已合入上游的 kernel
- FlashInfer 的 pull request #4262(KDA prefill)、#4279(KDA decode)、#4274(TinyGEMM)和 #4287(Alpha-MoE),均已合入:都是生成的 CUDA 代码,不需要 CAKE 也能直接阅读。CAKE kernel 跟踪 issue 里还列着更多。
- FlashLib 的 pull request #15、#16 和 #18:§14 中带分发器的 KNN 与 KMeans kernel。
GPU 与 kernel 背景知识
- FlashAttention(Dao et al., NeurIPS 2022):最清楚的例证,说明让一个核心运算快上好几倍的,可以是更好的调度,而不是新的数学。它的后续版本一路到面向 Blackwell 的 FlashAttention-4,就是一部浓缩的 GPU 特性演进史。
- CUTLASS 与 CuTe DSL:NVIDIA 用来写专家级 kernel 的模板库,以及 CAKE 有意不用的布局代数。
- Triton(Tillet et al., MAPL 2019):以 tile 为单位编程的语言,降低了写 kernel 的门槛。
- TVM(Chen et al., OSDI 2018):把张量程序“算什么”和“怎么调度”分开,正是这个编译器让这种做法流行开来。
kernel 智能体与进化搜索
- KernelBench(arXiv:2502.10517):确立了“PyTorch 算子进、快速 kernel 出”这一任务形式的基准。
- KernelEvolve(arXiv:2512.23236):Meta 面向异构加速器的智能体式 kernel 编程(本博客作者是论文合著者之一)。
- AlphaEvolve 与 FunSearch(Nature, 2023):由 LLM 驱动的进化式程序搜索。
- AVO、K-Search、KernelBlaster、CUDA Agent 和 EvoEngineer:近期的几种 kernel 智能体,在搜索、规划、记忆和训练上各有思路。
工作负载
- Flash-KMeans(arXiv:2603.09229) 与 Sparse VideoGen2(arXiv:2505.18875):k-means 这个工作负载,以及视频生成为什么需要它。
- Kimi Linear(arXiv:2510.26692)、Kimi-K3 与 Gated DeltaNet(ICLR 2025):前沿 kernel 背后的线性注意力架构。
- Alpha-MoE 与 FlashInfer:CAKE 重写的融合 MoE 巨型 kernel(megakernel),以及接收这些成果的 kernel 库。
写作风格
- Operating Systems: Three Easy Pieces,Remzi H. Arpaci-Dusseau 与 Andrea C. Arpaci-Dusseau 著,可在线免费阅读(中文版书名《操作系统导论》)。本文的对话,以及“关键问题”“补充”“提示”这几种方框,都借鉴自这本书的体例,在此致谢。
本站相关文章
- 逐层拆解 GLM-5.3 与 GLM-5.3-Flash 和 逐层拆解 DeepSeek-V4 与 V4-Flash:§13 中那些 kernel 对应的注意力变体(MLA、闪电索引器、KDA),这两篇都有讲解。
- Miles v0.1 深度解析:另一篇系统论文精读,主题是强化学习基础设施。