Add experimental Apple Silicon GPU inference support

[SuhasSrinivasan]

Add experimental Apple Silicon GPU (Metal / MPS) inference support

Summary

This PR enables AlphaFold 3 inference on the Apple Silicon integrated GPU via the community jax-mps Metal plugin. It is an opt-in, experimental path: the default backend stays gpu (NVIDIA CUDA), and the CUDA code path is untouched. It complements the existing CPU-only macOS mode by giving Mac users a much faster GPU option for small/medium structures. All changes are additive and gated behind a new flag.

Motivation: Current docs state Mac users must fall back to CPU because "jax-metal is unfinished." jax-mps is a different, working Metal plugin — this PR wires it in and backs it with a 52-target validation campaign (below).

Native build: Docker Desktop on macOS runs a Linux VM that cannot access the Metal GPU, so the containerized recipe is CPU-only on Mac. This path builds AlphaFold 3 natively (the same scikit-build/CMake build already supported on darwin arm64, plus build_data) and installs the MPS backend directly — no Docker required. See docs/installation_apple_silicon_gpu.md.

What changed (5 files, +650 / −56)

File	Change
run_alphafold.py	New --jax_backend {gpu,mps} (default gpu); new xla_chunked attention option; device selection/validation refactored into _select_inference_device / _validate_inference_device
run_alphafold_device_test.py	New — 31 unit tests for backend selection + validation
requirements/apple-silicon-gpu.txt	New — pinned Metal backend stack
docs/installation_apple_silicon_gpu.md	New — native (non-Docker) macOS GPU install + run guide
docs/installation.md	+7 lines — pointer to the new guide from the macOS section

Backward compatibility (nothing NVIDIA changes)

Aspect	Status
Default backend	gpu (CUDA) — unchanged
CUDA compute-capability checks (≥6.0, 7.x XLA_FLAGS workaround)	Preserved, byte-for-byte behavior
Model / weights / data pipeline	Untouched
triton / cudnn attention	Unchanged; still required on NVIDIA
Unit tests	31/31 pass, incl. explicit CUDA-path-preservation cases

New flags / usage

Flag	Values	Notes
--jax_backend	gpu (default), mps	mps requires xla/xla_chunked; --use_cpu_only overrides
--flash_attention_implementation	triton, cudnn, xla, xla_chunked (new)	xla_chunked = portable, linear-memory, slower

MPS invocation: --jax_backend=mps --flash_attention_implementation=xla

Backend dependency stack

Installed as a post-install override of the base jax==0.9.1 pin (one expected pip check warning); no change to uv.lock or the default install.

Package	Version
jax	0.10.2
jaxlib	0.10.2
jax-mps	0.10.9
Python	3.11–3.13 (tested 3.13)

Validation campaign

Item	Value
Hardware	M2 Ultra Mac Studio — 76-GPU-core, 192 GB unified memory, macOS 15
Targets	52 (12 molecule types × 4 + 4 token-scaling 1.5k–2.2k)
Pipeline	Full HMMER MSA + templates → 5-seed MPS inference
Result	52/52 completed, 0 failures, 40.7 h wall-time, 0 timeouts / OOM-kills
Reference	Verified vs experimental PDB structures

Accuracy (best of 5 seeds, vs PDB)

Overall (n = 52)

Metric	Result
Per-chain Cα RMSD (fold quality)	median 1.23 Å — 21 sub-1 Å, 42 sub-2 Å, 46 sub-3 Å
Complex Cα RMSD (quaternary)	median 1.62 Å
Ligand RMSD (pocket-aligned)	median 0.32 Å
Nucleic-acid backbone RMSD	median 0.94 Å
AF3 has_clash	0 / 52
pTM / ipTM	median 0.88 / 0.87
Confidence calibration (pTM vs complex RMSD)	r = −0.72
Determinism (same-seed re-run)	bit-identical, 0.000 Å

By molecule type (median). Read per-chain as fold quality; a high complex value with low per-chain = quaternary arrangement of a flexible assembly, not a fold error.

Type	n	per-chain Cα (Å)	complex Cα (Å)	DockQ	Note
hetero-trimer	4	0.56	1.65	0.96
protein-DNA-ligand	4	0.70	0.79	0.95
polymer-ligand	4	0.81	0.93	0.88
protein-RNA	4	0.94	12.28	0.05	RNA placement harder (subunits fold well)
hetero-dimer	4	0.99	1.51	0.89
protein-protein-ligand-ion	4	1.18	7.38	0.41
monomer-ligand	4	1.23	1.23	–
homo-polymer	4	1.28	16.79	0.03	quaternary/oligomer arrangement
protein-RNA-ligand	4	1.33	1.32	0.50
hetero-tetramer	4	1.48	16.34	0.79	quaternary
protein-DNA	4	1.93	1.93	–
large (1.5–2.2k tok)	4	2.10	9.48	0.86
monomer	4	11.18	11.18	–	2/4 genuine misses, both low-pTM (AF3-flagged)

Highlight: nitrogenase 3U7Q (2,177 tokens, FeMoco metal cluster) — per-chain 0.42 Å, complex 0.47 Å, DockQ 0.98, ligand 0.13 Å, pTM 0.96.

