Running AlphaFold3 (single seq mode) 3 times faster on TPUs

Background

Recently I’ve been training some protein design flow matching models on TPUs. To evaluate them protein folding models play an important role. So to keep things simpler (GCP buckets etc), I wanted to use AF3 on TPUs as well. However, the current AF3 codebase didn’t seem to support inference on TPUs.

Running AF3 with XLA

Unsupported doesn’t mean we cannot run it on TPUs.

AF3 is written in JAX and it compiles through XLA, which runs natively on TPUs. So what’s the issue? For TPUs it does not have an optimized kernel for pair-biased attention.³³ AF3 actually has four pair-biased attentions: triangle attention, the single attention in the pairformer, the diffusion transformer and the atom transformers. Only triangle attention goes through tokamax. Because of this, when you run AF3 on a v6e chip it runs slow even compared to a 5090 ⁴⁴ This is mostly not the attention’s fault: one unfused triangle-attention call takes 0.71 ms against 0.32 ms fused, which over the 1,056 calls in a fold adds about 0.4 s, while the TPU VM’s slower CPU costs about 1.1 s..

To dig a little deeper: AF3 calls tokamax, an OpenXLA library of fused JAX kernels, and --flash_attention_implementation picks which tokamax backend runs.

However, on a TPU we hit two dead ends:

1. tokamax’s fused GPU kernels (triton, the default, and cudnn) don’t run on TPUs.
2. tokamax’s TPU kernel (mosaic) doesn’t accept an attention bias at all ⁵⁵ The check is in pallas_mosaic_tpu_common.py [3]: passing any bias raises this error – still the case in tokamax 0.0.17, and AF3′s triangle attention needs one: the pair bias.

If we don’t care about performance, we can just run the standard AF3 inference code with this flag ⁶⁶ Defined in run_alphafold.py [4], whose help text calls xla “no flash attention” and “portable across GPU devices”; TPUs aren’t mentioned..

--flash_attention_implementation=xla

With it, AF3 computes attention with ordinary JAX operations instead of a hand-written GPU kernel. XLA compiles those for the TPU like the rest of AF3′s JAX code. But since this automatic compilation doesn’t result in a fused kernel, it writes the full attention-logit tensor to memory and reads it back ⁷⁷ Flash attention (Dao et al., 2022 [5]) computes softmax(QKᵀ + bias)·V in tiles held in fast on-chip memory, keeping a running maximum and sum for the softmax, so the N × N matrix of logits is never written to the slower main memory (HBM). That removes the largest memory traffic in attention. The TPU kernel later in this post uses the same idea in a simpler form: at these sizes a whole row of logits fits on-chip, so it never has to tile across keys.. How much does this hurt us?

Optimizations in a nutshell

If we look at the timing for a full end to end run on a v6e chip, it runs in about 5.8 seconds. For context I measured this on a 5090 as well.

How much do precision errors propagate?

Here are all the optimizations in one place, with what each one does to AF3′s predictions. Two of them just reorganize when and where the work happens, so the outputs stay the same as the original AF3 inference code’s.

The other four change the floating-point arithmetic. The kernels round the same way as AF3′s own code on more than 99% of values, but through 200 denoising steps even that can change the final structure.

Do the predictions change?

Pipelining and hoisting cannot change anything: the outputs are bit for bit identical. The kernels can, through rounding, so I checked them in two ways.

First, against real structures. I folded 40 proteins from the PDB from their deposited sequences (with their ligands, no MSA), once with stock AF3 and once with the kernels, with the same padding and seed, and scored each prediction against the deposited structure: lDDT over all the protein’s heavy atoms, where 1 is a perfect match. If the kernels changed accuracy, the points below would move off the diagonal. For scale, the other two panels make changes that have nothing to do with the kernels: a new random seed, and running stock AF3 on an RTX 5090 instead of the TPU.

The kernels don’t shift accuracy: the mean change is −0.002 lDDT. Individual proteins do move, by 0.027 on average, about as much as running stock AF3 on an RTX 5090 instead of the TPU (0.025) and a bit more than a new random seed (0.015). Single-sequence AF3 is unsure about most natural proteins (mean pLDDT around 50, and most points sit low on both axes), and for an unsure prediction any small numerical difference, ours, a new seed or another chip, sends the diffusion somewhere slightly different.

Second, on the inputs I actually fold: 40 designed proteins with ligands, where AF3 is mostly confident. There is no deposited structure, so I compared top-ranked structures directly, over three seeds. With the kernels they move a median of 0.52 Å from stock AF3′s, against 1.29 Å for changing the seed, and the confident ones (pLDDT of 90 or more) move at most 0.84 Å.

Related implementations

After I worked on this with claude + codex on optimizing AF3 for TPUs, I realized that I didn’t do much of a existing literature search. So we revisited Anthropic’s FlashPairformer code [6] that include three close comparisons, all targeting GPUs through Pallas’s Triton backend:

• Triangle attention: tiled flash attention with the pair bias shared across triangle rows.
• Transition: fused LayerNorm, SwiGLU and output projection.
• Triangle multiplication: fused input processing, a contraction, and fused output processing. Its source explicitly notes that fusion changes BF16 rounding and is not bitwise identical to stock execution.

However, Anthropic’s repo mainly targets NVIDIA GPUs through the Pallas kernels. With only their compiler settings switched to TPU ones, the triangle-attention and triangle-multiplication kernels do not compile on a TPU: their memory layouts break the TPU’s block-shape rules.

Apart from Anthropic’s repo, JAX’s own Pallas FlashAttention for TPUs [2] does accept an additive attention bias, but only as a separate copy for every row. On AF3′s triangle attention it is 1.6–4.4× slower than the kernel in this post (at 256 tokens 0.54 ms for the kernel alone and 0.84 ms with the bias copy and transposes it needs, against 0.32 ms), and no faster than plain XLA up to 384 tokens. Most of the gap is its handling of AF3′s 32-wide heads; building the per-row bias costs the rest. Separately, JAXBench: Benchmarking Autonomous TPU Kernel Optimization [7] includes an AlphaFold2-style triangle-multiplication workload in its TPU v6e benchmark; its reference workload [8] differs from the exact AF3 module.