arXiv:2608.12629cs.LG2026-08

让编程代理与编译器协同设计,自动生成高效GPU内核。

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

  • 代理直接编写硬件显式的调度中间表示CAKE IR,控制线程组角色和内存操作。
  • 在8000万次评估下,性能达基线的1.144倍,超越手动调优的CUDA/PTX。
  • 适合研究编译优化、自动化内核生成或追求极致性能的工程师。

GPU内核代理与编程语言发展脱节,导致专家级内核难以复现。现有代理将编译器视为黑盒,仅获错误、正确性结果和运行时间;而传统领域特定语言或隐藏关键调度决策,或以复杂布局抽象暴露。本文提出CAKE:一种编译器-代理协同设计框架,代理编写CAKE IR——一种带类型的、硬件显式调度表示,可显式表达线程组角色、内存移动、同步及流水线。该框架支持验证、成本建模与局部诊断,其测试环境也持续进化:重复失败转化为验证规则、IR原语、模型校准与可复用优化策略。在匹配实现但隐藏细节的B200上,8000万次评估预算下,最佳CAKE IR达到基线1.144倍性能,远超直接使用CUDA/PTX的0.928倍。此外,代理生成的Kimi Delta Attention相比官方FlashKDA实现2.05倍几何平均加速,并通过端到端服务验证;基于调度器的KNN与KMeans在超过400种形状上提升1.42至2.12倍,四个内核变更已提交上游PR。CAKE适用于NVIDIA Ampere至Blackwell系列GPU,区分单形状演化与库级泛化及分派。

原文摘要 · Abstract (English)

GPU kernel agents and GPU programming languages have advanced separately, leaving expert kernels difficult to reproduce. Agents usually treat the compiler as a fixed black box and receive only errors, correctness outcomes, and timing, while existing DSLs either hide critical scheduling decisions or expose them through difficult layout abstractions. We present CAKE, a compiler-agent co-design in which agents author CAKE IR, a typed, hardware-explicit schedule representation. CAKE exposes warp roles, memory movement, synchronization, and pipelines while supporting verification, cost modeling, and localized diagnostics. The harness itself evolves: recurring failures become verifier rules, IR primitives, model calibrations, and reusable optimization tactics. In matched implementation-hidden Flash-KMeans clean starts on B200, the best CAKE IR candidate at an 80-million-token budget runs at 1.144x the tuned FlashML baseline, compared with 0.928x for direct CUDA/PTX. Beyond this benchmark, agent-generated Kimi Delta Attention achieves a 2.05x geometric-mean speedup over official FlashKDA and passes end-to-end serving validation. Dispatcher-backed KNN and KMeans improve performance by 1.42x to 2.12x across more than 400 shapes, and four kernel changes are available as upstream PRs. CAKE targets NVIDIA GPUs from Ampere through Blackwell and separates single-shape evolution from library generalization and dispatch.

编译器GPU优化自动化内核生成

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。