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 与卡内基梅隆大学的一篇论文。它认为,光改进智能体还不够。kernel 智能体和 GPU 编程语言这两个方向的研究,一直各自独立推进。用论文的话说,“专家级 kernel,正是丢在了两者之间的空白里”。这片空白就是智能体所处的世界:它们写代码用的语言,以及它们收到的反馈。今天的智能体写的要么是原始 CUDA,要么是某种高层 kernel 语言,而环境的回应,不外乎“崩溃了”“输出不对”或者“耗时 0.94 ms”。CAKE 把这两头都改了。智能体改用一种新的、可以检查的语言来写 GPU 调度;编译器则给出诊断,直接指出受影响的资源、角色或流水线级。(论文没有公开诊断消息的格式;下面这个例子是我们按同样的思路自己编的:这个流水线级里的这个屏障,有人在等,却始终没人发信号。)CAKE 还往前多走了一步:如果智能体总在同一类问题上栽跟头,就去改编译器本身,把这种失败变成一项新的检查、一个新的语言特性,或是一个修正过的性能模型。

用论文自己的数字来说,核心结果如下:

1.144× 对 0.928×从零起步的 Flash-KMeans耗用 80M token 后的最佳 kernel,CAKE IR 对比直接写 CUDA/PTX,以调优过的基线为参照(3 次运行取中位数,B200)
2.05×Kimi Delta Attention 的 prefill处理提示词(prompt)的 kernel,对比官方 FlashKDA kernel,在 6 种输入尺寸上取平均;已在端到端推理服务中验证
1.42×–2.12×KNN 与 KMeans 库三个 kernel 家族,各自按输入尺寸挑选 kernel;在 112、198 和 124 种尺寸上的平均值分别为 1.418×、2.116× 和 1.803×(GB200);单个尺寸上的加速比波动更大
4 个 PR提交到 FlashInfer 上游KDA 的 prefill 与 decode、TinyGEMM,以及一个 Alpha-MoE 巨型 kernel(megakernel);据 GitHub 记录,四个 PR 截至 2026 年 9 月 13 日均已合入

这些标签现在看不太懂也很正常,每张卡片后面都有专门的一节来讲(第一张见 §12,第二张和第四张见 §13,第三张见 §14)。有两个词全文随处可见。形状(shape)指一组具体的输入尺寸,比如 32 个批次、每批 65,536 个点。如果一个结果涵盖许多形状,各个形状上的加速比就用几何平均(geometric mean)来汇总。这是对比值求平均的标准做法:一次 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 的 pull request,以及 Kimi-K3 的模型卡(model card)。这些资料在用到的地方都附有链接。除非另有说明,数字都来自论文。从图上读出的数值标为“近似值”;我们自己算出的数字标为“推算”;我们自己对证据的解读标为“我们的理解”。论文没有交代的地方,我们会直接指出,并在 §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。我用得着 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 上。
关于 NVIDIA B200 的四个面板。一个 warp 是 32 个线程,执行同一条指令,各自处理自己的数据。warp 组成线程块,一次 kernel 启动的线程块由硬件分散到各个流式多处理器(SM)上;B200 有 148 个 SM。每个 SM 内部有寄存器(256 KB)、张量内存(256 KB)、共享内存(最多 228 KB)和 Tensor Core,也就是矩阵乘法单元。内存阶梯从寄存器和张量内存(每个 SM 各 256 KB)开始,经过共享内存(每个 SM 最多 228 KB)和全芯片共享的 126 MB L2 缓存,一直到带宽最高 8 TB/s 的 180 GB HBM:片上每个 SM 只有几百 KB 的高速存储,往外则更大也更慢。
第一部分的词汇浓缩在一张图里:32 个线程组成 warp,warp 组成线程块,硬件把线程块分散到 B200 的 148 个 SM 上;各级存储构成一架阶梯,从片上的高速存储(每个 SM 几百 KB)一直到更大也更慢的 L2 和 HBM。打开原尺寸 SVG。

1.2 为什么矩阵乘法是最要紧的 kernel

向量加法很简单,因为每个输出恰好只需要两个输入。矩阵乘法则不同。要计算输出矩阵 C = A × B,每个输出元素都需要 A 的一整行和 B 的一整列,而每个输入元素又会被许多个输出用到。这种复用既是机会,也是难点:好的 kernel 把每块输入只加载一次,然后反复使用;差的 kernel 则一遍又一遍地从远处取回同样的数。

Transformer 做的几乎每件事都建立在矩阵乘法之上:注意力里的投影、前馈层,以及混合专家(MoE)模型里的专家。现代 GPU 为矩阵乘法配备了专门的硬件单元,叫作 Tensor Core(张量核心),一条指令就能完成一次小矩阵块的乘法。写 kernel 的人把矩阵乘法叫作 GEMM(general matrix multiply,通用矩阵乘法),把 Tensor Core 执行的这种运算叫作 MMA(matrix multiply-accumulate,矩阵乘加,即 C += A × B)。CAKE 论文里的大多数 kernel,说到底都是在精心设法让 Tensor Core 始终“吃饱”。

2. 快的 kernel,拼的是数据搬运

先说一个让新手意外的事实:在现代 GPU 上,算术本身很少是难点,难的是及时把数送到运算单元手里。

2.1 内存阶梯

GPU 上有好几种内存,像梯子的横档一样一级一级排开。离运算单元越近的那一级,速度越快,大体上容量也越小:SM 内部的几级,每个 SM 只有几百 KB;SM 之外的几级,则以 MB 乃至 GB 计。

梯级 位置 容量(B200) 用途
寄存器 每个 SM 内部 每个 SM 256 KB,分给 SM 上的各个线程 每个线程正在使用的数值
张量内存(TMEM) 每个 SM 内部,数据中心版 Blackwell(B200、GB200)新增 每个 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 运算。(BF16 即 bfloat16,是一种 16 位浮点格式。深度学习爱用它,因为它比普通的 32 位浮点数(FP32)少一半内存流量,代价是损失一些精度,而这点损失神经网络承受得住。FP8 和 FP4 则进一步降到 8 位和 4 位。“稠密”指的是不借助稀疏技巧时的算力;NVIDIA 宣传的标称数字,正是靠这种技巧翻了一倍。)两者一除就知道:要让 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 的一整行或一整列复用:用 128 × 128 的 tile 时,每个数会被用上 128 次,而不是只用一次。

分块也是 FlashAttention 的核心。FlashAttention 算的注意力和别人完全一样,但它一块一块地处理输入,并把各步运算融合在一起,于是那张巨大的注意力分数中间表(每一对 token 对应一个数)根本不必写入 HBM。

分块又带来一个新问题:线程块等待下一个 tile 送达时,Tensor Core 只能闲着。解决办法是准备两个或更多的 tile 缓冲区,在计算当前 tile 的同时提前去取下一个 tile。这叫作双缓冲(double buffering),推而广之就是软件流水线(software pipelining)。它的作用是掩盖内存延迟,而不是减少数据流量。后来的注意力 kernel 大量依赖这一技术,比如为 Hopper 编写的 FlashAttention-3,整个设计就是围绕“让复制与乘法重叠进行”展开的。

3. GPU 变成了一条装配线

近几代 NVIDIA GPU 加入的新硬件,让分块和流水线快了许多,也让编程难了许多。本节梳理这一变化,因为 CAKE 的设计正是紧扣着它展开的。

3.1 从“每个线程什么都干”到各司其职

在较早的 GPU kernel 里,每个 warp 什么活都干:加载 tile 的一部分,等待线程块里的其他线程,计算,如此反复。此后,每一代 GPU 都加入了新硬件,能在线程忙别的事情时,异步地独立完成其中某一项工作。下面这些名字不必记住,要紧的是其中的趋势:每一代都把又一项工作交给一个在后台运转的专用单元,于是线程不再亲自干活,而是转为协调这些工作。后文还会再出现的名字是:TMA(复制引擎)、TMEM(存放累加器的内存),以及宣告复制完成的屏障(barrier)。

  • Ampere(A100,2020 年)新增了 cp.async,它在后台把数据从全局显存复制到共享内存。这一代 GPU 的 Tensor Core 由 warp 级的 mma.sync 指令驱动,所需数据由 ldmatrix 从共享内存载入。
  • Hopper(H100,2022 年)新增了:Tensor Memory Accelerator(TMA),一个复制引擎,一条指令就能搬运一整个多维 tile;WGMMA,一种由 warpgroup(4 个 warp)共同发射的 Tensor Core 指令;异步事务屏障(一种 mbarrier,复制的数据陆续送达时,它会逐一清点字节数),用来在 tile 到齐时发出信号;以及线程块集群(cluster),同一集群里的线程块可以读取彼此的共享内存。
  • Blackwell(B200,2024–25 年)新增了:tcgen05.mma,一种由单个线程代表整个线程块发射的 Tensor Core 指令;以及张量内存(TMEM),即每个 SM 上一块 256 KB 的专用存储,专门存放累加器,累加器从此不再挤占寄存器。两个相邻的 SM 甚至可以协作完成一次更大的乘法(即“2-CTA”模式)。

