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

2026-08-12

AI Summary

Existing GPU kernel optimization approaches are hindered by compiler black-boxing or abstraction in domain-specific languages, impeding efficient co-scheduling and verification. This work proposes CAKE, a compiler–agent co-design framework featuring a typed, hardware-explicit intermediate representation (IR) that explicitly models warp roles, memory movement, synchronization, and pipelining, thereby enabling formal verification, cost modeling, and localized diagnostics. By encoding failure experiences into verification rules, IR primitives, and reusable strategies, CAKE decouples per-shape evolution from library-level generalization. Experiments on the B200 show that Flash-KMeans achieves 1.144× the performance of the baseline (versus only 0.928× with CUDA/PTX), Kimi Delta Attention obtains a 2.05× geometric mean speedup with end-to-end verification, and KNN/KMeans accelerate across over 400 shapes by 1.42–2.12×; four optimizations have been upstreamed.

Abstract

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.

Problem

• GPU kernel
• compiler-agent co-design
• schedule representation
• DSL
• kernel reproducibility

Innovation

• compiler-agent co-design
• CAKE IR
• hardware-explicit scheduling
• kernel evolution
• GPU code generation