# Accelerator Choice Is Not Enough: AlphaFold2 Inference on Cloud TPUs

arXiv: https://arxiv.org/abs/2609.34818 · HTML: https://mlsanity.com/paper/2609.34818/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 10 of its 27 references that are in the map.
- First posted on arXiv 28 Sep 2026.

## Builds on (most cited first, 10 of 10)

- [TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings](https://mlsanity.com/paper/2304.01433/) (2023)
- [GSPMD: General and Scalable Parallelization for ML Computation Graphs](https://mlsanity.com/paper/2105.04663/) (2021)
- [FastFold: Reducing AlphaFold Training Time from 11 Days to 67 Hours](https://mlsanity.com/paper/2203.00854/) (2022)
- [HelixFold: An Efficient Implementation of AlphaFold2 using PaddlePaddle](https://mlsanity.com/paper/2207.05477/) (2022)
- [DrJAX: Scalable and Differentiable MapReduce Primitives in JAX](https://mlsanity.com/paper/2403.07128/) (2024)
- [APACE: AlphaFold2 and advanced computing as a service for accelerated discovery in biophysics](https://mlsanity.com/paper/2308.07954/) (2023)
- [Compiler-First State Space Duality and Portable O(1) Autoregressive Caching for Inference](https://mlsanity.com/paper/2603.09555/) (2026)
- [CCL-Bench 1.0: A Trace-Based Benchmark for LLM Infrastructure](https://mlsanity.com/paper/2605.06544/) (2026)
- [Duet Benchmarking: Improving Measurement Accuracy in the Cloud](https://mlsanity.com/paper/2001.05811/) (2020)
- [Fine-Tuning and Serving Gemma 4 31B on Google Cloud TPU: A Technical Comparison with GPU Baselines](https://mlsanity.com/paper/2605.25645/) (2026)

## Abstract (arXiv)

AlphaFold2 is written in JAX, so the same inference code compiles and runs unchanged on CPUs, GPUs and Google Cloud TPUs. That portability makes the accelerator look like the main decision a user has to make. We show that it is not. Running one AlphaFold2 inference workload across a Colab CPU runtime, an NVIDIA T4 GPU and a dedicated eight-chip Cloud TPU v5e slice, we find a large hardware advantage for the TPU, 0.47 s per call in steady state on a single chip against 13.1 s on the T4 in the same measurement campaign, and three ways in which the software layer decides how much of it a user actually gets. The default execution path uses one chip of the eight, and at list prices the idle capacity makes the slice cost about as much per prediction as the GPU. Batching with JAX's vmap never exceeds single-query throughput, while mapping queries across chips with JAX's pmap gives eight chips 6.5-7.9x the throughput of one on a matched grid; automatic sharding leaves the per-chip footprint unchanged, consistent with replication, most plausibly because AlphaFold2 carries no sharding annotations. Our retained trace analysis of a first call at a new input shape reports about three quarters of the traced span in JAX tracing and compilation rather than execution. Reruns five weeks later reproduced neither cloud baseline, the GPU one off by roughly a factor of two, so the hardware ratio above is specific to one campaign.