CAKE Co-Designs Compiler-Agent for GPU Kernel Optimization

Optimizing GPU kernels for peak performance is a complex task, often involving a disconnect between AI agents that generate code and the compilers that execute it. Existing approaches either treat the compiler as a black box or expose scheduling decisions through difficult abstractions. This research introduces CAKE (Compiler-Agent Co-Design for Frontier Kernel Evolution), a novel framework that bridges this gap. CAKE enables AI agents to author CAKE IR, a typed, hardware-explicit schedule representation. This intermediate representation exposes critical GPU features like warp roles, memory movement, synchronization, and pipelines, while also supporting verification, cost modeling, and localized diagnostics. The system itself is designed to evolve, with recurring failures being converted into verifier rules, IR primitives, model calibrations, and reusable optimization tactics. In benchmarks, CAKE demonstrated significant performance gains. For Flash-KMeans on B200 GPUs, the best CAKE IR candidate achieved 1.144x the speed of a tuned FlashML baseline, outperforming direct CUDA/PTX (0.928x). Beyond benchmarks, agent-generated Kimi Delta Attention achieved a 2.05x geometric-mean speedup over official FlashKDA. CAKE also improved KNN and KMeans performance by 1.42x to 2.12x across over 400 shapes and has led to four kernel changes being proposed as upstream PRs. CAKE targets NVIDIA GPUs from Ampere through Blackwell, separating single-shape evolution from library generalization and dispatch.