MaxKernel: agents that out-optimize hand-tuned TPU kernels

Writing a fast custom kernel for a TPU is not a coding problem so much as a hardware negotiation. You are managing memory hierarchies by hand — deciding what lives in HBM versus VMEM, orchestrating DMA pipelining, choosing tiling strategies that will only reveal whether they were right after a compile-and-profile cycle. Very few engineers do this well, which is exactly why it's a good test of whether an agentic system can do real, load-bearing systems work rather than boilerplate.

A Google team just published MaxKernel, a multi-agent framework for generating JAX/Pallas kernels for TPUs, and the results are specific enough to take seriously: on a 50-task benchmark it goes from a 1.08x geometric-mean speedup with brute-force sampling to 1.58x with a coordinated agent loop, and on eight real production kernels its best configuration hits a 2.32x geometric-mean speedup against XLA — edging out the hand-tuned reference implementations, which averaged 2.02x. The code is open-sourced.

Three ways to run the same agents

MaxKernel's core idea is a shared pool of specialized sub-agents — planning, implementation, self-debugging, testing, and hardware profiling — that gets orchestrated three different ways depending on how much control you want to keep.

Human-in-the-loop (HITL) runs on a "one agent, then wait" principle: an orchestration agent routes work to a sub-agent, execution halts, and control returns to the developer to inspect the optimization plan or code draft before the next phase runs. This is the mode for kernels where expert judgment needs to steer the search.

Autonomous (Auto) is a closed loop with no human in it. It has a distinct preparation stage: before any optimization begins, a test-generation agent freezes a test suite against the reference implementation — deliberately locked so that later agents can't loosen the correctness bar to make failing code pass. Then it cycles: plan, implement, validate, profile, and feed the profiling data back into the next planning step. It also keeps a snapshot of every valid iteration and rolls back to the lowest-latency one at the end, so a bad late-stage mutation can't regress the final result.

Graph-based autonomous search wraps the Auto agent inside a formal search over a graph of kernel states, each node carrying its code, its optimization plan, and its measured correctness and speedup. MaxKernel implements two search strategies here: parallel search, which runs multiple independent optimization trajectories with a generous iteration budget per path, and beam search, which keeps a narrow frontier of the top-3 candidates and prunes aggressively between depths.

One design choice stands out: the retrieval-augmented knowledge base backing these agents is restricted to static framework documentation — Pallas, Mosaic, and XLA references and performance handbooks. The authors explicitly exclude hand-tuned kernel code from that corpus. That's a deliberate constraint on the experiment: it forces the agents to rediscover optimizations rather than retrieve someone else's answer, which makes the eventual comparison against human-written kernels a fairer one.

What the benchmark actually shows

The team evaluates on JaxBench, a 50-task suite covering attention variants, dense and sparse linear algebra, and fused operator chains, measured on TPU v6e hardware with XProf capturing real on-device execution time. Four methods are compared: a zero-shot best-of-100 sample, single-trajectory Auto, parallel search, and beam search.

The zero-shot baseline compiles only 10 of 50 kernels and produces a geometric-mean speedup near 1.0x — LLMs without compiler feedback mostly don't produce working accelerator code. Auto alone reaches 48-49 of 50 kernels compiling and correct, with a median 1.39x speedup, though single runs show real variance since a single trajectory can get stuck in a local optimum. Running five Auto trajectories in parallel and keeping the best gets every kernel compiling and correct, at 1.58x geometric-mean speedup. Beam search reaches the same 50/50 reliability at a slightly lower 1.49x, but at a fraction of the compute — it caps each candidate at two iterations instead of five.

That depth-versus-breadth tradeoff is the most useful practical takeaway. Parallel search's long, uninterrupted horizon per trajectory lets it work through multi-step compilation errors and memory-layout fixes that a shallow budget would abandon too early. Beam search's narrower per-candidate budget is cheaper and still competitive when the design space has many easily reachable local improvements — it plateaus when a kernel needs sustained, multi-step debugging to get past a single rigid constraint.

Where it beat the people who wrote the reference kernels

The more interesting comparison is against eight production kernels where hand-tuned Pallas implementations already exist — flash attention, GQA and MLA attention, sparse and paged attention, GEMM, and Megablox GMM. MaxKernel's parallel search matches or beats the hand-tuned version on seven of the eight, including a 6.74x speedup on paged attention against a 2.41x hand-tuned baseline, and 5.03x on sparse attention against 2.45x. On MLA attention, the human-written kernel actually underperformed the XLA baseline (0.69x) while both agent methods found real gains (1.21–1.23x) — a case where expert intuition guessed wrong and empirical search didn't. The one loss was ragged paged attention, where the hand-tuned kernel hit 4.65x against the agent's 1.42x — a reminder that matching experts on average doesn't mean matching them everywhere.

Beyond JaxBench, the team ran MaxKernel against kernels from current open-source models: a 4.70x speedup on the fused forward-backward training step for Qwen3-Next's gated DeltaNet, up to 7.85x on prefill for DeepSeek-V4's sparse attention, and a 9.5% throughput gain on multi-head latent attention over an already-optimized Pallas baseline. It also used its self-debugging agent to fix a real crash — a padding-related negative slice-size bug in ragged paged attention's prefill mode — by inserting protective clamp instructions, not just chasing latency numbers.

My take

What I find most credible about this paper is the eval design, not just the speedup numbers. Excluding hand-tuned kernels from the RAG corpus is the kind of constraint that's easy to skip and that quietly inflates results when skipped. Keeping it in makes the head-to-head against human engineers mean something. The other lesson worth carrying into your own agentic systems: closed-loop correctness gating — freezing the test suite before optimization starts, and rolling back to the best-known-good snapshot rather than trusting the last iteration — is doing as much work here as the LLM itself. The model proposes; the harness is what keeps it honest. That pattern generalizes well past kernel generation to any agentic loop where the model is allowed to grade its own homework.