license	cc-by-nc-4.0
pretty_name	TPU v5e single-device Pallas attention benchmark
tags	tpu jax pallas kernels attention grouped-query-attention benchmark performance kaggle
size_categories	n<1K
extra_gated_prompt	These records were measured by Laela Zorana on a Kaggle TPU v5e worker. Use them for non-commercial work, keep the attribution with them, and tell me what you are building.
extra_gated_fields	Name Organisation or affiliation What are you planning to use this for I will keep the attribution with the data text text text checkbox
extra_gated_button_content	Request access
configs	config_name default data_files comparisons true split path train dataset/comparisons.jsonl config_name data_files benchmarks split path train dataset/benchmarks.jsonl config_name data_files correctness split path train dataset/correctness.jsonl

JAX/Pallas attention measured on one TPU v5e device

I measured a Pallas masked-softmax kernel against the JAX/XLA reference inside grouped-query attention. The Kaggle worker exposed eight TPU devices. The unsharded arrays ran on JAX's default single device, so these numbers make no multi-device scaling claim.

All seven correctness checks passed. The benchmark contains twenty cases, with five warmups and twenty synchronized timing samples in each case. XLA was faster in all ten paired comparisons. At sequence length 2048, full attention measured 0.8063 ms for XLA and 1.9225 ms for Pallas.

Scope	Sequence	XLA reference	Pallas	Pallas slowdown
softmax	128	0.1761 ms	0.2166 ms	1.23x
softmax	256	0.1805 ms	0.2526 ms	1.40x
softmax	512	0.2445 ms	0.3516 ms	1.44x
softmax	1024	0.4080 ms	0.5200 ms	1.27x
softmax	2048	1.0967 ms	1.5886 ms	1.45x
attention	128	0.1811 ms	0.2353 ms	1.30x
attention	256	0.1801 ms	0.2462 ms	1.37x
attention	512	0.1840 ms	0.3155 ms	1.71x
attention	1024	0.2354 ms	0.4715 ms	2.00x
attention	2048	0.8063 ms	1.9225 ms	2.38x

Inputs were bfloat16. Logits, normalization, and output accumulation used fp32. The Pallas result matched the reference in all seven checks, including causal and padding masks, an all-masked row, a non-128 key width, and gradients with respect to Q, K, and V.

Created and measured by Laela Zorana. The project is on Hugging Face and GitHub.

The two engineering layers

Layer	What is tested
Kernel and compiler	8x128 Pallas tiles, TPU padding, fp32 reductions, StableHLO, compile time, synchronized device time
Model and evaluation	Grouped-query head mapping, causal decoding, padding masks, all-masked rows, Q/K/V gradients, adversarial candidates

The custom kernel owns masked softmax. XLA owns the QK and probability-times-V matrix multiplications. The StableHLO and raw timings record that boundary. Fusion and data movement are the next measurements.

Correctness and grader checks

Risk	Check
Wrong grouped-query mapping	4 query heads mapped to 2 distinct KV heads
Future-token leakage	Causal masks, including K longer than Q for cached decoding
Padding leakage	Non-power-of-two widths padded to TPU's 128-lane shape
Undefined all-masked rows	Exact zero output, no NaN or Inf
Precision loss	fp32 logits and reduction with bfloat16 inputs
Broken backward semantics	Q, K, and V gradient parity against pure JAX
Unsynchronized timing	Every timed result completes with block_until_ready()
Weak evaluation	Known-bad implementations must lose points

Repository contents

src/tpu_kernel_lab/
  attention.py        pure-JAX oracle and Pallas-backed attention
  pallas_softmax.py   8x128-tiled kernel with an explicit backward rule
  adversarial.py      missing-mask, low-precision, and wrong-GQA candidates
  rubric.py           correctness and gradient scorecard
  benchmark.py        synchronized JSON benchmark
tests/                parity, edge-case, gradient, and rubric tests
kaggle/               TPU measurement script
evidence/             hardware reports and StableHLO
dataset/              flat tables, raw JSON, data notes, and checksums

Reproduce the checks

python -m venv .venv
. .venv/bin/activate
python -m pip install -e ".[dev]"
pytest
tpu-kernel-rubric --candidate reference
tpu-kernel-rubric --candidate ignores-causal

On a CPU, Pallas runs in interpret mode for correctness. TPU_RUNBOOK.md carries the hardware commands. The checked hardware evidence is under evidence/kaggle-v5e-8/, and the viewer tables are under dataset/.

Rebuild the data release

After downloading a completed Kaggle run into kaggle-output/:

python scripts/build_publication.py
python scripts/verify_publication.py

The builder accepts only a run with eight TPU devices, seven passing correctness checks, and twenty successful benchmark cases. It keeps transfer, compilation, and synchronized steady-state time in separate fields and retains every timing sample.

The completed run first paired JAX 0.10.2 with the Kaggle image's June 2025 libtpu. Pallas rejected that version gap. The measured run used JAX and JAXLIB 0.6.2 with the bundled TPU runtime, and the bootstrap record stays in the data.

Licences

Source available under the CodeZorana Source Available License, Version 1.0, in the LICENSE file. Reading, study and independent reproduction of the reported measurements are permitted. Production use, commercial use, redistribution and derivative works require a separate written licence. No patent licence is granted, and patent rights are reserved. Data files under dataset/ keep the Creative Commons terms stated on their own card. Licence enquiries go to the copyright holder. available through the Hugging Face profile. The notice file carries the measurement attribution.