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

发表时间: 2026-08 · arXiv:2608.12629 (NVIDIA, CMU)

原文: https://arxiv.org/abs/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, Carnegie Mellon University

速读

一句话结论 本文提出了一个编译器与大模型代理协同设计的系统 Cake，通过让代理编写带有显式硬件调度且无需复杂布局演算的中间表示（IR），并利用编译器提供局部化诊断与自身演进，成功在缺乏底层参考的情况下合成了超越现有基线的高性能 GPU 算子。

要解决什么问题 现有大模型代理编写 GPU 算子的方案卡在反馈机制与抽象层级上。一方面，代理通常将编译器视为黑盒，环境只返回编译报错或端到端延迟，无法指出具体是哪个程序决策导致了同步失败、硬件契约违规或流水线停滞。另一方面，现有的领域特定语言并不适合代理：Triton 等 Tile 级别的语言隐藏了区分专家级算子与普通算子的关键硬件抽象（如 warp 特化、屏障编排和内存分层放置）；而 CuTe 等底层的语言虽然暴露了硬件控制权，却要求代理掌握复杂的布局代数，这使得代理极易出错且难以定位问题。这种反馈黑盒化与抽象两难的卡点，阻碍了代理独立探索和生成前沿架构下的高性能算子。

怎么做的 核心思路是将代理与编译器进行协同设计，让编译器从被动的测试黑盒变成与代理共同演进的合作者。这套机制既提供了细粒度的硬件控制，又通过结构化反馈避免了代理陷入复杂代数，其关键设计由三个部件构成。首先是类型化且硬件显式的中间表示 Cake IR。代理直接用它编写调度，通过四个属性发挥作用：使用固定的类型化词汇表处理计算与同步；通过资源声明让 IR 掌握数据形状与生命周期；显式命名 warp 组并让跨角色交接保持可见；由降级过程自动推导屏障地址、相位位和 TMEM 偏移量等物理细节。它故意不把内存布局作为一等抽象，代理只需记录存储和访问决策，由编译器检查兼容性，从而免去了布局演算。其次是局部化诊断的测试环境。在代码真正消耗 GPU 时间前，编译器会对 Cake IR 进行静态分析，检查同步、内存安全和数据流违规，并用校准过的成本模型估算性能。若出错，环境会精准指出受影响的程序区域和违规类别，在早期过滤掉无效候选者。最后是编译器自身的演进循环。系统利用代理反复出现的失败反馈，将其提炼为新的验证器规则、IR 原语和静态检查。在工作流中，代理生成的候选 IR 经硬门控和成本模型过滤后，幸存者才进行基准测试，诊断证据会再次路由回系统。泛化阶段则将强大的单一形状算子分组到形状桶中，生成特化或共享变体，构建出支持任意形状调用的算子库。

效果如何 实验在从 Ampere 到 Blackwell 架构的 NVIDIA GPU 上进行（主要基准测试在 B200，泛化评估在 GB200），代理模型使用 GPT-5.6-sol。对比基线涵盖了代表当前算子库与服务框架路线的 SGLang、vLLM、FlashInfer、Triton、CUTLASS、DeepGEMM 和 FlashAttention-4。在 B200 的 Flash-KMeans 冷启动测试中，代理直接写 CUDA 或 PTX 的成功率为零（最高仅达调优基线的 0.928 倍）；而使用 Cake IR 在消耗 5500 万 tokens 时便超越基线，最终性能达 1.144 倍。在无底层参考的前沿算子合成任务中，生成的 Kimi Delta Attention 预填充算子比官方 FlashKDA 快 2.05 倍，解码路径比 FlashInfer 快 1.14 倍；小形状 TinyGEMM 算子比 FlashInfer 的 TensorRT-LLM 派生内核减少 18% 到 23% 耗时；将 Alpha-MoE 的 W8A8 算子从 Hopper 重写到 Blackwell 后，相比 TensorRT-LLM 派生基线实现 4.025 倍到 6.204 倍的端到端加速。复现已知算子时，11 个测试中有 10 个达到或超越最先进参考标准，且物理代码行数更短。在 GB200 的库级泛化评估中，KNN 构建、KNN 搜索和 KMeans 任务在一百多个形状上分别实现了 1.418 倍、2.116 倍和 1.803 倍的几何平均加速比。该方法的局限性在于：目前的静态分析有意保持不完整，无法证明全局 GPU 的正确性或捕获所有微架构行为；此外，系统目前仅支持 NVIDIA 硬件，尚未探索对其他目标硬件的支持。

