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

GPU内核智能体（kernel agents）与GPU编程语言长期以来各自独立发展，导致专家级内核难以复现。现有智能体通常将编译器视为一个不可更改的“黑盒”，仅能接收编译错误、正确性验证结果及执行耗时等有限反馈；而当前各类领域特定语言（DSL）则要么完全隐藏关键的调度决策，要么仅通过晦涩难懂的内存布局抽象来暴露这些决策。本文提出CAKE——一种编译器与智能体协同设计的新范式：智能体直接编写CAKE中间表示（CAKE IR），这是一种具备类型安全、显式刻画硬件特性的调度表达形式。CAKE完整暴露了线程束（warp）角色划分、数据搬移路径、同步机制以及流水线结构，同时原生支持形式化验证、性能代价建模和局部化诊断能力。更重要的是，整个运行环境（harness）本身具备自演化能力：反复出现的失败案例会自动转化为新的验证规则、IR原语、代价模型校准参数，以及可复用的优化策略。在B200 GPU上、实现细节完全隐藏的Flash-KMeans基准测试中（即“clean start”模式），CAKE在8000万token预算下生成的最优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，且明确区分了面向单一形状的迭代优化与面向通用库的泛化设计及动态分发机制。

GPU kernel optimization suffers from a deep disconnect between AI agents (which treat compilers as opaque black boxes) and GPU programming languages/DSLs (which either oversimplify or over-complicate hardware-aware scheduling). This makes expert-level, performance-portable kernels difficult to reproduce, verify, or iteratively improve—especially across evolving GPU architectures and diverse workloads.

CAKE introduces compiler-agent co-design: agents generate CAKE IR—a typed, hardware-explicit intermediate representation that exposes warp roles, memory movement, synchronization, and pipelining—while the compiler harness evolves *with* the agent via feedback loops (e.g., recurring failures → new verifier rules, IR primitives, cost model calibrations, and reusable tactics). Unlike static DSLs or black-box agent tuning, CAKE enables *localized, interpretable, and verifiable* kernel synthesis grounded in real hardware semantics.

On NVIDIA B200: (1) Flash-KMeans at 80M tokens achieves 1.144× speedup over hand-tuned FlashML baseline (vs. 0.928× for direct CUDA/PTX); (2) Agent-generated Kimi Delta Attention delivers 2.05× geometric-mean speedup over official FlashKDA and passes end-to-end serving validation; (3) Dispatcher-backed KNN/KMeans show 1.42×–2.12× gains across >400 tensor shapes; (4) Four CAKE-driven kernel improvements submitted as upstream PRs; (5) Supports NVIDIA Ampere through Blackwell; separates single-shape evolution from library generalization and dispatch. No public code release mentioned; evaluation uses real hardware timing, correctness verification, and production-serving validation—not synthetic proxies.

FlashAttention (Dao et al., 2022); Triton (Tillet et al., 2019); Halide (Ragan-Kelley et al., 2013); TACO (Kjolstad et al., 2017); AccelWare (Zhang et al., 2021); SODA (Wang et al., 2023); GPT-Code (Chen et al., 2024); LLM-based CUDA generation (Li et al., 2023); AlphaDev (DeepMind, 2023); CompilerGym (Buchlovsky et al., 2022)