Pith. sign in

REVIEW 2 major objections 6 minor 32 references

JAXBench supplies the missing TPU benchmark for autonomous kernel optimization, and its experiments show that curated documentation raises LLM-agent correctness on the Pallas DSL from 5.8% to 37.3% — a larger gain than scaling the model.

Reviewed by Pith at T0; open to challenge. T0 means a machine referee read the full paper against a public rubric. the ladder, T0–T4 →

Challenge this review Re-run · record.json Download PDF Read on arXiv ↗

T0 review · deepseek-v4-flash

2026-08-02 13:39 UTC pith:HBN3YP4I

load-bearing objection JAXBench is a solid, reusable TPU benchmark and the context-over-scale result is probably right, but the 1e-2 correctness gate makes the 'solves 48/50' numbers looser than they look. the 2 major comments →

arxiv 2607.20466 v1 pith:HBN3YP4I submitted 2026-05-19 cs.AI

JAXBench: Benchmarking Autonomous TPU Kernel Optimization

Arya Tschand , Charles Hong , Julian Walker , Nina Cai , Shangkun Wang , Suvinay Subramanian , Sundar Dev , Vijay Janapa Reddi
show 2 more authors
Amir Yazdanbakhsh Sethu Sankaran
This is my paper
classification cs.AI
keywords JAXBenchTPU kernel optimizationPallasXLALLM agentskernel generation benchmarkMXU utilizationdocumentation conditioning
verification ladder T0 review T1 audit T2 compute T3 formal T4 reserved

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

This paper builds the missing TPU equivalent of KernelBench: a shared, reproducible benchmark that autonomous kernel-optimization agents can be measured against. JAXBench comprises 50 JAX workloads sized so the XLA baseline already approaches peak matrix-unit utilization, leaving headroom only for genuine algorithmic and scheduling improvements — 17 operators drawn from production LLM architectures and 33 fused operators translated from KernelBench, with 8 hand-tuned Pallas reference kernels providing an expert ceiling. The central experimental claim, measured with Gemini 3 Flash across the suite (and a 5-kernel subset with a stronger model), is that on TPUs the dominant lever is target-specific documentation, not model scale: conditioning agents on curated TPU documentation raises per-sample correctness from 5.8% to 37.3% and solves 48 of 50 workloads at 1.28x geomean speedup over XLA, while a beam-search pipeline that separates translation from optimization reaches 1.36x and recovers most of the 2.08x expert bound on the 8 hand-tuned kernels. If correct, the suite becomes the shared hill for TPU kernel-agent research, and 'context before scale' becomes a design principle for any niche, sparsely-documented accelerator.

Core claim

JAXBench is a benchmark of 50 JAX workloads sized to saturate TPU v6e matrix units, built from 17 production operators extracted from MaxText architectures (Llama-3.1, DeepSeek-V3, Mixtral, Mamba-2, AlphaFold2) and 33 fusion workloads translated from KernelBench and re-sized to at least 60% MXU utilization under XLA. The central discovery: for LLM agents writing Pallas kernels, target-specific documentation beats model scale. Distilled TPU context (hardware summary, Pallas API reference, examples, rules) raises per-sample correctness from 5.8% to 37.3% and solves 48 of 50 benchmarks at 1.28x geomean speedup over XLA; a beam-search variant reaches 1.36x, and 1.60x on the 8 kernels with hand-t

What carries the argument

The load-bearing machinery is the benchmark harness plus the documentation-conditioning pipeline. Each workload exposes a standard interface (CONFIG, create_inputs, workload), is sized so the XLA baseline reaches at least 60% MXU utilization, and is timed device-side via jax.profiler Perfetto traces to exclude host overhead. Correctness is gated by jnp.allclose atol=rtol=1e-2 on randomized bf16 inputs; aggregate speedups are floored at 1x so incorrect kernels contribute no gain. The documentation pipeline distills public Pallas and Cloud TPU docs into four prompt artifacts — hardware summary, per-benchmark-selected API reference, annotated code examples, rules block — prepended to every agen

