JAXBench: Benchmarking Autonomous TPU Kernel Optimization

# JAXBench: Benchmarking Autonomous TPU Kernel Optimization
Source: [https://arxiv.org/html/2607.20466](https://arxiv.org/html/2607.20466)
\\paperurl\\uselogo\\correspondingauthor

aryatschand@g\.harvard\.edu, sethusankaran@google\.com

Charles HongEqual contributions, work done while at GoogleUC BerkeleyGoogleJulian WalkerGoogle DeepMindNina CaiGoogleShangkun WangGoogleSuvinay SubramanianGoogleSundar DevGoogleVijay Janapa ReddiHarvard UniversityAmir YazdanbakhshGoogle DeepMindSethu SankaranGoogle

###### Abstract

Code:[https://github\.com/AI\-Hypercomputer/accelerator\-agents/tree/main/JAXBench](https://github.com/AI-Hypercomputer/accelerator-agents/tree/main/JAXBench) 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 presentJAXBench, a TPU\-native benchmark suite for AI\-generated kernel optimization on Google Cloud TPUs\.JAXBenchcomprises 50 JAX workloads that are bothrelevantand provideheadroomfor 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\. 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 forJAXBench\. 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 a1\.28×1\.28\\timesgeomean speedup\. Search structure yields significant gains once correctness is achieved, with Autocomp’s beam\-search pipeline reaching a1\.36×1\.36\\timesgeomean speedup over XLA\. On the 8 hand\-tuned kernels, Autocomp reaches1\.60×1\.60\\timesgeomean over XLA, recovering most of the2\.08×2\.08\\timesTokamax upper bound but trailing on the specialized paged and ragged attention operators\. High\-quality TPU kernel optimization remains a challenging task, and we release theJAXBenchbenchmark, evaluation harness, and baseline results to support open source contributions\.

## 1Introduction

Efficient kernel implementations are a bottleneck in achieving the full potential of hardware accelerators for machine learning\. Libraries such as cuBLAS, CUTLASS, and Triton\[cutlass,tillet2019triton\]and specialized kernels for GPU\[dao2022flashattention\]and TPU\[jiang2026ragged\]have unlocked generation\-over\-generation performance gains, but each new model architecture, quantization scheme, and hardware revision demands fresh low\-level implementations\. To realize the performance gains enabled by architecture, kernels must be rapidly co\-designed between the hardware\-specific features and emerging workload characteristics\.

![Refer to caption](https://arxiv.org/html/2607.20466v1/Figures/methodology.png)Figure 1:JAXBench overview: a TPU\-native benchmark for evaluating AI\-generated Pallas kernels against hand\-tuned baselines on TPU v6e\.We construct50 JAX reference workloads \(blue\)with 17 priority kernels from MaxText\[maxtext2022\]LLMs and 33 fused operators adapted from KernelBench L2, all sized to saturate TPU v6e MXUs\. We also extract8 priority kernels of hand\-tuned Pallas implementations \(green\)from Tokamax as expert upper bounds\. Theagent evaluation harness \(purple\)enables evaluation of any LLM\-driven method to generate candidate Pallas kernels, which are compiled, correctness\-checked on bf16 inputs, and profiled viajax\.profiler, with feedback fed back and final kernels compared against both baselines\.This gap has motivated a fast\-growing body of work on using language models to generate and optimize kernels automatically, ranging from one\-shot completions to iterative agents with compile and profile feedback to evolutionary search over program space\[liao2025kernelevolve,tschand2025swizzleperf,lange2025towards,cao2026k\]\. Rigorous benchmarks have been central to this progress\. KernelBench\[ouyang2025kernelbench\]standardized the evaluation protocol with 250 PyTorch workloads and has since been extended to Triton, CuTe, and other DSLs\. TritonBench\[li2025tritonbench\]specifically targets Triton generation with hardware\-aware performance measurement on NVIDIA and AMD GPUs\. FlashInfer\-Bench\[xing2026flashinfer\]grounds evaluation in real LLM serving traces with expert\-written FlashInfer kernels as references\. Together, these benchmarks have driven rapid community progress on correctness and speedup against GPU baselines and have enabled meaningful comparison of methods from best\-of\-NNsampling to reinforcement\-learned coding agents\.

However,*no analogous benchmark exists for TPU kernel optimization\.*Google’s Tensor Processing Units\[jouppi2017datacenter\]differ fundamentally from GPUs\. TPUs are sequential machines with wide SIMD vector registers \(8×\\times128 for 32\-bit values on v6e\) and dedicated 256×\\times256 systolic Matrix Multiply Units \(MXUs\), rather than the massively parallel SIMT execution model of GPUs\. Programming TPUs also requires a distinct software stack\. JAX\[bradbury2021jax\]programs compile through XLA, and low\-level kernel authoring goes through Pallas, which lowers to the Mosaic backend rather than Triton\. Pallas kernels must reason about TPU\-specific concerns including VMEM/SMEM/HBM memory hierarchies, software pipelining with prefetch scheduling, block shape constraints, and the lexicographic grid traversal order that Mosaic enforces\[jax\_pallas\_tpu\]\. Pallas also appears in LLM training data orders of magnitude less frequently than CUDA or even Triton, so models that fluently write GPU kernels routinely hallucinate Pallas APIs, emit memory\-space annotations that do not type\-check, or violate systolic tiling constraints in ways that no amount of generic compile feedback can resolve\. GPU benchmarks therefore cannot be directly applied because workloads, problem sizes, programming abstractions, and optimization strategies are all GPU\-specific\.

Recent concurrent work has addressed multi\-platform evaluation\. MultiKernelBench\[wen2025multikernelbench\]extends KernelBench to CUDA, Huawei AscendC, and Google Pallas, finding that the best model achieves only 8\.4–10\.5% Pass@1 on Pallas tasks\. However, MultiKernelBench does not target JAX, evaluates on TPU v2\-8, translates PyTorch workloads at problem sizes too small to saturate the MXU, and does not include production LLM operators or optimized reference kernels\. At small shapes, workloads are dominated by memory traffic and launch overhead rather than compute, so any measured speedup reflects bookkeeping rather than genuine algorithmic or scheduling improvement\. A useful TPU benchmark must instead push workloads into the compute\-bound regime where the MXU is binding and tiling, pipelining, and layout choices actually affect throughput\. Table[1](https://arxiv.org/html/2607.20466#S1.T1)summarizes how existing benchmarks position on these axes\. There remains a need for a TPU\-native benchmark that \(i\) uses contemporary TPU hardware, \(ii\) includes workloads representative of modern LLM training and inference, \(iii\) provides strong baselines through both XLA compilation and hand\-optimized Pallas kernels, and \(iv\) sizes problems to be compute\-bound so optimization headroom is meaningful\.

We introduceJAXBench, a benchmark suite of 50 JAX workloads designed specifically for TPU kernel optimization\. Our contributions are as follows\.

1. 1\.TPU\-native benchmark suite with 50 relevant workloads\.JAXBenchis a benchmark of 17 production operators from real LLM architectures \(Llama\-3\.1, DeepSeek\-V3, Mixtral, Mamba\-2, AlphaFold2\) extracted from MaxText, and 33 fused operator sequences from KernelBench translated from PyTorch to JAX and set with problem sizes adjusted to achieve high TPU MXU utilization\. For 8 priority kernels, we additionally provide expert\-optimized Pallas TPU kernels from the upstream Tokamax library with tuned block sizes\.
2. 2\.Rigorous and reproducible evaluation harness\.Device\-side profiling viajax\.profilerPerfetto traces eliminates host dispatch overhead, yielding reproducible per\-iteration kernel timings that can be used to easily evaluate any agent framework onJAXBench\.
3. 3\.Agent evaluation and failure modes\.We deployJAXBenchon TPU v6e \(Trillium\) and evaluate best\-of\-NN, iterative feedback agent loops, TPU context\-conditioned agent loops, and Autocomp\[hong2025autocomp\]augmented with TPU architecture documentation\. We also provide insights into failure modes and opportunities for further contributions in TPU\-centric kernel optimization\.

Table 1:Comparison of LLM kernel\-generation benchmarks\.JAXBenchis the only suite targeting contemporary TPUs with production\-scale workloads, hardware\-saturating problem sizes, and expert\-written Pallas reference kernels for a subset\.BenchmarkTarget\# TasksWorkload SourceExpert RefsSat\. SizesKernelBench\[ouyang2025kernelbench\]CUDA / GPU250PyTorch modules—YesTritonBench\[li2025tritonbench\]Triton / GPU184Real reposTritonPartialFlashInfer\-Bench\[xing2026flashinfer\]CUDA / GPU–LLM serving tracesFlashInferYesMultiKernelBench\[wen2025multikernelbench\]Pallas / TPU v2\-8285PyTorch modules—NoJAXBench\(ours\)Pallas / TPU v6e50Production LLMs \+Standard JAXPallasYes
## 2Methodology

### 2\.1Design Principles

JAXBenchis built around three principles that together make the benchmark both practically relevant and technically rigorous for TPU kernel optimization\.

#### Relevant workloads\.

The benchmark must evaluate performance on a diverse set of operators representative of real\-world TPU customers\. We include both well\-studied operators where expert\-optimized Pallas kernels already exist \(e\.g\., flash attention, grouped\-query attention\) and emerging operators where no such Pallas kernels have been developed \(e\.g\., Mamba\-2 state space duality, AlphaFold2 triangle multiplication\)\. This combination tests automated optimization methods across varying levels of prior human effort\.

#### Headroom to optimize\.

Workloads and problem sizes are chosen so that the XLA\-compiled baseline already achieves high TPU utilization, leaving kernel optimization to compete on genuine algorithmic and scheduling improvements rather than on closing artificial gaps\. This choice reflects how TPUs are actually used\. Customers running production workloads tune problem dimensions to maximally utilize the hardware, and FLOPs utilization at the workload level is tightly correlated with FLOPs utilization of the underlying operations\. A benchmark that evaluates kernels at undersized, memory\-bound, or launch\-overhead\-bound shapes would not reflect the regime in which TPU users \(and therefore TPU kernel optimizers\) actually operate\. We therefore size each workload to push the relevant operations toward the compute\-bound regime, where the MXU is the binding resource and optimization headroom is meaningful\. As shown in Figure[2\(a\)](https://arxiv.org/html/2607.20466#S2.F2.sf1), matmul\-heavy fused operators use dimensions such as\(4096,8192\)×\(8192,8192\)\(4096,8192\)\\times\(8192,8192\)in bf16, achieving 60–95% MXU utilization, and priority kernels use production\-scale dimensions from their source architectures\. The Llama\-3\.1\-70B GEMM workload, for example, runs at8192×8192×286728192\\times 8192\\times 28672and achieves∼\\sim79% MXU utilization\. Problem size is itself an axis of optimization, since different shapes expose different opportunities and sufficiently small problems offer no headroom at all\.

#### Reproducible baselines\.

Every workload is implemented in idiomatic JAX, compiled via XLA’sjax\.jit\. This is the strongest reproducible baseline for TPU because XLA already performs aggressive whole\-program optimization including operator fusion, buffer assignment, and layout optimization\. For priority kernels, when available, we additionally provide Pallas\-optimized variants that represent the current state of the art in hand\-tuned TPU kernel performance, establishing an upper bound that automated methods should aspire to match\.

### 2\.2Workload Construction

#### Priority kernels \(17 workloads\)\.

We extract 17 compute\-critical operators from production LLM architectures as implemented in MaxText\[maxtext2022\]\. The set spans attention variants \(flash\[dao2022flashattention\], GQA\[grattafiori2024llama\], MLA\[liu2024deepseek\], sparse/splash\[jiang2024mixtral\], flex, paged\[kwon2023efficient\], and ragged paged\), dense and sparse linear algebra \(GEMM, SwiGLU MLP\[shazeer2020glu\], sparse MoE\[shazeer2017outrageously\], Megablox GMM, ragged dot\), normalization and loss \(RMSNorm\[zhang2019root\], cross\-entropy\), and emerging architectures \(RetNet retention\[sun2023retentive\], Mamba\-2 SSD\[gu2023mamba\], AlphaFold2 triangle multiplication\[jumper2021highly\]\)\. Each workload uses production\-scale dimensions from its source model\. The GQA attention workload, for example, uses Llama\-3\.1\-405B dimensions with 128 query heads and 8 key\-value heads at sequence length 4096\.

Whenever available, we source optimized implementations by importing equivalent Pallas kernels from the upstream JAX Pallas ops library \(Tokamax\)\[tokamax2024\], then tune their block size parameters via exhaustive grid search on TPU v6e\. Eight such high quality kernels were available in the library\. The tuning procedure evaluated 203 configurations total, yielding speedups of up to 2\.79×\\times\(Megablox GMM\) over the Pallas default parameters\.

#### KernelBench fused operators \(33 workloads\)\.

We translate a curated subset of 33 workloads from KernelBench Level 2\[ouyang2025kernelbench\]from PyTorch to equivalent JAX\. We focus on Level 2 because Level 1 consists of isolated single operators with no fusion structure, whereas Level 2 combines a matrix multiplication or convolution with elementwise operations \(activations, normalization, pooling\), directly representing the kernel fusion opportunities that are most relevant for TPU optimization\. To select a non\-redundant subset, we group all Level 2 tasks by their fusion signature, namely the \(matmul or convolution\)×\\times\(activation, normalization, pooling\) operator class they instantiate, and retain one representative per class\. This procedure deduplicates structurally identical tasks and yields the 33 workloads we include\.

Translation is performed by prompting Gemini with each reference PyTorch module and asking for an idiomatic JAX equivalent\. To confirm numerical correctness, we execute the original PyTorch reference and the generated JAX implementation on TPU with matched bf16 inputs and require their outputs to agree underjnp\.allclosewithatol=rtol=10−2\\mathrm\{atol\}=\\mathrm\{rtol\}=10^\{\-2\}, the standard bf16 tolerance\. Any workload that fails this check is regenerated or repaired by hand before inclusion\.

Many of the original KernelBench problem sizes are too small to saturate TPU MXUs, since the 256×\\times256 systolic arrays require large matrix dimensions to achieve high utilization\. Rather than apply a uniform scaling factor, we tune each workload’s free dimensions independently by sweeping over candidate shapes and freezing the smallest configuration whose XLA baseline reaches at least 60% MXU utilization on TPU v6e\. This per\-operator procedure yields a heterogeneous set of shapes \(e\.g\., batch size 4096 with feature dimensions8192×81928192\\times 8192in bf16 for matmul\-fused operators\), all of which are compute\-bound rather than launch\-overhead\-bound on TPU\.

#### Workload interface\.

Every workload module exposes a standardized interface composed of aCONFIGdictionary of hyperparameters, acreate\_inputs\(dtype\)function that generates input tensors in bf16, and aworkload\(\*inputs\)function containing the computation to benchmark, which is annotated withjax\.named\_scopeby the profiler harness for visibility\. The benchmark runner JIT\-compiles each workload viajax\.jitand executes it under controlled timing \(Section[2\.3](https://arxiv.org/html/2607.20466#S2.SS3)\)\.

### 2\.3Timing and Profiling

![Refer to caption](https://arxiv.org/html/2607.20466v1/x1.png)\(a\)TPU v6e roofline\[williams2009roofline\]with all 50 workloads\.
![Refer to caption](https://arxiv.org/html/2607.20466v1/x2.png)\(b\)Tokamax Pallas kernel speedup over XLA\.

Figure 2:JAXBenchworkloads are predominantly compute\-bound on the TPU v6e roofline, and show headroom where existing Pallas kernel optimization can achieve up to a 16×\\timesspeedup over the XLA baseline\.Achieved performance vs\. operational intensity on a single TPU v6e chip \(918 TFLOPS bf16, 1640 GB/s HBM\) reveals that matmul\-fused and linear/MoE kernels already reach 60–95% MXU utilization under XLA, while attention and elementwise kernels remain memory\-bound\. Hand\-optimized Pallas kernels close this gap for the 8 priority workloads by tiling computation to reduce redundant HBM traffic, achieving speedups of up to 16\.3×\\timeson Ragged Paged Attention and 3\.5×\\timeson Flash Attention\. Compute\-bound workloads like GEMM show no improvement, confirming that XLA already saturates the MXU for high\-OI operations and that optimization effort is best directed at memory\-bound kernels\.Accurate kernel\-level timing on TPUs requires device\-side measurement\. Wall\-clock timing viatime\.perf\_counter\(\)withblock\_until\_ready\(\)includes host\-side overhead from Python dispatch, JAX runtime scheduling, and synchronization\. For short kernels \(<1<\\\!1ms\), this overhead can constitute 10–20% of measured time\.

We usejax\.profiler\.trace\(\)to capture Perfetto\-compatible traces and extract per\-iteration execution times\. For each workload, we \(1\) create inputs in bf16, \(2\) JIT\-compile, \(3\) run 5 warmup iterations, \(4\) execute 50 timed iterations under the profiler, and \(5\) parse the trace to extractjit\_\*\(\)\-level device events\. These events capture the total device\-side execution time per iteration, regardless of how XLA or Pallas decomposes the computation internally \(single fusion, multiple fusions, or a Pallas kernel call\)\. We report the median over 50 iterations as the primary metric\.

#### FLOP counting and utilization\.

We measure MXU utilization on TPU v6e and operational intensity of ragged paged attention\[jiang2026ragged\]as

MXU utilv6e=FLOPs/tmedian918×1012,OIragged\-attn=FLOPsv6eMemBWv6e=G​∑i=1BLiG​B\+∑i=1BLi\\text\{MXU util\}\_\{\\text\{v6e\}\}=\\frac\{\\text\{FLOPs\}/t\_\{\\text\{median\}\}\}\{918\\times 10^\{12\}\},\\qquad\\text\{OI\}\_\{\\text\{ragged\-attn\}\}=\\frac\{\\text\{FLOPs\}\_\{\\text\{v6e\}\}\}\{\\text\{MemBW\}\_\{\\text\{v6e\}\}\}=\\frac\{G\\sum\_\{i=1\}^\{B\}L\_\{i\}\}\{GB\+\\sum\_\{i=1\}^\{B\}L\_\{i\}\}\(1\)where 918 TFLOPS is the TPU v6e single\-chip bf16 peak,G=Hq/Hk​vG=H\_\{q\}/H\_\{kv\}is the GQA grouping factor,BBis the number of sequences, and\{Li\}\\\{L\_\{i\}\\\}are the per\-sequence KV lengths\. The numerator counts attention FLOPs, which scale asHq​∑iLiH\_\{q\}\\sum\_\{i\}L\_\{i\}since every query head attends over its full KV sequence\. The denominator counts bytes read, withHq​BH\_\{q\}Bfrom the per\-sequence query tensors andHk​v​∑iLiH\_\{kv\}\\sum\_\{i\}L\_\{i\}from the shared K and V cache \(with constants fordheadd\_\{\\text\{head\}\}and bytes\-per\-element absorbed into the proportionality\)\. Dividing numerator and denominator byHk​vH\_\{kv\}yields the form above\. In the long\-context limit∑iLi≫G​B\\sum\_\{i\}L\_\{i\}\\gg GBthe operational intensity therefore approachesGG, reflecting that each KV token is reused by exactlyGGquery heads and that this reuse caps achievable OI regardless of sequence length\. For Llama\-3\.1\-405B \(G=16G=16\), this upper bound of 16 FLOP/byte sits far below the v6e ridge point ofOI≈560\\mathrm\{OI\}\\approx 560, confirming that ragged paged attention is memory\-bandwidth\-bound at any context length\. For the workloads, FLOPs are extracted from XLA’scost\_analysis\(\)on the compiled HLO program and bytes accessed are computed analogously for each operator class\.

## 3Results

### 3\.1Evaluation Protocol

We evaluate generated kernels along three axes\. The first two are gating criteria, and the third is our reported performance metric\. We additionally report one derived budget metric \(fast1@NN\) to summarize how quickly each method produces useful kernels\.

#### 1\. Compilability\.

Does the generated code compile to a valid JAX/Pallas program? For Pallas kernels, this requires correct use ofBlockSpec,PrefetchScalarGridSpec, memory space annotations, and block shape constraints specific to TPU\.

#### 2\. Correctness\.

Does the generated kernel produce outputs that match the reference implementation? We check numerical agreement on randomized bf16 inputs underjnp\.allclosewithatol=rtol=10−2\\mathrm\{atol\}=\\mathrm\{rtol\}=10^\{\-2\}, the standard bf16 tolerance\. Kernels that fail either gate contribute1×1\\timesspeedup to the aggregate\.

#### 3\. Speedup over XLA\.

Our primary performance metric is wall\-clock speedup of the generated kernel over the XLA\-compiled JAX baseline, measured with device\-side profiler timing \(Section[2\.3](https://arxiv.org/html/2607.20466#S2.SS3)\)\. For the 8 priority kernels with hand\-optimized Pallas implementations, we report performance against those references separately as an upper\-bound comparison rather than as the denominator of the aggregate speedup\.

#### fast1@NN\.

As a sample\-efficiency metric, fast1@NNreports the fraction of benchmarks whose best\-so\-far generated kernel already beats the XLA baseline afterNNcumulative samples emitted by the method\. It complements geomean speedup by measuring breadth of coverage rather than magnitude of improvement at a given budget\.

#### Evaluated methods\.

We evaluate four approaches\.

\(1\)Best\-of\-NNgeneration\.Each sample is an independent one\-shot completion conditioned on a TPU v6e preamble and the JAX source, and the best correct sample is selected\. See Appendix[B](https://arxiv.org/html/2607.20466#A2)\.

\(2\)Iterative refinement\.Agent loops with access to compilation errors, correctness results, and profiler summaries as feedback between turns, following the paradigm of KernelBench\[ouyang2025kernelbench\]but adapted for JAX/Pallas\. We run 18 independent chains of 8 turns each\.

\(3\)Iterative refinement with Autocomp context\.Identical to iterative refinement except that every turn receives Autocomp’s agent context \(hardware architecture summary, per\-benchmark\-selected Pallas API reference, selected code examples, and rules block\) prepended to the preamble\. This ablation isolates the effect of curated documentation from the effect of structured search\.

\(4\)Autocomp\[hong2025autocomp\]\. An open\-source LLM\-driven kernel optimizer that ingests publicly available JAX Pallas and Cloud TPU documentation and distills it into four artifacts prepended to every agent prompt: a hardware\-architecture summary, a curated Pallas API reference, annotated code examples, and correctness rules\. We run a two\-phase pipeline: a*translation phase*\(4 beam\-search iterations\) rewrites the XLA baseline into a Pallas kernel, and an*optimization phase*\(4 iterations\) refines it for performance\. Both phases use beam size 3 with 6 new candidates per beam element per iteration\. We select Autocomp as our primary agent baseline because it is the most hardware\-aware open method available\. Other methods such as AlphaEvolve\[novikov2025alphaevolve\]and KernelEvolve\[liao2025kernelevolve\]are not designed with TPU\-specific context, and we leave their evaluation to future work\.

#### Setup\.

We evaluate all four methods on TPU v6e \(Trillium\), starting from the XLA baseline\. Best\-of\-NNand both iterative variants use a fixed budget of 144 samples per benchmark\. Autocomp uses a beam search of comparable total budget split across translation and optimization phases, and early\-stops once its beam stops improving, so its actual sample count per benchmark can be lower than 144\. For each method and benchmark we track best\-so\-far speedup over XLA at every sample, and aggregate across benchmarks by geometric mean\.111We set the speedup floor to1×1\\timesto prevent correct but slow kernels from counting as worse than incorrect kernels\.All methods use Gemini 3 Flash\[team2023gemini\]for the main evaluation on all 50 benchmarks \(Section[3\.2](https://arxiv.org/html/2607.20466#S3.SS2)\)\. We also show an ablation on a 5\-kernel subset with Gemini 3\.1 Pro \(Section[3\.4](https://arxiv.org/html/2607.20466#S3.SS4)\)\.

### 3\.2Baseline Comparison \(Gemini 3 Flash\)

We evaluate all four methods across the 50\-workload JAXBench suite with Gemini 3 Flash \(Table[2](https://arxiv.org/html/2607.20466#S3.T2)\)\.

Table 2:Cross\-method comparison on the full 50\-benchmark JAXBench suite with Gemini 3 Flash\. Per\-benchmark best speedups are floored at1×1\\timesbefore aggregating; missing \(incorrect\) benchmarks count as1×1\\times\. Per\-sample correctness is reported separately in Figure[4](https://arxiv.org/html/2607.20466#S3.F4)\.MethodGeomean speedupMean speedupPer\-benchmark correctnessBest\-of\-NN1\.01×\\times1\.02×\\times13/50Iterative refinement1\.18×\\times1\.50×\\times32/50Iterative \+ context1\.28×\\times1\.55×\\times48/50Autocomp1\.36×\\times1\.70×\\times45/50![Refer to caption](https://arxiv.org/html/2607.20466v1/x3.png)Figure 3:Speedup and solved\-rate trajectories on the full JAXBench suite \(Gemini 3 Flash\) as a function of cumulative samples\.Left:geomean speedup over XLA \(per\-kernel best speedups floored at1×1\\timesbefore the geomean, so curves are monotonically non\-decreasing\)\.Right:fast1@NN, the fraction of benchmarks whose best\-so\-far sample already beats XLA at budgetNN\.#### Speedup trajectory\.

Figure[3](https://arxiv.org/html/2607.20466#S3.F3)shows best\-so\-far geomean speedup \(left\) and fast1@NN\(right\) as samples accumulate\. The method ranking is stable almost from the first sample\. Autocomp pulls ahead within the first dozen samples and holds the lead through the full budget, finishing at1\.36×1\.36\\timesgeomean with76%76\\%of benchmarks beating XLA\. Iterative\+context is the runner\-up at1\.28×1\.28\\timesand32%32\\%fast1@NN, a clear improvement over plain iterative \(1\.18×1\.18\\times,18%18\\%\)\. Best\-of\-NNnever climbs meaningfully, producing correct kernels for 13 of 50 benchmarks but only beating XLA on one\.

Iterative\+context lands*more*correct kernels than Autocomp \(48/50 vs\. 45/50\) yet reaches a lower geomean speedup\. The two methods spend their 144\-sample budget differently: iterative\+context debugs until nearly every chain is correct, so it eventually solves 3 of the 5 benchmarks on which Autocomp generates no correct Pallas\. Autocomp commits 4 translation cycles up front and then spends its remaining budget on optimization rather than debugging, which enables higher speedups\. Per\-benchmark best speedups are reported in Figure[6](https://arxiv.org/html/2607.20466#A1.F6)\(Appendix[A](https://arxiv.org/html/2607.20466#A1)\)\.

![Refer to caption](https://arxiv.org/html/2607.20466v1/x4.png)Figure 4:Per\-sample outcome breakdown onJAXBench\(Gemini 3 Flash\), pooling all samples across the 50 benchmarks\. “Wrong output” denotes samples that compile and run but fall outside numerical tolerance\. “API/runtime” covers unsupported Pallas API syntax and compile/runtime exceptions\.
#### Failure analysis\.

To understand*where*each method spends its sample budget, we classify every sample by outcome \(Figure[4](https://arxiv.org/html/2607.20466#S3.F4)\)\. Adding TPU\-specific documentation sharply reduces API misuse\. Iterative refinement with Autocomp’s context raises per\-sample correctness from5\.8%5\.8\\%to37\.3%37\.3\\%with no change to the search algorithm\. Pallas API misuse dominates context\-free failures, with99\.7%99\.7\\%of best\-of\-NNand93\.8%93\.8\\%of iterative samples failing at compile or first execution by hallucinating Pallas APIs or misusingpallas\_callarguments andBlockSpectypes\. API misuse remains a challenge even with context, comprising59\.8%59\.8\\%of iterative\+ctx and55\.8%55\.8\\%of Autocomp samples\.

### 3\.3Comparison Against Hand\-Tuned Pallas

Eight of the 17 priority operators ship with hand\-optimized Pallas implementations from the upstream JAX Pallas ops library \(Tokamax\), with block sizes tuned via grid search on TPU v6e\. We measured these kernels through the same JAXBench harness and report per\-kernel speedups in Table[3](https://arxiv.org/html/2607.20466#S3.T3)\. Tokamax reaches a2\.08×2\.08\\timesfloor\-1 geomean over the 8 kernels \(the one sub\-1×1\\timesentry is a Splash Attention variant tuned for a different sparsity pattern than our dense\-mask workload\)\. Autocomp’s1\.60×1\.60\\timesgeomean on the same subset reaches roughly77%77\\%of the hand\-tuned geomean, exceeding Tokamax on 2 kernels and landing within6868–91%91\\%on 4 more, but failing correctness on the paged\- and ragged\-attention operators where Tokamax’s hand\-written scheduling matters most\. Iterative\+context with the same sample budget reaches1\.34×1\.34\\times\.

Table 3:Speedups on the 8 priority kernels with hand\-tuned Pallas references\. Agents use a 144\-sample Gemini 3 Flash budget\. “—” denotes no correct kernel\. Bold marks the best method per row\.BenchmarkXLA \(ms\)HT \(ms\)Hand\-tunedAutocompIter\+ctxIterFlash Attention25\.216\.453\.91×\\times2\.67×\\times0\.46×\\times—GQA Attention51\.1314\.593\.50×\\times2\.63×\\times2\.63×\\times—MLA Attention20\.3221\.720\.94×\\times0\.85×\\times1\.54×\\times—Sparse Attention5\.956\.890\.86×\\times2\.81×\\times2\.54×\\times—Paged Attention7\.823\.412\.29×\\times—0\.61×\\times—Ragged Paged Attention18\.702\.716\.91×\\times———GEMM5\.365\.590\.96×\\times0\.75×\\times0\.97×\\times0\.59×\\timesMegablox GMM3\.262\.011\.62×\\times2\.21×\\times0\.27×\\times0\.52×\\timesGeomean \(floor 1×\\times\)2\.08×\\times1\.60×\\times1\.34×\\times1\.00×\\times
### 3\.4Model Capacity Ablation \(Gemini 3\.1 Pro\)

To measure the effect of model capacity, we rerun all four methods on a 5\-kernel subset using Gemini 3\.1 Pro, limited in scope due to Pro’s substantially higher per\-sample cost\. The Flash curves in this ablation are remeasured on the same subset so that Flash and Pro are compared on identical benchmarks and sample budgets\. Pro lifts every method, with the largest gains on iterative refinement, iterative\+context, and Autocomp\. Iterative jumps from1\.07×1\.07\\timesto2\.43×2\.43\\times, iterative\+context from1\.59×1\.59\\timesto3\.82×3\.82\\times, and Autocomp from2\.35×2\.35\\timesto3\.79×3\.79\\times\(Figure[5](https://arxiv.org/html/2607.20466#S3.F5)\)\. Pro iterative\+context and Autocomp land correct Pallas on all 5 benchmarks, whereas Flash Autocomp still fails on Mamba\-2 SSD\.

![Refer to caption](https://arxiv.org/html/2607.20466v1/x5.png)Figure 5:Model\-capacity ablation on the 5\-kernel subset \(RMSNorm, Flex Attention, RetNet Retention, Mamba\-2 SSD, Flash Attention\): speedup over XLA \(left\) and fast1@NN\(right\) as a function of cumulative samples\. Solid lines are Gemini 3\.1 Pro; dashed lines are Gemini 3 Flash\.Plain iterative refinement benefits disproportionately from model scaling \(1\.07×→2\.43×1\.07\\times\\to 2\.43\\times\) purely from better Pallas priors, and once both context and capacity are engaged, iterative\+context and Autocomp converge to essentially the same geomean \(3\.82×3\.82\\timesvs\.3\.79×3\.79\\times\)\. This is unlike the full Flash suite where Autocomp leads by a wide margin, since once Pro produces correct seeds without extensive debugging, iterative\+context’s samples go toward optimization rather than correctness recovery\. On the two kernels where Autocomp Pro trailed iterative\+context \(Mamba\-2 SSD and Flash Attention\), Autocomp’s best latency was still improving monotonically through the final iteration, so the per\-kernel ordering reflects exploration depth at a fixed budget rather than beam\-search convergence\. We treat this narrowing as a preliminary observation given the 5\-kernel scope\.

## 4Discussion and Future Work

Compiler, runtime, and profiler feedback alone are insufficient to drive performance\. Iterative refinement without TPU\-specific documentation achieves only16\.9%16\.9\\%per\-sample correctness on the 5\-kernel subset with Gemini 3\.1 Pro, and1\.2%1\.2\\%with Flash\. The Pallas constraints that govern correctness, including lexicographic grid traversal, VMEM/SMEM/HBM placement,\(8,128\)\(8,128\)block divisibility, and prefetch scheduling, are absent from error messages, leaving feedback loops nothing to iterate against\.

Using a larger model helps, but not as much as adding useful context\. Autocomp\-generated documentation injected into the iterative loop raises per\-sample correctness from5\.8%5\.8\\%to37\.3%37\.3\\%across the full Flash suite \(Section[3\.2](https://arxiv.org/html/2607.20466#S3.SS2)\), a much larger jump than the Flash\-to\-Pro gain at fixed context\. For Pallas and other languages underrepresented in training data, the correctness bottleneck is information rather than reasoning, and written documentation supplies most of what is missing\.

Once correctness is achieved, performance depends on how the sample budget is spent\. Iterative refinement with only documentation solves*more*benchmarks than Autocomp on the full Flash suite \(48/50 vs\. 45/50\) but rarely pushes past the first correct kernel, finishing at1\.28×1\.28\\timesgeomean\. Autocomp’s translate/optimize split shifts the beam into optimization as soon as a correct seed exists, reaching1\.36×1\.36\\times\. Search structure, not just documentation, converts correctness into performance, though the Pro ablation shows this gap narrows once the model is strong enough to produce correct seeds without extensive debugging \(Section[3\.4](https://arxiv.org/html/2607.20466#S3.SS4)\)\.

#### Limitations\.

All 50 workloads run on a single TPU v6e chip, and multi\-chip sharding and collective communication are out of scope\. Hand\-tuned Pallas references are available for 8 of the 17 priority operators and serve as an aspirational upper bound\. The Gemini 3\.1 Pro evaluation is restricted to a 5\-kernel subset and is reported as a controlled model\-capacity ablation rather than a primary result\.

#### Future work\.

The most consequential extension is supporting multi\-TPU kernel evaluation and optimization\. Production training and inference workloads almost universally run across many chips, where end\-to\-end performance depends on sharding strategy, collective scheduling, and the overlap of communication with compute\. Future work will extendJAXBenchalong the parallelism axis, covering tensor\-parallel, pipeline\-parallel, and expert\-parallel sharding regimes on multi\-chip TPU pods\. A benchmark with these characteristics would evaluate whether the documentation\-versus\-search separation we observe generalizes to collective primitives \(shard\_map,psum\_scatter,all\_to\_all, asynchronous collectives, and overlapping pipelines\), which are even more sparsely represented in pretraining data than Pallas itself\. More broadly, scaling AI\-for\-systems methods from single\-device benchmarks to real\-world deployment surfaces recurring challenges around topology\-aware scheduling, communication overhead, and reproducibility of measured speedups\[tschand2026genaisystemsrecurringchallenges\]\.

Beyond multi\-chip, the information\-gap result motivates fine\-tuning or reinforcement learning from execution feedback on curated Pallas corpora, for whichJAXBenchcorrectness and latency signals are usable as rewards\. The tier structure in our breakdown also suggests hybrid controllers that route between iterative refinement and beam\-search agents on a per\-benchmark basis\. We release the benchmark, baselines, and harness to support these directions\.

## 5Conclusion

Hand\-optimized accelerator kernels remain a persistent bottleneck for ML systems, and the pool of engineers fluent in both ML and low\-level systems programming cannot keep pace with new architectures and hardware revisions\. Recent benchmarks show that LLMs can begin to fill this gap on CUDA and Triton, but the same machinery fails on TPUs because Pallas is underrepresented in training data\.JAXBenchhelps close the measurement gap with 50 production\-scale JAX workloads, 8 hand\-tuned Pallas references with grid\-searched block sizes, and a profiler\-grounded evaluation harness on TPU v6e\.

Across the suite, automated methods make real progress but still trail expert kernels\. Methods conditioned on documentation, like iterative refinement and Autocomp’s beam search, beat the baseline XLA performance on many workloadse\. The dominant lever on niche backends is target\-specific context rather than model scale or feedback alone, and search structure is what converts correctness into performance once that information gap is closed\. We release the suite, baselines, and harness to support evaluation for current methods and set a foundation for the multi\-chip and training\-time work to come\.

## Acknowledgements

We thank Alain Hamel, Newfel Harrat, Steven Ingram, Mehrdad Khani, Gerson Kroiz, and Tomas Pfister for their helpful feedback and support throughout this work\.

## References

## Appendix APer\-Benchmark Results

Figure[6](https://arxiv.org/html/2607.20466#A1.F6)shows the best speedup achieved by each method on every benchmark in the 50\-workload JAXBench suite with Gemini 3 Flash\. Benchmarks are sorted by the best speedup achieved by any method, and missing markers indicate the method produced no correct implementation within its sample budget\.

![Refer to caption](https://arxiv.org/html/2607.20466v1/x6.png)Figure 6:Best speedup over the XLA baseline for each benchmark and method\. Benchmarks without a marker had no correct implementation within the sample budget\.Table[4](https://arxiv.org/html/2607.20466#A1.T4)reports the corresponding per\-benchmark best speedups for the Gemini 3\.1 Pro ablation on the 5\-kernel subset \(Section[3\.4](https://arxiv.org/html/2607.20466#S3.SS4)\)\.

Table 4:Per\-benchmark best speedup over XLA on the 5\-kernel subset with Gemini 3\.1 Pro\. “—” denotes no correct kernel within the 144\-sample budget\. Bold marks the best method per row\. Geomean floors per\-kernel speedups at1×1\\timesbefore aggregating\.BenchmarkXLA \(ms\)Best\-of\-NNIterIter\+ctxAutocompRMSNorm1\.241\.42×\\times1\.42×\\times1\.43×\\times1\.44×\\timesFlex Attention36\.81—3\.99×\\times3\.46×\\times4\.01×\\timesRetNet Retention13\.04—4\.85×\\times6\.40×\\times9\.62×\\timesMamba\-2 SSD29\.59——5\.66×\\times4\.04×\\timesFlash Attention25\.38—3\.07×\\times4\.54×\\times3\.51×\\timesGeomean \(floor 1×\\times\)1\.07×\\times2\.43×\\times3\.82×\\times3\.79×\\times
## Appendix BBaseline Prompts and Agent Context

This appendix documents the exact prompt construction for each evaluated method so the setup can be reproduced from the paper text alone\.

### B\.1Best\-of\-NNand Iterative Refinement

Both baselines use the same short system preamble:

> You are optimizing a JAX kernel for TPU v6e \(Trillium\) using the Pallas programming model \(jax\.experimental\.pallas\)\. Target hardware: TPU v6e\. VMEM = 128 MiB, 8 TensorCores per chip, peak≈\\approx918 TFLOPS bf16, HBM≈\\approx1640 GB/s\. Rules for every turn: \(1\) keep the public workload\(\*inputs\) signature unchanged; \(2\) output a single complete Python file inside one‘‘‘python \.\.\. ‘‘‘block; \(3\) do not include explanatory prose outside the code block\. Strategy: if you do not yet have a correct Pallas implementation, your first priority is to produce a correct, straightforward translation even if it isn’t faster than the XLA baseline\. Only after you have a correct Pallas kernel should you focus on optimizing it for speed\.

#### Best\-of\-NN\.

Each of theNNsamples is an independent completion of a prompt consisting of the preamble above, a single lineBenchmark: <prob\_id\>, a one\-sentence task description \(“Below is the XLA/JAX reference implementation\. Rewrite it as a Pallas kernel that produces identical outputs and runs faster\.”\), and the JAX source wrapped in a Python code block\. No feedback is provided\.

#### Iterative refinement\.

The first turn of each chain uses the same prompt as Best\-of\-NN\. Each subsequent turn within a chain appends the following feedback block, then asks the model for a new complete Pallas implementation:

- •If the previous attempt was correct: “Previous attempt was correct\. Latency:<<x\>\>ms \(XLA baseline:<<y\>\>ms, speedup:<<s\>×\>\\times\)\. Best\-so\-far in this chain:<<z\>\>ms\.”
- •If the previous attempt was incorrect: “Previous attempt was incorrect or failed to run\.” followed by the last 40 lines of the runner’s stdout\.
- •A profiler summary \(top operators by time\) from the previous run when available\.
- •The previous implementation’s full source, wrapped in a Python code block\.

Chains run independently and never share context, and 18 chains of 8 turns give a total of 144 samples per benchmark\.

#### Iterative refinement with Autocomp context\.

The*iterative\+context*variant uses the same chain/turn/feedback layout as plain iterative refinement \(18 chains×\\times8 turns, same feedback block on correctness or failure, same final\-code echo\) with one change\. Before every prompt we prepend the static portion of Autocomp’s agent context, namely the hardware architecture summary, the per\-benchmark\-selected Pallas API reference, the selected code examples, and the rules block \(see next subsection\)\. The per\-benchmark category and example selections are made by the same LLM yes/no filter Autocomp uses and are cached per benchmark, so a chain’s eight turns see the same context menu\. The variant does not use Autocomp’s beam search or translate/optimize phase split, and is purely a context injection into the iterative baseline\.

### B\.2Autocomp Context Artifacts

We passed four JAX and Cloud TPU documentation sources into Autocomp’s Agent Builder pipeline:

- •
- •
- •
- •

Each source is crawled withmax\_depth=2andmax\_pages=250\. An LLM\-driven routing\-and\-synthesis pass then distills the crawled content into four textual artifacts that together form the static portion of every agent prompt:

- •Hardware architecture summary\(≈\\approx8 KB\): a single\-pass LLM summary of TPU v6e characteristics \(memory hierarchy, compute units, bandwidth and capacity numbers, and programming\-model constraints\)\.
- •Pallas API / ISA reference\(≈\\approx115 KB\): per\-API entries \(signature, description, parameter semantics, inline usage snippets\) extracted verbatim from the source documentation and grouped into functional categories\. At prompt time, an LLM yes/no filter selects which categories are relevant to the current kernel and caches the selection per\-benchmark\.
- •Annotated code examples\(≈\\approx30 KB\): representative Pallas kernels drawn from the ingested documentation, each accompanied by a one\-line summary\. Examples are selected per\-benchmark via a parallel LLM yes/no filter\.
- •Rules\(≈\\approx6 KB\): correctness constraints \(e\.g\., block\-shape divisibility, memory\-space annotations, aliasing rules\) extracted as imperative bullet points and appended to every planning and coding prompt\.

During beam search, each candidate\-generation prompt concatenates the architecture summary, the per\-benchmark\-selected API reference, any selected code examples, the candidate’s current code \(plus its latest score and any hardware\-counter feedback\), and the rules block, before asking the model for either a plan or a complete optimized kernel\. The translate and optimize phases use the same scaffold but draw from separate menus of strategies\.