MaxKernel: Agentic Kernel Generation for TPUs

MaxKernel. Agentic Kernel Generation for TPUs. By Shangkun Wang and nine others, from Google and Google DeepMind. Published on the third of September. Its code repository is worked out: about four hundred source files, with a readme and a licence.

Designing and authoring high-performance custom kernels for accelerators is a complex task, and what keeps it complex is that it requires deep hardware-level expertise. Large language models can be brought to bear on it, particularly when paired with real-time compiler feedback, which turns kernel generation into something an agentic system can attempt at all.

MaxKernel is a multi-agent system that declines to commit to one way of using that. It implements three distinct paradigms for TPU kernel development, and they differ in how much of the work a person does. A human-in-the-loop agent supports collaborative, step-by-step design. An autonomous agent executes a fully automated optimization loop driven by metrics and traces. A graph-based autonomous search scales that autonomous agent up to global exploration of the design space.

Underneath, the three share the same machinery. A shared pool of specialized sub-agents handles planning, implementation, self-debugging, testing and hardware profiling, and all three paradigms draw on it. What changes between them is the control structure above that pool rather than the capabilities inside it, which is why one system can offer step-by-step collaboration and unattended global search without being two systems.

Evaluation runs on JaxBench, a comprehensive suite of fifty diverse kernel tasks, and alongside it on complex real-world workloads drawn from state-of-the-art open-source models. Across those, MaxKernel consistently generates highly optimized implementations, matching expert hand-tuned baselines and delivering significant performance. Matching the hand-tuned baseline is the bar that counts, because that baseline is the expertise the task was said to require.

The agent is open-sourced. What the three paradigms amount to is a question about where the human belongs, answered by declining to settle it: the same sub-agents serve a person working step by step and a search running unattended, so the choice stays with whoever is doing the work.

MaxKernel. Agentic Kernel Generation for TPUs. By Shangkun Wang and nine others, from Google and Google DeepMind. Published on the third of September. This is the long version: the whole paper rather than the abstract.

To get the full compute and memory bandwidth out of an accelerator, engineers routinely bypass the standard compiler pipeline and write custom kernels by hand. Doing that well means managing memory hierarchies explicitly, deciding what lives in high-bandwidth memory and what lives in the fast on-chip memory, orchestrating direct memory access pipelining, and deriving tiling strategies across several dimensions. It is expensive expertise, it is specific to each accelerator family, and the APIs change. That combination is what makes kernel authoring an appealing target for automation and a hard one.

A language model asked to write a kernel cold does not do well, and the reasons are specific rather than general. Accelerator APIs are rigid, memory constraints are strict, and the compiler errors that come back are low-level and opaque. So generation has to be augmented with real-time compiler feedback simply to reach code that builds. Even then, functional correctness is not the goal. Reaching good hardware performance requires an iterative loop of empirical profiling and configuration tuning, which is why this problem attracts agentic treatment rather than a better prompt. Traditional compiler approaches attacked the same complexity by decoupling algorithm definition from scheduling, and still demanded either significant domain knowledge or an expensive evolutionary search. The authors also note a failure mode in earlier agentic work: abandoning complex ideas before the code has had a chance to compile, which discards the optimizations most likely to matter.

MaxKernel decomposes the job into specialized sub-agents that all three of its operating modes share. A planning agent produces a high-level optimization plan from the hardware specification and the reference implementation, and an implementation agent turns that plan into kernel code, with a validation loop that compiles, reads the compiler feedback, and fixes errors before anything executes. Testing is split between an agent that synthesises test suites and one that runs them on the target hardware, enforcing numerical tolerances specified per hardware generation. An autotuning agent explores block sizes and tile dimensions. A profiling agent captures low-level trace data and extracts execution latency, memory bandwidth utilisation and compute density, feeding those back as the signal for the next round of optimization. Documentation reaches the agents through retrieval rather than being loaded wholesale, because putting comprehensive hardware specifications into the context window is prohibitively expensive.

That restraint pays off in the numbers. Across the benchmark suite of fifty diverse kernel tasks, the system reaches a geometric mean speedup of about one and a half times. The sharper result is on eight production kernels where a human-written version exists for comparison. There the agent reaches a geometric mean speedup over the reference code of two point three two times, against two point zero two for the human-written kernels. The abstract describes this as matching expert hand-tuned baselines, which undersells it, and it lands with the expert code held out of retrieval entirely.

The evaluation is one accelerator family and one programming interface, which the authors acknowledge while arguing the agentic principles carry over. The benchmark is fifty tasks plus a handful of real workloads, and the human-comparison result rests on eight kernels, which is a small set to hang a claim about beating experts on. The retrieval corpus is static, so the system does not learn across runs. And the headline figures are geometric means over speedups, which compress the distribution: they say nothing about whether the agent ever produced a kernel slower than the reference it started from.

The design point to take away is the separation between the pool and the orchestration. All three modes draw on the same sub-agents for planning, implementation, self-debugging, testing and profiling; what differs is only the control structure above them. That is why one system can serve a developer working step by step and a search running unattended without becoming two systems, and it is why the question of where the human belongs can be left to whoever is doing the work rather than settled by the architecture.