主要贡献

当前编写GPU程序的代码代理（Agents）通常将编译器视为一个固定的黑盒：代理提出代码、编译、运行测试并获取延迟，但环境只返回编译器错误或端到端的时间信号，无法指出是哪个程序决策导致了同步失败、硬件契约违规或流水线停滞。同时，现有的领域特定语言（DSL）并不适合代理：Tile级别的DSL隐藏了区分专家级内核与普通内核的warp特化、屏障编排和内存分层放置；而低级别的DSL虽然暴露了控制权，但要求复杂的布局演算，这使得代理极易出错且难以定位。

为了解决这一问题，本文提出了Cake，它将代理与编译器进行了协同设计。其核心创新点包括：
1. 类型化、硬件显式的IR：代理编写Cake IR，这是一种类型化、硬件显式的调度表示，无需布局代数即可提供细粒度的控制，并携带足够的信息供验证器和成本模型在编译前对程序进行推理。
2. 局部化诊断：编译器测试环境（harness）不再仅仅返回通过/失败，而是返回局部化的正确性和性能诊断，从而在候选者消耗GPU时间之前对其进行过滤。
3. 编译器自身的演进：测试环境本身也是演进的目标。代理反复出现的失败会被转化为新的验证器规则、IR原语、成本模型校准和可重用的策略，而不是一次性的变通方法。

背景知识与设计原则

专家级内核程序员在工作时，通常会保持一个紧凑的工作负载模型，对显式硬件资源进行推理，并在不同内核之间传递可重用的规则。Cake旨在将编译器中已有的机制（结构化操作词汇、资源模型、合法性检查、静态分析等）面向代理开放。为此，Cake IR遵循以下八个核心设计原则：
* 符合人体工程学（P1）：保持NumPy/PyTorch用户熟悉的编辑模式，避免不必要的目标传递或网格簿记。
* 性能透明（P2）：保持与性能相关的硬件决策可见，并使降级（lowering）行为可检查。
* 规范化（P3）：每个操作倾向于使用一种规范形式，而不是等效的替代拼写。
* 静态类型检查（P4）：使用类型规则约束操作降级，并在构建期间拒绝类型错误的程序。
* 分析友好（P5）：暴露支持的静态分析所需的信息。
* 测试门控（P6）：针对静态分析和编译的内核矩阵测试来评估IR更改。
* 分析一致性（P7）：伴随IR数据模型的更改，进行相应的分析更新。
* 硬件接地（P8）：记录每个操作的预期硬件行为。

方法细节

自底向上的IR演进
Cake并没有从一个预定义的Cake IR词汇表开始。它的起始材料是生产级内核的语料库以及硬件设计原则。代理首先从这些内核中识别出重复的调度模式或缺失的能力；接着修改IR及其编译器支持；最后将内核移植并在修改后的系统中进行验证。当遇到下一个内核家族或验证暴露出新的差距时，这个三步循环将再次启动。这种持续的循环既产生了当前的IR，也推动着它不断扩展。