既然复制引擎和 Tensor Core 都能自己运转,最自然的设计就是一条装配线:一些 warp 只负责发起复制(即生产者);一个 warp 只负责发射矩阵乘法(即 MMA warp);其余的 warp 接过算好的结果,做完最后的逐元素处理(比如缩放或激活函数),再把结果写出去。最后这个阶段叫作 epilogue(收尾阶段)。这种设计叫作 warp 专用化(warp specialization),每一项工作就是一个 warp 角色。

一个 SM 上 warp 专用化的矩阵乘主循环动画。HBM 中的第 k、k+1、k+2 块由 TMA 拷进共享内存里的三级环形缓冲区,复制由生产者 warp 发起。每个 stage 有一个 full 屏障(数据已到,可以读)和一个 empty 屏障(数据已用完,可以覆盖)。第 1 拍:生产者把第 k 块拷进 stage 0,并在它的 full 屏障上 arrive。第 2 拍:MMA warp 等待 full,把 stage 0 乘进张量内存中的累加器,然后在 empty 上 arrive;与此同时第 k+1 块正被拷进 stage 1。第 3 拍:稳态,第 k+2 块正被拷进 stage 2、第 k+1 块在计算、stage 0 重新空出;最后一块算完后,epilogue warp 把结果写回 HBM。第 4 拍:stage 0 已装入第 k+3 块;如果生产者跳过对 empty 屏障的等待,第 k+4 块就会在 MMA warp 还在读取第 k+1 块时覆盖 stage 1,这种数据竞争会悄无声息地算错,或者让程序挂死。
把 warp 专用化的主循环看成装配线:生产者填充环形缓冲区的各个 stage,full 屏障表示数据已到;MMA warp 消费数据,empty 屏障把 stage 腾出来;少一次 wait,第 k+4 块就会覆盖一个还在被读的 stage。打开原尺寸 SVG。

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

装配线上的工位之间,得有办法交接工作。在 warp 专用化的 kernel 中,共享内存里的 tile 缓冲区组成一个分为若干级(stage)的环形缓冲区。以三级为例:生产者在填充第 2 级的同时,Tensor Core 在处理第 1 级,而刚刚腾出来的第 0 级正等着下一个 tile。

这些交接靠的是屏障。屏障是一个碰头点,设立时就规定好要等多少次到达,这个次数叫作它的到达计数(arrival count)。在屏障上等待的 warp 会一直被拦住,直到到达的次数凑够为止(对 TMA 复制来说,就是直到复制的字节全部送达);随后,屏障完成当前相位,并重置,准备进入下一个相位。NVIDIA 把“向屏障发信号”叫作在屏障上到达(arrive),本文的图里用的也是 arrive 这个词。

每一级都有两个屏障。生产者的复制一送达,就在这一级的 full 屏障上到达;消费者读取之前,先在 full 上等待。消费者用完这一级后,在这一级的 empty 屏障上到达;生产者覆盖这一级之前,先在 empty 上等待。由于环形缓冲区会绕回起点,每个屏障每转一圈就被复用一次,所以等待方必须说明自己等的是哪一圈。硬件用一个 1 比特的相位(phase)来记录这一点,它每转一圈翻转一次。如果等错了相位,等待要么过早返回,要么永远不返回。

还有两个细节后面会用到。GPU 访问内存有两条相互独立的路径:一条是线程发出的普通加载和存储,叫作通用代理(generic proxy);另一条是复制引擎和 Tensor Core,叫作异步代理(async proxy)。不妨把它们想象成往同一张办公桌送件的两个快递员:两人各送各的,谁也不能保证,一个人看到的对方送件顺序,就是对方实际送件的顺序。所以,如果线程用普通存储写下一些数据,接下来又要由某个异步操作来读(比如共享内存里的数值,随后要由 TMA 写回全局显存,或者交给 Tensor Core 做乘法),kernel 就必须显式插入一个代理栅栏(proxy fence),否则硬件可能会打乱这些访问的顺序。反方向的情况则不同:只要通过屏障观察到异步复制已经完成,数据就已经可见了。另外,每次 TMA 复制都需要一个描述符(descriptor),它编码了 tile 的形状、步长(stride,即相邻的行、相邻的列在内存中相隔多远)以及布局。

3.3 为什么这么容易出错

这些细节,每一处都可能出错,而且出起错来格外折磨人:

  • 漏掉一次 wait,消费者就会在数据送达之前读取这一级。程序不会崩溃,结果却是错的,而且有时只在部分输入上出错。
  • 屏障初始化时把到达计数设得过高,或者在一个永远不会完成的相位上等待,都会让某个 warp 永远等下去。kernel 就此挂死,唯一的症状是 GPU 不再响应。(反过来,计数设得过低,或者等的是一个过时的相位,等待就会提前返回,bug 于是变成悄无声息的数据竞争。)
  • 漏掉代理栅栏,只会在特定的时序下读到过期数据,于是 bug 这次运行出现,下次运行又不见了。
  • TMA 描述符里的 tile 形状如果与它要填充的共享内存缓冲区对不上,就会破坏属于别人的内存。

也有工具能帮上忙。NVIDIA 的 compute-sanitizer 可以在运行时检测出部分数据竞争和越界访问,但只限于你碰巧运行的那些输入,而且它不会告诉你问题出自哪一个设计决策。请记住:崩溃或挂死只告诉你出了问题,不告诉你为什么。这个缺口,正是第二部分的主题。

4. 专家级 kernel 与正确的 kernel 差在哪里

把这些拼在一起,就能看懂 CAKE 论文所说的专家级 kernel 是什么意思。论文摘要点出了三项决策,正是它们“把专家级 kernel 与仅仅正确的 kernel 区分开来”:

  1. warp 专用化。哪些 warp 担任哪些角色:几个生产者、几个消费者,两组消费者是否轮流上阵(“乒乓”,ping-pong),一组计算时,另一组写出结果。
  2. 屏障编排。流水线分几级,哪个屏障把守哪一次交接,每次等待等的是哪个相位,栅栏插在哪里。
  3. 内存层级放置。哪些数据放在寄存器里,哪些放在共享内存里,哪些放在张量内存里,以什么布局存放,好让复制引擎、Tensor Core 和输出路径对每个字节的位置看法一致。

这些决策合起来,构成 kernel 的调度(schedule):它规定怎样驱动这台机器,而不是计算什么。矩阵乘法的数学一行就能写完;一个快速的 Blackwell 矩阵乘法,调度却要写上几百行,而且改动其中一项决策,通常会牵动其他决策一起改。

第一部分小结

kernel 是由成千上万个 GPU 线程并行运行的函数。它的速度主要取决于数据搬运:数据以 tile 为单位,沿着内存阶梯一级级移动,并且在处理当前 tile 的同时就加载下一个。现代 NVIDIA GPU 把这一切变成了一条装配线,由各司其职的 warp 通过屏障相互传递 tile。专家的功夫在于调度(角色、屏障、内存放置)。调度一旦出错,换来的往往是悄无声息的错误结果或挂死,而不是一条有用的报错信息。


第二部分 · 语言 —— 一段对话

学生 也就是说,智能体写一个 kernel,跑一跑,看看结果。这有什么问题?我自己调试也是这么做的。

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

学生 看调用栈,看行号,看哪个变量成了 None。

教授 现在想象一下,遇到最棘手的那些 bug 时,你唯一能得到的反馈是“卡住了”“结果不对”或者“用时 0.94 毫秒”。没有行号,也没有变量。

学生 那我只能乱改一通,碰碰运气。

教授 面对最难的那些 kernel,kernel 智能体大致就是这个处境。所以有两个问题:智能体该用什么语言来写?出了错,编译器又该告诉它什么?

学生 像专家那样直接写 CUDA 不就行了?

教授 CUDA 什么都写得出来,包括每一种错误。更高层的语言会保护你,可专家写得出来的东西,它不让你写。论文的答案试图两者兼得:一门语言,既能写出专家要写的东西,写出来的形式又能让编译器在任何代码运行之前先检查一遍。我们来看看它是怎么做的。

5. 今天的智能体如何写 kernel

写 kernel 的智能体出现得很快。本节介绍它们共用的那个循环、这个循环擅长什么,以及 CAKE 的出发点:反馈问题。

5.1 “提出–测试–测量”循环