Load-bearing premise

The central claim rests on the correctness gate: a kernel counts as 'solving' a benchmark if its outputs match the reference under jnp.allclose(atol=rtol=1e-2) on randomized bf16 inputs, and kernels that pass this gate could still diverge on real activation distributions or ragged shapes — so the strength of 48/50, 37.3%, and the speedup geomeans is exactly the strength of that tolerance.

What would settle it

Re-run the evaluation with a stricter or distributional correctness gate — for instance, validate generated kernels on real activation distributions, ragged or masked shapes, or at atol=rtol=1e-3 against fp32 references — and check whether the 48/50 solved count and the 1.28x–1.36x geomeans survive. If the counts drop substantially, the documentation-beats-scale result holds only for the loose bf16 gate. A complementary check: repeat the documentation-injection ablation on a DSL that is well represented in training data (such as Triton for GPUs); if context then no longer beats scale, the resu

Watch this falsifier. Get emailed when new claim-graph text bears on it.

If this is right

  • The suite becomes the shared evaluation target for TPU kernel-agent research — a TPU analogue of what KernelBench did for GPUs — with a reproducible harness any method can be measured against.
  • For sparsely-documented accelerator DSLs, the cheap lever is curated context: the 5.8% to 37.3% correctness jump from documentation exceeds the gain from moving Flash to Pro at fixed context.
  • Once correctness is achieved, search structure converts it into speed: the translation/optimization beam split reaches 1.36x geomean versus 1.28x for context-conditioned iterative refinement on the full suite.
  • Autonomous agents recover roughly three-quarters of the hand-tuned ceiling (1.60x vs 2.08x geomean on 8 kernels), with paged and ragged attention still out of reach — a concrete map of where expert effort still beats automation.
  • The claims are scoped to single-chip TPU v6e kernels; multi-chip sharding and collective communication are explicitly out of scope, so the documentation-vs-search separation is untested on distributed kernels.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The context-beats-scale result is measured on a single correctness gate; if that gate were tightened to distributional or higher-precision checks, the absolute numbers (48/50, 37.3%) would likely fall even though the relative ordering of methods might survive.
  • The suite's single-chip scope makes its main design lesson untested in multi-chip settings; collectives like all_to_all and psum_scatter are even sparser in pretraining data than Pallas, so documentation conditioning could matter even more there — a direct extension the paper itself flags.
  • A testable corollary of the information-gap claim: fine-tuning an agent on a curated Pallas corpus with JAXBench correctness and speedup signals as rewards should outperform further model-scaling at fixed budget.
  • The per-benchmark tier structure suggests a cheap routing heuristic: use iterative refinement on kernels where the first correct seed is easy and reserve beam-search budgets for kernels where correctness is hard — the paper's own per-benchmark results already separate these regimes.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit.

Referee Report

2 major / 6 minor

Summary. The paper introduces JAXBench, a TPU-native benchmark suite of 50 JAX workloads (17 production operators from MaxText, 33 translated from KernelBench Level 2) for evaluating LLM-generated Pallas kernels on TPU v6e. It provides an evaluation harness using device-side Perfetto profiling with median-of-50 timing, a correctness gate at atol=rtol=1e-2, and, for 8 priority kernels, hand-tuned Tokamax Pallas baselines with grid-searched block sizes. Four methods are compared with Gemini 3 Flash: best-of-N, iterative refinement, iterative refinement with Autocomp context, and Autocomp beam search. Results show that adding curated TPU documentation raises per-sample correctness from 5.8% to 37.3%, iterative+context solves 48/50 benchmarks at 1.28x geomean speedup over XLA, and Autocomp reaches 45/50 at 1.36x. On the 8 hand-tuned kernels, Autocomp reaches 1.60x geomean versus the Tokamax 2.08x upper bound. A 5-kernel Gemini 3.1 Pro ablation is presented as a preliminary model-capacity comparison.

