JAXBench: Bridging the TPU Performance Gap with AI-Generated Kernels
JAXBench: Benchmarking Autonomous TPU Kernel Optimization
JAXBench is a TPU-native benchmark suite comprising 50 JAX workloads (17 production ML operators like Llama-3.1 and Mamba-2, and 33 fused operators) designed to evaluate AI-generated kernel optimization using the Pallas DSL. It establishes a rigorous evaluation harness on Google Cloud TPU v6e, leveraging the Autocomp framework to achieve a 1.36x geomean speedup over XLA baselines and recovering 77% of hand-tuned performance.
Executive Summary
TL;DR: Researchers from Google, Harvard, and UC Berkeley have released JAXBench, the first TPU-native benchmark for AI-generated kernel optimization. By providing 50 production-scale workloads and a rigorous evaluation harness, they demonstrate that while LLMs typically struggle with TPU-specific coding (Pallas), providing curated hardware context and structured search allows AI agents to outperform XLA compilers by a geomean of 1.36x.
This work marks a shift from GPU-centric benchmarking to a more heterogeneous AI infrastructure landscape, proving that "intelligence" in code generation is often a function of the specific documentation provided to the model.
Problem & Motivation: The Pallas Data Desert
While software like Triton has democratized GPU kernel writing, Google's TPU (Tensor Processing Unit) remains a "specialist" domain. TPUs rely on a sequential SIMD execution model with systolic Matrix Multiply Units (MXUs), fundamentally different from the SIMT (Single Instruction, Multiple Threads) model of GPUs.
The primary hurdle isn't just hardware complexity; it's data scarcity. The Pallas DSL (JAX’s low-level kernel language) appears in training sets orders of magnitude less than CUDA or Triton. Consequently, even the most capable LLMs "hallucinate" Pallas APIs, failing to account for VMEM/SMEM memory boundaries or systolic tiling constraints. Previous benchmarks often used workloads that were too small, measuring launch overhead rather than true compute efficiency.
Methodology: Pushing the MXU to the Limit
JAXBench ensures every workload is "compute-bound," meaning the hardware is actually working hard enough that optimization matters.
- Workload Selection: 17 production operators (FlashAttention, Mamba-2 SSD, DeepSeek-V3 MLA) and 33 fused operators from KernelBench translated to JAX.
- Saturation-Sizing: Instead of small test cases, dimensions are scaled (e.g., 8192x8192 matmuls) to reach 60-95% MXU utilization.
- The "Autocomp" Pipeline: The authors utilized a specific agentic flow:
- Translation Phase: Converting idiomatic JAX/XLA code into a functional Pallas kernel.
- Optimization Phase: Refining the kernel (tiling, prefetching) for maximum throughput.
Figure 1: The JAXBench evaluation harness on TPU v6e.
The Power of Context over Scale
The most striking insight from the paper is the Information-Gap. Naive models fail. However, when the authors injected ~150KB of curated TPU documentation and Pallas API references into the prompt:
- Correctness jumped from 5.8% to 37.3%.
- Gemini 3 Flash (a smaller model) with context outperformed Gemini 3.1 Pro without it.
This suggests that for niche backends, the bottleneck is knowledge, not reasoning capability.
Experiments & Results
JAXBench evaluates kernels on TPU v6e (Trillium). The results show that AI can now "hill-climb" toward expert-level performance.
- SOTA Comparison: Autocomp achieved a 1.36x geomean speedup over the standard XLA baseline.
- Closing the Gap: On priority kernels, Autocomp reached 1.60x speedup, recovering roughly 77% of the performance found in hand-tuned "Tokamax" libraries.
- Roofline Analysis: The benchmark successfully pushes workloads into the compute-bound regime of the TPU v6e roofline, where tiling and layout choices significantly impact TFLOPS.
Figure 2: TPU v6e Roofline showing kernels moving from memory-bound to compute-bound zones.
Key Performance Table (Select Kernels)
| Benchmark | XLA (ms) | Hand-Tuned Speedup | Autocomp Speedup |
|---|---|---|---|
| Flash Attention | 25.21 | 3.91x | 2.67x |
| GQA Attention | 51.13 | 3.50x | 2.63x |
| Megablox GMM | 3.26 | 1.62x | 2.21x |
Critical Analysis & Conclusion
Takeaway: JAXBench proves that LLMs can optimize TPU kernels, provided they aren't flying blind. The "Search vs. Knowledge" trade-off is clear: search structure (beam search) converts correctness into speed, but curated context is what enables correctness in the first place.
Limitations: Currently, JAXBench is limited to single-chip workloads. The "final boss" of TPU optimization—multi-chip sharding and collective communication (e.g., all_to_all at scale)—remains future work.
For AI system engineers, JAXBench provides the first rigorous "scoreboard" for the TPU ecosystem, paving the way for fully autonomous hardware-software co-design.
For more details, the code and benchmark suite are available at: https://github.com/AI-Hypercomputer/accelerator-agents