In 2022, FlashAttention made the attention layer of a transformer much faster on the same GPU, and it did so without changing a single number the math produced. It did not find a new formula. It found a better schedule: which piece of data sits in which memory, which part of the chip works on it, and in what order. Three more versions followed, each mostly a better schedule, the last two written for new GPU generations.

Programs like these are called GPU kernels, and they are where much of the money in AI is spent: every token a large language model (LLM) reads or writes passes through dozens of them. Lately, AI coding agents have started to write kernels too. They propose code, compile it, test it, time it, and try again. On many benchmark tasks this loop now produces correct, fast code. On the hardest kernels, the ones experts write by hand for each new GPU, it still falls short.

CAKE (Compiler–Agent co-design for frontier Kernel Evolution), a paper from NVIDIA and Carnegie Mellon University, argues that the agents are no longer the main bottleneck. The bottleneck is the world the agents work in: the language they write and the feedback they get back. Today’s agents write raw CUDA or a high-level kernel language, and the environment answers with little more than “it crashed” or “it took 0.94 ms.” CAKE changes both sides. Agents write a new, checkable language for GPU schedules, and the compiler answers with specific diagnostics: this barrier, in this pipeline stage, is waited on but never signalled. Then CAKE goes one step further. When agents keep failing in the same way, the compiler itself is changed: the failure becomes a new check, a new language feature, or a corrected performance model.

The headline results, in the paper’s own numbers:

1.144× vs 0.928×Flash-KMeans from a clean startbest kernel after 80M tokens, CAKE IR vs direct CUDA/PTX, relative to a tuned baseline (median of 3 runs, B200)
2.05×Kimi Delta Attention prefillgeometric mean over the official FlashKDA kernels, 6 shapes; validated in end-to-end serving
1.42×–2.12×KNN and KMeans librariesdispatcher-backed kernel families, more than 400 shapes in total (GB200)
4 PRsmerged into FlashInferKDA prefill and decode, TinyGEMM, and an Alpha-MoE megakernel

This post explains CAKE from the ground up. It is built from the paper (“CAKE: Compiler–Agent Co-Design for Frontier Kernel Evolution,” arXiv 2608.12629, submitted August 12, 2026), its LaTeX source and figures, and the public sources it cites. Numbers come from the paper unless we say otherwise. Values we read off a plot are marked “approximate”; our own readings of the evidence are marked as interpretation. Where the paper is silent, we say so, and §16 collects those gaps.

Who this post is for

Most papers about GPU kernels assume you already write them. This post does not. If you know some Python and roughly what a neural network does, you can follow it. We borrow the shape of a much-loved operating-systems textbook, Operating Systems: Three Easy Pieces by Remzi and Andrea Arpaci-Dusseau, and use three of its habits:

  • Three pieces. The post has three parts, and each opens with a short dialogue between a Professor and a Student. The dialogues only ask questions; the main text answers them.
  • The crux. Each chapter names its central problem in a box labelled The crux, so you always know what we are trying to solve.
  • Asides and tips. Aside boxes hold background you can skip. Tip boxes state a general lesson that outlives this paper.

If you already know GPUs well, skip Part I and start at §5.

The road map

  • Part I · The Machine (§1–4): what a GPU kernel is, why speed is mostly about moving data, how modern GPUs turned into asynchronous assembly lines, and what separates an expert’s kernel from a merely correct one.
  • Part II · The Language (§5–8): how agents write kernels today and what feedback they get; the menu of GPU languages and why none fits an agent well; CAKE IR, read line by line; and where it came from.
  • Part III · The Loop (§9–16): the checks CAKE runs before a kernel ever touches the GPU; one round of evolution; how the compiler itself evolves; the controlled experiment; the production results; turning a tuned kernel into a library; related work; and the limits.
  • Epilogue: a closing dialogue, a summary, questions to think about, and annotated sources.

Part I · The Machine — a dialogue

Professor Welcome back. Today: a paper about AI agents that write GPU kernels.

Student I use PyTorch every day and I've never written a kernel. Do I need one?

Professor You use them constantly. Every torch.matmul ends up as a kernel someone wrote. The question is who writes the next ones, for the next GPU.

Student Can't the compiler just generate them? That's what compilers are for.

Professor For simple operations, yes. For the important ones, the fastest code still comes from experts. They know things about the chip that compilers don't decide well yet.

Student Like what?

Professor Like which group of threads loads data while another group multiplies, and how they tell each other "your data is ready." Get that wrong and the program doesn't crash politely. It hangs, or it quietly produces garbage.

Student That sounds miserable to debug. And an AI is supposed to do it?

Professor That's the paper's question. But first you need to see the machine. Let's start with what a kernel actually is.

1. What a GPU Kernel Is

A GPU is a chip built to do the same simple operation on a huge amount of data at once. This section introduces the vocabulary the rest of the post needs: threads, warps, blocks, streaming multiprocessors, and the word kernel itself.

1.1 Many small workers instead of a few fast ones

A CPU has a handful of powerful cores, each built to finish one sequence of instructions as quickly as possible. A GPU makes the opposite bet. It has thousands of simpler lanes, and it wins by keeping all of them busy at once. Nobody cares how long one multiplication takes; what matters is how many finish per second. This is the difference between optimizing for latency and optimizing for throughput.

A kernel is a function that runs on the GPU. You write it once, and the GPU runs one copy per thread, often millions of threads, each working on a different piece of the data. Adding two vectors of a million numbers is the classic first kernel: thread i reads a[i] and b[i], adds them, and writes c[i].

Threads are organized in a hierarchy, and the hierarchy matches the hardware:

  • Warp. NVIDIA GPUs run threads in groups of 32 called warps. The 32 threads of a warp execute the same instruction at the same time, each on its own data.
  • Thread block. Warps are grouped into thread blocks (also called CTAs, cooperative thread arrays). The threads of one block run on the same part of the chip and can share a fast scratchpad memory.
  • Streaming multiprocessor (SM). The chip is made of many identical SMs. Each SM runs one or more thread blocks at a time, with its own registers, scratchpad memory, and math units. An NVIDIA B200 has 148 SMs.
  • Grid. One kernel launch creates a grid of thread blocks, and the hardware spreads them across the SMs.

1.2 Why matrix multiplication is the kernel that matters

Vector addition is easy because every output needs exactly two inputs. Matrix multiplication is different. To compute an output matrix C = A × B, every output element needs a whole row of A and a whole column of B, and every input element is used by many outputs. That reuse is both the opportunity and the difficulty: a good kernel loads each piece of input once and reuses it many times, and a bad one keeps fetching the same numbers from far away.

Almost everything a transformer does is built on matrix multiplications: the projections in attention, the feed-forward layers, the experts in a mixture-of-experts (MoE) model. Modern GPUs have dedicated units for them, called tensor cores, which multiply small matrix tiles in a single instruction. Most of the kernels in the CAKE paper are, at heart, careful ways to keep tensor cores fed.

2. Fast Kernels Are About Moving Data

Here is the first surprise for newcomers: on a modern GPU, doing the arithmetic is rarely the hard part. Getting the numbers to the arithmetic units in time is.

2.1 The memory ladder

A GPU has several kinds of memory, arranged like the rungs of a ladder. The closer a rung is to the math units, the faster and smaller it is:

Rung Where it lives Size (B200) What it is for
Registers inside each SM 256 KB per SM each thread’s working values
Tensor memory (TMEM) inside each SM, new in Blackwell 256 KB per SM the accumulators of tensor-core matrix multiplies
Shared memory (SMEM) inside each SM up to 228 KB per SM a scratchpad that the threads of a block share, programmed by hand
L2 cache shared by all SMs 126 MB a hardware-managed cache
HBM (global memory) stacks of DRAM next to the chip 180 GB at about 8 TB/s where tensors live
Where these hardware numbers come from
  • Registers (64K 32-bit registers per SM), shared memory (228 KB per SM, at most 227 KB per block) and the 126 MB L2: NVIDIA’s Blackwell tuning guide and the compute-capability tables of the CUDA programming guide. The tuning guide states the L2 size for GB200; Chips and Cheese measured the same 126 MB on a B200.
  • Tensor memory: the PTX ISA describes 128 rows × 512 columns of 32-bit cells per CTA, which is 256 KB.
  • HBM: 180 GB at up to 8 TB/s per B200 (NVIDIA reference architecture). Dense BF16: an 8-GPU HGX B200 is rated at 36 PFLOPS with sparsity, and dense is half of that, so 2.25 PFLOPS per GPU.
  • SM count: NVIDIA’s datasheets do not list it. The 148 comes from the FlashAttention-4 paper and independent measurements. The Blackwell GPU in a GB200 is a slightly larger configuration (186 GB of HBM, 2.5 PFLOPS dense BF16).