Significance. If the results hold, JAXBench fills a clear gap: no equivalent of KernelBench exists for TPUs, and the benchmark plus harness could become a shared target for TPU kernel-agent research. The methodology has real strengths: device-side timing, median-of-50 aggregation, roofline checks that shapes are compute-bound, PyTorch-reference validation of translations, fully documented prompts in Appendix B, and disclosed floor-1 aggregation. The paper also reports a result that is unfavorable to its preferred method (iterative+ctx solves 48/50 while Autocomp solves 45/50), which increases confidence. The context-vs-scale finding is striking and, if confirmed, would be a useful design principle for agents targeting niche, sparsely documented DSLs. The main threat to the central claims is the single, untested correctness gate and the unmatched comparison underlying the context-vs-scale conclusion.

major comments (2)
  1. [3.1, 2.2] All headline results pass through a single numerical gate: jnp.allclose(atol=rtol=1e-2) on randomized bf16 inputs. Section 2.2 uses the same gate to validate the PyTorch-to-JAX references and Section 3.1 uses it to certify generated kernels. The manuscript does not state whether correctness is evaluated on one random draw or multiple draws, whether a random seed is fixed, or how results would change under stricter tolerances, distributional inputs (zeros, clipped activations, large dynamic range), or ragged/masked shapes. Attention and normalization workloads are scale- and shape-sensitive, and a kernel that passes this gate can diverge on real activation distributions. Because correctness is a precondition for any speedup, the 48/50 solve count, the 37.3% per-sample correctness, and the 1.28x-1.36x geomeans all inherit the strength of this gate. The paper should add a sensitivity analys
  2. [3.4, 4] The paper's second headline claim, 'target-specific context matters more than model scale' (abstract, Section 4), is supported by comparing two effects on different benchmark sets. The context effect (5.8% to 37.3% per-sample correctness) is measured on the full 50-workload suite (Section 3.2), while the model-scale effect (1.2% to 16.9%) is measured on the 5-kernel subset (Section 3.4). These numbers are not commensurable. The conclusion would be supported by running the scale ablation on the full suite, or by recomputing the context effect on the same 5-kernel subset for both Flash and Pro. As written, the evidence is suggestive but not controlled, and the abstract states the conclusion categorically.
minor comments (6)
  1. [5] Typo: 'many workloadse' should be 'many workloads'.
  2. [Table 3] Table header 'HT (ms)' is undefined; use 'Hand-tuned (ms)' or spell out the abbreviation in the caption.
  3. [Figure 3] Caption and axis use 'fast@N' while the text defines 'fast1@N'; make the notation consistent.
  4. [2.2] The sentence 'we execute the original PyTorch reference and the generated JAX implementation on TPU' is unclear, since PyTorch does not run natively on TPU. Clarify the actual execution environments (e.g., PyTorch on CPU/GPU and JAX on TPU, with matched inputs).
  5. [2.2, 3.1] The phrase 'standard bf16 tolerance' is used without citation or derivation. bf16 has limited precision, but atol=rtol=1e-2 is not a universally standard choice; justify or qualify it.
  6. [Table 2] Footnote 1 defines the speedup floor; consider using 'floor-1 geomean' in table and figure captions for consistency with Section 3.3.

Circularity Check

0 steps flagged

No significant circularity; headline results are external measurements against XLA and Tokamax baselines.

full rationale

JAXBench is an empirical benchmark paper. The headline numbers (48/50 solved, 37.3% per-sample correctness, 1.28x/1.36x geomean speedups) are measured by compiling and running generated kernels against XLA-compiled JAX references, with hand-tuned Tokamax kernels as external upper bounds; there is no fitted parameter that is later renamed as a prediction. The only formal derivation, Eq. 1 for ragged-paged-attention operational intensity, is a FLOP/byte counting bound and is not an input to any headline result. The use of jnp.allclose(atol=rtol=1e-2) in both reference validation (Section 2.2) and agent evaluation (Section 3.1) is a stated tolerance convention, not a reduction of a predicted quantity to a fitted one; concerns about gate strictness on real distributions are correctness/robustness risks, not circularity. The paper includes self-citations: Autocomp [7] is co-authored by one of the present authors and is used as the primary agent baseline, and [27]/[28] are cited for context. This is not load-bearing: Autocomp's results are measured through the same harness, the iterative+context ablation does not depend on Autocomp's search, and no central claim is justified solely by these citations. The ablations even show results that do not favor the authors' preferred method (iterative+context solves 48/50 vs Autocomp 45/50), consistent with external measurement rather than construction. No equation or definition makes a claimed result true by construction.

