CAKE: Compiler-Agent Co-Design Lets AI Write Faster GPU Kernels
CAKE pairs AI coding agents with a hardware-explicit compiler IR, beating tuned GPU kernel baselines on B200 with verifiable, evolving toolchains.
CAKE introduces a co-designed compiler and IR built specifically for AI coding agents generating GPU kernels, rather than treating the compiler as an opaque black box. The typed, hardware-explicit CAKE IR exposes warp roles, memory movement, synchronization and pipeline structure directly to agents, alongside verification, cost modeling and localized diagnostics — turning cryptic error and timing feedback into actionable signals.
The system's harness itself evolves over time: recurring failures get converted into new verifier rules, IR primitives, calibrated cost models and reusable optimization tactics, effectively building institutional knowledge for future kernel generation tasks.
On B200 GPUs, CAKE-generated kernels for Flash-KMeans reached 1.144x the tuned FlashML baseline at an 80-million-token budget, versus only 0.928x when agents wrote raw CUDA/PTX directly. Agent-authored Kimi Delta Attention delivered a 2.05x geometric-mean speedup over the official FlashKDA implementation and passed end-to-end serving checks, while dispatcher-backed KNN/KMeans kernels gained 1.42x to 2.12x across over 400 input shapes — with four of the resulting kernel changes merged as upstream PRs.
For engineers, this signals a shift from prompting agents to write raw kernel code toward giving them a structured, verifiable intermediate representation — potentially closing the gap between hand-tuned expert kernels and what LLM-driven pipelines can produce across Ampere through Blackwell hardware.