Replace two-einsum OuterProductMean with fused three-way einsum

[mooreneural]

Replaces the two-einsum OuterProductMean chunk computation with a single
three-way einsum, letting XLA contract the two C_outer dimensions before
summing over the MSA axis. Single file change, always-on.

The change

Standard path (~256 MB intermediate at N=1024, C_outer=32, chunk=128, bfloat16):

act = jnp.einsum('acb,ade->dceb', left_act, right_act)
act = jnp.einsum('dceb,cef->dbf', act, output_w) + output_b

Fused path (~72 MB peak):

act = jnp.einsum('acb,ade,cef->dbf', left_act, right_act, output_w)
return jnp.transpose(act, [1, 0, 2]) + output_b

Numerical testing

Verified on RTX 5080, JAX 0.10.2, bfloat16, N=1024 residues, MSA depth=512:

• Max absolute difference between standard and fused path: 0.0
• Memory reduction: 256 MB -> 72 MB (3.56x)
• Throughput: no regression (0.97x, within noise)

Files changed

• src/alphafold3/model/network/modules.py -- fused three-way einsum in compute_chunk

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[mooreneural]

@googlebot rescan

[mooreneural]

Manual testing: serialization bug fix (templateIndices)

Ran manual_test.py on the Alphafold4 branch (Python 3.11.9, Windows)

PASS Test 1: empty template map serializes as [] and round-trips
PASS Test 2: non-empty template map round-trips correctly
PASS Test 3: buggy null templateIndices fails on deserialize: 'NoneType' object is not iterable
PASS Test 4: folding_input.py contains the PR #674 fix (no 'or None')
SKIP Test 5: would just do the same thing as Tests 1–4 but using the real ProteinChain class instead of mirroring the logic inline. It's redundant, and it can't run on Windows anyway without building the C++ extension.

All manual serialization tests passed.

Test 3 reproduces the original crash, zip([], None) raises TypeError on the pre-fix serialization. Tests 1 and 2 confirm the fix: empty query_to_template_map now serializes as [] instead of null, so round-trip deserialization no longer crashes.

CLA is signed.

[mooreneural]

GPU memory optimization benchmark results on RTX 5080 (16 GB VRAM), JAX 0.10.2, bfloat16

OuterProductMean at N=1024 residues, MSA depth=512:

Path	Peak intermediate memory	Time (ms)
Standard (two einsums)	256 MB	3.37 ± 0.32
Fused (PR #674, 3-way einsum)	72 MB	3.46 ± 0.06

3.56x memory reduction with no throughput regression (0.97x, within noise).
Numerical equivalence confirmed: max absolute difference = 0.0 in bfloat16.

At larger sequence lengths (N=2048+) where the standard path's intermediate grows to ~1 GB+,
this reduction is what allows sequences to fit in GPU memory at all.

[mooreneural]

Thank you for pulling in the docstring fixes, @Augustin-Zidek I appreciate the acknowledgment. I'll keep the remaining items in this PR (serialization fix + GPU optimization) available for review whenever the team has bandwidth.

[Augustin-Zidek]

Hi, thanks for submitting this PR. What worries is the size: 1.7k lines code is it hard to review dilligently and also it would increas our maintenance burden (for comparison, the entire AlphaFold 3 source code is 24k lines of code, so this PR increases its size by 7%).

Could you update the PR to contain only the outer product change and nothing more (i.e. remove the typo fixes, unrelated parallel code and the unrelated XLA cache code)?

If I didn't miss anything, this change should be just a few line change in modules.py.

Moreover, how did you test this? Did you test that this doesn't change numerics?

[mooreneural]

Hi @Augustin-Zidek, I've stripped the PR down to just the OuterProductMean change; parallel.py, xla_cache.py, the benchmark script, and the serialization fix are all removed. The test files are also removed since AF3 doesn't use that pattern.

What remains (160 lines total):

• fused_ops.py; the three-way einsum implementation
• modules.py; lazy dispatch behind global_config.use_fused_outer_product_scan (default False)
• model_config.py + golden file; two new config fields

Numerical testing: verified on RTX 5080, JAX 0.10.2, bfloat16, N=1024 residues, MSA depth=512:

• Max absolute difference between standard and fused path: 0.0
• Memory reduction: 256 MB → 72 MB intermediate (3.56x)
• Throughput: no regression (0.97x, within noise)

[Augustin-Zidek]

Thanks, could you squash and rebase this on top of HEAD? There are some license changes and changes from other commits in this PR that should not be there.

Moreover, is there any reason to put this behind a flag and in a separate module? I.e. if it is the same speed and uses less memory, why not use that by default? I would also inline the implementation, so the change truly becomes just a single file change over a few lines of code.

[mooreneural]

@Augustin-Zidek, totally agree on both points, will squash to a single commit and rebase onto HEAD, and making it the default (no flag, inlined directly in modules.py) is the right call. The change is a strict improvement so there's no reason to gate it. Making it default should meaningfully extend what’s feasible on consumer and mid-range GPUs without any tradeoff.

I'll have the cleaned-up push ready in a few hours. Thanks for the continued engagement on this. I look forward to contributing to AF3.

[mooreneural]

@Augustin-Zidek
Done, squashed to a single commit on top of HEAD, removed the flag and separate module, inlined directly into compute_chunk in modules.py.

Re-ran numerical verification after inlining:

Output shape: (128, 1024, 128)
Max absolute difference: 0.0
Numerically equivalent: True

Single file change, 5 insertions, 3 deletions.

[Augustin-Zidek]

Hello Clay, I've now done this in 7b4c2a5, sorry for the delay. I got the code even simpler by removing the unnecessary transposes before and after the einsum. I will make sure to credit you in the release notes for v3.0.4.

Thanks again for this optimisation!

[mooreneural]

@Augustin-Zidek Wow! How exciting. Thank you so much for taking my suggestions into the next release.

I wanted to ask if you'd connect with me on LinkedIn https://www.linkedin.com/in/claynovel

Best,
Clay