Axiom & Free-Parameter Ledger

5 free parameters · 5 axioms · 0 invented entities

The central claims rest on five hand-chosen conventions (MXU-utilization shape floor, 1x speedup floor, 1e-2 correctness tolerance, 144-sample budget with Autocomp early-stopping, best-of-203 Tokamax tuning) and on domain assumptions about XLA being the right denominator, profiler events equaling kernel time, and workload representativeness. No invented entities are introduced. The free parameters are disclosed design choices; they do not make the results unfalsifiable, but each one shifts the headline numbers, so they are counted here.

free parameters (5)
  • MXU utilization threshold (shape selection) = 60% (chosen floor)
    Workload shapes for the 33 translated operators were frozen at the smallest configuration whose XLA baseline reaches at least 60% MXU utilization (Sec. 2.2). This hand-chosen constant determines which shapes enter the suite and underwrites the 'compute-bound headroom' design claim; a different threshold would change headroom and measured speedups.
  • Speedup floor (floor-1 aggregation) = 1.0x
    All aggregate geomeans floor per-kernel speedups at 1x, so incorrect kernels and regressions (e.g., Autocomp at 0.75x on GEMM, Table 3) count as 1x (Sec. 3.1). This disclosed convention materially inflates the headline 1.28x/1.36x geomeans.
  • Correctness tolerance (bf16 gate) = atol=rtol=1e-2
    The gating check behind 'solves 48 of 50' and the 37.3% per-sample correctness number (Sec. 3.1). Chosen as 'standard bf16 tolerance', but loose enough to admit approximate kernels, and validated only on randomized inputs, never on distributional workloads.
  • Sample budget = 144 samples/benchmark (Autocomp early-stops)
    Best-of-N and iterative variants use 18 chains x 8 turns = 144 samples; Autocomp uses 'comparable' budget but early-stops once its beam stops improving (Sec. 3.1). The method comparison is therefore not exactly budget-matched.
  • Tokamax block-size grid search = best of 203 configurations per kernel
    The 8 'expert upper-bound' references are community kernels tuned by exhaustive grid search on TPU v6e (Sec. 2.2). The 2.08x expert geomean is deliberately optimistic, and the 'recovering 77% of expert bound' framing inherits this selection.
axioms (5)
  • domain assumption XLA is the right competitive denominator for TPU kernel quality
    The entire metric is speedup over the XLA-compiled jax.jit baseline (Sec. 2.1, 3.1). The paper justifies XLA as 'the strongest reproducible baseline', but the suitability of that denominator for every shape class is assumed, not verified against independent hand-written kernels beyond the 8 Tokamax references.
  • domain assumption High MXU utilization implies workloads are in the meaningful optimization regime
    The design principle sizes workloads to reach 60-95% MXU utilization (Sec. 2.1), but the paper's own Figure 2a shows attention and elementwise kernels are memory-bound below the ridge point - and those are exactly where the largest speedups (16.3x ragged paged attention) occur. The mapping 'high utilization = meaningful headroom' is therefore not uniform across the suite.
  • domain assumption jax.profiler device-side events equal true per-iteration kernel time
    The timing protocol (Sec. 2.3) assumes that jit_*()-level events extracted from Perfetto traces capture device-side execution time without host dispatch overhead. This is standard practice but is not independently validated against hardware counters in the paper.
  • domain assumption bf16 allclose at 1e-2 is a valid correctness criterion
    Used both for workload translation validation (against the PyTorch reference) and for agent-kernel gating (Sec. 3.1). Precision-tolerant by design, but may admit kernels that are incorrect on real value distributions, attention masks, or ragged shapes.
  • domain assumption The 50 workloads are representative of production TPU kernel demand
    Selection was made by the authors: 17 MaxText operators plus a one-per-fusion-signature subset of KernelBench Level 2 (Sec. 2.2). Representativeness is asserted rather than demonstrated, and the 8-slot 'priority' tier plus 33 translated ops mix two different provenance pools.

