MaxKernel: Agentic Kernel Generation for TPUs
Shangkun Wang, Nina Cai, Charles Hoong, Julian Walker, Gerson Kroiz, George Vanica, Deepak Patil, Andi Gavrilescu, Hassan Sipra, Sethu Sankaran
cs.AI, cs.PF, cs.PL
2026-09-04
MaxKernel's agent loop writes Pallas TPU kernels. Parallel search hits 1.58× geomean on JaxBench's 50 tasks and 2.32× on eight production kernels, beating humans at 2.02×.
Custom TPU kernels stall in the same place every time. XLA cannot schedule attention variants, sparse ops, and heavily fused kernels well, so engineers drop into JAX Pallas and manage HBM versus VMEM, DMA pipelines, and multi-dimensional tiling by hand. CUDA and Triton are already hard on GPUs. TPU APIs are stricter and the compiler errors are worse.
Zero-shot sampling does not get past compilation. On JaxBench's 50 tasks, 100 independent draws of the fastest correct sample compile and match numerically on only 10/50, for a 1.08× geometric mean over XLA. Compiling and matching the reference still does not mean the kernel is fast. Speed needs hardware traces, retiled memory layouts, and another compile. Halide, TVM, and Ansor split algorithm from schedule and still demand expert search. Google's MaxKernel wires compiler feedback, XProf traces, a frozen test suite, and a search graph into one multi-agent system.
Shared sub-agents plan, emit Pallas, patch from compiler errors, synthesize tests, run on device, autotune block and tile sizes, and pull latency, bandwidth, and compute density from XProf. Retrieval is limited to Pallas, Mosaic, and XLA docs and handbooks. Hand-tuned kernels are kept out of the corpus so the agent cannot copy expert code.
Three orchestrators share that pool.
HITL runs one sub-agent and stops for a human to inspect the plan or the draft. The paper reports no numbers for this path.
Auto freezes a test suite from the reference implementation before the loop, so the implementer cannot edit the checks. The cycle is plan, implement, compile, test, autotune, profile, then feed the trace into the next plan. Compile or numeric failure short-circuits back to planning. Snapshots are kept and the run rolls back to the lowest-latency valid kernel. Each trajectory is capped at 5 iterations; each task is run 5 times and the median is reported.
Graph search treats every Auto output as a node holding code, plan, and metrics. Each expansion starts a fresh session so context does not overflow. Parallel search runs 5 unpruned Auto trajectories of 5 iterations and keeps the fastest correct kernel. Beam search uses width 3, depth 3, 2 branches per node, and only 2 inner iterations, trading horizon for breadth.
All LLM calls use Gemini 3.1 Pro on TPU v6e. Correctness is jnp.allclose, mostly atol=rtol=1e-2, relaxed to 1e-1 on some bf16 tasks. Latency is on-device time with host dispatch and compile stripped. Geometric mean floors regressions at 1.0×.
JaxBench mixes 17 LLM-style operators with 33 fused kernels adapted from KernelBench.
| Method | Compile | Correct | Geomean speedup | fast1 |
| Best-of-100 zero-shot | 10/50 | 10/50 | 1.08 | 6/50 |
| Auto median | 49/50 | 48/50 | 1.39 | 22/50 |
| Parallel search | 50/50 | 50/50 | 1.58 | 34/50 |
| Beam | 50/50 | 50/50 | 1.49 | 31/50 |
Auto's speedup band is [1.19, 1.42]; a single trajectory can stick in a compilable but slow local state. Parallel search reaches fast1 of 34/50 and still clears 24% of tasks at a 2.0× threshold.
On eight production kernels that have hand-tuned Pallas, geometric mean over XLA (floor 1×) is 2.02× for humans, 2.32× for parallel search, 1.78× for beam. Agents win 7 of 8. Parallel search hits 6.74× on Paged Attention versus 2.41× human, and 5.03× versus 2.45× on Sparse Attention. The miss is Ragged Paged Attention: humans 4.65×, agent 1.42×. On MLA, the human kernel is 0.69×, slower than XLA; the agent reaches 1.21–1.23×. GEMM sits at 1.02–1.03× for everyone. XLA already owns that op.
On recent model kernels versus JAX or existing Pallas:
| Workload | Baseline | MaxKernel | Gain |
| MLA v1 latency | 3.127 ms (Pallas) | 2.856 ms | −8.68% |
| Qwen3-Next GDN forward | 17.09 ms (JAX) | 10.47 ms | 1.63× |
| Qwen3-Next training step | 84.51 ms (JAX) | 17.99 ms | 4.70× |
| DeepSeek-V4 Sparse prefill | 27.08 ms (JAX) | 3.449 ms | 7.85× |
| DeepSeek-V4 Sparse decode | 0.047 ms (JAX) | 0.02 ms | 2.36× |
| Mamba v2 SSD | 0.211 ms (JAX) | 0.192 ms | 1.10× |
| StreamIndextopk | 0.415 ms (JAX) | 0.250 ms | 1.66× |
MLA v1 throughput moves from 116.7 to 127.8 TFLOPS. The same loop patched a Ragged Page Attention v3 prefill crash by inserting ALU clamps for left padding, so negative slices and dirty DMA prefetches stop.
For teams already writing Pallas, this is a runnable closed loop: a JAX reference in, compiler errors and XProf traces back, a numerically gated kernel out. JaxBench going from 10/50 compiling under 100 blind samples to 50/50 correct at 1.58× geomean is mostly hardware feedback, not a bigger base model.
Beating the human geomean on eight kernels does not mean replacing the expert. Ragged Paged Attention still has the human more than 3× ahead. GEMM has no headroom. The 1.58× parallel number is best-of-5 Auto runs, roughly 5× the compute of one Auto pass, and token or wall-clock cost is unreported. HITL is a UI, not a measured baseline.
The usable setting is narrow: a JAX reference exists, TPU v6e is available, and extra search trajectories are acceptable. CUDA, Triton, NKI, and MTIA appear only in the introduction. Every number is Pallas on TPU.
The authors flag follow-ups: evolutionary or greedy search, a knowledge base that grows with past runs, and a stateful mix of exploration and exploitation. Only Gemini 3.1 Pro was tested.
Flooring the geomean at 1.0× inflates the headline. Kernels slower than XLA count as "not slower." Tolerances of 1e-2 to 1e-1 are loose for training kernels with backward passes. The 6.74× Paged Attention result sits in the Parallel column; the prose once attributes it to beam search. Trust the table.
Parallel search's edge over beam is entangled with 5 iterations versus 2, so it is not a clean algorithm comparison. Keeping hand-tuned kernels out of RAG is honest for the study and may be the wrong production default. There is no head-to-head with AutoComp on the same TPU tasks, and search cost is missing.