这个任务设定因 KernelBench(2025 年)而流行开来:给模型一个 PyTorch 算子,让它写出一个计算结果相同、但更快的 GPU kernel。此后,许多系统搭起了类似的循环,有的面向 GPU,有的面向其他 AI 加速器。

其中很多系统借用了生物学里进化(evolution)的思想,CAKE 标题里的“进化”也是这个意思。具体做法是:维护一组候选 kernel,让 LLM 改动(“变异”)其中有希望的那些,对每个新候选都做测试和计时,留下正确且最快的几个,如此循环往复。各个系统的区别,在于改进的是这个循环的哪一部分:

  • 循环本身。由 LLM 根据编译错误、测试结果或性能分析器(profiler)的输出修改 kernel(AccelOpt,面向 AWS Trainium;Autocomp,面向张量加速器;KernelAgent,面向 GPU)。
  • 记忆。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 指令的发射速率?是流水线太浅,还是某个屏障让两个角色串行执行了?单凭一个数字,回答不了。
  • 性能计数器。有些循环还会把性能分析器的输出一并交给智能体。它能说明 kernel 是在等显存还是在等 Tensor Core,却说不出是哪个程序决策(流水线级数、角色划分、缓冲区放在哪里)造成了这种局面。

还有第二个更隐蔽的问题:这些信号“在前沿工作负载暴露出能力缺失时,无法随之扩充”。如果语言表达不了新工作负载需要的调度,每次尝试都会因为同一个结构性原因失败,而环境始终无法从中吸取教训。

5.3 CAKE 的三个核心主张

论文的观察是:专家装在脑子里的那套机制,编译器里大多已经有了,即“结构化的操作词汇表、资源模型、合法性检查、静态分析、代价模型、lowering 规则”。(lowering 就是编译器把程序逐层向下翻译成 GPU 自身代码的过程,见 §7.1。)CAKE 要问的是:“如何让这套机制面向智能体,以及在前沿工作负载暴露出缺口时如何改进它。”它的回答分三部分:

  1. 智能体编辑的是带类型的中间表示(IR),而不是原始 CUDA。IR 是一种专为编译器分析而设计的程序格式。CAKE 的 IR 是一门调度语言,硬件决策在其中都是显式的,因此在生成任何代码之前就能检查这些决策。
  2. 编译器返回精确定位的诊断,而不是一个“通过/失败”比特。正确性和安全性方面的诊断会指向受影响的资源、角色或流水线级,并指出违反的是哪一类规则;性能报告则估算开销,并指出瓶颈大致属于哪一类。轻量级的分析会在候选 kernel 消耗 GPU 时间之前,先把它们筛一遍。
  3. harness(测评环境)本身也会进化。所谓 harness,就是智能体周围负责评判其工作的一切:各项检查、测试和基准测试(见 §9)。其中有两个部分在这里很重要。验证器(verifier)读取调度,在编译任何代码之前就把不合法的调度拒之门外,作用很像代码检查工具(linter)。代价模型(cost model)不用真正运行候选 kernel,就能预测它会跑多快。这有点像天气预报:成本低,通常八九不离十,但最终还得到真实的 GPU 上核实。如果智能体总以同一种方式失败,harness 就会随之改变:可能是一条新的验证器规则、一个新的原语(语言内置的操作)、一个重新校准过的代价模型(用新的实测数据重新拟合),也可能是一条可复用的优化策略。每一项改动都要经过 kernel 语料库上的测试把关。
两个面板。第一个是今天常见的只进化 kernel 的做法:智能体写出一个 CUDA 候选,经过编译、测试和计时;收到的反馈只有编译错误、挂死或结果错误,或者一个时间(比如 0.94 ms),也就是通过/不通过外加一个数字,只说明出了问题,不说明原因。第二个是 CAKE:智能体写出 CAKE IR 候选,harness(静态检查、代价模型、代码生成和 GPU 运行)给出定位到具体位置的诊断,例如屏障 q_full 在 stage 2 上有人等待却没人 arrive,或者估计的瓶颈是 Tensor Core 的发射速率。外层循环把能力缺口变成编译器进化,并更新 harness 的 IR、分析和 lowering。内层循环进化 kernel,外层循环进化编译器。诊断文字仅为示意。
今天的 kernel 智能体只收到通过/不通过和一个时间;CAKE 的 harness 返回定位到具体位置的诊断,外层循环再把能力缺口变成对编译器的修改。图中的诊断文字仅为示意。打开原尺寸 SVG。

论文把这种做法称为编译器–智能体协同设计(compiler–agent co-design)。智能体所处的环境(语言、分析、代价模型)围绕智能体的需要来设计,主要由智能体维护,并随着它们遇到的问题不断变化。智能体本身,也就是它的模型和脚手架,则保持不变(见 §10.2)。接下来三节讲第一个主张:语言。

6. 智能体该用哪种语言?

写 GPU kernel 的方式有很多,它们最大的区别在于:调度由谁来定,程序员还是编译器?本节按 CAKE 论文自己的分类,把这份“菜单”逐项过一遍。阅读时,请带着这样一个问题:

6.1 高层:tile 级语言

Triton(2019 年)让程序员以 tile 为单位思考,从而降低了写 kernel 的门槛:你只写矩阵的一个块上要做什么,至于线程、各级内存和指令具体怎么执行,交给编译器决定。更新的语言,比如 Helion、TileLang 和 NVIDIA 的 cuTile,在不同的抽象层次上沿用了类似的思路。

这对人类来说非常好用,写出来的代码正确,而且往往很快。CAKE 论文的异议很具体:tile 级语言(tile DSL)“隐藏了 warp 专用化、屏障编排和内存层级放置,而正是这些把专家级 kernel 与仅仅正确的 kernel 区分开来”。用 Triton 写代码的智能体可以要求流水线分几级(num_stages=3),在 Blackwell 上还能打开自动 warp 专用化。但它没法说明哪些 warp 担任哪个角色、哪个屏障把守哪次交接,也没法要求累加器放在 TMEM 里。这些选择都由编译器来做,智能体只能指望编译器选得和专家一样。

6.2 低层:CUTLASS、CuTe DSL 与原始 CUDA

另一端是 NVIDIA 的 CUTLASS 库及其 Python 前端 CuTe DSL。它们把硬件提供的一切都暴露出来,代价是要掌握一套布局代数(layout algebra),也就是一个描述张量的逻辑坐标如何映射到内存和线程上的数学框架。再往下是带内联 PTX(NVIDIA 的一种类似汇编的虚拟指令集)的 CUDA C++,地址、屏障相位位和描述符编码都要你亲手管理。

论文认为,CuTe DSL 这类低层 DSL“暴露了这种控制,却要求掌握一套布局演算,使得智能体的错误既容易发生,又难以定位”。很多布局错误照样能通过编译。CuTe 会拒绝维数(rank)和尺寸不匹配的布局,但如果布局的形状对了,步长或 swizzle(地址交错)却错了,正确的数就会落到错误的位置上,直到后面才暴露出来:要么输出错误,要么 kernel 很慢。

6.3 两者之间:Gluon

Gluon 介于两大阵营之间。它复用 Triton 的编译器栈,但在布局、数据搬运和异步执行上开放了更底层的控制。它把刻度盘往“控制”一侧拨了拨,也就一并拨向了控制带来的负担。

6.4 CAKE 瞄准的空白

把这些选项并排一放,空白就显现出来了。能表达专家级调度的语言,要你自己管理布局和各种底层细节,而且很多最要命的错误(等错了相位、漏了一个栅栏、布局形状对了映射却错了)要到运行时才会暴露。能保护你的语言,则把调度的大部分留在编译器手里,最多让你给些提示,比如流水线分几级。

一张简化表格(沿用 CAKE 论文的分类)向四类 GPU kernel 语言(自上而下从高层到底层)以及单独列在下方、以蓝色突出的 CAKE IR 提出三个问题:warp 角色、屏障和存储层级是否由你来写;是否需要布局代数;程序运行前是否会被检查。Tile DSL(Triton、Helion、TileLang、cuTile):否,由编译器决定;不需要;部分,有类型检查。Gluon(复用 Triton 的编译器栈):是;部分,显式布局;部分,编译器检查。CuTe DSL / CUTLASS(Python DSL 与 C++ 模板):是;需要,CuTe 布局;部分,形状会提前检查,同步问题只在运行时暴露。CUDA C++ + 内联 PTX(PTX 即 GPU 汇编):是;不需要,但要算裸地址;否,跑了才知道。CAKE IR(带类型的调度语言):是;不需要,直接写具体偏移;是,编译之前就由验证器和代价模型进行分析。高层语言写得对,却写不出专家级调度;底层语言什么都写得出,什么错也犯得了;Gluon 介于两者之间;CAKE IR 调度显式、无需布局代数,并且运行前先检查。前两列沿用 CAKE 论文的分类,第三列是我们的简化;每个工具的实际情况都比表中更细致。
谁来决定调度?各类 GPU 语言在控制权、布局代数和运行前检查上的位置(前两列沿用 CAKE 论文的分类,第三列是我们的简化)。打开原尺寸 SVG。