pith-pipeline@v1.3.0-alltime-deepseek · 13703 in / 22047 out tokens · 225782 ms · 2026-08-02T13:39:48.407744+00:00 · methodology

0 comments
read the original 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. 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. 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. 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. 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. 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.

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Reference graph

Works this paper leans on

32 extracted references · 13 linked inside Pith

  1. [1]

    Jax: Autograd and xla.Astrophysics Source Code Library, pages ascl–2111, 2021

    James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, et al. Jax: Autograd and xla.Astrophysics Source Code Library, pages ascl–2111, 2021

    2021

  2. [2]

    K-search: Llm kernel generation via co-evolving intrinsic world model.arXiv preprint arXiv:2602.19128, 2026

    Shiyi Cao, Ziming Mao, Joseph E Gonzalez, and Ion Stoica. K-search: Llm kernel generation via co-evolving intrinsic world model.arXiv preprint arXiv:2602.19128, 2026

    arXiv 2026

  3. [3]

    Flashattention: Fast and memory-efficient exact attention with io-awareness.Advances in neural information processing systems, 35:16344–16359, 2022

    Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness.Advances in neural information processing systems, 35:16344–16359, 2022

    2022

  4. [4]

    MaxText: A high performance, scalable, open-source LLM in pure Python/JAX

    Google AI Hypercomputer. MaxText: A high performance, scalable, open-source LLM in pure Python/JAX. https://github.com/AI-Hypercomputer/maxtext, 2022. Accessed: 2026- 04-26

    2022

  5. [5]

    The llama 3 herd of models.arXiv preprint arXiv:2407.21783, 2024

    Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models.arXiv preprint arXiv:2407.21783, 2024

    Pith/arXiv arXiv 2024

  6. [6]

    Mamba: Linear-time sequence modeling with selective state spaces

    Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023

    Pith/arXiv arXiv 2023

  7. [7]

    Autocomp: Llm-driven code optimization for tensor accelerators

    Charles Hong, Sahil Bhatia, Alvin Cheung, and Sophia Shao. Autocomp: Llm-driven code optimization for tensor accelerators. InMachine Learning for Computer Architecture and Systems 2025, 2025

    2025

  8. [8]

    Writing TPU kernels with pallas.https://docs.jax.dev/en/ latest/pallas/tpu/details.html, 2024

    JAX Development Team. Writing TPU kernels with pallas.https://docs.jax.dev/en/ latest/pallas/tpu/details.html, 2024. Accessed: 2026-04-26

    2024

  9. [9]

    Mixtral of experts.arXiv preprint arXiv:2401.04088, 2024

    Albert Q Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, et al. Mixtral of experts.arXiv preprint arXiv:2401.04088, 2024

    Pith/arXiv arXiv 2024

  10. [10]

    Ragged paged attention: A high-performance and flexible llm inference kernel for tpu.arXiv preprint arXiv:2604.15464, 2026

    Jevin Jiang, Ying Chen, Blake A Hechtman, Fenghui Zhang, and Yarong Mu. Ragged paged attention: A high-performance and flexible llm inference kernel for tpu.arXiv preprint arXiv:2604.15464, 2026

    Pith/arXiv arXiv 2026

  11. [11]

    In-datacenter performance analysis of a tensor processing unit

    Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. In-datacenter performance analysis of a tensor processing unit. InProceedings of the 44th annual international symposium on computer architecture, pages 1–12, 2017

    2017

  12. [12]

    Highly accurate protein structure prediction with alphafold.nature, 596(7873):583–589, 2021

    John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold.nature, 596(7873):583–589, 2021

    2021

  13. [13]

    Efficient memory management for large language model serving with pagedattention

    Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. InProceedings of the 29th symposium on operating systems principles, pages 611–626, 2023. 12 JAXBench: Benchmarking Autonomous TPU Kernel Optimization

    2023

  14. [14]

    Towards robust agentic cuda kernel benchmarking, verification, and optimization.arXiv preprint arXiv:2509.14279, 2025

    Robert Tjarko Lange, Qi Sun, Aaditya Prasad, Maxence Faldor, Yujin Tang, and David Ha. Towards robust agentic cuda kernel benchmarking, verification, and optimization.arXiv preprint arXiv:2509.14279, 2025

    arXiv 2025

  15. [15]

    Tritonbench: Benchmarking large language model capabilities for generating triton operators

    Jianling Li, Shangzhan Li, Zhenye Gao, Qi Shi, Yuxuan Li, Zefan Wang, Jiacheng Huang, WangHaojie WangHaojie, Jianrong Wang, Xu Han, et al. Tritonbench: Benchmarking large language model capabilities for generating triton operators. InFindings of the Association for Computational Linguistics: ACL 2025, pages 23053–23066, 2025

    2025

  16. [16]

    Kernelevolve: Scaling agentic kernel coding for heterogeneous ai accelerators at meta.arXiv preprint arXiv:2512.23236, 2025

    Gang Liao, Hongsen Qin, Ying Wang, Alicia Golden, Michael Kuchnik, Yavuz Yetim, Jia Jiunn Ang, Chunli Fu, Yihan He, Samuel Hsia, et al. Kernelevolve: Scaling agentic kernel coding for heterogeneous ai accelerators at meta.arXiv preprint arXiv:2512.23236, 2025

    Pith/arXiv arXiv 2025

  17. [17]

    Deepseek-v3 technical report.arXiv preprint arXiv:2412.19437, 2024

    Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report.arXiv preprint arXiv:2412.19437, 2024

    Pith/arXiv arXiv 2024

  18. [18]

    AlphaEvolve: A coding agent for scientific and algorithmic discovery.arXiv preprint arXiv:2506.13131, 2025

    Alexander Novikov, Ngân V˜u, Marvin Eisenberger, Emilien Dupont, Po-Sen Huang, Adam Zsolt Wagner, SergeyShirobokov, BorislavKozlovskii, FranciscoJ.R.Ruiz, AbbasMehrabian, M.Pawan Kumar, Abigail See, Swarat Chaudhuri, George Holland, Alex Davies, Sebastian Nowozin, Pushmeet Kohli, and Matej Balog. AlphaEvolve: A coding agent for scientific and algorithmic ...

    Pith/arXiv arXiv 2025

  19. [19]

    Tokamax: A GPU and TPU kernel library.https://github.com/ openxla/tokamax, 2024

    OpenXLA Contributors. Tokamax: A GPU and TPU kernel library.https://github.com/ openxla/tokamax, 2024. Accessed: 2026-04-26

    2024

  20. [20]

    Kernelbench: Can llms write efficient gpu kernels?arXiv preprint arXiv:2502.10517, 2025

    Anne Ouyang, Simon Guo, Simran Arora, Alex L Zhang, William Hu, Christopher Ré, and Azalia Mirhoseini. Kernelbench: Can llms write efficient gpu kernels?arXiv preprint arXiv:2502.10517, 2025

    Pith/arXiv arXiv 2025

  21. [21]

    Glu variants improve transformer.arXiv preprint arXiv:2002.05202, 2020

    Noam Shazeer. Glu variants improve transformer.arXiv preprint arXiv:2002.05202, 2020

    Pith/arXiv arXiv 2002

  22. [22]

    Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.arXiv preprint arXiv:1701.06538, 2017

    Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.arXiv preprint arXiv:1701.06538, 2017

    Pith/arXiv arXiv 2017

  23. [23]

    Retentive network: A successor to transformer for large language models.arXiv preprint arXiv:2307.08621, 2023

    Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models.arXiv preprint arXiv:2307.08621, 2023

    Pith/arXiv arXiv 2023

  24. [24]

    Gemini: a family of highly capable multimodal models.arXiv preprint arXiv:2312.11805, 2023

    Gemini Team, Rohan Anil, Sebastian Borgeaud, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, Katie Millican, et al. Gemini: a family of highly capable multimodal models.arXiv preprint arXiv:2312.11805, 2023

    Pith/arXiv arXiv 2023

  25. [25]

    CUTLASS.https://github.com/NVIDIA/cutlass, 2023

    Vijay Thakkar, Pradeep Ramani, Cris Cecka, Aniket Shivam, Honghao Lu, Ethan Yan, Jack Kosaian, Mark Hoemmen, Haicheng Wu, Andrew Kerr, Matt Nicely, Duane Merrill, Dustyn Blasig, Aditya Atluri, Fengqi Qiao, Piotr Majcher, Paul Springer, Markus Hohnerbach, Jin Wang, and Manish Gupta. CUTLASS.https://github.com/NVIDIA/cutlass, 2023. Accessed: 2026-04-26

    2023

  26. [26]

    Triton: an intermediate language and compiler for tiled neural network computations

    Philippe Tillet, Hsiang-Tsung Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. InProceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages, pages 10–19, 2019. 13 JAXBench: Benchmarking Autonomous TPU Kernel Optimization

    2019

  27. [27]

    Swizzleperf: Hardware-aware llms for gpu kernel performance optimization.arXiv preprint arXiv:2508.20258, 2025

    Arya Tschand, Muhammad Awad, Ryan Swann, Kesavan Ramakrishnan, Jeffrey Ma, Keith Lowery, Ganesh Dasika, and Vijay Janapa Reddi. Swizzleperf: Hardware-aware llms for gpu kernel performance optimization.arXiv preprint arXiv:2508.20258, 2025

    Pith/arXiv arXiv 2025

  28. [28]

    Genai for systems: Recurring challenges and design principles from software to silicon, 2026

    Arya Tschand, Chenyu Wang, Zishen Wan, Andrew Cheng, Ioana Cristescu, Kevin He, Howard Huang, Alexander Ingare, Akseli Kangaslahti, Sara Kangaslahti, Theo Lebryk, Hongjin Lin, Jeffrey Jian Ma, Alexandru Meterez, Clara Mohri, Depen Morwani, Sunny Qin, Roy Rinberg, Paula Rodriguez-Diaz, Alyssa Mia Taliotis, Pernille Undrum Fathi, Rosie Zhao, Todd Zhou, and ...

    arXiv 2026

  29. [29]

    Multiker- nelbench: A multi-platform benchmark for kernel generation.arXiv e-prints, pp

    Zhongzhen Wen, Yinghui Zhang, Zhong Li, Zhongxin Liu, Linna Xie, and Tian Zhang. Multiker- nelbench: A multi-platform benchmark for kernel generation.arXiv e-prints, pp. arXiv–2507, 2025

    2025

  30. [30]

    Roofline: an insightful visual performance model for multicore architectures.Communications of the ACM, 52(4):65–76, 2009

    Samuel Williams, Andrew Waterman, and David Patterson. Roofline: an insightful visual performance model for multicore architectures.Communications of the ACM, 52(4):65–76, 2009

    2009

  31. [31]

    Flashinfer-bench: Building the virtuous cycle for ai-driven llm systems.arXiv preprint arXiv:2601.00227, 2026

    Shanli Xing, Yiyan Zhai, Alexander Jiang, Yixin Dong, Yong Wu, Zihao Ye, Charlie Ruan, Yingyi Huang, Yineng Zhang, Liangsheng Yin, et al. Flashinfer-bench: Building the virtuous cycle for ai-driven llm systems.arXiv preprint arXiv:2601.00227, 2026

    arXiv 2026

  32. [32]

    Below is the XLA/JAX reference implementation. Rewrite it as a Pallas kernel that produces identical outputs and runs faster

    Biao Zhang and Rico Sennrich. Root mean square layer normalization.Advances in neural information processing systems, 32, 2019. A. Per-Benchmark Results Figure 6 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 miss...

    2019