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

Enable AlphaFold 3 inference on the Apple Silicon integrated GPU via the
community jax-mps Metal plugin for JAX. The NVIDIA CUDA path is unchanged
and remains the reference; MPS is an opt-in, experimental feasibility path.

run_alphafold.py:
  * Add --jax_backend {gpu,mps} to explicitly select the accelerator backend
    (default 'gpu' preserves existing behavior).
  * Add the portable, linear-memory 'xla_chunked' flash-attention option
    alongside 'xla'; both are the only implementations valid on CPU/MPS
    (triton/cudnn remain NVIDIA-only).
  * Refactor device selection/validation into _select_inference_device and
    _validate_inference_device: CPU and MPS require xla/xla_chunked; the CUDA
    compute-capability checks (>=6.0, 7.x XLA_FLAGS workaround) are preserved
    unchanged; unvalidated accelerators are rejected with actionable errors.

run_alphafold_device_test.py:
  * New unit tests (31 cases) covering backend selection, index bounds,
    attention validation, and explicit CUDA-path-preservation checks.

requirements/apple-silicon-gpu.txt:
  * Pinned Metal backend (jax==0.10.2, jaxlib==0.10.2, jax-mps==0.10.9),
    installed as a post-install override of the base jax==0.9.1 pin.

docs/installation_apple_silicon_gpu.md:
  * Native (non-Docker) macOS build and run guide, including HMMER build,
    build_data, the MPS backend install, run flags, and measured
    memory/performance guidance from a 52-target validation campaign.

Author: SuhasSrinivasan

docs/installation.md

@@ -496,6 +496,13 @@ Since JAX doesn't support running natively on Mac GPU as of 2026, you have to
 resort to running AlphaFold 3 in the slow CPU-only mode even though it has a GPU
 (`jax-metal` is unfinished as of July 2026).
 
