NVIDIA与CMU提出CAKE：Agent与编译器协同设计，前沿Kernel提速2.05倍！

CAKE: Compiler-Agent Co-Design for Frontier Kernel Evolution

在过去几年中，利用大型语言模型（Agent）来编写和优化底层 GPU 代码（Kernel）成为了提升 AI 计算效率的热门路径。然而，当前的大部分尝试都面临着一个难以逾越的瓶颈：Agent 通常将编译器视为一个僵化的黑盒。它们提交代码、等待编译、运行测试，最后只能获得一个冰冷的报错信息或一个总体的执行延迟数字。系统从不告诉 Agent 究竟是哪一次同步引发了死锁、哪一种内存布局打破了硬件契约，或者是哪一步流水线导致了计算停滞。

ArXiv URL：https://arxiv.org/abs/2608.12629v1

另一方面，现存的编程语言也并不是为 AI 智能体设计的。如果让 Agent 使用像 Triton 这样的高层领域特定语言（DSL），它会发现那些区分“能跑的代码”和“专家级极致性能代码”的关键细节——例如 Warp 专门化、细粒度的内存层级调度和屏障同步编排——都被编译器完全隐藏了。而如果让 Agent 直接编写极度底层的 CUDA 或 PTX 代码，由于涉及极其复杂的布局代数计算，Agent 极易犯错且很难从海量崩溃中定位问题。

为了打破这种局面，Carnegie Mellon University 与 NVIDIA 的研究人员提出了 CAKE 框架（Compiler-Agent Co-Design）。该工作不仅为 Agent 量身打造了一种硬件显式的新型中间表示 Cake IR，更重要的是，它实现了编译器环境与大模型 Agent 的双向协同进化。基于 CAKE 框架，Agent 在毫无底层代码参考的“纯净启动”条件下，于 B200 GPU 上一举超越了经过人类极致微调的底层算子基准，并在合成新兴大模型的前沿 Kernel 时取得了高达 2.05 倍的执行提速。

这篇论文揭示了 AI 编写高性能底层基础设施的全新范式。

告别黑盒瞎猜：为智能体量身定制的 Cake IR

以往的代码 Agent 就像是蒙着眼睛的微调工人。它们只能通过“突变代码-运行耗时测试-保留较快版本”这种极低效率的盲搜来逼近最优解。要让专家级的算子调优知识固化下来，就必须将编译器内部丰富的静态分析、资源模型和合法性检查工具直接暴露给 Agent。

CAKE 选择的切入点是程序的表示形式。它放弃了直接生成底层 CUDA 代码，转而让 Agent 编写一种名为 Cake IR 的新型中间表示。Cake IR 是一种强类型的、直接表达硬件调度的语言。在 Cake IR 中，程序清晰地记录了机器的驱动方式：哪些 Warp（线程束）承担了什么具体的角色？哪些缓冲区被分阶段推入了多深的流水线？哪一道屏障（Barrier）正在管控数据的交接？

这种设计的精妙之处在于“声明式调度”与“细节隐藏”的平衡。Agent 只需要在 Cake IR 中显式规定“机器应当发生什么”（例如角色的划分与同步控制），而底层的编译器框架（Harness）则会自动推导出“如何实现”（例如具体的屏障内存地址、相位位元、TMEM 的偏移量以及底层的指令集编码）。更为关键的是，CAKE 刻意避免了让 Agent 去操作复杂的布局代数（Layout Algebra），存储和访问的决策直接依附于调度流程上，由编译器去验证生产者与消费者的数据表示是否与当前硬件兼容。

得益于这种结构，编译器现在能够在生成可执行代码前，就对程序进行验证并生成详细的性能预估模型。一旦代码中存在潜在的内存不安全或是同步违规，编译器抛给 Agent 的不再是单纯的“失败”，而是精确到具体资源、角色和代码阶段的诊断报告，使得大模型能够快速修正路径，极大节省了昂贵的 GPU 试错算力。

不仅仅是代码，环境本身也在进化

CAKE 最具颠覆性的机制，在于它并未把编译器的检查规则和 IR 原语视为一成不变的圣旨。在真正的专家级编程中，当遇到全新的前沿负载时，现有的抽象工具和检查规则往往会暴露出漏洞或表达能力的缺失。

在这套系统中，不仅 Agent 在迭代生成高性能代码，CAKE 的编译器外壳本身也成为了进化的目标对象。在整个底层的验证循环中：

沉淀新的规则：当 Agent 在某一种算子的推导中遭遇反复出现的失败模式时，这些反馈不会被简单丢弃或临时绕过，而是作为证据驱动验证器规则的更新。
丰富原语与成本模型：当发现某种物理调度模式无法被当前的 Cake IR 表达时，它会促发新 IR 原语的诞生；而性能预测失准的案例，则直接用于校准内部的成本模型（Cost Model）。

