[Steboss]

Summary

This change optimises the GridSelfAttention by:

1. Removing the outer mapping.inference_subbatch loop;
2. Passing the complete row batch to Tokamax in one logical attention call;
3. Replacing the dense boolean attention mask with key_value_seq_lengths.

Tokamax FlashAttention already tiles the operation internally. In the measured kernels it uses 64 x 64 query/key tiles, so the additional AlphaFold-level row subbatching fragmented an operation that was already memory-aware.

The change preserves the attention calculation for valid tokens, while substantially reducing launches, dynamic slicing, device-to-device copies, and host-side scheduling overhead.

results were tested with the following input tokens: 512, 1024, 2048, 3072.
We have used:

• For the Evoformer model we have used random weights, script
• For the inputs we have generated a simple amminoacidic sequence of size 512, 1024, 2048, 3072, script
• This is the inference script we have used to extract GPU time (s), which was inspired by the AF README. In the final outcome we have not considered the compilation time, and metrics were computed form 5 runs, 10 recycles, triton flash attention.

Explanation

Before Alphafold divided the first residue dimension into small row subbatches:

  GridSelfAttention
      └── inference_subbatch
            ├── rows   0–127  → Tokamax → internal 64×64 tiles
            ├── rows 128–255  → Tokamax → internal 64×64 tiles
            ├── ...
            └── rows 896–1023 → Tokamax → internal 64×64 tiles

The changed simplified this view as:

  GridSelfAttention
      └── all 1024 rows → Tokamax → internal 64×64 tiles

At 1024 tokens, the original chunk size was 128, resulting in eight separate outer attention calls, at 2048 the chunk size fell to 32. Each subbatch processes different query rows, so AF was not mathematically recomputing the same attention outputs.
HOWEVER, The redundancy was in the execution machinery around those calculations. Analysing the Before/After profiles we could see:

• Repeated attention dispatches.
• Fragmented Q/K/V projections.
• Dynamic slicing of inputs.
• Dynamic updates used to assemble the output.
• Repeatedly supplying the shared pair-bias operand.
• Additional device copies and command-buffer executions.
• More host-side scheduling work.

Implementation changes

Previously we had

  pair_mask = pair_mask[:, None, None, :]

  act = mapping.inference_subbatch(
      self._attention,
      chunk_size,
      batched_args=[act, pair_mask],
      nonbatched_args=[nonbatched_bias],
  )

now we're exploiting directly tokamax and identifying with key_lengths the valid keys end and paddings:

  # pair_mask still has shape (N, N) here.
  key_lengths = pair_mask.astype(jnp.int32).sum(axis=-1)

  act = self._attention(
      act,
      key_lengths,
      nonbatched_bias,
  )

The complete act tensor is passed to Tokamax. Tokamax remains responsible for the lower-level GPU tiling, online softmax and memory movement within the FlashAttention kernel.

This removes only the redundant outer chunking layer—it does not disable attention tiling. key_lengths is independent of subbatching. It does not control how Tokamax partitions the work; it describes which keys are valid for each attention row.

In particular, for GridSelfAttention, Q/K/V have shape:

  [B=N, T=N, H, D]

AlphaFold pads tokens contiguously at the end. Consequently, every row’s valid keys form a prefix:

  valid keys = [0, key_length)

For a bucket of eight tokens containing five real tokens:

  pair_mask =
  [
    1 1 1 1 1 0 0 0
    1 1 1 1 1 0 0 0
    1 1 1 1 1 0 0 0
    1 1 1 1 1 0 0 0
    1 1 1 1 1 0 0 0
    0 0 0 0 0 0 0 0
    0 0 0 0 0 0 0 0
    0 0 0 0 0 0 0 0
  ]

  key_lengths = [5, 5, 5, 5, 5, 0, 0, 0]

The dense attention-side mask therefore contains no more information than one integer per row.

Passing key_value_seq_lengths allows Tokamax to:

• Run fully valid key blocks without loading a boolean mask.
• Mask only the block containing the padding boundary.
• Skip blocks that contain only padding.
• Represent the attention constraint using O(N) integers instead of an O(N^2) attention-mask operand.

Even for a bucket-aligned input, where every length equals N, this avoids dense-mask operations. For padded inputs, it can additionally avoid computing key blocks beyond the real sequence length.

Results

Here are the results on A100 machine, measured as GPU time in seconds, or warm time, given as average +/- stdev

Input tokens	AF3	AF3 + optimisation	Improvement
1024	54.37 +/- 0.09	36.14 +/- 0.05	-33.5% (1.5x)
2048	300.58 +/- 0.47	194.08 +/- 0.13	-35.4% (1.5x)

on a B300 machine

Input tokens	AF3	AF3 + optimisation	Improvement
1024	15.07 +/- 0.01	12.59 +/- 0.01	-16.3% (1.2x)
2048	76.99 +/- 0.01	65.33 +/- 0.01	-15.1% (1.2x)
3072	461.22 +/- 0.06	443.5 +/- 0.06	-3.9% (1.1x)

Profile comparison

The 1024-token traces provide direct evidence that XLA was executing the outer subbatch loop and was not optimizing it away.
Across 5 warm runs we had:

Profile metric	Change
Flash-attention launches	31872 --> 3984 (8x less)
GPU events	-25%
Host events	-71.3%
MemcpyD2D events	-91%
Slice-fusion events	-71.5%
Dynamic-Slice events	68784 --> 0
Dynamic-update-slice events	-97%
Command-buffer executions	-61.2%

[Steboss]

CC: @Augustin-Zidek

[Augustin-Zidek]

Thanks for the PR. How was this tested?

[Steboss]

Hi @Augustin-Zidek, thanks for your question and I am sorry for my late reply, I was on leave.

This was tested on:

• end-to-end correctness
• operator equivalence

For end-to-end correctness we verified the inference results through this script.
The output results were tested with respect to the not-optimized results for:

• Structure RMSD after Kabsch alignment.
• Maximum absolute pLDDT difference.
• Maximum absolute PAE difference.
• Ranking-score difference.

using this script.

Identical-code GPU runs showed an approximately 0.20–0.22 Å RMSD drift because of diffusion and nondeterministic GPU reductions. The end-to-end acceptance threshold was therefore calibrated to 0.30 Å, with:

• max |ΔpLDDT| ≤ 1.0
• max |ΔPAE| ≤ 0.5
• |Δranking| ≤ 0.01

The complete candidate passed at both ends of the tested range:

Input tokens	RMSD	max ΔpLDDT	max ΔPAE	Δranking
512	0.21 Å	0.74	0.19	1.7e-5
2048	0.20 Å	0.69	0.27	2.4e-4

The key_lengths change was additionally tested against the immediately preceding optimized implementation and passed the stricter 0.1 Å threshold:

• Bucket-aligned input: 0.0295 Å
• 400 real tokens in a 512 bucket: 0.0293 Å

The de-subbatching change was tested to ensure operator equivalence, using shared parameters and identical inputs.
Chunked and full-batch GridSelfAttention produced max|Δ| = 0. with this script.

[Steboss]

@Augustin-Zidek is there anything we should consider in the setup to have this better tested? thank you

[Augustin-Zidek]

Hey, sorry for the late reply...

This looks great, we will do some internal testing on this and merge if things look good.

[Steboss]

Hey, sorry for the late reply...

This looks great, we will do some internal testing on this and merge if things look good.

No worries Augustin. Thanks very much indeed