7. 逐行读懂 CAKE IR

CAKE IR 就是智能体所写的语言。要理解它,最好的办法是直接读一段代码。

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

这一设计建立在一种分工之上。用论文的话说:“调度声明要发生什么,lowering 推导如何发生。”调度记录的是:哪些 warp 担任哪些角色,哪些缓冲区分几级,哪个屏障把守哪次交接,哪种指令形式消费哪个操作数。lowering 是编译器把程序向下翻译成 CUDA 和 PTX 的过程,它负责算出这些决策的机械性后果:“屏障地址、相位位、TMEM 偏移、描述符编码和 warp 身份,全都根据声明计算得出,而不是由智能体逐一写出。”

下面是论文用来说明这一思路的片段。论文只称它为“一个 CAKE IR 调度片段”。从名字 fmha_fwd(fused multi-head attention forward,即融合多头注意力的前向计算)和它对张量内存的使用来看,它仿照的应该是某个 Blackwell 注意力 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))

分四步来读:

  • 函数签名。@cake.schedule() 把这个函数标记为一个 CAKE 调度。lm 是构建器(builder),所有资源、角色和操作都通过它来声明,就像你写 np. 或 torch. 那样。从名字推断,Q: LM.tma3d 是一个由 TMA 复制引擎读取的 3 维张量,O: LM.tma2d 是一个 2 维输出,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 角色:一个负责加载,一个负责发射 Tensor Core 指令。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: 中,加载 warp 为每一级发起一次 TMA 复制;每次复制送达时,由硬件在 q_full 上到达。在 with mma: 中,Tensor Core warp 逐级在 q_full 上等待,调用 lm.fence_proxy(),然后发射矩阵乘法,结果写入 TMEM 累加器,并在第一级时将累加器初始化。从名字看,fence_proxy 是一个代理栅栏(见 §3.2)。论文没有说明它 lowering 成什么指令;而 TMA 的写入和 Tensor Core 的读取走的都是异步代理,所以单看这个片段,看不出这里为什么需要它。

注意这里少了什么:没有屏障地址,没有相位位,没有描述符编码,也不用计算哪个线程属于哪个 warp;这些全由编译器根据声明推导出来。不过,缓冲区确实带着几项简单明确的承诺,比如视图在内存池中的偏移,以及累加器在张量内存中的起始列。§7.3 会解释 CAKE 为什么要求直接写出这些值;其余的寻址信息,则由编译器以它们为起点推导出来。这段代码只是教学示例,不是完整的 kernel:它把 Q tile 和它自己相乘,而真正的注意力 kernel 还要加载 K 和 V,并计算 softmax。

你可能还注意到少了一个屏障。§3.2 说过,每一级都需要一个 full 屏障和一个 empty 屏障,可这个片段只声明了 q_full。它的三级各只填充一次,从不绕回环形缓冲区的起点,所以没有哪个槽位会被覆盖。真正的主循环要让许多 tile 依次流过环形缓冲区,因此还得再声明一个 empty 屏障,由 mma 负责到达、load 负责等待:§3 那张图里的数据竞争,正是缺了这一次等待造成的。

论文图 3 中的 CAKE IR 调度片段,分成三块。资源:96 KB 共享内存池、三级的 128×128 BF16 缓冲区 smem_q、张量内存中 128×128 的 FP32 累加器。角色与同步:load 在 warp 0,mma 在 warp 1,外加把数据块从 load 交给 mma 的三级屏障 q_full。各角色的工作:load 发起三次 TMA 复制;mma 等待 q_full、发出代理栅栏、再做矩阵乘。另有一栏列出由 lowering 推导的内容:屏障地址、相位位、TMEM 偏移、描述符编码和 warp 身份,最终生成 CUDA/PTX。
论文示例调度的三个部分:资源、角色及其间的屏障、各角色的工作;机械性的细节由 lowering 推导(据论文图 3 简化,省略了类型标注和注释)。打开原尺寸 SVG。

7.2 让它可以检查的四个特性

论文指出,真正“起作用”的是四个特性:

  • 经过类型检查的词汇表。计算、数据搬运、同步、数学运算和 warp 控制,用的都是一组固定的 IR 操作,而不是嵌入的 C 或 PTX 字符串。类型错误的程序,在构建阶段就会被拒绝。
  • 预先声明的资源。内存区域、同步对象和流水线都只声明一次,因此“IR 知道每个缓冲区的形状、数据类型(dtype)和生命周期”。
  • 显式的角色。warp 组都有名字,而且“每一次跨角色的交接都是可见的,而不是一种隐含的约定”。
  • 自动推导的元数据。声明带来的机械性后果“由 lowering 生成,而不是由作者编写”。

这样做的回报是:分析可以在代码生成之前就对调度决策进行推理,harness 也能把诊断结果定位到受影响的资源、角色或流水线级,“而不是只返回一个后端错误或一次挂死”。

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 却按另一种 swizzle 去读,编译器会在任何代码运行之前,就指出是哪条声明对不上。(这个例子是我们自己举的;论文只按类别描述了这些检查。)

这和 CuTe 的取舍正好相反。CuTe 要求作者精通布局,回报是强大的表达能力;CAKE 只要求作者写明简单的事实,揪出其中的矛盾则是编译器的责任。至于布局验证器,论文只描述了它的契约和覆盖范围,没有介绍内部实现。

7.4 一种语言,多款 GPU

同一种调度语言覆盖了从 Ampere 到 Blackwell 的 NVIDIA GPU。“角色–屏障–流水线”式调度的结构可以沿用;至于允许使用哪些指令、这些指令如何 lowering,仍然因目标而异。(sm_ 开头的名字是 NVIDIA 给各代 GPU 起的代号:sm_80 是 Ampere,sm_90a 是 Hopper,sm_100a 是 Blackwell;末尾的 a 标记的是该芯片特有的功能。)下面这些细节不必深究;最右边一列只是让你看到:调度结构不变,底下的指令集却在变。

CAKE 面向哪些 GPU,各代之间有什么变化
目标 代表 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 语料库,外加一组硬件设计原则。论文附录描述了五个步骤:

  1. 收集语料。从生产级库中收集高质量的 CUDA kernel:Alpha-MoE、CUTLASS、cuTile、DeepGEMM、FlashAttention-4、Flash-KMeans、FlashInfer、SonicMoE 和 TileLang。用其他语言写的 kernel,先由编程智能体翻译成带内联 PTX 的 CUDA。
  2. 提炼抽象。智能体分析语料,把反复出现的模式(屏障编排、流水线分级、warp 角色划分、TMA 描述符设置、TMEM 累加器的生命周期)归纳为候选抽象。
  3. 依据硬件做设计。人类专家凭经验把这些抽象引向 Blackwell 的编程模型:把 TMEM 当作一等资源,把 warp 专用化作为主要的并行形式,把异步屏障作为同步原语,并用集群范围的操作来协调多个 SM。
  4. 按原则迭代。每个候选抽象都要对照八条设计原则(见下文)检查;违反原则的,要么改进,要么淘汰。
  5. 靠移植扩展。新的 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 给了它们一门语言:专家的决策(角色、缓冲区、屏障、流水线)在其中是显式的、带类型的,容易出错的簿记工作(屏障地址、相位位、描述符编码、warp 身份)则由编译器推导。它刻意避开布局代数,转而让编译器检查具体的布局承诺。这门语言是从生产级 kernel 中长出来的,八条原则让它在成长中始终保持可分析。


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

学生 好,智能体就用这门调度语言来写。然后呢?代码最后总还是得上 GPU 跑吧。

教授 最后是要跑。可是靠上 GPU 跑一趟来发现漏掉的屏障,代价太高了。有没有更便宜的办法?

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

教授 没错。先做检查:在编译之前就运行,还能指出是哪个屏障、哪一级出了问题。接着用代价模型低成本地估算每个过关的候选能跑多快,挑出值得拿去计时的那些。最后才轮到 GPU。

学生 要是检查本身漏掉了一整类 bug 呢?

教授 那你就找到这篇论文最有意思的想法了。这个 bug 会变成一条新的检查。编译器会长记性。

学生 这一套真能做出更快的 kernel 吗?还是只换来更好看的报错信息?

教授 很好,保持这份怀疑。我们会去看证据,也会看看这些证据说明不了什么。