Eight terabytes per second sounds enormous until you compare it with the math. The tensor cores of a B200 can do about 2.25 thousand trillion (2.25 × 1015) dense BF16 operations per second. Divide one by the other: to keep the tensor cores busy, a kernel must do roughly 280 operations for every byte it reads from HBM. A kernel that does less work per byte spends its time waiting for memory, no matter how clever its math is.

This ratio of operations to bytes is called arithmetic intensity, and the idea that a kernel is limited either by memory bandwidth or by compute, whichever runs out first, is called the roofline model. Adding two vectors does one addition per 12 bytes moved (two 4-byte reads, one write); it will always be memory-bound. A large matrix multiplication can reach thousands of operations per byte, but only if the kernel reuses data well.

2.2 Tiling: the trick behind every fast kernel

The standard way to raise reuse is tiling. Instead of fetching inputs element by element, a thread block copies a tile of A and a tile of B from HBM into shared memory, multiplies them using the tensor cores, accumulates the partial result, and then moves on to the next pair of tiles. Each number fetched from HBM is reused by every output in the tile.

Tiling raises a new problem: while the block waits for the next tile to arrive, the tensor cores sit idle. The fix is to fetch the next tile while computing on the current one, using two or more tile buffers. This is called double buffering, or more generally software pipelining, and it is exactly the kind of scheduling decision that separated FlashAttention from the attention code before it. FlashAttention computed the same attention as everyone else; it just arranged the tiles so that the large intermediate matrix never had to be written to HBM at all.

3. The GPU Becomes an Assembly Line

Recent NVIDIA GPUs added hardware that makes tiling and pipelining much faster, and much harder to program. This section walks through that change, because CAKE’s design follows it closely.

3.1 From “every thread does everything” to specialized roles

In older GPU kernels, every warp did every job: load a slice of the tile, wait at a barrier for the rest of the block, compute, repeat. Each GPU generation since then has added hardware that does one of these jobs on its own, asynchronously, while the threads do something else:

  • Ampere (A100, 2020) added cp.async, which copies data from global memory into shared memory in the background. Its tensor cores are driven by warp-level mma.sync instructions, fed from shared memory by ldmatrix.
  • Hopper (H100, 2022) added the Tensor Memory Accelerator (TMA), a copy engine that moves an entire multi-dimensional tile with one instruction; WGMMA, a tensor-core instruction issued by a warpgroup of four warps; asynchronous transaction barriers (mbarriers that count the bytes of a copy as they land) to signal when a tile has arrived; and thread-block clusters whose blocks can read each other’s shared memory.
  • Blackwell (B200, 2024–25) added tcgen05.mma, a tensor-core instruction that a single thread issues on behalf of the whole block, and tensor memory (TMEM), a dedicated 256 KB per SM that holds the accumulators so they no longer clog the registers. Two neighbouring SMs can even cooperate on one larger multiply (the “2-CTA” mode).

With copy engines and tensor cores that run on their own, the natural design is an assembly line. Some warps do nothing but issue copies; one warp does nothing but issue matrix multiplies; others take the finished results and write them out. This is called warp specialization, and each job is a warp role.

3.2 The handoff: ring buffers and barriers

An assembly line needs a way for stations to hand work to each other. In a warp-specialized kernel, the shared-memory tile buffers form a ring buffer of stages: with three stages, the producer can be filling stage 2 while the tensor cores work on stage 1 and stage 0 waits to be freed.

Each stage has two barriers. The producer signals the stage’s full barrier when its copy lands; the consumer waits on full before reading. The consumer signals the empty barrier when it has finished with the stage; the producer waits on empty before overwriting it. Because the ring wraps around, each barrier is reused once per lap, and the hardware tracks laps with a one-bit phase: a waiter must say which phase it expects, and waiting on the wrong one either returns too early or never returns.

Two more details matter later. Copies issued by the TMA and multiplies issued by the tensor cores belong to the hardware’s async proxy, while ordinary loads and stores belong to the generic proxy; when data crosses from one to the other, the kernel needs an explicit proxy fence, or the hardware may reorder the accesses. And every TMA copy needs a descriptor that encodes the tile’s shape, strides, and layout in memory.

3.3 Why this is so easy to get wrong

Every one of these details is a way to fail, and the failures are cruel:

  • A missing wait lets the consumer read a stage before its data arrives. Nothing crashes; the results are simply wrong, sometimes only for some inputs.
  • A barrier initialized with the wrong arrival count, or a wait on the wrong phase, makes a warp wait forever. The kernel hangs, and the only symptom is a GPU that stops responding.
  • A missing proxy fence produces stale data only under particular timing, so the bug appears on one run and vanishes on the next.
  • A TMA descriptor whose tile shape disagrees with the shared-memory buffer it fills corrupts memory that belongs to someone else.

Tools exist to help. NVIDIA’s compute-sanitizer can detect some races and out-of-bounds accesses at run time, but only on the inputs you happen to run, and it does not explain which design decision caused the problem. Remember this: a crash or a hang says that something broke, not why. That gap is the subject of Part II.

4. What Separates an Expert Kernel from a Correct One

Put the pieces together and you can see what the CAKE paper means by an expert kernel. Its abstract names three decisions that “separate expert kernels from merely correct ones”:

  1. Warp specialization. Which warps take which roles: how many producers, how many consumers, whether two groups of consumers alternate (“ping-pong”) so one computes while the other writes results.
  2. Barrier choreography. How many pipeline stages, which barrier gates which handoff, which phase each wait expects, where the fences go.
  3. Memory-tier placement. Which data lives in registers, which in shared memory, which in tensor memory, and in what layout, so that the copy engine, the tensor cores, and the output path all agree on where every byte sits.

Together these decisions form the kernel’s schedule: the plan for how the machine is driven, as opposed to what is computed. The math of a matrix multiply fits on one line. The schedule of a fast Blackwell matrix multiply runs to hundreds of lines, and changing one decision usually forces changes in the others.

Summary of Part I

A kernel is a function that thousands of GPU threads run in parallel. Its speed depends mostly on moving data: through a ladder of memories, in tiles, with the next tile loading while the current one is processed. Modern NVIDIA GPUs turned this into an assembly line of specialized warps that hand tiles to each other through barriers. The expert’s art is the schedule (roles, barriers, memory placement), and mistakes in it tend to cause silent wrong answers or hangs rather than helpful error messages.


Part II · The Language — a dialogue

Student So an agent writes a kernel, runs it, and looks at what happened. What's wrong with that? It's how I debug.

Professor What do you look at when your code breaks?

Student The stack trace. The line number. The variable that was None.

Professor Now imagine your only feedback is "it froze" or "it took 0.94 milliseconds." No line, no variable.

Student I'd change things at random and hope.

Professor That's roughly where kernel agents are today. So, two questions. What language should the agent write in? And what should the compiler tell it when something is wrong?

Student Shouldn't it just write CUDA, like the experts?

Professor CUDA lets you say everything, including every mistake. A higher-level language protects you, but it won't let you say what the expert says. The paper's answer sits in between. Let's see where.

5. How Agents Write Kernels Today

Kernel-writing agents arrived quickly. This section describes the loop they share, what that loop is good at, and the feedback problem that CAKE starts from.

5.1 The propose–test–measure loop

