One-sentence Summary

Researchers from NVIDIA and Carnegie Mellon University propose Cake, a compiler-agent co-design for GPU kernel evolution whose typed, hardware-explicit Cake IR provides fine-grained control without a layout algebra and enables localized verifier and cost-model diagnostics that evolve with recurring failures, achieving on B200 a $1.144\times$1.144× speedup over the tuned FlashML baseline for FlashKMeans versus $0.928\times$0.928× for direct CUDA/PTX and a $2.05\times$2.05× geometric-mean speedup over FlashKDA for Kimi Delta Attention.

Key Contributions

• Cake provides a typed, hardware-explicit IR and co-designed compiler harness for GPU kernel agents, giving fine-grained control without a layout algebra and enabling verifier and cost-model reasoning before compilation.
• The harness is an evolvable target: recurring kernel failures become new verifier rules, IR primitives, cost-model calibrations, and reusable tactics, with localized correctness and performance diagnostics returned to the agent.
• On implementation-hidden Flash-KMeans clean starts on B200, the best Cake IR candidate at an 80-million-token budget reaches a median 1.144× the tuned FlashML baseline versus 0.928× for direct CUDA/PTX; agent-generated Kimi Delta Attention achieves a 2.05× geometric-mean speedup over official FlashKDA and is validated in end-to-end serving, and dispatcher-backed KNN and KMeans families improve performance by 1.42× to 2.12× across more than 400 shapes.

Introduction

Coding agents increasingly write and revise GPU kernels, but most systems treat the programming environment as a fixed black box: they compile, test, measure latency, and edit, so a crash does not identify the violated safety or hardware condition, and a single latency number does not explain which program decision limits performance. Existing GPU DSLs are also awkward for agent-driven kernel development because high-level tile DSLs hide the warp specialization, barrier choreography, and memory-tier placement that expert kernels need, while low-level DSLs require layout algebra expertise and can produce brittle code. The authors introduce Cake, a system that makes the compiler an evolving collaborator for kernel agents. Cake has agents edit a typed IR rather than raw CUDA, returns localized correctness and performance diagnostics instead of pass/fail feedback, and turns repeated failures into verifier rules, calibration tasks, or new primitives under corpus tests and human merge gates.

Method

The authors leverage a bottom-up approach to design the Cake IR, rather than starting with a predefined vocabulary. The process begins with a corpus of production kernels and hardware design principles. Agents identify recurring schedules or missing capabilities, revise the IR and its compiler support, and then port and validate kernels against the revised system. This cycle repeats for new kernel families or gaps exposed by validation, continuously growing the IR.

Refer to the framework diagram:

The Cake IR records explicit machine schedules, detailing how the machine should be driven, including warp roles, buffer staging, barrier gates, and instruction forms. A program combines explicit operations, declared resources, warp roles, and grid configuration. Key properties include a type-checked vocabulary for compute and memory movement, declared resources for memory and synchronization state, explicit warp group roles, and auto-derived metadata where mechanical consequences are lowered rather than authored. This design allows analyses to reason from explicit schedule decisions before code generation. The system targets NVIDIA GPUs from Ampere through Blackwell, mapping the attached GPU exactly and emitting performance estimates where calibration is available.

The compiler harness serves as the agent-facing environment around the Cake IR. Humans provide high-level descriptions of intended analyses, while agents implement, maintain, and refine them under validation. Before compilation, the harness checks the typed schedule for synchronization, memory-safety, data-flow, and resource violations. It also verifies numerical correctness against reference outputs and uses a calibrated cost model to estimate performance and rank candidates.

Compiler evolution follows two complementary paths to improve the system alongside kernel development.

As shown in the figure below:

