Give an AI model a TPU kernel to write and it will fail 94 times out of 100. Hand it the right documentation first and that failure rate collapses: per-sample correctness jumps from 5.8% to 37.3% S¹. That six-fold leap is the buried finding from JAXBench, a new open-source benchmark for AI-generated TPU kernels released on arXiv this week S¹, and it raises a question the entire AI infrastructure field has been dancing around: what if the bottleneck is not model size, but the manual?

My read: This is the first TPU benchmark I have seen that measures whether AI can write fast kernels and what it needs to get there. The 5.8% baseline is brutal. Without curated TPU docs, Gemini 3 Flash essentially cannot write correct Pallas kernels. The fact that documentation, not model scale, is the lever tells me the gap between AI-written and human-written infrastructure code is closing faster than people realise, but only for teams willing to invest in the context layer. I am skeptical of the self-reported speedup numbers until independent teams reproduce them, but the correctness finding is the one to watch.

Why TPUs needed their own test

GPU kernel generation has had benchmarks for a while. KernelBench exists for CUDA. TPUs, Google's custom AI accelerators, had nothing equivalent S¹. That gap matters because TPUs run on a different software stack: JAX and Pallas, a domain-specific language for writing low-level TPU code. Pallas is sparsely documented, which means AI models trained on general code have very little to draw from when asked to write TPU kernels.

JAXBench fills that gap with 50 JAX workloads S¹. Seventeen are production ML operators pulled from the MaxText library, drawn from real architectures including Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2, and AlphaFold2 S¹. The remaining 33 are translated from KernelBench, validated for correctness and resized for high utilization on TPU v6e hardware S¹.

Eight of the 17 production operators come with hand-optimized Pallas kernels from Google's Tokamax library, block-size tuned to set an expert upper bound S¹. That gives the benchmark a ceiling to measure against: how close can AI get to what a human specialist writes?

The documentation finding that matters more than the speedup

The authors evaluated four feedback-driven methods for generating Pallas kernels using Gemini 3 Flash S¹. The headline finding across the full suite: target-specific context matters more than model scale on a sparsely-documented DSL like Pallas S¹.

The numbers tell the story. Without curated TPU documentation, the model produces correct kernels 5.8% of the time. With it, that rises to 37.3% S¹. The same documentation conditioning lets the model solve 48 of 50 benchmarks at a 1.28x geomean speedup over XLA, Google's default compiler backend S¹.

How close AI gets to human experts

Once correctness is achieved, search structure adds real gains S¹. Autocomp, a beam-search pipeline, reaches 1.36x geomean speedup over XLA across the full suite S¹. On the eight hand-tuned kernels specifically, Autocomp hits 1.60x S¹.

The Tokamax expert upper bound on those same eight kernels is 2.08x S¹. So AI recovers most of the gap, but not all of it. Autocomp trails specifically on paged and ragged attention operators, the specialised memory-management patterns used in serving large language models S¹.

Geomean speedup over XLA baseline

The benchmark, evaluation harness, and baseline results are released as open source on GitHub under the AI-Hypercomputer organisation [P3, P4].

What to do about it

If you run ML infrastructure on Google Cloud TPUs, the practical takeaway is concrete. The single biggest investment you can make in AI-assisted kernel optimisation is not a bigger model. It is curated TPU documentation, fed as context to whatever model you are already using.

Consider a team training custom recommendation models on TPU v6e pods. Their bottleneck is a fused embedding-and-attention kernel that XLA compiles but does not optimise well. Rather than waiting for a Pallas specialist to hand-tune it, they could assemble a context pack from the Pallas documentation, the Tokamax kernel patterns, and their own profiling data, then run an automated search loop. JAXBench's results suggest this approach would get them most of the way to a hand-tuned kernel, faster than hiring for the role.

This week: check the JAXBench repository P³, run the evaluation harness against one of your own kernels, and measure whether your model produces correct output with and without curated TPU docs. The 5.8% to 37.3% gap is the baseline to beat.

What we don't know yet

All performance figures in the paper are self-reported by the authors, and the paper is an arXiv preprint that has not undergone peer review S¹. The speedup numbers are relative to XLA and Tokamax baselines the authors chose; absolute performance on different TPU generations or software versions may vary.

The claim that no equivalent TPU benchmark exists is the authors' own framing and may be incomplete S¹. The finding that context matters more than model scale was tested with Gemini 3 Flash specifically, and the paper does not claim this model is uniquely optimal for the task S¹.

The specialised attention operators where Autocomp trails, paged and ragged attention, are exactly the kernels that matter most for LLM serving at scale. Whether AI can close that gap with better search strategies, or whether those patterns require human expertise for now, is the open question.

The next signal: community contributions to the JAXBench GitHub repository over the coming weeks, particularly any independent reproduction of the 1.60x speedup claim. We will check the repository's issue tracker and pull requests for verification. If this work matters to you, subscribe to follow that audit.

Sources: S1 — JAXBench: Benchmarking Autonomous TPU Kernel Optimization · P2 — JAXBench: Benchmarking Autonomous TPU Kernel Optimization · P3 — accelerator-agents/JAXBench at main · AI-Hypercomputer/accelerator-age · P4 — JAXBench/README.md · P5 — 81be73a Add JAXBench: TPU kernel benchmark suite (#34)

More from Not A Tech Guy

Generated from an audited evidence pack with primary-source research. Social-media items are discussion signals, not verified facts. Nothing here is financial, legal or medical advice.