JAXBench: Benchmarking Autonomous TPU Kernel Optimization
Arya Tschand, Charles Hong, Julian Walker, Nina Cai, Shangkun Wang, Suvinay Subramanian, Sundar Dev, Vijay Janapa Reddi, Amir Yazdanbakhsh, Sethu Sankaran
cs.AI
2026-05-19
JAXBench adds 50 TPU-native Pallas kernel tasks. For a sparsely documented DSL, curated docs matter more than model scale, lifting per-sample correctness from 5.8% to 37.3%.
Autonomous GPU kernel optimization has KernelBench as a shared target to hill-climb on; TPU has had nothing. A TPU is a sequential machine with matrix-multiply units (MXUs) that uses not CUDA but the Pallas DSL, and Pallas is sparse in training corpora and thin on docs. Without a benchmark there is no way to measure how close LLM-written Pallas gets to human experts and no common hill for the community to climb. Google fills the gap and asks a sharper question: for an underrepresented DSL, is the bottleneck the model's reasoning, or the fact that it has never seen the documentation?
JAXBench is 50 tasks, all on a single Google Cloud TPU v6e chip. Composition: 17 production operators pulled from real architectures in MaxText (Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2, AlphaFold2), spanning attention variants, GEMM, RMSNorm, MoE, RetNet and the like; 33 translated from KernelBench Level 2, sized to at least 60% MXU utilization; 8 priority operators ship hand-tuned Pallas kernels from the Tokamax library as an expert upper bound.
Evaluation runs four checks: compiles, correctness (jnp.allclose, bf16), speedup over XLA, and fast1@NN (fraction beating XLA within N samples). Speedups take the per-task best, floor at 1x, then geomean; wrong kernels count as 1x.
Four feedback-driven methods are tested, all with Gemini 3 Flash on a 144-sample budget: Best-of-NN, iterative refinement (with compile errors, correctness, and profiler feedback), iterative refinement plus curated TPU documentation, and Autocomp (a two-phase beam search, 4 translation rounds plus 4 optimization rounds). The docs are targeted: a hardware-architecture summary, a Pallas API reference, code examples, and a rules block. The motivation is direct: without docs, 99.7% of Best-of-NN samples and 93.8% of iterative samples fail at compile or first execution, mostly on Pallas API misuse.
Full 50, Gemini 3 Flash:
| Method | Geomean speedup | Correct |
| Best-of-NN | 1.01x | 13/50 |
| Iterative | 1.18x | 32/50 |
| Iterative + docs | 1.28x | 48/50 |
| Autocomp | 1.36x | 45/50 |
Adding docs lifts per-sample correctness from 5.8% to 37.3% and solves 48 of 50. The cleanest contrast: upgrading the model from Flash to Pro (5-kernel subset) takes iterative geomean from 1.18x to 2.43x, while adding docs on Flash lifts correctness by 31 points. In the authors' words, a bigger model helps, but not as much as useful docs. The conclusion: for a sparse DSL the bottleneck is information, not reasoning.
On the 8 kernels with hand-tuned references, Autocomp reaches a 1.60x geomean, about 77% of the 2.08x Tokamax upper bound, and even beats hand-tuned on Sparse Attention (2.81x vs 0.86x) and Megablox GMM (2.21x vs 1.62x). It fails outright on paged and ragged paged attention, where hand-tuned scheduling matters most (ragged reaches 6.91x by hand, nothing correct from any agent).
For anyone in compilers or autonomous kernel optimization the practical signal is hard: for a DSL that is rare in training data, stuffing targeted docs into the prompt beats upgrading the model, and probably generalizes to other cold hardware DSLs. JAXBench is a usable hill, released with the benchmark, eval harness, and baselines. A subtler point: solving more tasks does not mean being faster. Iterative plus docs solves the most (48/50) yet posts a lower geomean (1.28x) than Autocomp (1.36x), because spending the budget on debugging versus optimization yields different returns.
The authors concede: everything runs on a single TPU v6e chip, with no multi-chip sharding or collectives; hand-tuned references cover only 8 of 17 priority operators, so the upper bound is partial; Pro is run only on a 5-kernel subset (cost); and on paged and ragged attention agents solve nothing, which the authors call a challenging task. Other gaps: