Experimenting with Threefry and Pallas on TPU

[pascalMontmory]

Hi everyone,

I've been working on my own PRNG, Montmory, for several months. While porting it to TPU, I also started experimenting with compiling Threefry myself using Pallas and exploring optimizations to improve performance.

Here are some results from a single session on one TPU v5 lite, with JAX/jaxlib 0.10.2. These measure output-buffer generation; GB/s uses decimal bytes, while MiB denotes binary output sizes.

Implementation	64 MiB output	128 MiB output
Threefry2x32, measured JAX path	31.954 GB/s	33.364 GB/s
Threefry2x32, optimized Pallas	65.249 GB/s	93.006 GB/s
Montmory, previous Pallas kernel	54.200 GB/s	73.498 GB/s
Montmory, retained Pallas kernels	62.458 GB/s	89.488 GB/s
JAX/XLA RBG	86.811 GB/s	90.791 GB/s

Confirmation used seven new seeds, with nine measurements per seed, after a separate search phase. Reported throughput is derived from the median of the per-seed median times. Timing includes the host call, allocation and completion wait; compilation, key preparation, host transfers and output checks are excluded.

The Threefry Pallas kernel retains all 20 rounds. Its tested outputs match the JAX Threefry reference, and Montmory's match its CPU reference. Hashes of timed outputs were checked outside the timing window; RBG was checked only for repeatability on the same backend.

Threefry Pallas remains faster than Montmory in these measurements. The RBG row is the measured JAX/XLA path, not a direct measurement of the hardware PRNG's maximum throughput. These are generation benchmarks, not application speedups or results across multiple TPU generations.

Public research record (code, protocols, measurements, limitations):
https://github.com/pascalMontmory/montmory-prng-research

Documentation is available in English and French; see the repository license and third-party notices for usage terms.

• Main results and methodology
• Raw measurements (CSV)
• Confirmation summary (CSV)
• TPU optimization notebook (English)

The result links above point to the exact published revision. My GitHub profile: pascalMontmory.

Pascal