NVIDIA's CAKE has agents author GPU IR directly, hitting 2.05x on Kimi Delta Attention kernels

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

Zihao Ye, Yingyi Huang, Hongyi Jin, Bohan Hou, Junru Shao, Zhongming Yu, Jinqi Chen, Meghan Cowan, Shiyi Cao, Shanli Xing, Hanfeng Chen, Vinod Grover, Tianqi Chen, Luis Ceze

cs.LG

2026-08-13

CAKE lets agents write a typed, hardware-explicit IR while the compiler turns recurring failures into verifier rules and tactics; on B200 it reaches 2.05x over official FlashKDA on Kimi Delta Attention.

What problem this solves

GPU kernel agents and GPU programming languages have advanced on separate tracks, and expert kernels get lost in the gap. Agents treat the compiler as a fixed black box: the environment returns compiler errors, pass/fail verdicts, and end-to-end timing, but never says which program decision caused a synchronization failure, a hardware-contract violation, or a pipeline stall. When a frontier workload exposes a missing capability, those signals cannot grow either.

The writing side is equally awkward. Tile-level DSLs like Triton hide warp specialization, barrier choreography, and memory-tier placement, which is exactly what separates expert kernels from merely correct ones. Low-level DSLs like CuTe expose that control but demand a layout calculus that makes agent errors likely and hard to localize.

Method

CAKE, from NVIDIA and CMU, co-designs both sides. Agents author Cake IR, a typed, hardware-explicit schedule representation:

• Explicit roles: warp groups are named; lm.role(warps=[0]) assigns who loads and who computes, with every cross-role handoff visible;
• Declared resources: memory regions, synchronization state, and pipelines are declared once, fixing shape and dtype;
• Derived mechanics: barrier addresses, phase bits, and TMEM offsets are computed by the compiler instead of authored;
• No layout algebra: the compiler only verifies that producer and consumer representations are compatible.

The distinctive part is that the harness itself evolves. Recurring failures get distilled into four things: new verifier rules (an opaque runtime crash becomes a legality check), new IR primitives, cost-model calibration targets, and reusable optimization tactics. The object of evolution is the compiler environment; the coding agent and foundation model stay fixed. Changes are test-gated against a kernel corpus, with humans still guarding the merge gates.

Results

In the KDA experiment, FlashKDA source was never shown to the agent. Of 11 known-kernel reproductions, 10 meet or exceed the reference and one reaches 96.5%. Four kernel changes landed as upstream PRs (KDA prefill, KDA decode, TinyGEMM2, Alpha-MoE).

Why it matters

This is a clean demonstration that environment design caps agent performance: same agents, same token budget, but a compiler that returns localized diagnostics and accumulates rules turns "never converges from scratch" into "converges three out of three". The 1.144x versus 0.928x gap comes entirely from the environment, not from a stronger model.

It is most practical for new-architecture bring-up. When an attention variant like Kimi Delta Attention ships without a hand-tuned reference, the CAKE route lets agents close that gap at the DSL level, and the output can go straight into upstream PRs. Kernel-library and inference-engine teams get reusable capacity; agent-infrastructure teams get an environment-design template worth copying.

Limitations

The authors are unusually direct: most performance evidence is on B200, the timing model is calibrated only for B200 and H100 and declines to predict elsewhere; static analyses and performance models are intentionally incomplete, used only to rank and filter while GPU execution remains ground truth; compiler evolution is still human-guided at merge gates. Transfer cost to non-NVIDIA targets is unmeasured and the backend lowering would need a rebuild. Blackwell-only features (TMEM, TMA, clusters) are absent on Ampere.

One caveat worth flagging: the headline multipliers cluster on workloads with weak references (early official KDA, the megakernel MoE), while the best result across 11 mature-kernel reproductions is about 1.27x. For kernels that experts have already tuned by hand, the headroom here is visibly narrower, and the flashiest numbers are not the general expectation.

Terms

• Cake IR: the intermediate representation agents write directly, typed and explicit about hardware scheduling
• warp specialization: assigning warp groups fixed roles such as loading or compute, standard in expert Blackwell kernels
• TMEM: Blackwell's tensor memory, a high-speed scratchpad dedicated to tensor cores
• TMA: tensor memory accelerator, hardware that moves data blocks asynchronously instead of manual copies
• FlashKDA: the official FlashML implementation of Kimi Delta Attention, used here only as a black-box timing baseline
• megakernel: a single device program fusing routing, expert compute, and output accumulation

Source

• arXiv 2608.12629v1
• PDF
• Hugging Face

What people are saying

• CAKE Paper: Compiler-Agent Co-Design Achieves 2.05x Speedup on Blackwell Kernels hsu_byron
• CAKE paper proposes compiler-agent co-design Tim_Dettmers
• CAKE Paper: Evolving Compiler Harness via AI for SOTA Kernels lazowska

Related papers

• ICML position: monad algebras in Para recover GDL equivariance and RNNs as one theory
• Any non-trivial joint confidence forces conformal sets to grow as e^{nH(Y|X)}
• Calibrated Clipping Fixes FP8 RL Over-Clipping and Recovers BF16 Scores at up to 1.5× Throughput
• Trimmed instance norm stops MCL head collapse; TimePre wins Distortion on all six GluonTS sets
• IER Shows 0.1%–1% of Tokens Can Match Full On-Policy Distillation
• 4-bit models that keep classification accuracy still swap 14–46% of retrieval top-1s

All paper explainers