Cake IR中的显式机器调度
Cake IR详细记录了应该如何驱动机器：哪些warp承担哪些角色、哪些缓冲区被分级到何种深度、哪个屏障（barrier）控制哪次交接，以及哪种指令形式消耗哪个操作数。程序结合了显式操作、声明的资源、warp角色以及网格和流水线配置。它通过四个属性发挥作用：首先是类型检查的词汇表，计算、内存移动、同步等使用固定的IR词汇而不是嵌入的C或PTX；其次是声明的资源，内存区域、同步状态和流水线只需声明一次，IR就能知道每个缓冲区的形状、数据类型和生命周期；接着是显式角色，为warp组命名，并且每次跨角色的交接都是可见的；最后是自动派生的元数据，诸如屏障地址、相位位、TMEM偏移量和warp身份等物理细节由降级过程推导，而不是由代理编写。这种设计的收益在于，分析工具可以在代码生成之前基于显式的调度决策进行推理，将发现的问题准确定位到受影响的资源或阶段。

布局（Layout）处理
Cake故意不将布局作为一等抽象。Cake IR不要求代理操作复杂的布局代数，而是直接在调度中记录存储和访问决策。随后，编译器会检查生产者和消费者的表示是否与目标硬件兼容。

架构与降级（Lowering）
相同的调度语言针对从Ampere到Blackwell的NVIDIA GPU，因此角色-屏障-流水线调度的结构具有可移植性，而指令的准入和降级则保持特定于目标。Cake会精确映射连接的GPU，明确报告不支持的目标而不是默默替换架构，并且仅在有特定目标校准可用的情况下才发出性能估计。

分析与验证
在编译之前，测试环境（harness）会对类型化调度进行广泛类别的检查，包括同步、内存安全、数据流、资源、指令和数据表示违规。这些检查通过稳定的分析接口，拒绝那些在数学上合理但与目标执行模型不兼容的候选者，并指出受影响的程序区域和违规的契约类别。对于数值正确性，系统会在不同形状和输入分布下比较内核与参考输出，并在目标框架中进行端到端评估。此外，经过校准的成本模型会估计候选者的性能，返回高级别的瓶颈归因和优化建议，用于在消耗GPU时间之前对候选者进行排序和过滤。

编译器演进
Cake在内核演进的同时推进编译器的演进。候选内核、验证结果、基准测试和失败报告构成了提出和验证编译器更改的证据。编译器演进遵循两条互补的路径：在第一条路径中，代理检查生产内核和硬件文档以寻找缺失的Blackwell模式（如新指令形式、资源类型、同步习惯用法），并提出编译器更改建议，每项建议在实施前都会根据Cake IR设计原则进行检查；在第二条路径中，代理利用失败候选者的反馈（如消毒剂报告、正确性不匹配、调试日志），将反复出现或高成本的失败模式提炼为新的分析规则，例如将不透明的运行时崩溃转化为验证器规则，将非法的降级模式转化为静态检查。这两条路径相互耦合，新原语为编译器暴露了更多硬件事实从而增强分析能力，而新分析反过来又限制了未来原语的设计空间。所有编译器更改都必须在内核语料库中进行测试门控验证。

代理工作流
一次内核演进运行包含四个阶段：首先生成结构上不同的Cake IR候选者；接着在消耗GPU时间之前，利用IR构建检查、验证器硬门控和成本模型排序对它们进行过滤；然后利用基准测试和分析器证据，根据外部预言机（oracle）评估幸存者；最后根据诊断结果，将产生的证据路由回候选者、验证器、成本模型或IR词汇表中。所有报告的代理任务均使用推理工作量设置为xhigh的GPT-5.6-sol【10，GPT-5.6 Sol Model + 2026 + OpenAI API Documentation + https://developers.openai.com/api/docs/models/gpt-5.6-sol】 。

从调优形状到库的泛化
将针对单一形状调优的内核转化为支持任意形状调用的库是一个独立的阶段。首先，泛化仅在存在强大的单形状种子后才开始，并根据固定工作负载下包含调度器（dispatcher-inclusive）的性能进行评分，不正确或缓慢的种子会退回内循环。接着，泛化阶段将测量的种子分组到形状桶中，生成特化或共享的变体，并在显式回退（fallback）之后对它们的守卫（guards）进行排序。调优可以改变实现参数但不能改变输入形状。最后，为防止评估泄漏，有效的形状域在调优前声明，调度器谓词可以划分该域但不能引入新的评估行。系统策略要求尽可能在形状域中重用单个物理调度，只有当域需要实质性的调度更改时才引入新的调度路由。