The setting was popularized by KernelBench (2025): give a model a PyTorch operator and ask it for a faster GPU kernel that computes the same thing. Since then, many systems have built loops around this task. Some use a human-designed loop in which an LLM revises a kernel from compiler errors, correctness results, or profiler output (AccelOpt, Autocomp, KernelAgent). KernelBlaster adds a persistent knowledge base of past optimizations. KernelEvolve and EvoEngineer run evolutionary search over populations of candidate kernels. AVO, from several of CAKE’s authors, replaces fixed mutation rules with coding agents. K-Search separates planning from implementation. AutoTriton and CUDA Agent train the model itself with reinforcement learning.

Underneath the variety, the environment is the same. The agent proposes code, the code is compiled, a numerical test compares its output with a reference, a benchmark measures its latency, and the agent picks the next edit. The CAKE paper’s summary: these methods “evolve the search process, accumulated memory, or model weights while retaining a chosen DSL and evaluation environment.”

5.2 What the environment says back

This loop works well for local tuning: try a different tile size, unroll a loop, fuse two operations. It struggles with the decisions from §4, for a reason the paper puts in one sentence: compiler errors, correctness outcomes, and end-to-end timing “never say which program decision caused a synchronization failure, a hardware-contract violation, or a pipeline stall.”

Think about what each signal tells an agent that just rewrote a pipelined kernel:

  • A hang. Some warp is waiting on a barrier that never completes. Which barrier? Which stage? Which phase? The process just stops.
  • Wrong output. A race, a bad layout, a missing fence, or a plain indexing bug. The test can say which elements differ, not which decision caused it.
  • A time. 0.94 ms. Is the kernel limited by memory bandwidth, by tensor-core issue, by a pipeline that is too shallow, by a barrier that serializes two roles? One number cannot say.

There is a second, quieter problem. The signals “cannot grow when a frontier workload exposes a missing capability.” If the language cannot express a schedule the new workload needs, every attempt fails for the same structural reason, and the environment never learns from it.

5.3 CAKE’s three commitments

The paper’s observation is that a compiler already contains most of the machinery an expert keeps in their head: “structured operation vocabularies, resource models, legality checks, static analyses, cost models, lowering rules.” The question CAKE asks is “how to make that machinery agent-facing, and how to improve it when a frontier workload exposes a gap.” Its answer has three parts:

  1. Agents edit a typed IR, not raw CUDA. An IR (intermediate representation) is a program format designed for a compiler to analyze. CAKE’s is a schedule language in which the hardware decisions are explicit, so they can be inspected before any code is generated.
  2. The compiler returns localized diagnostics, not a pass/fail bit. Correctness and performance findings point at the role, stage, buffer, or instruction at fault, and cheap analyses filter candidates before they spend GPU time.
  3. The harness itself evolves. Repeated failures become new verifier rules, new IR primitives, recalibrated cost models, and reusable tactics, each change gated by tests over a corpus of kernels.

The paper calls the result compiler–agent co-design: the agent and its environment are designed together, and the environment keeps changing in response to what the agent runs into. The next three sections cover the first commitment, the language.

6. Which Language Should an Agent Write?

There are many ways to write a GPU kernel. They differ in one question above all: who decides the schedule, the programmer or the compiler? This section walks the menu, using the CAKE paper’s own classification.

6.1 High-level: tile languages

Triton (2019) made kernel writing accessible by letting programmers think in tiles: you write what happens to a block of a matrix, and the compiler decides how threads, memories, and instructions carry it out. Newer languages such as Helion, TileLang, and NVIDIA’s cuTile follow similar ideas at different levels of abstraction.

This is excellent for humans and gives correct, often fast code. The CAKE paper’s objection is specific: tile languages “hide the warp specialization, barrier choreography, and memory-tier placement that separate expert kernels from merely correct ones.” An agent writing Triton cannot ask for “one producer warp, one MMA warp, a three-stage ring, accumulators in TMEM” if the language has no way to say it; it can only hope the compiler chooses that.

6.2 Low-level: CUTLASS, CuTe DSL, and raw CUDA

At the other end is NVIDIA’s CUTLASS library and its Python front end, CuTe DSL. These expose everything the hardware offers. The price is a layout algebra: a mathematical framework for describing how a tensor’s logical coordinates map onto memory and onto threads. Below that is CUDA C++ with inline PTX (NVIDIA’s assembly-like virtual instruction set), where you manage addresses, barrier phase bits, and descriptor encodings by hand.

According to the paper, low-level languages “expose that control but demand a layout calculus that makes agent errors both likely and hard to localize.” A wrong layout does not cause a compile error. It causes the right numbers to land in the wrong places, which shows up later as wrong output.

6.3 In between: Gluon

Gluon sits between the two camps. It reuses Triton’s compiler stack but exposes lower-level control over layouts, memory movement, and asynchrony. It moves the dial toward control, and with it toward the burdens of control.

6.4 The gap CAKE aims at

Lay the options side by side and a gap appears. The languages that let you express an expert schedule make you manage layouts and raw details yourself, and they find your mistakes only at run time. The languages that protect you keep the schedule out of your hands.