这种由真实算子的失败证据所驱动的“双轨进化”，使得 CAKE 系统具备了极强的适应性。它不但向后兼容从 Ampere 到 Blackwell 架构的各类 NVIDIA GPU，还在不断吸纳诸如注意力机制、稀疏与稠密 GEMM、量化运算、混合专家系统（MoE）以及 KNN 图计算等不同领域的硬件经验。

无参考启动下的惊艳成绩：超越人工极限

为了严格验证 CAKE 生成算子的真实水平，本文设计了一组极其严苛的基准测试。最核心的验证是基于 B200 显卡的 Flash-KMeans 算子“纯净启动（Clean-start）”实验。在这个场景下，Agent 可以阅读相关的数学规范、评估契约和高层 Python 描述，但被完全禁止访问任何底层的 CUDA、PTX 或 SASS 实现。这种设置强迫 Agent 从零推导出正确的物理调度。

实验对比了直接编写 CUDA/PTX 和编写 Cake IR 的差别。在规定消耗不超过 8000 万 Token 的预算范围内，直接编写 CUDA/PTX 的控制组没有任何一次能够达到预期的性能饱和点，其最终中位数性能仅为人类微调优化的 FlashML 库基准性能的 0.928 倍。

与此形成强烈反差的是，使用 Cake IR 的模型组展现出了碾压级的效率与性能：不仅所有独立运行都在 8000 万 Token 预算内突破了瓶颈，其中位数表现更是在 B200 上达到了手工微调基准的 1.144 倍。并且由于大量错误被提前拦截在了静态检查阶段，模型组达成最优解所需的活动演进时间仅为 1.89 小时，几乎是直接编写 CUDA 组（3.73 小时）的一半。这充分证明了显式机器调度带来的引导优势。

突进前沿架构：提速高达 2.05 倍

更令人振奋的是 CAKE 在实际前沿大模型算子中的表现。对于像 Kimi Delta Attention (KDA) 这样新涌现的模型架构，业内往往缺乏成熟度极高的底层实现作为参考。

在这种最具挑战性的“前沿合成”任务中，Agent 通过 Cake IR 成功重写了 KDA 的预填充（Prefill）阶段算子。新合成的算子在处理固定长度、变长压缩以及尾部输入时表现出了惊人的效率，跨六种不同的 B200 BF16 张量形状（Shape）测试中，相较于官方的 FlashKDA 基线实现了高达 2.05 倍的几何平均提速。该成果已经通过了严格的端到端服务验证，并作为上游合并请求（PR）提交至开源算子库 FlashInfer。

同样令人瞩目的还有对通信密集型超大型算子的重构。在处理 Alpha-MoE 的 W8A8 混合专家模型算子时，Agent 成功实现了令人咋舌的整合优化，将原本需要五次独立 GPU 活动的流程（包括路由收集、投影、激活、重量化和累加）完美融合进一个极其精简的巨型内核（Megakernel）中。在端到端 API 级别的对比中，这种整合削减了海量的调度空隙和工作区开销，于不同规模下取得了 4.025 倍到 6.204 倍的极限加速。

从单点跑分到构建生产级算子库

在学术界的许多研究中，模型在一个特定的张量形状下跑出高分往往就是终点。但 CAKE 的设计初衷是为了生产环境，因此论文在最后花大量篇幅阐述了如何从一个高度特化（Tuned Shape）的内核跨越到泛化的生产级算子分发库（Library）。

对于一个能接受任意维度输入的库而言，盲目追求某一维度的绝对优化反而会导致整体覆盖率崩塌。因此，CAKE 建立了一个独立的泛化阶段：

分桶与路由编排：它将先前挖掘出的不同维度下的高性能内核“种子”按规模进行聚类分组（Shape Buckets），自动生成一个前置的分发器（Dispatcher），并为一个庞大库设置统一的逻辑入口。
严格防止评估泄露：为了确保泛化能力的真实性，分发器的边界和评估样本集被强行隔离，Agent 无法通过切割特定输入形状来粉饰覆盖率数据。

在经过这一阶段的组装后，像 KNN 和 KMeans 这样的底层算子族能够在超过 400 种不同的输入形状配置下，均实现 1.42 到 2.12 倍的稳定性能提升，这标志着由大模型生成的底层代码真正具备了商用泛化水平。

总结

CAKE 证明了在利用 AI 接管底层计算基础设施的道路上，仅仅提升模型本身的推理能力是不够的。现有的高级 DSL 无法释放硬件底层的全部潜力，而传统底层语言的报错黑盒又让智能体深陷泥潭。通过将编译器转化为一个能够输出细粒度结构化反馈、并且自身能够随新负载进化的动态系统，CAKE 重塑了大模型编写高性能算子的范式。可以预见，随着该类协同设计框架的成熟，未来前沿 AI 模型算子的极致优化，将越来越脱离对顶尖底层专家的重度依赖，转向由智能体与编译器共同主导的自动化纪元。