实验环境

• 硬件配置：评估涵盖从Ampere到Blackwell的NVIDIA GPU。主要基准测试和正确性检查在B200 GPU上进行（每次定时采样前刷新L2缓存），泛化投资组合评估在GB200上进行。
• 软件配置：代理使用GPT-5.6-sol大语言模型。基准测试和对比框架包括SGLang、vLLM、FlashInfer、Triton、CUTLASS、DeepGEMM和FlashAttention-4。
• 数据集与工作负载：
  • Flash-KMeans：用于Sparse VideoGen2中语义感知token排列的精确$k$-means工作负载（测试形状如$B=32$, $N=65536$, $K=1024$, $D=128$）。
  • 前沿内核合成：Kimi Delta Attention (KDA)、Gated DeltaNet、稀疏注意力、TinyGEMM、Alpha-MoE（W8A8融合MoE巨型内核）。
  • 已知内核复现：注意力前向/后向和解码、低精度GEMM、MQA/MLA logits和解码内核。

实验结果

1. Flash-KMeans冷启动演进

• 实验内容：在隐藏底层实现的B200冷启动环境中，对比代理编写Cake IR与直接编写CUDA C++/PTX的能力（8000万token预算，各运行3次）。
• 实验结果：编写Cake IR的组在3/3次运行中均达到了预设的平台期，中位数最佳性能达到调优后FlashML基准的1.144倍；而直接编写CUDA/PTX的组为0/3次，中位数为0.928倍。Cake IR的平均性能在消耗5500万token时就超越了基准线。
• 分析结论：Cake IR及其结构化反馈显著提高了代理在未见过底层实现的复杂工作负载上的探索效率和最终性能。
  

2. 前沿内核合成 (Frontier-Kernel Synthesis)
* 实验内容：在不提供底层目标实现作为参考的情况下，合成新兴模型架构的内核。
* 实验结果：
* Kimi Delta Attention (KDA)：在六个B200 BF16形状上，生成的prefill内核比官方FlashKDA基准的几何平均速度提高了2.05倍，并在SGLang端到端服务中得到验证。单独的decode路径比上游FlashInfer提高了1.14倍（图6）。
* TinyGEMM：针对小形状，代理生成了自适应深浅流水线家族，在35个规范形状上比FlashInfer的TensorRT-LLM派生内核实现了18-23%的几何平均内核时间减少（图7）。
* Alpha-MoE：代理成功将原针对Hopper的W8A8融合MoE巨型内核重写为Blackwell版本。相比TensorRT-LLM派生的预路由API，在$N=256$和$N=512$时端到端API级别加速比分别达到6.204倍和4.025倍（图8）。

• 分析结论：Cake的IR和演进系统能够在缺乏已知良好设计锚点的情况下，发现并表达高度复杂、通信密集的物理调度。

3. 已知内核复现 (Known-Kernel Reproduction)
* 实验内容：针对LLM关键执行路径上的已知算子家族（如注意力、低精度GEMM、MQA/MLA解码），将Cake生成的内核与TensorRT-LLM、CUTLASS、DeepGEMM、FlashAttention-4等SOTA基准进行比较。
* 实验结果：在11个固定的对比中，有10个达到或超过了列出的参考标准（如MQA索引器达到了约1.27倍的性能），剩余1个达到了参考的96.5%。所有Cake IR实现的物理代码行数均短于其审计的参考设备核心代码（见表4）。
* 分析结论：Cake IR能够紧凑地表示评估的硬件调度，同时保留分析和降级所需的决策，证明了其在生产质量输出上的能力。