9. Harness:上 GPU 之前的检查

所谓 harness(测评环境),就是围绕智能体的一整套东西:各项分析、测试、基准测试,以及判定什么结果才算数的规则。本节介绍 CAKE 的 harness 检查什么,又会向智能体反馈什么。

9.1 七类检查

论文是按功能来介绍这套分析的:看它能为智能体做什么,而不是看编译器内部有哪些 pass。检查一共七类,分属三种处置方式:拒绝候选的关卡(gate)、描述候选的报告(report),以及建议如何改进的提示(hint)。

类别 处置方式 作用 示例(我们自己举的,仅作说明)
程序安全 编译前关卡 找出同步、顺序和内存使用方面的隐患 消费者在等待某个屏障,却永远不会有生产者到达那里(§3.3)
硬件合规 编译前关卡 强制遵守受支持的资源、指令与架构契约 缓冲区需要的共享内存超出了 SM 的容量,或者在面向 Hopper 的代码里用了 Blackwell 独有的指令
数据一致性 编译前关卡 检查数据流,以及生产者与消费者对数据表示的约定是否一致 TMA 按一种 swizzle 写入 tile,MMA 读取时却假定是另一种
调度语义 编译前关卡 检查所声明调度的结构不变量 流水线只有三级,循环却访问了第四级
数值验证 执行关卡 把编译后的输出与权威的外部参考结果对比 输出与参考结果的偏差超出了容差
性能分析 报告 估算开销,指出瓶颈大致属于哪一类 “很可能受限于显存带宽”
优化指导 提示 建议值得尝试的修改,但不阻止编译 “更深的流水线也许能隐藏复制延迟”

四道编译前关卡会拒掉“许多在数学上说得通、却与目标执行模型不兼容的候选”。每条诊断结果都会“指出受影响的程序区域,以及违反的是哪一类契约”。拿它和 §3.3 里的挂死比一比:智能体面对的不再是一块毫无反应的 GPU,而是类似这样一句话:消费者在这个屏障的第 2 级上等待,但永远不会有生产者到达那里。这个例子是我们自己举的,论文没有公开诊断消息的格式。

9.2 正确性与性能

检查数值正确性的方法,是在不同的形状和输入分布下与参考实现对比。最终验收还要求在目标框架里做端到端评测(例如用 SGLang 部署完整的模型并提供服务;SGLang 是一个流行的开源 LLM 推理服务系统)。

性能评估分为两层。经过校准的代价模型先估计候选能跑多快,并指出它的瓶颈可能在哪里,这样 harness 不必把每个候选都跑一遍,就能给它们排序、筛选。但这个模型只是一道筛子:“设备上的实测和性能剖析(profiling)仍是最终的评判依据。”在论文的评测中,kernel 的计时都是用 CUPTI(NVIDIA 的性能剖析接口)在 GPU 上实测的;每次计时采样之前都会清空 L2 缓存,免得上一次运行替下一次把缓存预热好。(§13.2 中有少数几个数字的测法不同:Alpha-MoE 在 API 层测得的加速比,以及 SGLang 推理服务的吞吐。)

对于静态分析能保证什么,论文说得很谨慎:它们只“在其建模范围之内”充当关卡;它们“并不证明 GPU 上的全局正确性,也无法刻画所有微架构行为”;“误报和漏报都可能发生”。缺失的信息会被“明确报告出来,而不是当作检查通过”。

10. 一轮进化

语言和 harness 都就位之后,一次 kernel 进化运行的流程就很简单了。本节跟着走完一轮。

10.1 四个阶段

论文把一次运行分为四个阶段:

  1. 生成结构各异的 CAKE IR 候选:角色划分、流水线深度、内存放置各不相同,而不只是换几个常数。
  2. 筛选这些候选:在花费任何 GPU 时间之前,先用 IR 构建检查、验证器的硬性关卡和代价模型排序把它们过一遍。
  3. 评估过关的候选:对照外部的正确性判据(oracle,即一份可信的参考实现,它的输出就是标准答案)进行检验,并收集基准测试和性能剖析证据。
  4. 分流证据。一条诊断结果送往哪里,要看诊断出的是什么问题:可能送回候选本身(修复它),也可能送给验证器(新增一条规则)、代价模型(重新校准)或 IR 词汇表(新增一个原语)。
一个漏斗(数量仅为示意)。先生成 12 个结构各不相同的 CAKE IR 候选。构建与类型检查剔除类型错误的候选,剩下 9 个;验证器的硬性关卡(程序安全、硬件合规、数据一致性、调度语义)带着原因剔除调度,比如缺少 wait 或张量内存越界,剩下 5 个;校准过的代价模型给它们排序,保留 3 个。这些步骤都不花 GPU 时间。幸存者在 GPU 上对照正确性判据运行,并用 CUPTI 计时和性能分析,选出一个最佳 kernel。随后各阶段的诊断都被送到该从中学习的地方:候选本身(修复它)、验证器(加一条规则)、代价模型(重新校准)或 IR(加一个原语)。
一轮进化:便宜的静态检查带着原因剔除不合法的调度,代价模型给幸存者排序,只有最好的几个上 GPU 运行;各阶段的诊断都被送到系统中该从中学习的那一部分(数量仅为示意)。打开原尺寸 SVG。

正是第四个阶段,让 CAKE 有别于一个只是报错信息更好的搜索循环。诊断结果不会用完就扔,而是被送到系统里最该从中吸取教训的那个部分。

10.2 哪些保持不变

每次运行都有两个锚点。第一个是工作负载契约(workload contract),也就是论文所说的“稳定的权威”。它固定了以下几项:形状;正确性判据及其容差(输出与参考结果要接近到什么程度才算对,就像 torch.allclose 那样);硬件;以及智能体可以查看哪些参考资料。运行结果都会保留下来,所以决策可以审计,反复出现的诊断结果也能复用。

第二个是模型。论文里所有智能体任务用的都是同一个模型:GPT-5.6-sol,推理强度(reasoning effort)设为 xhigh。“把模型和智能体脚手架(scaffold)固定下来,这些比较……就可以归因于环境,而不是模型能力。”(脚手架是包在 LLM 外面的那层程序:它为 LLM 提供编辑文件、运行代码的工具,并驱动整个循环。)一项实验的两组之间无论出现什么差别,都不会是因为其中一组用了更聪明的模型。

11. 编译器也在进化

到目前为止,智能体都是在一个固定的环境里进化 kernel,只不过这个环境恰好设计得很好。CAKE 的最后一个主张是:环境本身并不固定。

11.1 进入编译器的两条路径

候选 kernel、验证结果、基准测试和失败报告,都可以成为修改编译器的证据。论文描述了两条路径:

  • 路径 1:挖掘专家 kernel。智能体研读生产级 kernel 和硬件文档,找出 IR 尚未覆盖的 Blackwell 模式(“新的指令形式、资源类型、描述符变体、同步惯用法”),并撰写编译器修改提案。每份提案在实现之前,都要对照 §8.2 的设计原则检查,其中包括性能透明和易于验证。
  • 路径 2:从失败中提炼。智能体利用失败候选留下的反馈(“sanitizer 报告、失败用例、正确性比对不一致、调试日志”),把反复出现或代价高昂的失败模式提炼成新的分析:“一次原因不明的运行时崩溃变成一条验证器规则,一种反复出现的非法 lowering 模式变成一项静态检查,一种系统性的预测偏差变成一个校准目标。”
CAKE 编译器的进化流程图。路径 1 从生产级 kernel 和硬件文档中找出 IR 缺失的模式:新的指令形式、资源类型、描述符变体、同步惯用法。路径 2 从失败候选中提炼检查:原因不明的崩溃变成验证器规则,非法 lowering 变成静态检查,预测偏差变成校准目标。两条路径汇成编译器修改提案,先对照 IR 设计原则检查,再把原语与其分析一起实现,并须通过语料库测试(400 多个静态/编译用例、399 个 GPU 正确性用例)、经人类把关才能合入;更新后的 harness 用于下一个 kernel 家族。
外层循环:智能体从专家 kernel 中挖掘 IR 还表达不了的模式,并把反复出现的失败提炼成检查;每项修改都必须通过整个 kernel 语料库的测试,并经人类把关才能合入。打开原尺寸 SVG。

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

