CAKE: Compiler-Agent Co-Design for Frontier Kernel Evolution
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
cs.LG
2026-08-13
CAKE 让 Agent 直接编写带类型、显式描述硬件的中间表示,编译器持续把常见失败变成校验规则和优化战术;在 B200 上把 Kimi Delta Attention 内核做到官方实现的 2.05 倍。
GPU kernel agent 和 GPU 编程语言是两个各自在进步、互不来往的领域,专家级内核就丢在这个缝里。现有 agent 把编译器当黑盒:环境只回编译错误、对错判定和总耗时,从来不说哪个程序决策引发了同步失败、硬件契约违规或流水线停顿。新架构冒出缺失能力时,这套反馈信号也没法跟着长。
Agent 能写的那侧同样别扭。Triton 这类 tile 级 DSL 把 warp 特化、barrier 编排、显存分层全藏起来,而这些恰恰是专家内核与「 merely 正确」内核的差距所在;CuTe 这类低层 DSL 倒是全开放,但要求 agent 摆弄一套 layout 代数,写错概率高、错了还难定位。
CAKE(NVIDIA 与 CMU 合作)的答案是两边一起改。Agent 直接编写 Cake IR,一种带类型、显式描述硬件的调度表示:
关键是 harness 本身在进化。反复出现的失败会被沉淀成四样东西:新的校验规则(比如运行时不透明崩溃变成一条合法性检查)、新的 IR 原语、cost model 的校准目标、可复用的优化战术。进化对象是编译器环境,coding agent 和底座模型保持不动。改动要在内核语料上过测试门槛,合并点仍有人工把关。
| 实验 | 基线 | 结果 |
| Flash-KMeans 冷启动(B200,8000 万 token 预算) | 调优过的 FlashML | Cake IR 中位 1.144×,直接写 CUDA/PTX 只有 0.928× |
| 同上,收敛率 | — | Cake IR 三次全收敛,直接 CUDA/PTX 零次;进化耗时 1.89 小时 vs 3.73 小时 |
| Kimi Delta Attention prefill(六个 B200 BF16 形状) | 官方 FlashKDA 黑盒计时 | 几何均值 2.05× |
| KDA decode(30 个公开 API 形状) | 上游 FlashInfer | 1.14× |
| Alpha-MoE megakernel | FlashInfer 的 TensorRT-LLM 派生 API | N=256 时 6.204×,N=512 时 4.025× |
| KNN/KMeans 分发器组合(GB200,434 个形状) | 各自参考 | KNN 构建 1.418×,KNN 检索 2.116×,KMeans 1.803× |
KDA 那组实验里 FlashKDA 只当黑盒计时基线,源码不给 agent。11 组已知内核复现中 10 组达到或超过参考实现,剩 1 组到 96.5%。四项内核改动已作为 PR 进入上游(KDA prefill、KDA decode、TinyGEMM2、Alpha-MoE)。
这是「环境设计决定 agent 上限」的一个干净论证:同一批 agent、同样的 token 预算,换一个能给出局部诊断、能积累规则的编译器环境,从零收敛都做不到变成三次全收敛。1.144× 对 0.928× 的差距全来自环境,不来自模型变强。
对新架构落地尤其实用。Kimi Delta Attention 这类新注意力变体出来时往往没有手工调优的参考实现,CAKE 路线意味着 agent 可以在 DSL 层补齐这段空窗,产物还能直接变成上游 PR。对做推理引擎和 kernel 库的团队,这是可复用的产能;对做 agent 基建的团队,这是个值得抄的环境设计范本。
作者自己列得相当直白:性能证据大头在 B200,计时模型只校准了 B200 和 H100,其他架构拒绝外推;静态分析和性能模型刻意不完整,进化中只用来排序过滤,真值仍靠 GPU 实测;编译器进化在合并点仍是人工导引。非 NVIDIA 目标的迁移成本未测量,后端 lowering 得重写。TMEM、TMA、cluster 这些 Blackwell 特有能力在 Ampere 上没有,老卡用户拿不到完整体验。
一个需要留意的点:2.05×、6.2× 这些高倍数集中在参考实现薄弱的新工作负载上(KDA 官方实现早期、megakernel 化的 MoE),而 11 组成熟内核复现里最好也只有约 1.27×。对已经有专家手工调优的内核,这套方法的增益空间明显更窄,宣传里最亮的数字不代表通用预期。