+> **Experimental Apple Silicon GPU (Metal / MPS) inference.** As an alternative
+> to CPU-only mode, AlphaFold 3 inference can run on the Apple Silicon GPU
+> through the community `jax-mps` Metal plugin. This is an unofficial,
+> experimental feasibility path (not numerically certified), but it is
+> substantially faster than CPU for small and medium structures. See
+> [installation_apple_silicon_gpu.md](installation_apple_silicon_gpu.md).
+
 1.  Download all required databases and AlphaFold 3 weights (see above).
 2.  Install the [HMMER Suite](http://hmmer.org/). See
     http://hmmer.org/documentation.html for installation instructions.

docs/installation_apple_silicon_gpu.md

@@ -0,0 +1,243 @@
+# Running AlphaFold 3 on Apple Silicon GPU (Metal / MPS)
+
+> **Status: experimental feasibility path.** This mode runs AlphaFold 3
+> inference natively on the Apple Silicon integrated GPU through the community
+> [`jax-mps`](https://pypi.org/project/jax-mps/) Metal plugin for JAX. It is
+> **not** an officially supported or numerically certified configuration. The
+> NVIDIA CUDA path is unchanged and remains the reference. Use this for research,
+> local development, and prototyping on Macs — not for production or benchmarking
+> against the published CUDA results.
+
+This guide describes a **native build with no Docker**. Docker Desktop on macOS
+runs a Linux VM that cannot see the Mac GPU, so the containerized recipe in
+[`installation.md`](installation.md) can only ever reach the CPU. To use the
+Apple GPU you must build AlphaFold 3 natively, as described below.
+
+If you only want CPU-only inference on a Mac (much slower, but officially
+described), follow the "Running AlphaFold 3 without a GPU" section of
+[`installation.md`](installation.md) instead.
+
+## What works and what to expect
+
+Validated in July 2026 across a 52-target campaign (all 12 molecule-type classes:
+monomers, homo/hetero oligomers, protein–DNA, protein–RNA, and ligand/ion
+complexes, plus token-scaling targets up to ~2,200 tokens) on an **M2 Ultra Mac
+Studio (24-core CPU, 76-core GPU, 192 GB unified memory)**, macOS 15.
+
+- **Correctness:** the full data pipeline (HMMER MSA + templates) and inference
+  run end to end; predictions match the CUDA reference closely for the cases
+  tested. Inference is **bit-reproducible run-to-run** with fixed seeds.
+- **Attention:** only the portable `xla` and `xla_chunked` implementations work
+  on Metal. `triton` and `cudnn` require NVIDIA GPUs and are rejected early.
+- **Memory is the main limit** (see [Performance and memory](#performance-and-memory)).
+  Peak GPU (unified) memory grows with the model's padded token bucket and
+  reaches ~167 GiB at ~2,200 tokens on a 192 GB machine — that is the practical
+  single-structure ceiling on this hardware.
+- **Large targets are slow.** Per-seed inference scales roughly cubically with
+  tokens and hits a steep Metal-specific penalty above ~1,500 tokens.
+
+## Prerequisites
+
+- Apple Silicon Mac (M1/M2/M3/M4 family). More unified memory = larger tractable
+  structures; 64 GB is a reasonable floor, 128–192 GB recommended for large
+  complexes.
+- macOS 14 or newer.
+- Xcode Command Line Tools (provides `clang`, `make`, `patch`):
+
+  ```sh
+  xcode-select --install
+  ```
+
+- A recent CMake (≥ 3.28) and a Python 3.11–3.13 toolchain. The steps below use
+  [`uv`](https://docs.astral.sh/uv/); a Conda/Miniforge environment works
+  equally well (see [Alternative: Conda](#alternative-conda-environment)).
+
+## Step 1 — Install the HMMER suite (for the MSA data pipeline)
+
+HMMER is only needed if you run the data pipeline (i.e. not with
+`--norun_data_pipeline`). Prefer a user-local source build of HMMER 3.4 with the
+sequence-truncation patch that ships in this repository
+(`docker/jackhmmer_seq_limit.patch`); a plain HMMER build also works but is not
+the preferred configuration.
+
+```sh
+# Pick a location OUTSIDE the git checkout.
+export AF3_REPO="$PWD"                    # path to your alphafold3 clone
+export HMMER_PREFIX="$HOME/af3-tools/hmmer-3.4"
+mkdir -p "$HOME/af3-tools/hmmer-source" && cd "$HOME/af3-tools/hmmer-source"
+
+curl -LO http://eddylab.org/software/hmmer/hmmer-3.4.tar.gz
+echo "ca70d94fd0cf271bd7063423aabb116d42de533117343a9b27a65c17ff06fbf3  hmmer-3.4.tar.gz" | shasum -a 256 -c
+tar -xzf hmmer-3.4.tar.gz
+cd hmmer-3.4
+patch -p0 < "$AF3_REPO/docker/jackhmmer_seq_limit.patch"
+./configure --prefix="$HMMER_PREFIX"
+make -j"$(sysctl -n hw.perflevel0.logicalcpu)"
+make install
+"$HMMER_PREFIX/bin/jackhmmer" -h | head -1   # sanity check
+```
+
+You will pass the binaries to `run_alphafold.py` in Step 6 via the
+`--jackhmmer_binary_path`, `--nhmmer_binary_path`, `--hmmalign_binary_path`,
+`--hmmsearch_binary_path`, and `--hmmbuild_binary_path` flags (or put
+`$HMMER_PREFIX/bin` first on your `PATH`).
+
+## Step 2 — Get the code, model weights, and databases
+
+1. Clone the repository (if you have not already):
+
+   ```sh
+   git clone https://github.com/google-deepmind/alphafold3.git
+   cd alphafold3
+   ```
+
+2. Obtain the model weights by following the process in the
+   [main README](https://github.com/google-deepmind/alphafold3). Keep the
+   weights **outside** the git checkout.
+3. Download the genetic databases (see [`installation.md`](installation.md) and
+   `fetch_databases.py`). These are large (hundreds of GB); place them on fast
+   local storage.
+
+## Step 3 — Build and install AlphaFold 3 natively
+
+This compiles the AlphaFold 3 C++/CMake extension and installs the Python
+package. Upstream already supports the `darwin arm64` build environment.
+
+```sh
+uv venv --python 3.12
+source .venv/bin/activate
+uv sync            # builds and installs alphafold3 (scikit-build-core + cmake)
+uv run build_data  # builds the CCD / chemical-component data used at runtime
+```
+
+`uv sync` installs the base dependencies, which pin `jax==0.9.1` (CPU wheel on
+macOS). That JAX cannot reach the Apple GPU — Step 4 replaces it.
+
+Verify the base install with the data-pipeline test:
+
+```sh
+uv run python run_alphafold_data_test.py
+```
+
+## Step 4 — Install the Apple Silicon GPU backend
+
+Replace the base JAX with the Metal-capable stack pinned in
+[`requirements/apple-silicon-gpu.txt`](../requirements/apple-silicon-gpu.txt):
+
+```sh
+uv pip install -r requirements/apple-silicon-gpu.txt
+# equivalently: uv pip install jax==0.10.2 jaxlib==0.10.2 jax-mps==0.10.9
+```
+
+Confirm JAX now sees the Metal device:
+
+```sh
+python -c "import jax; print([d.platform for d in jax.local_devices()])"
+# -> ['mps']   (a one-time 'Platform mps is experimental' warning is expected)
+```
+
+`pip check` will report exactly one warning that AlphaFold 3 pins `jax==0.9.1`.
+That single mismatch is expected; any other mismatch means the environment is
+inconsistent.
+
+## Step 5 — Verify device selection
+
+Run the launcher's device-selection unit tests (fast, no GPU needed):
+
+```sh
+uv run python run_alphafold_device_test.py
+```
+
+These confirm the MPS backend is selectable and — importantly — that the CUDA
+path is unchanged.
+
+## Step 6 — Run inference on the Apple GPU
+
+Select the Metal backend with `--jax_backend mps` and a portable attention
+implementation. `xla` is the validated choice; `xla_chunked` uses less memory
+but is slower.
+
+```sh
+uv run python run_alphafold.py \
+  --json_path=/path/to/fold_input.json \
+  --model_dir=/path/to/weights \
+  --db_dir=/path/to/databases \
+  --output_dir=/path/to/output \
+  --jax_backend=mps \
+  --gpu_device=0 \
+  --flash_attention_implementation=xla \
+  --jackhmmer_binary_path="$HMMER_PREFIX/bin/jackhmmer" \
+  --nhmmer_binary_path="$HMMER_PREFIX/bin/nhmmer" \
+  --hmmalign_binary_path="$HMMER_PREFIX/bin/hmmalign" \
+  --hmmsearch_binary_path="$HMMER_PREFIX/bin/hmmsearch" \
+  --hmmbuild_binary_path="$HMMER_PREFIX/bin/hmmbuild"
+```
+
+For inference on a pre-computed input (skipping the MSA pipeline and HMMER), add
+`--norun_data_pipeline` and drop the HMMER flags.
+
+## Performance and memory
+
+Measured on the M2 Ultra (192 GB). Numbers are per diffusion sample ("seed") and
+scale with the **padded token bucket**, not the exact token count. Treat them as
+order-of-magnitude guidance; other Apple Silicon chips will differ.
+
+| Padded tokens | Peak GPU (unified) memory | Inference time / seed |
+| ------------: | ------------------------: | --------------------: |
+|           256 |                  ~36 GiB  |                 ~29 s |
+|           512 |                  ~50 GiB  |                 ~90 s |
+|           768 |                  ~73 GiB  |                ~3.9 min |
+|         1,024 |                 ~101 GiB  |                ~8.6 min |
+|         1,536 |                 ~155 GiB  |               ~26.6 min |
+|        ~2,180 |    ~167 GiB (+ ~11 GiB swap) |             ~3.9 hours |
+
+Key points:
+
+- **Unified memory is the ceiling.** GPU tensors live in unified memory and are
+  reported by macOS as "Memory Used"; they are **not** visible as process RSS
+  (the Python process itself stays near ~7 GiB). Do not rely on an RSS-based
+  memory limiter to cap GPU use. On a 192 GB machine, ~2,200 tokens
+  (~167 GiB + light swap) is the practical single-structure limit; larger
+  structures will run out of memory. Scale your expectation to your Mac's memory.
+- **Inference cost is roughly cubic** in tokens (triangle attention dominates),
+  with an additional steep Metal-specific penalty beyond ~1,500 tokens. Small to
+  medium structures (≤ ~1,000 tokens) are comfortable; multi-thousand-token
+  complexes are expensive. Use `xla_chunked` to trade speed for lower memory.
+- **The MSA / data pipeline runs on the CPU** and its cost depends on the number
+  of unique chains and whether RNA is present (≈ 7–35 min in the campaign), not
+  on token count. Set the HMMER thread counts so that
+  (concurrent searches × threads-per-search) ≈ the number of performance cores
+  to avoid oversubscription (`sysctl hw.perflevel0.logicalcpu`).
+- **No persistent JAX compilation cache on Metal.** `--jax_compilation_cache_dir`
+  does not persist across processes on the MPS backend, so every process pays the
+  compile cost for each new token bucket. Amortize it by folding many inputs in
+  one process (e.g. `--input_dir`) rather than launching one process per target.
+- **Thermals** were nominal on the Mac Studio throughout a ~40 h campaign (no
+  throttling). Laptops with less cooling headroom will behave differently.
+
+## Alternative: Conda environment
+
+If you prefer Conda/Miniforge over `uv`, the same result is achieved with:
+
+```sh
+conda create -n af3-mps python=3.13 cmake -y
+conda activate af3-mps
+pip install --no-build-isolation -e .   # or: pip install .
+build_data
+pip install -r requirements/apple-silicon-gpu.txt
+```
+
+Then run `run_alphafold.py` with the same MPS flags as in Step 6.
+
+## Troubleshooting
+
+- **`No local MPS device was found for --jax_backend=mps`** — `jax-mps` is not
+  installed or JAX was not upgraded. Re-run Step 4 and re-check
+  `jax.local_devices()`.
+- **`--flash_attention_implementation must be set to "xla" or "xla_chunked"`** —
+  `triton`/`cudnn` are NVIDIA-only; use `xla` (or `xla_chunked`) on Metal.
+- **Out-of-memory / heavy swapping on large targets** — reduce the structure
+  size, switch to `--flash_attention_implementation=xla_chunked`, or run on a Mac
+  with more unified memory. See the memory table above.
+- **`pip check` reports a jax version mismatch** — one warning about the
+  `jax==0.9.1` pin is expected (see Step 4); investigate any other mismatch.

requirements/apple-silicon-gpu.txt

@@ -0,0 +1,14 @@
+# Experimental AlphaFold 3 Apple Silicon GPU (Metal / MPS) backend.
+#
+# Install these *after* the base AlphaFold 3 install. The base project pins
+# jax==0.9.1 for the CUDA / CPU path; the experimental `jax-mps` Metal plugin
+# requires a newer JAX, so these pins upgrade JAX to 0.10.2. After installing,
+# `pip check` reports exactly one expected warning that AlphaFold 3 pins
+# jax==0.9.1 -- that single mismatch is expected and the MPS path is validated
+# with the versions below. Any *other* mismatch means something is wrong.
+#
+# Validated on macOS (Apple Silicon, M2 Ultra, 192 GB), Python 3.13, July 2026.
+# See docs/installation_apple_silicon_gpu.md for the full guide.
+jax==0.10.2
+jaxlib==0.10.2
+jax-mps==0.10.9