两条路径相辅相成。新原语让编译器更了解硬件,从而能做更强的分析;新分析反过来又约束了未来原语可以是什么样子。论文坚持认为,一个原语和它的分析“必须共同进化”:“只有语法而没有副作用和合法性规则,会让 IR 更难分析;而未经语料库验证的新验证器规则,可能会拒掉有效的 kernel。”通俗地说:往语言里加一个新操作,却不告诉检查器它读什么、写什么、要求什么,那么凡是用到它的程序,检查器都会两眼一抹黑。反过来,加一项新检查,却不先在已知正确的 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 是 k-means 的一个快速、精确的 GPU 实现。它之所以重要,原因可能出乎你的意料:视频生成。Sparse VideoGen2 用 k-means 把相似的 token 聚到一起,让稀疏注意力能在由相关 token 组成的连续块上工作(“语义感知重排”,semantic-aware permutation)。在这种用法里,k-means 运行在模型的前向传播之中,所以它花的时间会直接计入生成时间。(稀疏注意力只为最重要的那些 token 对计算注意力。)

按论文的说法,一次 Lloyd 迭代的时间主要花在两个 BF16 kernel 上,两者合计占端到端时间的 95% 以上:

  • assign 计算每个点到全部 K 个质心的平方距离,并返回最近那个质心的下标。它是计算受限(compute-bound)的:一次矩阵乘法,后接一次归约。
  • centroid_update 把每个簇里的点加起来,并统计点数。它受限于显存带宽,以及原子更新时的争用(许多线程同时往同一个簇的累加和里加数,而硬件只能一个接一个地执行这些更新)。

实验聚焦于 assign,它“考验 Tensor Core 流水线和标量 epilogue(收尾阶段)”,这里的 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 或等价的生成代码这类底层目标实现”。(SASS 是 GPU 真正执行的机器码,比 PTX 还低一层。)现成的实现可以作为黑盒计时基线拿来运行,但其内部细节始终对智能体隐藏。这条限制在隔离环境中强制执行,事后还做了审计。

两组实验的区别在于智能体写什么,因而它收到的编译器反馈也不同:

  • 实验组:智能体写带类型的 CAKE IR,配合第三部分介绍的 harness。
  • 对照组:智能体直接写 CUDA C++ 和内联 PTX,遵守同样的参考实现访问规则。

其余一切都固定不变:编程智能体及其脚手架、模型及推理强度、任务描述、正确性判据、基准测试 harness,以及目标形状。每组各跑三次,预算为 80M token。

12.3 结果

论文表 2 给出了用满 80M token 预算时的结果,每组三次运行:

表示形式 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 一次也没有。所谓一次运行进入平台期,是指它的最佳速度已经不再提升;具体的判定规则由论文事先定好,但没有公开。也就是说,到 80M token 时,CAKE 的三次运行都已趋于平稳,而 CUDA/PTX 的三次运行在预算用完时仍在提升,我们无从知道它们最终会停在哪里。CAKE 的运行用掉的有效时间也只有对方的一半左右。但论文没有把这段时间拆开来报告(它的停止与计时记录“保留在配套材料(artifact)中”),所以我们分不清 CAKE 的运行更早结束,是因为停在了平台期,还是因为每一步都更快。

图表展示截至当时经过验证的最佳速度(相对调优过的 FlashML KMeans 基线,0.938 ms)随累计 token 数(1000 万到 8000 万)的变化;对象是 B200 上一个固定形状的 assign kernel,每组 3 次运行。曲线是 3 次运行的均值,从论文图 5 上近似读出。CAKE IR 从 10M token 时约 0.34× 升到 20M 时约 0.8×,约在 55M token 时越过基线;直接写 CUDA/PTX 从约 0.14× 升到 80M 时约 0.97×,始终低于基线。80M token 时,CAKE IR 的最佳速度中位数为 1.144× [1.041, 1.205],3 次运行都进入平台期,有效进化时间 1.89 小时;直接写 CUDA/PTX 为 0.928× [0.852, 1.151],0 次进入平台期,用时 3.73 小时。
受控实验:模型、任务和判据都相同,CAKE IR 的均值大约在 55M token 时越过调优基线,直接写 CUDA/PTX 的均值到 80M 时仍低于基线(最好的一次达到 1.151×)。曲线是从论文图 5 读出的 3 次运行均值(近似值);卡片中是论文表 2 的数字。打开原尺寸 SVG。

随时间变化的轨迹,讲的也是同一个故事。每隔 5M token 有一个检查点;每次运行在检查点上计入的,是它截至当时、经过验证的最佳速度。CAKE IR 各次运行的均值“到 55M token 时越过调优过的 FlashML 基线,并继续提升;而直接写 CUDA/PTX 的均值到 80M token 截止时仍低于基线”。从论文的图上读数(近似值):到 20M token 时,CAKE IR 的均值已接近 0.8×,而 CUDA/PTX 的均值还在 0.2× 左右。

12.4 怎么解读

下面是我们的理解,以及这些数字本身附带的限定条件:

  • 在同样的预算内,环境改变了结果。用 CAKE IR 加上它的 harness,三次运行全部超过了调优基线;CUDA 的运行按中位数看没有做到,而且仍在往上爬。究竟是这套环境中的哪一部分起了作用(语言、编译前检查、代价模型,还是编译器进化),仍是一个开放问题(§16.2)。
  • 但直接写 CUDA 并非毫无希望。CUDA/PTX 最好的一次运行达到 1.151×,比 CAKE 的中位数还高。差别在于稳定性和进展的快慢:CUDA 也能达到这个水平,只是不那么可靠,也来得更慢。每组只有三次运行,结果的离散程度和中位数同样重要。
  • 这只是一个 kernel、一个形状、一个模型。这个实验很好地隔离出了环境的作用,但没有告诉我们,换成别的 kernel 或别的模型,这个作用有多大。

13. 基准测试之外

对照实验干净地展示了效果;其余评测则展示了这套系统在真实 kernel 上的表现。这部分评测还回答了另外两个问题:在不许查看任何现有底层实现(即便已有快速实现,也只能当黑盒计时基线)的情况下,CAKE 能不能发现 kernel 的调度?对于专家已经调优过的 kernel,它能不能复现,甚至超越?

13.1 前沿 kernel:新的模型架构,没有参考可抄

论文把一个 kernel 称为前沿 kernel,取的是“操作层面的含义,即智能体必须在不查看底层目标实现的情况下,发现它的物理调度”。协同进化的 IR 在这里理应最能帮上忙,但这里也最容易暴露它的短板:缺了某种能力,就会表现为一个“智能体根本无法表达”的调度。

Kimi Delta Attention。官方的 FlashKDA kernel 只用作黑盒计时基线,它们的源码和生成代码都没有提供给智能体。智能体生成的 prefill kernel 与 FlashKDA 兼容,覆盖定长、打包变长(packed variable-length)和尾部输入(分别指所有序列长度相同的批次、把长短不一的序列首尾相接打包而成的批次,以及序列末尾剩下的 token)。在六个 B200 BF16 形状上,它相对该基线取得了 2.05× 的几何平均加速比。论文称它“在其验证契约上逐位正确”(pull request 本身报告的是:输出与 FlashKDA 在 10−2 的容差内一致),它还在 SGLang 上通过了 Kimi-K3 端到端推理服务的验证。另有一组单独的 decode kernel,在 30 个公开 API 形状上相对上游 FlashInfer(即这个库目前官方提供的 kernel)取得 1.14× 的几何平均加速比。FlashInfer 是一个开源 kernel 库,SGLang、vLLM 等 LLM 推理服务引擎内部都在用它;这两类 kernel 都以生成的 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 中面向极小批次的 BF16 矩阵乘法(论文称之为小 M),源自 NVIDIA 自家的 LLM 推理库 TensorRT-LLM。服务端同时只为少数几个请求做 decode 时,批次通常就这么小。智能体产出了一个自适应的 kernel 家族,既有浅流水线也有深流水线,其中包括使用 PDL(programmatic dependent launch,允许下一个 kernel 在上一个结束之前就开始执行自己的前导部分)的变体,以及专为批大小不到 8 的情况准备的路径。FlashInfer PR #4274 报告,在 35 个标准形状和一套覆盖面更广的回归测试上,kernel 耗时的几何平均降低了 18–23%。贪心解码(每一步总是选概率最高的那个 token)时,GPT-OSS-20B 和 GPT-OSS-120B 的输出保持逐位一致,在 B200 和更新的 Blackwell 系统 GB300 上都是如此。一项用 GPT-OSS-120B 做的 SGLang 实验测得:在单卡上同时处理 128 个请求时,每秒输出的 token 数最多提高 7.6%;把模型拆到四张卡上运行(四路张量并行)时,差异则没有超出噪声范围。

