JAXBench: Benchmarking Autonomous TPU Kernel Optimization

Abstract: Rigorous benchmarks have driven progress in autonomous GPU kernel performance optimization by establishing a shared target to hillclimb on, but no equivalent exists for TPUs. We present JAXBench, a TPU-native benchmark suite for AI-generated kernel optimization on Google Cloud TPUs.

摘要： 严格的基准测试通过建立共同的优化目标，推动了自动 GPU 内核性能优化的进展，但 TPU 领域尚不存在类似的基准。我们提出了 JAXBench，这是一个专为 Google Cloud TPU 上 AI 生成内核优化而设计的原生 TPU 基准测试套件。

JAXBench comprises 50 JAX workloads that are both relevant and provide headroom for optimization. We extract 17 production ML operators from architectures in the public MaxText library such as Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2, and AlphaFold2, and translate 33 operators from KernelBench that are validated for correctness and set with new problem sizes that achieve high TPU v6e MXU utilization.

JAXBench 包含 50 个既具有实际意义又具备优化空间的 JAX 工作负载。我们从 MaxText 公共库的架构（如 Llama-3.1、DeepSeek-V3、Mixtral、Mamba-2 和 AlphaFold2）中提取了 17 个生产级机器学习算子，并从 KernelBench 中转换了 33 个算子。这些算子经过了正确性验证，并设定了新的问题规模，以实现高 TPU v6e MXU 利用率。

Eight of the 17 production operators ship with hand-optimized Pallas kernels from the public Tokamax library and block-size tuned to establish an expert upper-bound baseline. We evaluate four feedback-driven methods on generating candidate Pallas kernels for JAXBench.

在 17 个生产级算子中，有 8 个配备了来自公共 Tokamax 库的手动优化 Pallas 内核，并通过块大小调整建立了专家级的性能上限基准。我们评估了四种基于反馈的方法，用于为 JAXBench 生成候选 Pallas 内核。

Across the full suite with Gemini 3 Flash, we find that target-specific context matters more than model scale on a sparsely-documented DSL like Pallas. Conditioning on curated TPU documentation raises per-sample correctness from 5.8% to 37.3% and solves 48 of 50 benchmarks at a 1.28x geomean speedup.

在使用 Gemini 3 Flash 进行全套测试后，我们发现，对于像 Pallas 这样文档稀缺的领域特定语言（DSL），针对特定目标的上下文比模型规模更为重要。通过引入精选的 TPU 文档作为条件，单样本正确率从 5.8% 提升至 37.3%，并成功解决了 50 个基准测试中的 48 个，实现了 1.28 倍的几何平均加速比。

Search structure yields significant gains once correctness is achieved, with Autocomp’s beam-search pipeline reaching a 1.36x geomean speedup over XLA. On the 8 hand-tuned kernels, Autocomp reaches 1.60x geomean over XLA, recovering most of the 2.08x Tokamax upper bound but trailing on the specialized paged and ragged attention operators.

一旦实现了正确性，搜索结构将带来显著的性能提升；Autocomp 的束搜索（beam-search）流水线相比 XLA 实现了 1.36 倍的几何平均加速比。在 8 个手动调优的内核上，Autocomp 相比 XLA 达到了 1.60 倍的几何平均加速比，恢复了 2.08 倍 Tokamax 上限的大部分性能，但在专门的分页（paged）和不规则（ragged）注意力算子上仍有差距。

High-quality TPU kernel optimization remains a challenging task, and we release the JAXBench benchmark, evaluation harness, and baseline results to support open source contributions.

高质量的 TPU 内核优化仍然是一项具有挑战性的任务，我们发布了 JAXBench 基准测试、评估工具和基准结果，以支持开源贡献。