Command Palette
Search for a command to run...
CAKE:面向前沿内核演化的编译器–智能体协同设计
CAKE:面向前沿内核演化的编译器–智能体协同设计
摘要
GPU 内核智能体的研究与 GPU 编程语言的研究长期各自独立发展,二者之间的鸿沟正是专家级内核流失之处。内核智能体将编译器视为固定的黑盒:它们改进提案、变异和排序,但环境只返回编译器错误、正确性结果和端到端计时——这些信号既无法说明究竟是哪个程序决策导致了同步失败、硬件契约违规或流水线停顿,也无法在前沿工作负载暴露出缺失能力时随之增长。与此同时,智能体可能使用的语言也并非为其量身打造。Tile 级 DSL 隐藏了经线特化、屏障编排和内存层级放置,而这些正是专家级内核区别于仅仅正确内核的关键;低级 DSL 虽然暴露了这些控制能力,却要求一种布局演算,使智能体既容易出错,又难以定位错误。我们提出 Cake,对二者进行协同设计。智能体编写 Cake IR,这是一种带类型、硬件显式的调度表示,无需布局代数即可提供细粒度控制,并携带足够信息,使验证器和成本模型能够在编译之前对程序进行推理;测试框架以局部化的正确性与性能诊断作为回应,并且其本身也是演化的对象,因此反复出现的失败会转化为新的验证器规则、IR 原语、成本模型校准和可复用策略,而非一次性的变通方案。在 B200 上对实现隐藏的 Flash-KMeans 进行匹配的干净启动实验(每种表示运行三次)中,在 8000 万 token 预算下,使用 Cake IR 的最佳候选方案达到调优后 FlashML 基线中位性能的 1.144 倍,而直接使用 CUDA/PTX 仅为 0.928 倍。除干净启动基准之外,智能体生成的 Kimi Delta Attention 相比官方 FlashKDA 实现了 2.05 倍的几何平均加速,并在端到端服务中得到验证。基于 Dispatcher 的 KNN 和 KMeans 系列在 400 多种形状上性能提升 1.42 倍至 2.12 倍,且有四项内核改动已作为上游 PR 提供。Cake 面向从 Ampere 到 Blackwell 的 NVIDIA GPU,并将单形状演化与库集成所需的泛化和调度阶段相分离。
一句话总结
NVIDIA 和卡内基梅隆大学的研究人员提出 Cake,一种面向 GPU 内核演化的编译器与 Agent 协同设计,其类型化、硬件显式的 Cake IR 在没有布局代数的情况下提供细粒度控制,并支持随反复失败而演化的局部化验证器和成本模型诊断;在 B200 上,针对 FlashKMeans,相对调优 FlashML 基线实现 1.144× 加速,而直接 CUDA/PTX 为 0.928×;针对 Kimi Delta Attention,实现相对 FlashKDA 的 2.05× 几何平均加速。
核心贡献
- Cake 为 GPU 内核 Agent 提供类型化、硬件显式的 IR 和协同设计的编译器 harness,无需布局代数即可进行细粒度控制,并支持在编译前进行验证器和成本模型推理。
- harness 是一个可演化目标:反复出现的内核失败会变成新的验证器规则、IR 原语、成本模型校准和可复用策略,并向 Agent 返回局部化的正确性和性能诊断。
- 在 B200 上实现隐藏的 Flash-KMeans 干净启动中,8000 万 token 预算下最佳 Cake IR 候选的中位性能达到调优 FlashML 基线的 1.144×,而直接 CUDA/PTX 为 0.928×;Agent 生成的 Kimi Delta Attention 相对官方 FlashKDA 达到 2.05× 几何平均加速,并在端到端服务中得到验证;带 dispatcher 的 KNN 和 KMeans 家族在 400 多个 shape 上性能提升 1.42× 到 2.12×。
引言
编码 Agent 越来越频繁地编写和修改 GPU 内核,但大多数系统将编程环境视为固定黑盒:它们编译、测试、测量延迟并编辑,因此一次崩溃并不能指出违反了哪条安全或硬件条件,单一延迟数字也不能说明哪个程序决策限制了性能。现有 GPU DSL 对 Agent 驱动的内核开发也不方便,因为高层 tile DSL 隐藏了专家内核所需的 warp 特化、barrier 编排和内存层级放置,而低层 DSL 需要布局代数专业知识,并可能生成脆弱代码。作者提出 Cake,这个系统让编译器成为内核 Agent 的持续演化协作者。Cake 让 Agent 编辑类型化 IR 而非原始 CUDA,返回局部化的正确性和性能诊断而非通过/失败反馈,并将反复失败转化为验证器规则、校准任务或新原语,同时受语料测试和人工合并门控约束。
方法
作者采用自底向上的方式设计 Cake IR,而不是从预定义词汇开始。过程从生产内核语料库和硬件设计原则开始。Agent 识别反复出现的调度或缺失能力,修改 IR 及其编译器支持,然后将内核移植到修订后的系统并针对该系统进行验证。该循环针对新内核家族或验证暴露的缺口反复进行,持续扩展 IR。
参见框架图:
Cake IR 记录显式机器调度,详细描述应如何驱动机器,包括 warp 角色、缓冲区暂存、barrier 门控和指令形式。一个程序组合显式操作、声明资源、warp 角色和网格配置。关键属性包括用于计算和内存移动的类型检查词汇、用于内存和同步状态的声明资源、显式 warp group 角色,以及自动推导的元数据,其中机械性结果被下沉生成而非手工编写。这种设计使分析能够在代码生成前从显式调度决策进行推理。系统面向从 Ampere 到 Blackwell 的 NVIDIA GPU,精确映射所连接的 GPU,并在校准可用时输出性能估计。
编译器 harness 作为围绕 Cake IR 的面向 Agent 的环境。人类提供目标分析的高层描述,Agent 在验证下实现、维护和完善这些分析。编译前,harness 检查类型化调度是否存在同步、内存安全、数据流和资源违规。它还根据参考输出验证数值正确性,并使用校准后的成本模型估计性能和候选排序。
编译器演化沿两条互补路径进行,在内核开发的同时改进系统。
如下图所示:
在第一条路径中,Agent 检查生产内核和硬件文档,识别缺失模式,例如新指令形式或同步惯用法,并形成编译器变更提案。这些提案在实现前根据设计原则进行检查。在第二条路径中,Agent 利用来自失败候选的反馈,包括 sanitizer 报告和调试日志,将反复出现的失败模式提炼为新分析。例如,难以理解的运行时崩溃变成验证器规则,反复出现的非法 lowering 模式变成静态检查。这两条路径相互耦合;新原语暴露硬件事实,使分析更强,而新分析约束未来原语的设计空间。变更在内核语料库中进行测试门控,以确保原语和分析共同演化。
外部工作负载契约作为 Agent 工作流的稳定权威,由四个阶段组成。首先,Agent 生成结构上不同的 Cake IR 候选。其次,使用 IR 构造检查、验证器硬门控和成本模型排序过滤这些候选,以避免不必要的 GPU 使用。第三,幸存候选根据外部 oracle,通过基准测试和 profiler 证据进行评估。最后,根据诊断将所得证据路由到候选、验证器、成本模型或 IR 词汇表。
为了从调优后的 shape 过渡到库,系统采用单独的泛化阶段。该阶段将测量种子分组为 shape 桶,生成专用或共享变体,并将其 guard 排序在显式 fallback 之后。验证覆盖代表性输入、边界情况和 fallback 路径。系统尽可能在 shape 域中复用单一物理调度,只有需要实质性调度变更时才引入新路由,确保路由复杂度由实测工作负载收益证明。
实验
评估涵盖干净启动演化、前沿内核合成、已知内核复现和包含 dispatcher 的库泛化。干净启动 Flash-KMeans 实验表明 Cake IR 超过调优基线,而直接 CUDA/PTX 分支未超过;KDA、Gated DeltaNet、MiniMax 稀疏注意力、TinyGEMM 和 Alpha-MoE 等前沿内核在无低级参考的情况下被合成,并优于黑盒基线。已知内核复现通过在几乎所有固定比较中达到或超过高度优化的专家参考,验证了生产质量输出。经过验证的语料库和泛化组合进一步确认了广泛的架构支持、正确的输出以及跨大量 shape 的一致 dispatcher 级加速。
harness 暴露七类分析和验证,涵盖编译前门控、执行门控、报告和非阻塞提示。预编译检查在候选执行前拒绝同步、内存使用、资源、指令、数据流和调度结构违规,而数值验证将编译输出与权威外部参考进行比较。性能分析估计成本并识别大致瓶颈类别,但设备上测量和 profiling 仍是最终 ground truth。程序安全、硬件一致性、数据一致性和调度语义作为预编译门控,用局部化原因阻止无效候选。数值验证是执行门控,要求在最终接受前与权威外部参考一致。性能分析返回估计成本和大致瓶颈归因,但不替代设备上测量作为最终 ground truth。优化指导是建议性的,提出有希望的修订而不阻止编译。
在匹配的干净启动 Flash-KMeans 运行中,在 8000 万 token 预算下,CAKE IR 分支在全部三次运行中达到平台期,而直接 CUDA/PTX 分支在任何一次运行中均未达到平台期。CAKE IR 分支还使用较低的中位活跃演化时间,并在 80M 预算下以更高的最佳中位性能结束。直接 CUDA/PTX 的最佳中位性能低于参考水平。到 8000 万 token 预算时,CAKE IR 在每次运行中均达到平台期,而直接 CUDA/PTX 达到零次。CAKE IR 的中位活跃演化时间明显低于直接 CUDA/PTX。在 80M token 预算下,CAKE IR 的最佳中位性能高于参考,而直接 CUDA/PTX 低于参考。
评估将同步、内存、资源、指令、数据流和调度结构违规的预编译检查与针对权威参考的执行时数值门控相结合,而性能分析和优化指导仍为建议性。在匹配的干净启动 Flash-KMeans 运行中,在 8000 万 token 预算下,CAKE IR 分支在全部三次运行中达到性能平台期,使用较低的中位活跃演化时间,并最终高于参考,而直接 CUDA/PTX 一次也未达到平台期并低于参考。这些结果表明,在相同预算下,CAKE IR 路径带来更可靠的收敛和更好的最终性能。