Alpha-MoE。在混合专家(MoE)层里,一个小型的路由器(router)把每个 token 分派给众多“专家”前馈网络中的少数几个,再把这些专家的输出加权合并。原版 Alpha-MoE 用一个 GPU 程序完成整层计算,这样的程序叫作巨型 kernel(megakernel);它为 Hopper 编写,权重和激活都是 8 位(W8A8)。CAKE 的智能体把它改写到了 Blackwell 上。这个巨型 kernel 先收集路由到各个专家的 token,再执行该专家的两次矩阵乘法,中间夹着激活函数。接着,它把中间结果转换回 8 位(即重新量化,requantization),最后按路由器给出的分数加权,把每个专家的输出累加到对应 token 的结果上。与 FlashInfer 中源自 TensorRT-LLM 的预路由(pre-routed)API 相比,端到端的 API 层加速比在 N = 256 时为 6.204×,在 N = 512 时为 4.025×;若按 GPU 跨度(GPU span)衡量,则分别为 1.215× 和 1.170×。(论文没有说明这里的 N 指什么;它不是 k-means 实验里的那个 N。)这个 kernel 已于 2026 年 9 月合入 FlashInfer(PR #4287),但 pull request 里报告的是另一组对比:在 GB300 上,与 SGLang 自带的、由五个 Triton kernel 组成的 MoE 路径相比,融合后的 kernel 按 GPU 时间快 1.37×–1.75×。论文里的 6.204× 和 4.025× 在 pull request 中并没有出现。

13.3 复现专家 kernel

最后一个问题是:当专家设计的结构已经存在,而且智能体可以查看时,harness 能否帮助智能体既保住正确性,又追平那些高度优化的 kernel。参考实现来自 FlashAttention-4、TensorRT-LLM、DeepGEMM、CUTLASS 和 FlashInfer。所有运行都在 B200 上以固定形状进行,用 CUPTI 测得的 GPU 跨度中位数计时,并且都通过了各 kernel 专属的正确性关卡。变体名称不必细究:看“相对性能”这一列就行(高于 1.0× 表示 CAKE 的 kernel 比专家写的更快),再比较一下两边的行数(“设备端”代码指在 GPU 上运行的那部分代码)。

家族 变体 相对性能 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 路由和运行时常量下可达的源码。
  • GQA(分组查询注意力)让多个查询头共用同一个键/值(KV)头;“因果”注意力只允许一个 token 看它之前的 token。MQA(多查询注意力)索引器为历史 token 打分,决定一个查询应该关注哪些 token,就像 DeepSeek-V3.2 的闪电索引器那样;MLA 是 DeepSeek 的多头潜在注意力。“分页”kernel 读取的 KV cache 按固定大小的块存放。我们的 GLM-5.3 和 DeepSeek-V4 深度解读对两者都有讲解。

13.4 语料库的广度

只看前面几项重头实验,看不出 harness 如今撑起了多大的规模。经过验证的语料库包含 400 多个静态检查与编译用例,以及 399 个 GPU 正确性用例,覆盖约 28 个 kernel 家族,其中包括注意力、稠密与稀疏 GEMM、MoE、量化、归一化、状态空间模型、KNN(k 近邻搜索)和 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 构建分发器

分发器是放在几个 kernel 前面的一小段代码:它查看输入的形状,然后调用合适的那个 kernel。用 Python 写,它就是一串 if/elif/else 分支(下面只是示意,并非 CAKE 的代码):

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)

代码里的每个判断条件叫作守卫条件(guard),每个条件把守的那个 kernel 叫作一条路由(route),最后那个 else 就是兜底路径(fallback)。“Split-K”把很长的归约维度切开,分给多个线程块,再把它们的部分结果加起来;当其他维度太小、填不满 GPU 时,这种做法很有帮助。“尾部填充”(padded tail)路由则把不规整的尺寸向上补齐到一个方便处理的尺寸。

泛化阶段先把测过的种子按形状分桶,生成特化的或共享的变体,再把它们的守卫条件依次排好,最后接一条显式的兜底路径。守卫条件是对形状的判断,比如“D = 352”;兜底路径负责处理所有守卫条件都没匹配上的输入。调优可以改实现参数,但绝不能改输入形状。在报告任何汇总数字之前,验证要覆盖这几类情况:有代表性的输入和留出的(held-out)输入,边界情况和尾部情况,相互重叠或留有空隙的守卫条件,以及兜底路径本身。

CAKE 如何把只为单个形状调好的 kernel 变成一次库调用。进化只为一个确切的形状调出一个种子,也就是一个 CAKE IR 程序;可库调用(例如 kmeans(B, N, K, D))传进来的可能是任意形状。泛化阶段把种子按形状分桶,在一条显式的兜底路径前面依次排好针对形状的守卫条件:自上而下逐个检查,第一个匹配的生效。图中路由取自论文的 KMeans 组合(守卫条件的顺序为示意):D = 288 和 D = 352 走 split-K 路由(分别覆盖 14 个和 6 个形状),D 在 144 到 176 之间走尾部填充路由(8 个形状),D = 112 走特化路由(6 个形状),另外 6 条路由共覆盖 14 个形状,通用路由覆盖 64 个,兜底路径覆盖 12 个,合计 12 条路由、124 个形状。每条路由都是一个独立的 CAKE IR 程序,可以单独分析、单独做基准测试。报告任何汇总数字之前,验证要覆盖有代表性的输入和留出的输入、边界情况和尾部情况、相互重叠或缺失的守卫条件,以及兜底路径;形状集合在调优前就已固定,守卫条件可以划分它,但不能往里加评测形状。在 GB200 上,相对参考实现的 GPU 跨度几何平均加速比为:KNN 构建 1.418×(112 个形状,8 个外层族,80 条路由),KNN 搜索 2.116×(198 个形状),KMeans 1.803×(124 个形状,12 条路由),没有任何错误输出,KNN 召回率为 1.0。
从调好的单个形状到一次库调用:按顺序检查的守卫条件,把每次调用分发到若干条可单独检查的 CAKE IR 路由之一,或交给显式的兜底路径。KMeans 的路由与形状数来自论文,守卫条件是根据路由名称推断的。打开原尺寸 SVG。

有一条策略让路由不至于沦为大杂烩:这个阶段“在形状域中尽可能大的范围内复用同一个物理调度,只有当形状域要求调度发生实质性变化时,才引入另一个”。每条路由都是一个独立的 CAKE IR 程序,所以每个备选方案仍然可以单独分析、单独做基准测试。

14.3 不能自己出题自己考

这个阶段藏着一个不易察觉的危险:评测泄漏(evaluation leakage)。如果调优分发器和报告它的速度用的是同一批形状,分发器完全可以把这些形状直接背下来。CAKE 在调优之前就先声明有效的形状域。分发器的判定条件(predicate)“可以划分这个域,但不能引入讨巧的新评测条目”;覆盖范围也只能“通过同一来源中按确定性方式划出的、未见过的分片”来扩展。也就是说,新的测试形状要按一条固定的规则从同一来源中抽取,而且在整个调优过程中始终不可见。

14.4 在真实库上的结果

经典的机器学习工作负载给出的读数最干净,因为它们的形状组合(portfolio)很大,单靠某一个形状撑不起整体结果。KNN(k 近邻,k-nearest neighbours)要为每个查询点找出数据集中离它最近的那些点:“KNN 构建”通常指在同一个点集内部、对所有点做这件事(即构建一张近邻图),“KNN 搜索”则是针对另外给出的查询点来做。CAKE 为 FlashLib 贡献了泛化版的 KNN 构建、KNN 搜索和 KMeans kernel,已于 2026 年 8 月通过 pull request #15、#16 和 #18 合入(论文本身没有引用任何 PR 编号)。按论文的测量,这些 kernel 在 GB200 上的结果如下表所示。其中 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 测量。

15. CAKE 在相关工作中的位置

CAKE 涉及四个研究领域:GPU 语言(§6)、编译器、kernel 智能体,以及自我改进的系统。要给它定位,最清楚的办法是对每一条研究路线都问一句:进化的是什么,固定不变的又是什么?

研究路线 代表工作 变的是什么 不变的是什么
带搜索或记忆的 kernel 智能体 KernelEvolve、EvoEngineer、AVO、K-Search、KernelBlaster 搜索过程、积累下来的记忆 语言和评测环境
用强化学习训练的 kernel 模型 AutoTriton、CUDA Agent 模型权重 语言和评测环境
LLM 驱动的程序搜索 FunSearch、AlphaEvolve 程序本身(借助 LLM 写出的变异) 评估器
自我改进的智能体 ADAS(智能体系统自动设计,Automated Design of Agentic Systems)、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 运作、能够改变自身机制和抽象的系统。随后,论文按“进化的是什么”,把 CAKE 与所有这些会进化的系统区分开来:智能体和模型保持固定,变的只有一个面向特定领域的编译器 harness,而且每次改动都要接受语料库测试和人工合并关卡的约束。

16. 局限与开放问题

好论文会交代自己止步于何处。CAKE 交代了,细心的读者还能再补上几个问题。