In the first path, agents inspect production kernels and hardware documentation to identify missing patterns, such as new instruction forms or synchronization idioms, and formulate compiler change proposals. These proposals are checked against design principles before implementation. In the second path, agents use feedback from failed candidates, including sanitizer reports and debugging logs, to distill recurring failure modes into new analyses. For instance, an opaque runtime crash becomes a verifier rule, and a repeated illegal lowering pattern becomes a static check. These two paths are coupled; new primitives expose hardware facts enabling stronger analysis, while new analyses constrain the design space for future primitives. Changes are test-gated across the kernel corpus to ensure primitives and analyses evolve together.

The external workload contract acts as the stable authority for the agent workflow, which consists of four stages. First, agents generate structurally distinct Cake IR candidates. Second, they filter these candidates using IR construction checks, verifier hard gates, and cost-model ranking to avoid unnecessary GPU usage. Third, survivors are evaluated against an external oracle with benchmarking and profiler evidence. Finally, the resulting evidence is routed to the candidate, verifier, cost model, or IR vocabulary based on the diagnosis.

To transition from a tuned shape to a library, the system employs a separate generalization stage. This stage groups measured seeds into shape buckets, produces specialized or shared variants, and orders their guards behind an explicit fallback. Validation covers representative inputs, boundary cases, and fallback paths. The system reuses a single physical schedule across as much of the shape domain as possible, introducing new routes only when a material schedule change is required, ensuring routing complexity is justified by measured workload gain.

Experiment

The evaluation spans clean-start evolution, frontier-kernel synthesis, known-kernel reproduction, and dispatcher-inclusive library generalization. Clean-start Flash-KMeans experiments show Cake IR surpasses a tuned baseline while the direct CUDA/PTX arm does not, and frontier kernels such as KDA, Gated DeltaNet, MiniMax sparse attention, TinyGEMM, and Alpha-MoE are synthesized without low-level references and improve over black-box baselines. Known-kernel reproduction validates production-quality output by matching or exceeding highly optimized expert references in nearly all fixed comparisons. The validated corpus and generalized portfolios further confirm broad architecture support, correct outputs, and consistent dispatcher-level speedups across many shapes.

The harness exposes seven analysis and validation categories spanning pre-compile gates, an execution gate, reporting, and non-blocking hints. Pre-compile checks reject candidates for synchronization, memory-use, resource, instruction, data-flow, and schedule-structure violations before execution, while numerical validation compares compiled outputs against an authoritative external reference. Performance analysis estimates cost and identifies broad bottleneck classes, but on-device measurement and profiling remain the final ground truth. Program safety, hardware conformance, data consistency, and schedule semantics act as pre-compile gates that block invalid candidates with localized reasons. Numerical validation is an execution gate requiring agreement with an authoritative external reference before final acceptance. Performance analysis returns estimated cost and broad bottleneck attribution without replacing on-device measurement as the final ground truth. Optimization guidance is advisory and suggests promising revisions without blocking compilation.

In matched clean-start Flash-KMeans runs at the 80-million-token budget, the CAKE IR arm reached a plateau in all three runs, while the direct CUDA/PTX arm did not plateau in any run. The CAKE IR arm also used less median active evolve time and ended with a higher best median performance at the 80M budget. Direct CUDA/PTX ended with a best median performance below the reference level. CAKE IR reached the plateau in every run by the 80M-token budget, whereas direct CUDA/PTX reached it in none. Median active evolve time was substantially lower for CAKE IR than for direct CUDA/PTX. At the 80M-token budget, CAKE IR's best median performance was above the reference while direct CUDA/PTX's was below it.

The evaluation combines pre-compile checks for synchronization, memory, resource, instruction, data-flow, and schedule-structure violations with an execution-time numerical gate against an authoritative reference, while performance analysis and optimization guidance remain advisory. In matched clean-start Flash-KMeans runs at an 80-million-token budget, the CAKE IR arm reached a performance plateau in all three runs, used less median active evolve time, and finished above the reference, whereas direct CUDA/PTX plateaued in none and stayed below the reference. These results indicate that the CAKE IR path gives more reliable convergence and better final performance under the same budget.