JAXBench: Bridging the "Pallas Gap" in Autonomous TPU Kernel Optimization

TL;DR

While GPU kernel optimization has seen a surge in automated benchmarking, TPUs have remained a "black box" for AI agents. JAXBench introduces the first rigorous, TPU-native benchmark suite of 50 workloads (including Llama-3.1 and AlphaFold2 operators). The core finding? Providing LLMs with target-specific hardware context is more effective than simply increasing model size, enabling agents to outperform XLA-compiled code by 1.36× on average.

Problem: The GPU-Centric Bias of AI Systems

Most modern AI-for-Systems research gravitates toward NVIDIA's ecosystem. Tools like Triton and CUDA have massive footprints in training corpora, allowing LLMs to "hallucinate" functional kernels reasonably well.

TPUs, however, operate on a different physical intuition:

• Sequential vs. Parallel: Unlike the SIMT model of GPUs, TPUs are sequential machines with wide SIMD vectors and systolic Matrix Multiply Units (MXUs).
• The Pallas DSL: Pallas, the low-level language for TPUs, is rarely seen in public codebases compared to CUDA.
• Physical Constraints: Factors like VMEM/HBM hierarchies and lexicographic grid traversal make generic code generation fail at the compilation stage.

Methodology: High-Fidelity Benchmarking

JAXBench doesn't just evaluate "toy" problems. It targets the compute-bound regime where optimization actually matters for production.

1. Workload Composition

The suite consists of:

• 17 Priority Kernels: Extracted from MaxText (Google's production library), featuring FlashAttention, GQA, and Mamba-2.
• 33 Fused Operators: Translated from KernelBench, but resized to ensure 60-95% MXU utilization on TPU v6e (Trillium).

2. The Agent Framework (Autocomp)

The researchers evaluated several methods, but the standout was Autocomp, which utilizes a two-phase "Translate then Optimize" beam search. It injects a "Context Artifact" into the prompt:

• 8KB Hardware summary.
• 115KB Pallas API/ISA reference.
• Annotated code examples and correctness rules.

Figure 1: The JAXBench evaluation harness workflow, from reference JAX workload to profiled Pallas kernel.

Experiments & Results: Context is King

The study compared Best-of-N, Iterative Refinement, and Autocomp using Gemini 3 Flash and 3.1 Pro.

The Breakdown of Success

• Correctness Jump: Simply adding documentation (Iterative + Context) raised correctness from a dismal 5.8% to 37.3%.
• Performance vs. XLA: Autocomp achieved a 1.36× geomean speedup over XLA. On a subset of hand-tuned kernels from the Tokamax library, it recovered most of the performance gap, trailing only on highly specialized ragged/paged attention operators.

Figure 2: Performance trajectory showing how Autocomp's search strategy converts correctness into wall-clock speedups.

The "Model Scale" Illusion

In an ablation study using Gemini 3.1 Pro, the researchers found that while a larger model improves the "prior" knowledge of Pallas, it still cannot match a smaller model (Flash) that has been provided with the correct hardware context and a structured search algorithm.

Critical Analysis & Conclusion

Takeaway

JAXBench proves that for specialized hardware optimization, information is the bottleneck, not just reasoning. If the agent doesn't know the specific tiling constraints of a TPU v6e MXU, no amount of iterative feedback will help it "guess" the right memory-space annotation.

Limitations & Future Work

• Single-Chip Limit: The current benchmark focuses on single TPU chips. The next frontier is multi-chip collective communication (sharding and pipelining), which is even more sparsely documented.
• Expert Gap: Agents still struggle with complex control flows like PageAttention, where human-written scheduling is significantly more efficient.

Conclusion: JAXBench provides the necessary "hill-climbing" target for the next generation of TPU-resident AI agents. By open-sourcing the harness and baselines, the authors have set the stage for LLMs to move beyond GPU-centricity and into the heart of the AI Hypercomputer.