CAKE: Compiler-Agent Co-Design for Frontier Kernel Evolution
paperYour notes
NVIDIA and CMU (Zihao Ye and Yingyi Huang co-first; senior authors Vinod Grover, Tianqi Chen and Luis Ceze) argue that GPU-kernel agents and GPU DSLs have evolved separately: agents treat the compiler as a black box that returns only errors, pass/fail and timings, while tile-level DSLs hide the warp specialization, barrier choreography and memory-tier placement that make expert kernels fast, and low-level DSLs demand a layout calculus agents handle poorly. CAKE is a compiler–agent co-design: agents author Cake IR, a typed, hardware-explicit schedule representation exposing warp roles, memory movement, synchronization and pipelines, backed by verifier hard gates, a calibrated cost model and localized diagnostics. The harness itself evolves — recurring failures are promoted into new verifier rules, IR primitives, model calibrations and reusable optimization tactics — and a four-stage loop (generate structurally distinct candidates → filter by IR checks, verifier and cost model before spending GPU time → benchmark survivors with profiler evidence → route findings back to candidate, verifier, cost model or IR vocabulary) keeps the process auditable. All agent runs use GPT-5.6 Sol at xhigh reasoning effort so gains are attributable to the environment, not the model.
In matched implementation-hidden Flash-KMeans clean starts on B200 with an 80M-token budget, the best Cake IR candidate reaches 1.144× the tuned FlashML baseline versus 0.928× for agents writing CUDA/PTX directly. With FlashKDA available only as a black-box timing oracle, agent-generated Kimi Delta Attention prefill kernels hit a 2.05× geometric-mean speedup over the official implementation across six BF16 shapes, bitwise-correct and validated in end-to-end Kimi K3 serving under SGLang, with decode paths 1.14× over upstream FlashInfer; Gated DeltaNet and MiniMax sparse attention paths also improve. Dispatcher-backed KNN/KMeans libraries gain 1.42–2.12× across 400+ shapes, and four kernel changes ship as upstream FlashInfer PRs so users take no dependency on Cake. Targets Ampere through Blackwell (cost model calibrated only for B200/H100); compiler evolution is still human-gated at merge time. A compiler-side complement to NVIDIA's agentic-search AVO work.