A table compares GPU kernel languages on three questions: do you write the warp roles, barriers and memory tiers yourself, is a layout algebra needed, and is the program checked before it runs. Tile DSLs (Triton, Helion, TileLang, cuTile): no, the compiler decides; no; partly (type checks). Gluon: partly on all three. CuTe DSL / CUTLASS: yes; yes (CuTe layouts); mostly at run time. CUDA C++ with inline PTX: yes; no (raw addresses); run it and see. CAKE IR, set apart: yes; no (concrete offsets); yes (verifier and cost model).
Who decides the schedule? Tile languages hide it, low-level languages hand it to you without a safety net, and CAKE IR makes it explicit, without a layout algebra, and checks it before it runs (simplified, following the CAKE paper's classification). Open the full-size SVG.

7. CAKE IR, Line by Line

CAKE IR is the language the agents write. The best way to understand it is to read some.

7.1 What the schedule states, and what lowering derives

The design rests on a division of labour. In the paper’s words, “a schedule states what is to happen and lowering derives how.” The schedule records which warps take which roles, which buffers are staged how deeply, which barrier gates which handoff, and which instruction form consumes which operand. Lowering, the compiler’s translation down to CUDA and PTX, computes the mechanical consequences: “barrier addresses, phase bits, TMEM offsets, descriptor encodings, and warp identity are all computed from the declarations rather than written out by the agent.”

Here is the fragment the paper uses to show the idea, a piece of a fused multi-head attention (FMHA) forward kernel for Blackwell:

@cake.schedule()
def fmha_fwd(lm, Q: LM.tma3d, O: LM.tma2d, seqlen_q: LM.i32):
    # declare resources: named resources, not raw addresses
    pool     = lm.smem(98304)
    smem_q   = pool.view(offset=0, shape=(128,128), dtype=lm.bf16, stage=3)
    tmem_acc = lm.tmem(cols=0, width=128, shape=(128,128), dtype=lm.f32)

    # declare roles: warp groups with assigned work
    load = lm.role(warps=[0])
    mma  = lm.role(warps=[1])
    pipe = lm.pipeline(stages=3)
    q_full = lm.barrier(count=3, prod=[load], cons=[mma],
                        init_count=1, pipeline=pipe)

    with load:
        for stage in lm.range(0, 3):
            smem_q.tma_load(Q, coords=(0,0,stage), stage=stage, barrier=q_full)

    with mma:
        for stage in lm.range(0, 3):
            lm.wait(q_full, stage=stage)
            lm.fence_proxy()
            lm.mma(tmem_acc, smem_q[stage], smem_q[stage], init=(stage == 0))

Read it in three blocks:

  • Resources. lm.smem(98304) reserves a 96 KB pool of shared memory. pool.view(...) carves out a named buffer of 128 × 128 BF16 values with three stages: a three-slot ring buffer. lm.tmem(...) reserves a 128 × 128 FP32 accumulator in tensor memory. Nothing here is an address; every buffer has a name, a shape, a type, and a lifetime the compiler can see.
  • Roles and synchronization. lm.role(warps=[0]) and lm.role(warps=[1]) name two warp roles, a loader and a tensor-core issuer. lm.pipeline(stages=3) declares the pipeline, and lm.barrier(...) declares the handoff between them: a barrier with three stages whose producer is load and whose consumer is mma.
  • Work per role. Inside with load:, the loader issues one TMA copy per stage and tells the hardware to signal q_full when each copy lands. Inside with mma:, the tensor-core warp waits on q_full for each stage, issues a proxy fence (the ordering instruction between the hardware’s two memory paths, §3.2), and issues the matrix multiply into the TMEM accumulator, initializing it on the first stage.

Notice what is missing. There are no barrier addresses, no phase bits, no descriptor encodings, and no computation of which thread is in which warp; the compiler derives all of them from the declarations. The buffers do carry a few plain commitments, such as the view’s offset inside the pool and the accumulator’s starting column in tensor memory. §7.3 explains why CAKE asks for those directly, and the compiler derives the rest of the addressing from them. The fragment is a teaching example, not a complete kernel: it multiplies the Q tile by itself, where a real attention kernel would also load K and V and run a softmax.

The CAKE IR schedule fragment from the paper's Figure 3, split into three bracketed blocks. Resources: a 96 KB shared-memory pool, a three-stage 128×128 BF16 buffer smem_q, and a 128×128 FP32 accumulator in tensor memory. Roles and sync: load on warp 0, mma on warp 1, and a three-stage barrier q_full handing tiles from load to mma. Work per role: load issues three TMA copies; mma waits on q_full, issues a proxy fence, and multiplies. A side panel lists what lowering derives: barrier addresses, phase bits, TMEM offsets, descriptor encodings and warp identity, emitted as CUDA/PTX.
The paper's example schedule in three blocks: resources, roles with the barrier between them, and each role's work; lowering derives the mechanical details (simplified from the paper's Figure 3, with type annotations and comments omitted). Open the full-size SVG.

7.2 Four properties that make it checkable

The paper names four properties that “do the work”:

  • Type-checked vocabulary. Compute, memory movement, synchronization, math, and warp control use a fixed set of IR operations, not embedded C or PTX strings. An ill-typed program is rejected while it is being built.
  • Declared resources. Memory regions, synchronization objects, and pipelines are declared once, so “the IR knows the shape, dtype, and lifetime of every buffer.”
  • Explicit roles. Warp groups are named, and “every cross-role handoff is visible rather than an implicit convention.”
  • Auto-derived metadata. The mechanical consequences of the declarations are “lowered, not authored.”

The payoff is that analyses can reason about schedule decisions before code generation, and the harness can tie a finding to the affected resource, role, or stage, “rather than returning only a backend error or a hang.”

7.3 No layout algebra, on purpose

CAKE “deliberately does not make layout a first-class abstraction.” Instead of manipulating layouts as mathematical objects, the agent writes down concrete commitments: “an SMEM view offset, an operand byte offset, a TMEM column range, a swizzle tag, a TMA descriptor coordinate.” The compiler then “carries the burden of deciding whether those commitments are legal”: it checks that they stay consistent along the program’s data flow (what the producer writes is what the consumer expects to read) and that they satisfy the target instruction’s rules. A diagnostic names the IR decision at fault and a broad category of mismatch.

This is the reverse of CuTe’s bargain. CuTe asks the author to be an expert in layouts and rewards them with expressive power. CAKE asks the author for plain facts and makes the compiler responsible for catching contradictions. The paper describes its layout verifier only by its contract and coverage, not its internals.

7.4 One language, many GPUs

The same schedule language targets NVIDIA GPUs from Ampere through Blackwell. The structure of a role–barrier–pipeline schedule carries over; which instructions are allowed, and how they are lowered, remains specific to each target:

Target Example GPUs Tensor-core path and notable features
sm_80 A100 mma.sync, ldmatrix, cp.async; no TMA, clusters, or TMEM
sm_89 L40S, RTX 6000 Ada as sm_80, plus FP8 tensor cores; no TMA or TMEM
sm_90a H100, H200 WGMMA + mma.sync; TMA, clusters, async barriers, distributed shared memory
sm_100a B200 tcgen05.mma + TMEM; 2-CTA MMA; tcgen05.{ld,cp,shift}
sm_103a B300 adds tcgen05.ld.red and K = 96 block-scaled MMA
sm_120a RTX 5090, RTX PRO 6000 mma.sync + ldmatrix, TMA, clusters, distributed shared memory; no tcgen05 or TMEM
sm_121a DGX Spark (GB10) same instructions as sm_120a, separate binary, different special-function-unit rates

Two policies stand out. The compiler “requires an exact target match”: it reports a missing device or toolchain rather than quietly lowering a schedule to an older architecture. And performance estimates appear only where they are calibrated: B200 is the measured baseline, H100 is calibrated separately, and other targets report a coverage limitation instead of borrowing numbers. After the static checks, CAKE IR lowers deterministically to readable CUDA/PTX and then through NVIDIA’s standard toolchain. The generated source stays available as an escape hatch; by default, the decisions stay in the IR where they can be analyzed.

8. Where CAKE IR Came From

You might expect a language like this to be designed on a whiteboard and then implemented. CAKE IR was grown instead.

8.1 Grown from expert kernels

“CAKE did not begin with a predefined CAKE IR vocabulary.” Its starting material was a corpus of production kernels and a set of hardware design principles. The appendix describes five steps:

  1. Corpus collection. High-quality CUDA kernels from production libraries: Alpha-MoE, CUTLASS, cuTile, DeepGEMM, FlashAttention-4, Flash-KMeans, FlashInfer, SonicMoE, and TileLang. Kernels written in other languages were first translated to CUDA with inline PTX by coding agents.
  2. Abstraction extraction. Agents analyzed the corpus and summarized recurring patterns (barrier choreography, pipeline staging, warp-role partitioning, TMA descriptor setup, TMEM accumulator lifecycles) into candidate abstractions.
  3. Hardware-informed design. Human expertise steered the abstractions toward the Blackwell programming model: TMEM as a first-class resource, warp specialization as the main form of parallelism, asynchronous barriers as the synchronization primitive, and cluster-scoped operations for coordinating several SMs.
  4. Principle-driven iteration. Each candidate abstraction was checked against eight design principles (below); violations were refined or rejected.
  5. Port-driven expansion. New kernels are ported continuously. Each port either succeeds, which validates the abstractions, or reveals a gap, which triggers a proposal to extend the IR or its lowering.

The loop never closes: “every new kernel family stress-tests the IR and drives further evolution.” The design target throughout was demanding: the IR must “reproduce the physical schedules and performance of expert-written kernels.” According to the paper, the harness is likewise “maintained primarily by agents under human merge gates.”

8.2 Eight principles

Principle What it asks
P1 Ergonomic Feel familiar to NumPy and PyTorch users; avoid needless bookkeeping.
P2 Performance-transparent Keep performance-relevant hardware decisions visible and lowering inspectable.
P3 Canonical One canonical form per operation, not several equivalent spellings.
P4 Statically type-checked Typing rules constrain lowering and reject ill-typed programs during construction.
P5 Analysis-friendly Expose the information the static analyses need.
P6 Test-gated Evaluate every IR change against the kernel-matrix tests for analysis and compilation.
P7 Analysis-consistent Change the analyses whenever the IR’s data model changes.
P8 Hardware-grounded Document the intended hardware behaviour of every operation.

P5 and P7 keep new features analyzable, P6 guards against regressions, and P2 and P8 keep the mapping to hardware legible “to humans and agents alike.” The paper connects this to an observation by the mathematician Terence Tao about AI-generated proofs: when generating artifacts becomes cheap, the bottleneck shifts to verifying and understanding them. CAKE’s version: when agents can produce kernels in bulk, what matters is that each kernel can be checked and read.

Summary of Part II

Kernel agents today live in a black box that answers with pass/fail and a time. CAKE gives them a language in which the expert’s decisions (roles, buffers, barriers, pipelines) are explicit and typed, while the error-prone bookkeeping (addresses, phases, offsets, encodings) is derived by the compiler. It deliberately avoids a layout algebra and makes the compiler check concrete layout commitments instead. The language was grown from production kernels, under eight principles that keep it analyzable as it grows.


Part III · The Loop — a dialogue

Student OK, so the agent writes this schedule language. Then what? It still has to run on the GPU eventually.

Professor Eventually. But a GPU run is the expensive way to find a missing barrier. What's cheaper?

Student Reading the code first? Like a linter?

Professor Exactly: checks that run before compilation and say which barrier, which stage. Then a model that guesses which survivors are worth timing. Only then the GPU.

Student And when the checks themselves miss a whole kind of bug?

Professor Then you've found the paper's most interesting idea. The bug becomes a new check. The compiler learns.

Student Does all this actually make faster kernels? Or just nicer error messages?

Professor Good. Keep that skepticism. We'll look at the evidence, and at what it doesn't show.

9. The Harness: Checks Before the GPU

The harness is everything around the agent: the analyses, the tests, the benchmark, and the rules for what counts as a result. This section describes what CAKE’s harness checks and what it says back.

9.1 Seven kinds of checks

The paper describes its analysis suite by what it does for the agent, not by its internal compiler passes. There are seven categories, in three dispositions: gates that reject a candidate, reports that describe it, and hints that suggest improvements.

Category Disposition Purpose
Program safety pre-compile gate Find synchronization, ordering, and memory-use hazards
Hardware conformance pre-compile gate Enforce supported resource, instruction, and architecture contracts
Data consistency pre-compile gate Check data flow and that producer and consumer agree on representations
Schedule semantics pre-compile gate Check the structural invariants of the declared schedule
Numerical validation execution gate Compare compiled outputs with an authoritative external reference
Performance analysis report Estimate cost and identify broad bottleneck classes
Optimization guidance hint Suggest promising revisions without blocking compilation

The four pre-compile gates reject “many candidates that are mathematically plausible but incompatible with the target execution model.” A finding “identifies the affected program region and the class of violated contract.” Compare that with the hang from §3.3: instead of a frozen GPU, the agent hears something like the consumer waits on stage 2 of this barrier, but no producer ever arrives there. That example is ours; the paper does not publish its message format.

9.2 Correctness and performance

Numerical correctness is checked against a reference implementation across different shapes and input distributions, and final acceptance requires end-to-end evaluation in the target framework (for example, a full model served by SGLang).

Performance has two layers. A calibrated cost model estimates how fast a candidate will run and names its likely bottleneck, so the harness can rank and filter candidates without running all of them. But the model is only a filter: “on-device measurement and profiling remain the final ground truth.” In the evaluation, every reported number is measured on the GPU with CUPTI, NVIDIA’s profiling interface, with the L2 cache flushed before each timed sample so that one run cannot warm the cache for the next.

10. One Round of Evolution

With the language and the harness in place, a run of kernel evolution has a simple shape. This section follows one round.

10.1 Four stages

The paper describes a run as four stages:

  1. Generate structurally distinct CAKE IR candidates: different role splits, pipeline depths, memory placements, not just different constants.
  2. Filter them with IR construction checks, the verifier’s hard gates, and cost-model ranking, before any GPU time is spent.
  3. Evaluate the survivors against the external oracle, with benchmarking and profiler evidence.
  4. Route the evidence. Depending on the diagnosis, a finding goes back to the candidate (repair it), to the verifier (add a rule), to the cost model (recalibrate it), or to the IR vocabulary (add a primitive).

The fourth stage is what makes CAKE different from a search loop with better error messages. A finding is not just used and discarded; it is sent to whichever part of the system should learn from it.

10.2 What stays fixed

Two things anchor every run. The first is the workload contract, the “stable authority”: it fixes the shapes, the correctness oracle and its tolerances, the hardware, and which references the agent may see. Results are retained, so decisions can be audited and recurring findings reused.

The second is the model. All agent tasks in the paper use the same model, GPT-5.6-sol, at reasoning effort xhigh. “Holding the model and agent scaffold fixed makes the comparisons … attributable to the environment rather than to model capability.” Whatever difference shows up between two arms of an experiment, it is not because one had a smarter model.

11. The Compiler Evolves Too

So far, the agent evolves kernels inside a fixed environment that happens to be well designed. CAKE’s last commitment is that the environment is not fixed.

11.1 Two paths into the compiler

Kernel candidates, validation results, benchmarks, and failure reports are all evidence for changing the compiler. The paper describes two paths:

  • Path 1: mine expert kernels. Agents inspect production kernels and hardware documentation to find Blackwell patterns the IR is missing (“new instruction forms, resource types, descriptor variants, synchronization idioms”) and write compiler change proposals. Each proposal is checked against the design principles of §8.2, including performance transparency and friendliness to verification, before it is implemented.
  • Path 2: distill failures. Agents use feedback from failed candidates (“sanitizer reports, failure cases, correctness mismatches, debugging logs”) and turn recurring or costly failure modes into new analyses: “an opaque runtime crash becomes a verifier rule, a repeated illegal lowering pattern becomes a static check, a systematic misprediction becomes a calibration target.”
Flow diagram of how CAKE's compiler evolves. Path 1 mines production kernels and hardware documentation for patterns the IR is missing: new instruction forms, resource types, descriptor variants, synchronization idioms. Path 2 distills failed candidates: an opaque crash becomes a verifier rule, an illegal lowering a static check, a misprediction a calibration target. Both feed a compiler change proposal, which is checked against the IR design principles, implemented together with its analyses, and must pass corpus tests (400+ static and compile cases, 399 GPU correctness cases) at a human merge gate before the updated harness meets the next kernel family.
The outer loop: agents mine expert kernels for patterns the IR cannot yet express and distill recurring failures into checks; every change must pass the whole kernel corpus behind a human merge gate. Open the full-size SVG.

11.2 Why primitives and analyses must change together

The two paths feed each other. A new primitive tells the compiler more about the hardware, which makes stronger analyses possible; a new analysis constrains what future primitives may look like. The paper insists that a primitive and its analyses “must evolve together”: “syntax without effects and legality rules makes the IR less analyzable, and a new verifier rule without corpus validation can reject valid kernels.”

That is why every compiler change is test-gated across the kernel corpus. The corpus is large: more than 400 static and compile cases and 399 GPU correctness cases, across roughly 28 kernel families (§13.4). A new rule that rejects a known-good kernel does not merge.

Who makes these changes? Agents implement and maintain the analyses from high-level descriptions written by humans, and humans approve merges: the paper says compiler evolution “is still human-guided at merge gates.”

12. The Controlled Experiment: Flash-KMeans from a Clean Start

Descriptions of a system are cheap; a controlled comparison is not. This section covers the paper’s cleanest experiment: the same agent, the same model, the same task, written either in CAKE IR or directly in CUDA/PTX.

12.1 The workload: one step of k-means

k-means is one of the oldest algorithms in machine learning. You have N points and want K groups. Start with K guesses for the group centres (centroids), then repeat two steps: assign every point to its nearest centroid, and update every centroid to the average of the points assigned to it. Each repetition is called a Lloyd iteration.

Flash-KMeans is a fast, exact GPU implementation, and it matters for a reason you might not expect: video generation. Sparse VideoGen2 groups similar tokens with k-means so that sparse attention can work on contiguous blocks of related tokens (“semantic-aware permutation”). There, k-means runs inside the model, so its speed is the model’s speed.

According to the paper, a Lloyd iteration is dominated by two BF16 kernels that account for more than 95% of the end-to-end time:

  • assign computes the squared distance from every point to all K centroids and returns the index of the nearest one. It is compute-bound: a matrix multiplication followed by a reduction.
  • centroid_update sums the points of each cluster and counts them. It is limited by memory bandwidth and by contention on atomic updates.

The experiment focuses on assign, which “exercises the tensor-core pipeline and scalar epilogue,” at one fixed shape: B = 32 batches, N = 65,536 points, K = 1,024 centroids, D = 128 dimensions, BF16 inputs and FP32 accumulation. The bar is high: the baseline is a tuned Triton implementation from the FlashML project (whose library, FlashLib, collects fast GPU versions of classical machine-learning operators), measured at 0.938 ms.

12.2 The rules: what the agent may and may not see

The paper is careful about a subtle form of cheating. If the agent could read an expert’s CUDA for this kernel, it could copy the schedule instead of discovering it. So in these clean-start runs the agent may see “the mathematical specification, evaluation contract, correctness oracle, and high-level code, but not low-level target implementations such as CUDA, PTX, SASS, or equivalent generated source.” An existing implementation may be run as a black-box timing baseline, but its internals stay hidden. The restriction was enforced in isolated environments and audited afterward.

The two arms differ only in what the agent writes:

  • Treatment: the agent writes typed CAKE IR, with the harness of Part III.
  • Control: the agent writes CUDA C++ and inline PTX directly, under the same reference rules.

Everything else is fixed: the coding agent and its scaffold, the model and reasoning effort, the task statement, the correctness oracle, the benchmark harness, and the target shape. Each arm runs three times, with a budget of 80 million tokens.

12.3 The result

Representation Reached a plateau by 80M tokens Active evolve time (hours) Best speed at 80M tokens (× baseline)
CAKE IR 3 of 3 runs 1.89 [1.02, 2.33] 1.144 [1.041, 1.205]
Direct CUDA/PTX 0 of 3 runs 3.73 [3.59, 4.34] 0.928 [0.852, 1.151]

Entries are medians over three runs, with [minimum, maximum] in brackets. With CAKE IR, the median best kernel runs 14.4% faster than the tuned baseline (1.144×); with direct CUDA/PTX, it runs at 92.8% of the baseline’s speed. CAKE IR runs met the paper’s “prespecified plateau criterion” in all three runs; the CUDA/PTX runs met it in none. They also used about half the active time.

The trajectory tells the same story over time. At each 5-million-token checkpoint, every run contributes its best validated speed so far. The mean of the CAKE IR runs “crosses the tuned FlashML baseline by 55 million tokens and continues to improve, while the direct CUDA/PTX mean remains below baseline at the 80-million-token cutoff.” Reading the paper’s plot (approximate values), the CAKE IR mean is already near 0.8× by 20M tokens, when the CUDA/PTX mean is still around 0.2×.

12.4 How to read it

Our interpretation, with the caveats the numbers themselves carry:

  • The language changed the outcome, not just the speed of getting there. In the same budget, all three CAKE runs beat a tuned baseline and the CUDA runs, on median, did not.
  • But direct CUDA is not hopeless. The best CUDA/PTX run reached 1.151×, above the CAKE median. The difference is consistency and pace: CUDA can get there, but less reliably and more slowly. With three runs per arm, the spread matters as much as the median.
  • It is one kernel, at one shape, with one model. The experiment isolates the effect of the representation well, but it does not tell us how large the effect is for other kernels or other models. The paper does not define its plateau criterion in the text, and says the detailed stopping and timing accounting is “retained in the artifact.”

13. Beyond the Benchmark

The controlled experiment shows the effect cleanly; the rest of the evaluation shows the system at work on real kernels. It answers two more questions: can CAKE find the schedule of a kernel nobody has published a fast low-level version of, and can it reproduce, or beat, kernels that experts already tuned?

13.1 Frontier kernels: new architectures, no reference to copy

The paper calls a kernel frontier “in the operational sense that the agent must discover its physical schedule without inspecting a low-level target implementation.” This is where a co-evolving IR should help most, and also where it is most exposed: a missing capability shows up as a schedule “the agent cannot express at all.”

Kimi Delta Attention. The official FlashKDA kernels served only as a black-box timing baseline; their source and generated code were not given to the agent. The agent-generated prefill kernel, compatible with FlashKDA and covering fixed, packed variable-length, and tail inputs, reaches a 2.05× geometric-mean speedup over that baseline across six B200 BF16 shapes. The paper calls it “bitwise correct on its validation contract” (the pull request itself reports outputs that match FlashKDA within a tolerance of 10−2), and it was verified in end-to-end serving of Kimi-K3 under SGLang. Separate decode kernels reach 1.14× over upstream FlashInfer across 30 public-API shapes. Both went into FlashInfer as generated CUDA (pull requests #4262 and #4279, merged in August 2026), “so downstream users take on no dependency on CAKE.”

Gated DeltaNet and sparse attention. Against FlashInfer, CAKE’s Gated DeltaNet prefill and speculative-decoding paths “improve performance while preserving the model’s recurrent state,” and MiniMax sparse attention shows that the representation supports sparse-attention families across prefill and decode. The paper gives no speedup numbers for these two. They are “dispatch families rather than single kernels”: several physical schedules, each its own CAKE IR program, behind one entry point (§14).

13.2 Production kernels, evolved from a reference

TinyGEMM. Here the agent starts from an existing kernel: FlashInfer’s BF16 matrix multiply for tiny batches (the paper calls it small-M), derived from TensorRT-LLM. Tiny batches are typical when a server decodes a handful of requests at once. The agent produced an adaptive family of shallow and deep pipelines, including variants that use programmatic dependent launch (PDL, which lets the next kernel start its preamble before the previous one finishes) and paths for batch sizes below eight. FlashInfer PR #4274 reports an 18–23% geometric-mean reduction in kernel time across 35 canonical shapes and a broader regression suite. Greedy decoding stayed bitwise identical for GPT-OSS-20B and GPT-OSS-120B on B200 and GB300, and an SGLang experiment with GPT-OSS-120B measured up to 7.6% higher output throughput at concurrency 128 on one GPU (differences within noise with four-way tensor parallelism).

Alpha-MoE. The original Alpha-MoE is a fused mixture-of-experts “megakernel” written for Hopper, with 8-bit weights and activations (W8A8). CAKE agents rewrote it for Blackwell. One device program performs the routed gather, two projections, the activation, requantization, and route-weighted accumulation of the output. Against FlashInfer’s TensorRT-LLM-derived pre-routed API, the end-to-end API-level speedups are 6.204× at N = 256 and 4.025× at N = 512; measured as GPU span, they are 1.215× and 1.170×. The kernel was merged into FlashInfer in September 2026 (PR #4287), but the pull request reports a different comparison: against stock SGLang’s five-kernel Triton MoE path on GB300, its fused kernel is 1.37×–1.75× faster in GPU time. The paper’s 6.204× and 4.025× do not appear there.

13.3 Reproducing expert kernels

The last question is whether, when expert structure already exists and the agent may look at it, the harness helps agents preserve correctness while matching highly optimized kernels. The references come from FlashAttention-4, TensorRT-LLM, DeepGEMM, CUTLASS, and FlashInfer. All runs are on B200 at a fixed shape, use median CUPTI GPU span, and pass kernel-specific correctness gates.

Family Variant Relative performance CAKE IR lines Reference device lines
FlashAttention-4 forward, BF16, non-causal 1.0045× 430 2,369
FlashAttention-4 backward, BF16, non-causal 1.0470× 514 2,552
TensorRT-LLM GQA decode, FP16 1.043× 783 8,515
DeepGEMM 1D1D GEMM, FP8 1.0370× 221 516
DeepGEMM grouped GEMM, BF16, masked 1.0174× 401 442
DeepGEMM MQA indexer, FP8 1.2700× 480 704
DeepGEMM MQA indexer, FP4 1.2730× 392 704
DeepGEMM paged MQA indexer, FP4 1.0036× 395 779
CUTLASS MLA decode, BF16, TMA 1.2174× 845 1,860
DeepSeek-V4 sparse MLA decode, BF16 1.1297× 1,299 13,609
DeepSeek-V4 sparse MLA decode, FP8 0.9649× 1,393 8,942

Ten of the eleven comparisons meet or beat the reference; the last reaches 96.5%. The paper adds two honest qualifications. Entries below the reference “generally reflect compiler-integration maturity”: a feature the kernel wants is still being integrated into the compiler, so the submitted kernel uses the closest supported strategy. And the strongest wins, the two MQA indexers at about 1.27×, “are not faithful transcriptions of the reference kernels”: the agent explored optimizations the original lacked, so above-parity results “reflect search rather than transcription fidelity alone.” The line counts are descriptive only; the languages and counting scopes differ, and they show just that these schedules “can be represented compactly.”

Shapes and counting details
  • S1 (FA4): B = 4, H = 32, S = 8,192, D = 128. S2 (TRT-LLM GQA): B = 128, 64 query heads, 8 KV heads, KV length 4,096, D = 128, page size 16. S3 (FP8 GEMM): M = 4,096, N = 7,168, K = 4,096. S4 (grouped GEMM): 256 groups × M = 128, N = 4,096, K = 7,168. S5 (MQA indexers): H = 32, D = 128, query length 1,024, KV length 2,048. S6 (paged indexer): B = 256, H = 64, D = 128, average KV 4,096, block 64. S7 (CUTLASS MLA): DeepSeek-V3, B = 128, KV length 4,096, page 128. S8 (DeepSeek-V4 sparse MLA): B = 3, ragged query lengths [3, 4, 5], H = 128, Dqk = Dv = 512.
  • The paper’s table also lists “reference + support” line counts (for example 3,039 and 3,690 for the two FA4 kernels). For DeepSeek-V4, reference lines are restricted to source reachable under the fixed S8 route and runtime constants.
  • MQA (multi-query attention) indexers score which past tokens a query should attend to, as in DeepSeek-V3.2’s lightning indexer; MLA is DeepSeek’s multi-head latent attention. Our GLM-5.3 and DeepSeek-V4 deep dives explain both.

13.4 The breadth of the corpus

The headline experiments hide how much the harness now sustains. The validated corpus contains more than 400 static and compile cases and 399 GPU correctness cases across roughly 28 families: attention, dense and sparse GEMM, MoE, quantization, normalization, state-space models, KNN, and KMeans, with architecture-specific paths from Ampere through Blackwell. It includes more than 100 ports of TensorRT-LLM kernels for attention, MLA decode, and MoE.

One property stands out: composition. Because roles, barriers, and buffers are declared rather than implied, schedules that would normally be separate kernels can be written as one device program. BatchAttention combines decode and prefill work in one kernel; the Alpha-MoE and mega-MoE families fuse routing, expert computation, and output accumulation “without materializing intermediate results.” The Alpha-MoE rewrite also shows that “schedule structure can survive even when target instructions change,” from Hopper’s to Blackwell’s.

14. From a Tuned Shape to a Library

Everything so far optimizes a kernel for a particular shape. A library has to accept whatever shape the caller passes. This section covers the step that benchmark papers usually skip, and that CAKE treats as a stage of its own.

14.1 A different objective, so a separate stage

“Closing that gap is not a matter of running the inner loop on more shapes; it is a separate stage with a different objective, a different ranking signal, and a different failure mode.” An exact shape gives evolution a clean target and permits aggressive specialization; scoring the inner loop on broad coverage would blur that signal. So generalization begins only after strong per-shape seeds exist, and it is scored on performance including the dispatcher over a fixed workload. Incorrect or slow seeds go back to the inner loop “rather than being hidden behind routing.”

14.2 Building a dispatcher

The generalization stage groups the measured seeds into shape buckets, produces specialized or shared variants, and orders their guards (conditions on the shape, such as “D = 352”) behind an explicit fallback that handles anything the guards do not match. Tuning may change implementation parameters, but never the input shape. Before any aggregate number is reported, validation covers representative and held-out inputs, boundary and tail cases, guards that overlap or leave gaps, and the fallback path itself.

A library call kmeans(B, N, K, D) passes an ordered chain of guards on the dimension D (D = 288, D = 352, D = 144–176, D = 112, six more, a general case, else), each leading to a separate CAKE IR route that serves 14, 6, 8, 6, 14, 64 and 12 of the 124 KMeans shapes; an evolved seed tuned for one exact shape sits above the routes. A panel lists what is validated before reporting (held-out inputs, boundary and tail cases, overlapping or missing guards, the fallback path) and that the shape set is fixed before tuning. Results on GB200, GPU-span geometric mean versus the reference: KNN build 1.418× (112 shapes, 8 families, 80 routes), KNN search 2.116× (198 shapes), KMeans 1.803× (124 shapes), no incorrect outputs, KNN recall 1.0.
From one tuned shape to a library call: ordered guards send each call to one of several separately checkable CAKE IR routes or an explicit fallback. The KMeans routes and shape counts are from the paper; the guard conditions are inferred from the route names. Open the full-size SVG.

One policy keeps routing from turning into a zoo: the stage “reuses a single physical schedule across as much of the shape domain as it can, and introduces another only when the domain requires a material schedule change.” Each route is its own CAKE IR program, so every alternative can still be analyzed and benchmarked on its own.

14.3 Not grading your own homework

The subtle danger in this stage is evaluation leakage: if the same shapes are used to tune the dispatcher and to report its speed, the dispatcher can simply memorize them. CAKE declares the valid shape domain before tuning. Dispatcher predicates “may partition that domain, but they may not introduce convenient new evaluation rows,” and coverage expands only “through deterministic unseen shards of the same source.”

14.4 Results on real libraries

The classic machine-learning workloads give the cleanest reading, because their portfolios are too large for one shape to carry the result. On GB200, the generalized KNN build, KNN search, and KMeans kernels, contributed to FlashLib (pull requests #15, #16 and #18, merged in August 2026), reach these results, where Gspan is the unweighted geometric mean, over shapes, of the reference’s median CUPTI GPU span divided by CAKE’s:

Library function Shapes Gspan
KNN build 112 1.418×
KNN search 198 2.116×
KMeans 124 1.803×

There were no incorrect outputs, and KNN recall was 1.0 (every true nearest neighbour was found). The paper warns against comparing these numbers with the clean-start experiment of §12: the measurements “use different hosts, shape distributions, baselines, and protocols,” so the difference “is not, by itself, a measured cost of generalization.”

Route-level breakdown (the paper's appendix figure)

KNN build (baseline FlashLib 0.2.0): 8 outer families covering 112 shapes, with 80 distinct final routes observed.

Outer family Shapes / routes Gspan
D64 7 / 5 1.130×
D128 low-K BF16 38 / 17 1.258×
D128 low-K FP16 1 / 1 1.353×
D128 mid-K 32 / 30 1.425×
D128 large-K 10 / 5 1.909×
D192 1 / 1 1.432×
D256 8 / 8 1.386×
high-D 15 / 13 1.758×

Flash-KMeans (baseline: the authors’ tuned FlashLib implementation): 12 final routes covering 124 shapes.

Final route Shapes Gspan
gap-fused general 64 1.753×
D288 split-K 14 2.227×
aligned-shape fallback 12 1.439×
D144–176 padded tail 8 1.786×
D112 specialized 6 1.414×
D352 split-K 6 3.757×
D480 split-K 6 2.597×
micro-D hybrid 2 1.309×
micro-D pipelined 2 1.095×
large paired 2 1.217×
D128 specialized 1 1.036×
D224 TMEM 1 1.037×

Values are per-row geometric means from same-session CUPTI measurements. “Split-K” splits the reduction dimension across several thread blocks and combines their partial results, which helps when the other dimensions are too small to fill the GPU.

CAKE touches three research areas: GPU languages (§6), compilers, and self-improving systems. The cleanest way to place it is to ask of each line of work: what gets evolved, and what stays fixed?

Line of work Examples What changes What stays fixed
Kernel agents with search or memory KernelEvolve, EvoEngineer, AVO, K-Search, KernelBlaster the search process, the accumulated memory the language and the evaluation environment
Kernel models trained with RL AutoTriton, CUDA Agent the model’s weights the language and the evaluation environment
LLM-driven program search FunSearch, AlphaEvolve the programs, through LLM-written mutations the evaluator
Self-improving agents ADAS, Darwin Gödel Machine, Meta-Harness the agent’s own code, or the code around a fixed model the foundation model
CAKE   the compiler harness: IR vocabulary, analyses, cost calibration the coding agent and the model

In the paper’s words, CAKE “targets that complementary layer: it changes the representation being searched and the structured compiler evidence returned to the agent.” Nothing stops the two from combining: a stronger search procedure or a better-trained model could run inside CAKE’s environment.

On the compiler side, TVM, XLA, MLIR, TensorIR, Ansor, and MetaSchedule structure programs so they can be analyzed and optimized automatically, and Graphene, Twill, and Tawa model asynchronous GPU execution, pipelining, or warp specialization. CAKE “shares the principle that structured programs enable useful analysis, but places compiler findings inside an agent evolution loop and lets recurring kernel evidence evolve the harness.” More broadly, it is an instance of the “self-defining systems” agenda: AI-operated systems that can change their own mechanisms and abstractions.

16. Limits and Open Questions

A good paper says where it stops. CAKE does, and a careful reader can add a few more questions.

16.1 What the paper says about its own limits

  • NVIDIA only. CAKE targets Ampere through Blackwell. The schedule language, role model, and analyses carry across those generations, but instruction forms, legality rules, and cost anchors are architecture-specific, and “that cost is the honest measure of transfer.” For non-NVIDIA hardware it “remains unmeasured,” and the backend would have to be rebuilt.
  • Uneven evidence. Most performance results are on B200. The timing model is calibrated only for B200 and H100 and “declines to predict elsewhere.”
  • Incomplete analyses. The static analyses and performance models are intentionally incomplete; they rank and filter, and GPU execution stays the ground truth.
  • Humans in the loop. Compiler evolution is “still human-guided at merge gates.”

16.2 Questions the paper leaves open

These are ours, not the paper’s:

  1. Can others use it? The paper does not say that CAKE’s compiler, IR, or harness is publicly released, and we found no release as of October 2026 (only unofficial community re-implementations). The generated code is public: the four FlashInfer pull requests and three FlashLib pull requests are merged, and a FlashInfer tracker issue indexes many more CAKE-generated kernels.
  2. How general is the controlled result? The clean-start comparison is three runs per arm, one kernel, one shape, and one model (GPT-5.6-sol). We do not know how the gap changes with other kernels, other models, or a larger budget. (The FlashInfer tracker describes CAKE-generated kernels produced with Claude and Codex agents, which suggests the harness is not tied to one model, but no controlled comparison is reported.)
  3. Which part does the work? The language, the pre-compile checks, the cost model, and compiler evolution all change at once between the two arms. There is no ablation that, for example, keeps CAKE IR but removes the cost model, or freezes the compiler.
  4. Why not a tile-language arm? The control arm writes CUDA/PTX. A third arm writing Triton or CuTe DSL under the same rules would show where tile languages and layout algebras land.
  5. What does compiler evolution cost? The paper does not report how many compiler changes were made, how many came from each path, how often the corpus gate rejected a change, or how much human review the merge gates took.
  6. How accurate are the checks? False-positive and false-negative rates for the verifier, and the accuracy of the cost model’s predictions, are not reported.
  7. Missing numbers, and two mismatches. The Gated DeltaNet and MiniMax sparse-attention results are described without speedups, and the plateau criterion is not defined in the text. Two upstream records also differ from the paper: the merged Alpha-MoE pull request reports a different baseline and different speedups (§13.2), and the KDA prefill pull request reports agreement within a tolerance where the paper says “bitwise correct.”

None of these undercuts the main idea. They mark where the next paper, or an independent reproduction, could add the most.


Epilogue — a dialogue

Professor So. What did you learn?

Student That a fast kernel is mostly a good schedule: who loads, who multiplies, how they hand tiles to each other, and where every byte lives.

Professor And why agents struggled to write one?

Student They were writing in a language that either hid the schedule or made them manage every address by hand. And all they heard back was "crashed" or a number.

Professor And CAKE's fix?

Student A language where the decisions are explicit and the bookkeeping is derived, checks that name the broken barrier before the GPU ever runs, and a compiler that turns repeated failures into new checks. The agent gets better because its world gets better.

Professor And what would you still want to know?

Student Whether it holds with other models and other kernels, which piece matters most, and whether I can try it myself.

Professor Good. Now you're reading papers properly.

Summary

  • The problem. Expert GPU kernels are distinguished by their schedule (warp roles, barrier choreography, memory placement), and mistakes in a schedule cause silent wrong answers or hangs. Kernel agents have improved their search, memory, and models, but they still write in languages that either hide the schedule or expose it without guardrails, and they get back only errors, pass/fail results, and timings.
  • The idea. Co-design the agent’s environment with the agent. Agents write CAKE IR, a typed, hardware-explicit schedule language with no layout algebra, in which resources, roles, and handoffs are declared and the mechanical details are derived. A harness checks safety, hardware conformance, data consistency, and schedule semantics before compilation, ranks candidates with a calibrated cost model, and treats GPU measurement as ground truth. The compiler evolves: recurring failures and newly found expert patterns become verifier rules, primitives, and calibrations, gated by a corpus of more than 400 static and compile cases and 399 GPU correctness cases.
  • The evidence. With the model fixed, CAKE IR beats direct CUDA/PTX on a clean-start Flash-KMeans task (median 1.144× vs 0.928× of a tuned baseline at 80M tokens; 3/3 vs 0/3 runs plateaued; about half the active time). Without seeing reference implementations, it produced a KDA prefill kernel 2.05× faster than the official kernels (geometric mean over six shapes), validated in end-to-end serving. It matches or beats 10 of 11 expert kernels, and its dispatcher-backed KNN and KMeans families are 1.42×–2.12× faster across more than 400 shapes. Four kernels are merged into FlashInfer.
  • The limits. NVIDIA GPUs only, mostly B200 evidence, a small controlled study, incomplete analyses by design, human merge gates, and no public release of CAKE itself (its generated kernels are public).

Questions to Ponder

OSTEP ends its chapters with homework. Ours has no simulator, but these are worth a few minutes each:

  1. In §3.2, a ring buffer with three stages uses a one-bit phase per barrier. Why is one bit enough? What would go wrong if the producer could run two full laps ahead of the consumer?
  2. CAKE refuses to step a schedule down to an older GPU (§7.4). What could go wrong if a Blackwell schedule were silently lowered to Hopper instructions?
  3. The pre-compile checks are allowed to have false positives (§9). In an evolution loop, which is worse: a check that wrongly rejects 5% of valid kernels, or one that misses 5% of broken ones? Does your answer change for a production compiler?
  4. Design the missing ablation from §16.2: which arms would you run to separate the effect of the language from the effect of the harness and of compiler evolution, and what would each arm keep fixed?
  5. §14.3 forbids the dispatcher from adding its own evaluation shapes. Give a concrete example of how a dispatcher could look 2× faster than it is if this rule were broken.
  6. The paper borrows Terence Tao’s point that cheap generation shifts the bottleneck to verification and understanding (§8.2). Where else in machine-learning systems do you see the same shift happening?

Sources and Further Reading

The paper

  • CAKE: Compiler–Agent Co-Design for Frontier Kernel Evolution (arXiv:2608.12629), 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, and Luis Ceze (NVIDIA and Carnegie Mellon University). Every number in this post comes from here unless marked otherwise. Read §2–3 for the design and the appendix for the eight principles and the full kernel tables.

The upstream kernels

  • FlashInfer pull requests #4262 (KDA prefill), #4279 (KDA decode), #4274 (TinyGEMM), and #4287 (Alpha-MoE), all merged: generated CUDA you can read without CAKE. The CAKE kernel tracker lists many more.
  • FlashLib pull requests #15, #16 and #18: the dispatcher-backed KNN and KMeans kernels of §14.

Background on GPUs and kernels

Kernel agents and evolutionary search

The workloads

The writing style

  • Operating Systems: Three Easy Pieces by Remzi H. Arpaci-Dusseau and Andrea C. Arpaci-Dusseau, free online. The dialogues, crux boxes, asides, and tips in this post are borrowed, with gratitude, from its format.

On this blog