16.1 论文自述的局限

  • 只支持 NVIDIA。CAKE 面向从 Ampere 到 Blackwell 的各代架构。调度语言、角色模型和各项分析可以在这几代之间通用,但指令形式、合法性规则和代价锚点都与具体架构绑定,而“这部分成本才是衡量迁移的诚实尺度”。对于非 NVIDIA 硬件,这一成本“仍未测量”,而且后端得重新构建。
  • 证据分布不均。大部分性能结果来自 B200。耗时模型只针对 B200 和 H100 校准过,在其他硬件上则“拒绝预测”。
  • 分析并不完备。静态分析和性能模型是有意做得不完备的:它们负责排序和过滤,GPU 上的实际执行始终是最终依据。
  • 人仍在回路中。编译器进化“在合并关卡上仍由人类把关”。

16.2 论文留下的开放问题

以下问题是我们提出的,并非出自论文:

  1. 别人用得上吗?论文没有说 CAKE 的编译器、IR 或 harness 已经公开发布;截至 2026 年 10 月,我们也没找到任何发布版本(只有一些非官方的社区复现)。公开的是生成出来的代码:四个 FlashInfer pull request 和三个 FlashLib pull request 都已合入,FlashInfer 还有一个跟踪 issue,汇总了更多由 CAKE 生成的 kernel。
  2. 对照实验的结果有多普适?从零起步的对比,每组只跑了三次,而且只用了一个 kernel、一个形状和一个模型(GPT-5.6-sol)。换成别的 kernel、别的模型,或者给更大的预算,差距会怎样变化,我们不得而知。(FlashInfer 的跟踪 issue 提到,CAKE 生成的 kernel 也出自 Claude 和 Codex 智能体之手,这似乎说明 harness 并不绑定某一个模型;但没有报告任何对照比较。)
  3. 到底是哪一部分在起作用?两组之间,语言、编译前检查、代价模型和编译器进化是同时改变的。论文没有做这样的消融实验:比如保留 CAKE IR 但去掉代价模型,或者把编译器冻结住。
  4. 为什么没有 tile 级语言这一组?对照组写的是 CUDA/PTX。如果在同样的规则下再加第三组,改用 Triton 或 CuTe DSL 来写,就能看出 tile 级语言和布局代数会落在什么位置。
  5. 编译器进化的代价有多大?论文没有报告编译器一共改了多少次、每条路径各贡献了多少、语料库关卡拒绝改动的频率,以及合并关卡耗费了多少人工审查。
  6. 这些检查有多准?论文没有报告验证器的误报率和漏报率,也没有报告代价模型预测的准确度。
  7. 缺失的数字,以及两处出入。Gated DeltaNet 和 MiniMax 稀疏注意力的结果只有文字描述,没有给出加速比;平台期判据在正文中也没有定义。另有两份上游记录与论文的说法不一致:已合入的 Alpha-MoE pull request 报告的是另一个基线和另一组加速比(§13.2);KDA prefill 的 pull request 报告的是容差范围内一致,而论文说的是“逐位正确”。

这些问题都动摇不了论文的核心思想。它们标出的,正是下一篇论文或一次独立复现最能有所贡献的地方。

第三部分小结

CAKE 的 harness 在调度送上 GPU 之前先检查它(程序安全、硬件合规、数据一致性、调度语义),再用校准过的代价模型给过关的候选排序,并以设备上的实测为最终依据。每一轮的诊断结果都会流向候选本身、验证器、代价模型或 IR 中的某一处;反复出现的失败则变成新的检查或原语,它们必须通过整个 kernel 语料库的测试,还要经过人工合并关卡。在模型固定的前提下,CAKE IR 配合它的 harness 取得了这些成果:在一项小规模对照实验中胜过了直接写 CUDA/PTX;在看不到底层实现的情况下,为 KDA 等架构写出了新的 kernel;追平或超越了 11 个专家 kernel 中的 10 个;还把只针对单一形状的种子做成了有分发器支撑的库函数家族。这些开放问题关乎普适性、归因和代价,而不在于这个想法本身行不行得通。


尾声 —— 一段对话

教授 好了。你都学到了什么?

学生 快的 kernel,很大程度上就是一份好的调度:谁来加载,谁来做乘法,彼此之间怎么交接 tile,每个字节放在哪里。

教授 那智能体为什么写得这么吃力?

学生 因为它们用的语言,要么把调度藏起来,要么逼它们亲手管理每一个地址。得到的反馈,也只有一句“崩了”,或者一个数字。

教授 那 CAKE 怎么解决?

学生 一门语言,决策都写在明面上,琐碎的簿记交给编译器去推导;一套检查,GPU 还没开跑,就能点出是哪个屏障坏了;还有一个编译器,会把反复出现的失败变成新的检查。智能体变强,是因为它所处的世界变好了。

教授 那你还想知道些什么?

学生 换成别的模型、别的 kernel,结论还成不成立;哪一部分最关键;还有,我能不能自己动手试试。

教授 很好。现在你算是会读论文了。

全文小结

  • 问题。专家级 GPU kernel 胜在调度(warp 角色、屏障编排、内存放置),而调度一旦出错,就会悄无声息地算出错误结果,或者干脆挂死。kernel 智能体在搜索、记忆和模型上都有了改进,但它们写代码用的语言,仍然要么把调度藏起来,要么把调度暴露出来却不设护栏;它们得到的反馈也只有报错、通过与否和耗时。
  • 思路。智能体和它的工作环境要协同设计(co-design)。智能体编写 CAKE IR:一种带类型的调度语言,硬件决策都摆在明面上,也不用布局代数;资源、角色和交接都要显式声明,机械性的细节则自动推导出来。harness 在编译之前检查程序安全、硬件合规、数据一致性和调度语义,用校准过的代价模型给候选 kernel 排序,并以 GPU 实测为最终依据。编译器也在进化:反复出现的失败和新发现的专家模式,会变成验证器规则、原语和校准;每项改动都要经过语料库测试的把关,这个语料库包含 400 多个静态检查与编译用例,外加 399 个 GPU 正确性用例。
  • 证据。模型固定时,在从零起步的 Flash-KMeans 任务上,CAKE IR 配合它的 harness,胜过直接写 CUDA/PTX(80M token 时,速度中位数分别为调优基线的 1.144× 和 0.928×;进入平台期的运行分别为 3/3 和 0/3;有效进化时间约为后者的一半)。在看不到任何底层实现(CUDA、PTX、SASS)的情况下,它写出的 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)每一章都以作业收尾。我们的作业没有配套的模拟器,但下面每道题都值得花几分钟想一想:

  1. 在 §3.2 中,三级环形缓冲区的每个屏障只用 1 位来记录相位。为什么 1 位就够了?如果生产者能领先消费者整整两圈,会出什么问题?
  2. CAKE 拒绝把调度降级到更老的 GPU 上(§7.4)。如果编译器悄悄地把一份 Blackwell 调度 lowering 成 Hopper 指令,可能会出什么问题?
  3. 编译前检查允许出现误报(§9)。在进化循环里,一项检查会错误地拒掉 5% 的合法 kernel,另一项会漏掉 5% 有问题的 kernel,哪一项更糟?如果换成生产环境用的编译器,你的答案会变吗?
  4. 请把 §16.2 中缺失的那个消融实验设计出来:你会跑哪几组实验,才能把语言、harness 和编译器进化各自的作用区分开?每一组又要固定哪些条件?
  5. §14.3 禁止分发器自行添加评测形状。请举一个具体例子:如果打破这条规则,分发器怎么会看起来比实际快 2×?
  6. 论文借用了陶哲轩的观点:生成变得廉价之后,瓶颈会转移到验证和理解上(§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(程序表示与编译器 harness),八条原则和完整的 kernel 表格见论文附录。

已合入上游的 kernel

  • FlashInfer 的 pull request #4262(KDA prefill)、#4279(KDA decode)、#4274(TinyGEMM)和 #4287(Alpha-MoE),均已合入。它们都是生成的 CUDA 代码,不需要 CAKE 也能直接阅读。CAKE kernel 跟踪 issue 里还列出了更多 kernel。
  • FlashLib 的 pull request #15、#16 和 #18:§14 中有分发器支撑的 KNN 与 KMeans kernel。

GPU 与 kernel 背景知识

kernel 智能体与进化搜索

工作负载

写作风格

  • Operating Systems: Three Easy Pieces,Remzi H. Arpaci-Dusseau 与 Andrea C. Arpaci-Dusseau 著,可在线免费阅读(中文版书名《操作系统导论》)。本文的对话,以及“关键问题”“补充”“提示”这几种方框,都借鉴自这本书的体例,在此致谢。

本站相关文章