4. 调度器投资组合泛化 (Dispatcher Portfolio)
* 实验内容：在GB200上评估泛化后的KNN构建、KNN搜索和KMeans库的包含调度器的整体性能。
* 实验结果：在112、198和124个形状上，分别实现了1.418倍、2.116倍和1.803倍的无加权几何平均加速比（$G_{span}$），且没有错误输出（图9）。
* 分析结论：Cake不仅能优化单一形状，还能成功构建覆盖广泛形状域的高性能库级实现。

补充细节

相关工作对比
现有的GPU编程系统对代理驱动的内核开发存在两难：高级和中级Tile DSL（如Triton、TileLang）向代理隐藏了硬件抽象，阻碍了代理表达区分专家级内核的warp特化、屏障编排和内存分层放置；而低级DSL（如CuTe DSL）虽然暴露了硬件控制，但要求代理掌握特定领域的专业知识（如布局代数），这极易导致错误且难以定位。Cake通过共同设计IR解决了这一差距，提供了细粒度的硬件控制而无需布局演算。此外，与现出演进搜索过程或模型权重的内核代理系统不同，Cake改变的是被搜索的底层表示形式，以及返回给代理的结构化编译器证据。

结论

本文提出了Cake，一个将编译器环境作为内核代理不断演进的合作者的系统。Cake目前支持从Ampere到Blackwell架构的NVIDIA GPU。通过让代理编写类型化、硬件显式的Cake IR，系统能够将硬件决策暴露给结构化分析，并在代码生成前提供局部化的诊断。同时，编译器演进循环将代理反复出现的失败转化为可重用的验证器规则和新原语。实验证明，Cake不仅能在冷启动下超越调优基准，还能在未参考底层实现的情况下合成前沿内核，并以更紧凑的代码复现专家级内核。未来的工作需要探索非NVIDIA目标的支持，并进一步完善目前有意保持不完整的静态分析和性能模型。

附录细节

IR抽象演进过程（附录A）
Cake IR的演进是一个代理驱动的抽象发现过程：首先从生产库（如FlashInfer、CUTLASS等）收集高质量的CUDA内核作为语料库；接着，代理分析语料库并提取出重复的模式（如屏障编排、流水线分段、TMA描述符设置等）作为候选抽象；然后，人类专业知识将这些抽象偏向Blackwell编程模型的设计（如将TMEM作为一等资源、异步屏障等）；随后，每个候选抽象都会根据设计原则进行验证和迭代；最后，通过持续移植新内核来扩展IR，每个成功的移植验证了抽象，而暴露的差距则触发对IR或其降级过程的扩展提案。

能力类别与资源模型（附录B）
Cake IR将其操作词汇分为四大类：计算（矩阵、逐元素和归约操作）、内存移动（全局与片上内存层之间的显式传输，包括异步和集体传输模式）、同步（跨角色、流水线阶段和硬件范围的排序与协调）以及控制和调度（warp角色分配、流水线管理和多块协调）。程序声明五种资源：共享内存区域、张量内存区域、同步对象、warp角色和流水线。这些声明记录了降级和分析调度所需的信息（包括数据形状、所有权和生命周期），使得硬件敏感的选择在程序中保持可见。

架构覆盖细节（附录B.5）
编译器要求精确的目标匹配。例如，sm_80（A100）支持mma.sync、ldmatrix等，但没有TMA或TMEM；sm_90a（H100/H200）增加了WGMMA、TMA、线程块集群和异步屏障；sm_100a（B200）进一步增加了tcgen05.mma和TMEM支持。如果目标不支持，编译器会明确报告缺失的设备或工具链支持。

分析范围与验证（附录C）
测试环境在支持的构造范围内对类型化调度进行安全、一致性、语义和数据一致性类别的评估。静态分析仅在其建模域内作为编译前门控，它不能证明全局GPU的正确性，也无法捕获所有的微架构行为。缺失的信息或分析覆盖范围会被显式报告，而不是被默认视为检查成功。