Performance and memory (M2 Ultra, 192 GB)

Padded tokens	MSA time¹	Inference / seed²	Peak unified memory (GB)²
256	~7 min	~29 s	36
512	~13 min	~90 s	50
768	~13 min	~3.9 min	73
1,024	~7 min	~8.6 min	101
1,536	~19 min	~26.6 min	155
~2,180	~20 min	~3.9 h	167

¹ MSA / data pipeline is CPU-bound and driven by the number of unique chains and RNA presence, not token count — hence non-monotonic across sizes (observed range 7–35 min). ² Inference time and peak memory are set by the padded token bucket. Inference scales ~polynomially (exponent ~1.7 → ~2.8) with a steeper Metal-specific penalty above ~1,536 tokens; unified memory (not process RSS) is the capacity limit — practical ceiling ≈ 2,200 tokens on 192 GB.

Known limitations

Limitation	Detail
Experimental backend	jax-mps Metal plugin; not numerically certified
Attention	Only xla / xla_chunked on MPS (triton/cudnn are NVIDIA-only)
JAX version	MPS needs jax 0.10.2 vs the base 0.9.1 pin (documented override)
Compile cache	--jax_compilation_cache_dir does not persist on Metal
Memory / speed	~2,200-token ceiling on 192 GB; steep time cost above ~1,500 tokens

The inference times are expected to be lower on Apple M5 series, because of the Neural Accelerator in each GPU core, which can speed up compute-bound AI tasks up to 4x.

Developed with help from Claude!

[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.

[SuhasSrinivasan]

Signed the CLA

[Augustin-Zidek]

Hello, thanks for the PR. Could you simplify it so that it contains only the minimal change to make it work? There are quite a few unrelated changes, e.g. the xla_chunked seems to be redunant.

Also, could you use the dependencies in pyproject.toml and make the jax-mps conditional instead of shipping them in an additional file?

[SuhasSrinivasan]

Thanks for the review — both points addressed.
I've force-pushed a squashed, single clean commit on top of main.

1. Minimized the change / dropped xla_chunked. run_alphafold.py is now +18 / −3:

• Removed xla_chunked, the _select_inference_device/_validate_inference_device refactor, and the extra device-selection test.
• What remains is just: a new --jax_backend {gpu,mps} flag (default gpu), one elif requiring --flash_attention_implementation=xla for MPS, and selecting the device from the chosen backend. The NVIDIA CUDA path (compute-capability checks, triton/cudnn) is unchanged.

2. Moved the backend into pyproject.toml, conditionally. Deleted requirements/apple-silicon-gpu.txt and made the Metal stack a platform-conditional dependency:
   "jax==0.9.1; sys_platform != 'darwin'",
   "jax==0.10.2; sys_platform == 'darwin'",
   "jaxlib==0.10.2; sys_platform == 'darwin'",
   "jax-mps==0.10.9; sys_platform == 'darwin' and platform_machine == 'arm64'",
   jax-mps requires jax >= 0.10.0, and tokamax/dm-haiku only require jax >= 0.9.1, so this resolves without conflict. uv.lock is regenerated (Linux keeps jax 0.9.1 + CUDA; darwin arm64 gets jax 0.10.2 + jax-mps).

The net diff is now run_alphafold.py, pyproject.toml, uv.lock, and the macOS install doc. Verified on an M2 Ultra that MPS inference selects MpsDevice(id=0) and runs.

[Augustin-Zidek]

Neat! Could you rebase on top of a3b8355? That should simplify the change even further.

I am going to test things on my side and I will send a PR that is inspired by this, but I want to refactor a few things (I think I will get rid of the --use_cpu_only flag and replace it with a --device flag with options gpu/cpu/mps).

However, I will make sure to mention you in the release notes and this PR is a great proof of concept that this is a path worth pursuing -- many thanks for sending it.

[SuhasSrinivasan]

Thanks again!
Shall I change the device flags as well?
Since cpu/gpu/mps are JAX's own backend names and then create a deprecated alias: --use_cpu_only mapping to --device=cpu with a warning?
I can run the tests too for the devices?

[Augustin-Zidek]

Shall I change the device flags as well?

No need.

Since cpu/gpu/mps are JAX's own backend names and then create a deprecated alias: --use_cpu_only mapping to --device=cpu with a warning?

I introduced the flag only last week, so it is OK to rename it, I don't think anyone depends on it yet.

But feel free to try with the proposed changes, these things are easy with agents these days after all. :)

[SuhasSrinivasan]

Rebased onto a3b8355 and squashed to a single commit.
Net change is now run_alphafold.py, pyproject.toml, uv.lock, plus the macOS install doc.

[Augustin-Zidek]

Hello, I've done some preparatory work for this in 98a15dc. Could you rebase on top of this commit?

[SuhasSrinivasan]

Thank you.
Rebased onto 98a15dc.
The change is now just an mps enum value + its two elif branches (run_alphafold.py), the conditional jax-mps dep, and the install doc.

[SuhasSrinivasan]

Thank you!