Google's MaxKernel: an agentic framework that writes and tunes Pallas TPU kernels — 2.32× over reference code, open-sourced under Apache 2.0
A three-mode multi-agent system writes and tunes JAX/Pallas kernels, deliberately excluding hand-tuned code from its retrieval corpus — with the caveat that every speedup is vendor-reported.
Google has open-sourced MaxKernel, a multi-agent framework that writes and optimizes Pallas kernels for TPUs. It reports a 2.32× geometric-mean speedup over reference code on eight production kernels — beating the human-written versions (2.02×) — and a 1.62× speedup on the 50-task JaxBench suite (paper). The code is under Apache 2.0 in Google's AI-Hypercomputer/accelerator-agents repo (GitHub).
What shipped
MaxKernel is a multi-agent system that generates, debugs, tests, and profiles custom JAX/Pallas TPU kernels, using LLMs plus real-time compiler and hardware-trace feedback (paper). It first shipped in May 2026 as a human-in-the-loop-only agent (Google developer forum); the two automated modes arrived with the September 2026 paper (v1), and the current November 2, 2026 revision (v2) is the version evaluated here.
The framework runs three orchestration paradigms over a shared pool of specialized sub-agents:
- Human-in-the-Loop (HITL) — an orchestrator follows a "one agent, then wait" pattern, halting after each phase so a developer can review the optimization plan or code draft and steer the trajectory.
- Autonomous Loop — chains the sub-agents into a closed loop that iterates plan → implementation → compile validation → testing → profiling without human intervention, effectively a localized hill-climbing search.
- Graph-Based Autonomous Search — models the kernel design space as a formal search graph and runs parallel or beam search over it, trading speed for broader coverage to escape local optima.
The sub-agents are divided by responsibility: a Planning Agent produces an algorithmic optimization plan from hardware specs and the reference implementation; an Implementation Agent turns it into Pallas code; Testing and Verification Agents synthesize test suites and run them on the target hardware against numerical tolerances; an Autotuning Agent sweeps block sizes and tile dimensions; and a Profiling Agent extracts latency, memory-bandwidth utilization, and compute density from XProf traces to feed the next iteration (paper).
One design choice matters for interpreting the results: the RAG knowledge store is restricted to static sources — Pallas, Mosaic, and XLA documentation, memory-layout guides, and performance handbooks — and explicitly excludes hand-tuned kernel code (paper). The authors say this is intentional: they want to test whether the agents can rediscover TPU-specific optimizations from scratch rather than interpolate from a library of known-good solutions.
The headline numbers, all from Google's own evaluation:
- 2.32× geometric-mean speedup over reference code on the eight production kernels where human-written kernels exist, "outperforming human-written ones (2.02×)" (paper).
- 1.62× geometric-mean speedup across the 50-task JaxBench suite, whose baseline is XLA, JAX's default compiler (paper).
- A concrete internal dogfooding result on a DeepSeek MLA kernel on a v5p TPU: latency dropped from 3.12 ms to 2.856 ms (8.7%) and throughput rose from 116.73 to 127.82 TFLOPS (9%) against a human-written, optimized Pallas baseline (Google developer forum).
Why it matters
Pallas kernel expertise is scarce and concentrated; hand-optimizing a kernel is a weeks-long loop of writing code, compiling, reading hardware traces, and adjusting. MaxKernel automates most of that loop, and the DeepSeek MLA example shows a real inference win on a state-of-the-art open model's kernel, not a toy benchmark (Google developer forum).
For TPU engineers, the practical takeaway is that HITL mode is usable today as a kernel co-pilot: it converts CUDA, Triton, or PyTorch kernels into JAX/Pallas, generates pytest suites, runs compilation and numerical checks, and profiles DMA and compute-vs-memory ratios, all under scoped file-system permissions (Google developer forum). The autonomous and graph-search modes carry the "beats the experts" claim, and they are the part most worth watching.
Strategically, Google has converted an internal advantage into shared infrastructure, raising the bar for the CUDA and HIP ecosystems. If agentic kernel optimization becomes standard in the TPU stack, the same pressure lands on NVIDIA and AMD toolchains — and the underlying principles are claimed to generalize to CUDA/Triton, NKI, and MTIA (paper).
Caveats
Every headline number here is vendor-reported. No third party has independently benchmarked MaxKernel — that is an inference from the sources, not a stated fact — and the paper does not break out per-task wins and losses across the 50 JaxBench kernels, so the uniformity behind the geometric mean is unverified (AI Weekly).
Two details matter for anyone citing the numbers. First, the JaxBench figure changed between revisions: the v1 abstract said 1.58×, while the v2 abstract and body say 1.62× — AI Weekly still quotes 1.58×, but the current paper says 1.62× (v1, v2). Second, the paper's own framing is hedged: the abstract says MaxKernel produces code "matching expert hand-tuned baselines," while the body claims to "outperform human-written ones (2.02×)."
The paper also concedes that some complex kernels still need expert guidance to avoid local optima — the automation is not fully hands-off (paper). And the repo carries a "not an officially supported Google product" disclaimer (GitHub).