bios-labs/RFD4-Proteina-MLX-fp16

RFD4-Proteina on Apple silicon (MLX)

Design protein binders with RFD4-Proteina entirely on a MacBook: target structure in, ranked and co-fold-screened binders out. No CUDA and no cloud. After the one-time downloads, this repo's commands need no network (set HF_HUB_OFFLINE=1 to enforce it for Hugging Face lookups; python -m rfd4_mlx.prep --fetch can download a structure from RCSB/AFDB, and only if you pass --fetch).

This repo is an MLX/Metal implementation of the model's recurrent transformer trunk (the part that takes 96% of generation time), wired into the unmodified upstream sampler, plus a local pipeline around it: target preparation, generation, cheap geometric filters, Boltz-2 co-folding on the Apple GPU, ranking and clustering.

Unofficial community port, maintained by BIOS Labs. Not affiliated with, reviewed by, or endorsed by the Baker Lab / Institute for Protein Design (UW) or NVIDIA. All biology comes from their model; this repo changes where and how fast it runs. Everything here is in silico; nothing has been tested in a wet lab. The wet-lab results in the RFD4-Proteina preprint belong to the original CUDA implementation.

What you get

Feature	What you get
Faster generation	20 binders (80 aa, GLP-1R ECD target) in ≈ 6 min 55 s on an M5 Pro, 2.2× faster (95 % CI 1.8–2.7) than stock PyTorch-MPS at upstream's default precision (bf16-mixed) and 3.2× (CI 2.9–3.6) faster than stock fp32; one machine, see Results for what "stock" means
A complete local binder pipeline	python -m rfd4_mlx.binder: raw PDB → 20 designs → Boltz-2 co-fold on MPS → pass/fail, ranked, diversity-clustered top designs in 12 min 04 s (one run; needs no network after the one-time downloads; BINDER_PIPELINE.md)
43 single-sample feature tests	functional checks (not success rates) across monomers, PPI, motifs, enzymes, ligands, DNA/RNA, inverse folding, partial diffusion, beam search, LoRA: zero fallbacks to the torch trunk in every successful run (FEATURES.md)
Raw data and negative results	raw data, negative results (including a shortcut we tested and removed: skipping pair updates on recycles destroyed the designs) and every number's provenance (BENCHMARKS.md, benchmarks/, paper/preprint.pdf)

Stage times are from one idle-GPU run (generation includes startup). A separate idle-GPU 20-design run took 6 min 59 s wall-clock including startup, against 6 min 17 s (377 s) for the generation stage here; we treat the difference as run-to-run variability (thermal state) and did not separate it further. "13 passed" means 13 designs met the on-device criteria (an in-silico screen, not a binding rate); 10 diverse ones were exported.

Quick start

Requirements.

• Apple-silicon Mac; 16 GB+ recommended (about 6.5 GB peak process footprint measured during generation at N = 180, plus about 5 GB for a Boltz-2 co-fold, run one after the other). 8 GB Macs: not tested.
• macOS: MLX needs macOS 14+. Every measurement here is on macOS 26.4.1 and an M5 Pro; other macOS versions and chips were not tested. The only part gated on the GPU generation is the hand-written Metal tensor-op kernel: the code enables it when the GPU architecture is applegpu_g17 or newer (M5-class) and a startup self-test passes, and otherwise uses MLX's own GEMMs. The code does not check the macOS version.
• On M1-M4 (or if the self-test fails) the tensor-op kernel is disabled and the speed-up over PyTorch-MPS has not been measured.
• Disk: about 9 GB, measured in a clean-room install (Python environment 1.8 GB, Boltz-2 home ~/.rfd4_mlx 5.2 GB, Hugging Face cache 1.6 GB, this repo 0.36 GB).
• Time: first-time setup is about 30-45 min, mostly downloads, and depends on your bandwidth (clean-room install of this repo 3 min, upstream checkpoint 9 min, Boltz-2 setup 27.7 min on our connection). One 20-design run takes about 12 min end to end on an M5 Pro.
• Tools: install uv first (brew install uv or curl -LsSf https://astral.sh/uv/install.sh | sh). uv venv --python 3.12 uses a Python 3.12 if one is installed and otherwise downloads one (uv's documented behavior; we have not tested a Mac with no Python at all). Xcode Command Line Tools / git: a few dependencies (ihm, antlr4-python3-runtime, assertpy) were built from source during our install, which ran on a Mac that has the Command Line Tools; we did not test a Mac without them, and nothing in this repo's own steps invokes git.
• Network: setup needs it. setup_oracles boltz creates a second virtual environment (in ~/.rfd4_mlx) and downloads Boltz-2; uv pip install ./rfd4-mlx pulls rfd4-proteina==0.0.3 and other dependencies from PyPI. The upstream model nvidia/RFD4-Proteina is a public Hugging Face repository (not gated; no login or token needed, checked with the unauthenticated Hub API on 28 Oct 2026).

uv venv --python 3.12 .venv && source .venv/bin/activate   # uv: see Requirements
uv pip install huggingface_hub                   # provides the `hf` command
export HF_HUB_DISABLE_XET=1                      # in our tests the Xet backend crawled (~3 KB/s); plain HTTP ran at MB/s
hf download bios-labs/RFD4-Proteina-MLX-fp16 --local-dir rfd4-mlx
uv pip install ./rfd4-mlx                        # also installs rfd4-proteina==0.0.3 and mlx
hf download nvidia/RFD4-Proteina model.safetensors config.json LICENSE release-checkpoints.json provenance.json inference.yaml   # once: upstream checkpoint (1.7 GB) into the Hugging Face cache
python -m rfd4_mlx.setup_oracles boltz           # once: local Boltz-2 co-fold oracle (≈ 4 GB weights, MIT)

# your target -> ranked binders (residue numbers in the file's own numbering)
python -m rfd4_mlx.binder --target-file my_target.pdb --target-sel A:29-128 \
       --binder-length 80 --auto-hotspots --n-designs 20 --out out/

Notes on the command. Numbering: --target-sel uses the file's own (author) residue numbers. In the shipped example examples/glp1r_binder80.yaml the same residues appear as A10-109, because that spec uses label_seq_id numbering: author residues A:29-128 of 3IOL are label_seq_id 10-109. Hotspots: --auto-hotspots picks a surface-hydrophobic patch by solvent-accessible surface area; it is not an epitope or binding-site prediction. If you know the site, give it yourself with --hotspots A:56,A:60. Binder length: a single integer; 20-250 is accepted, only 60-100 has been tried (untested outside it). Other targets: the 66 % pass rate below is for the GLP-1R ECD only; on another target run --n-designs 20 first and read report.md before scaling up.

Outputs: ranked.csv (all designs, all metrics; file paths relative to --out), report.md, top/*_design.cif and top/*_cofold.cif. Runs are resumable (an existing run is reused; a different --n-designs than it holds stops with a message, see BINDER_PIPELINE.md) and co-folds are cached. With --oracle none the designs are not co-fold screened, and report.md says so. Only generation, no co-fold: add --oracle none, or, from rfd4-mlx/examples, run python -m rfd4_mlx.generate inputs=glp1r_binder80.yaml num_replicates=20 out_dir=./out; it accepts every rfd4-generate argument. The upstream checkpoint (1.7 GB, downloaded once in the setup above) is read from the local Hugging Face cache; if it is missing and the network is allowed, the first generate downloads it (this repo redistributes only the denoiser trunk, as an optional fp16 copy).

Working offline. After the one-time setup above, export HF_HUB_OFFLINE=1. The sampler looks for nvidia/RFD4-Proteina in the local Hugging Face cache first; with HF_HUB_OFFLINE=1 it never goes online and, if the checkpoint is missing, stops and prints the exact command to run once with network access: hf download nvidia/RFD4-Proteina model.safetensors config.json LICENSE release-checkpoints.json provenance.json inference.yaml. Boltz-2 is always run offline. (python -m rfd4_mlx.generate --help does not touch the checkpoint.) This is not a claim about every package in the environment, only about what this repo's commands do. A clean-room test ran the full pipeline on 2 GLP-1R designs under macOS sandbox-exec with outbound network denied, with and without HF_HUB_OFFLINE, and it completed (exit code 0); that is 2 designs on one machine.

Tests: cd rfd4-mlx && uv pip install pytest && pytest tests -q (the tests that need the upstream checkpoint or the Metal GPU are skipped with a reason when those are absent; optional overrides RFD4_BUNDLE = checkpoint directory, RFD4_EXAMPLE_CIF = example structure). Debug knob: RFD4_MLX_DTYPE=fp16|bf16|fp32 (default fp16, see BENCHMARKS.md).

Results (MacBook Pro, M5 Pro 20-core GPU, 48 GB, macOS 26.4.1)

Speed. 20 × 80-residue binders vs GLP-1R (PDB 3IOL), upstream's default sampler settings (60 steps, medium/front_loaded recycling), alternating-block and contention-filtered. "Stock" means upstream's own PyTorch code run unchanged on MPS at upstream's default precision (bf16-mixed; fp32 for the second column), with no torch.compile or tuning (torch attention and triangle-multiplication backends, batch 1, because cuEquivariance is CUDA-only). The port's torch-side parts (features, sampler, heads, decoder) run at fp32 in the rfd4_mlx.generate wrapper, not at bf16-mixed, so the port column is not "upstream defaults" in precision; the trunk itself runs in fp16 under MLX.

metric	stock PyTorch-MPS bf16-mixed	stock PyTorch-MPS fp32	this port
s / design (95 % CI)	46.5 (38.3–54.3)	66.8 (62.0–70.0)	20.8 (18.4–21.7)
20 designs	≈ 16 min	≈ 22 min	≈ 6 min 55 s
speed-up of this port (95 % CI)	2.2× (CI 1.8–2.7)	3.2× (CI 2.9–3.6)	–

Baselines rest on 6 and 5 intervals (bf16-mixed, fp32); one M5 Pro. We did not separate which design choices (compile, fp16, fusion, kernel) produce the 2.2×; only the Metal kernel was ablated.

The 20-design time is 20 × the unrounded per-design time (20 × 20.75 s = 415 s = 6 min 55 s; the table shows the rounded 20.8 s). A complete idle-GPU 20-design run measured 6 min 59 s wall-clock including startup. The bootstrap CIs treat intervals as independent and rest on only 5 / 6 / 14 intervals (fp32 / bf16-mixed / MLX), so they understate the uncertainty.

• Per-cycle speed-ups (cycles 1/2/3): 3.20/3.53/3.13× vs fp32 and 2.15/2.23/2.38× vs bf16-mixed. Cycle 1 overlapped another GPU job and is excluded from the headline; cycles 2 and 3 are the contention-free ones.
• The machine slows by ≈ 25 % under 15 minutes of sustained load, so only compare numbers inside a paired session.
• Protocol, thermal behavior, scaling by length and approximate times for other Macs (projections, not measurements: ≈ 24 min for 20 designs on an M4 Pro) are in BENCHMARKS.md.

For scale, the same upstream code on one NVIDIA L40S took 7.0 s/design (untuned, including startup, no cuEquivariance), so a datacenter GPU is faster than this Mac; the point is that the model now runs locally.

Quality against CUDA (same task, n = 100 designs per arm, in silico):

(n = 100 per arm)	CUDA bf16	CUDA fp32	MLX fp16
ESMFold designability (binder alone: Cα-RMSD < 2 Å and pLDDT > 70)	81%	82%	86%
Boltz-2 complex (target MSA, binder single-sequence; cloud GPU), designed pose reproduced (< 5 Å)	64%	68%	75%

MLX minus CUDA bf16: designability +5 points (95 % CI −5 to +15; Fisher p = 0.45); pose +11 points (95 % CI −2 to +23; Fisher p = 0.12), and Mann-Whitney p = 0.06 for the pose RMSD itself. Against fp32: designability +4 points (95 % CI −6 to +14; Fisher p = 0.56); pose +7 points (95 % CI −5 to +19; Fisher p = 0.35). n = 100 per arm resolves only gaps of ≈ 15 (designability) to ≈ 19 (pose) points, so the data cannot exclude a deficit of a few points: no detectable degradation, not equivalence. The trunk is a re-implementation in fp16, so it is not bit-identical to CUDA; against a float64 reference its CA-velocity error is ≈ 4–14× smaller per step (≈ 10× on the mean) than torch bf16 autocast on MPS (upstream's default precision) on 4 replayed steps. The error is not uniform: in a full 60-step trace, 1 step in 60 had a ~10× larger error (5.7e-03 against a median of 6.0e-04, mean absolute error of the single representation): on that step the relative error grew from 5.7e-04 after the first pass to 4.2e-02 after 6 recycles. The spike is not a catastrophic failure (no overflow or non-finite values) but single trajectories can diverge between precisions; details in BENCHMARKS.md.

The on-device co-fold oracle. Upstream scores designs with AF3/OpenFold3 containers that need CUDA. Here Boltz-2 runs on the Apple GPU with the target supplied as a structural template and no MSA (without the template it under-scores binders badly). On 100 designs plus 30 residue-shuffled controls (one target): 66/100 designs pass (ipTM ≥ 0.8, mean interface PAE ≤ 6.0 Å, pose within 5.0 Å; 95 % CI 56–75%) and 0/30 controls pass; ipTM separates the two with AUC 0.93 (local template protocol, n = 30 shuffled controls). The cloud protocol used for the table above (target MSA, no template) separates poorly (AUC 0.63, n = 20 shuffled per set), which is why the template matters. Per-design agreement with cloud Boltz-2 scores is moderate (Spearman 0.74 ipTM), so treat it as a screen whose thresholds were checked on one target, not a replica. About 19 s per complex in batches.

What these counts are. 66/100 (and 13/20 in the pipeline run above) are in-silico pass counts for one target under Boltz-2 with the target structure given as a template; they are not binding or hit rates. Because the co-fold is given the same target structure as a template, it cannot detect a wrong target conformation.

What runs on a Mac

Status	Features
Runs, tested	unconditional and conditional monomers, PPI binders, motif scaffolding (indexed, unindexed, internal-coordinate), enzymes, small-molecule and DNA binders, RNA, fold-class and sequence conditioning, inverse folding, folding, partial diffusion, seeds, ODE/SDE, trajectories, structure-only beam search, LoRA, other checkpoints, rfd4-evaluate with the local oracle
Not accelerated; falls back to the torch trunk automatically (untested)	training, gradient guidance through the trunk
Needs Linux/CUDA	upstream AF2/AF3/OpenFold3 containers, PyRosetta/TMol, bioir backend, cuEquivariance. FoldSeek: a macOS binary exists but is not installed or tested here
Upstream bugs found	classifier-free guidance assert (proteina.py:555), CA-only trajectory export crash; details and workarounds in FEATURES.md

How it works

Each of the 60 flow-matching steps runs the same trunk three times: 14 attention blocks with pair bias and 7 pair-representation updates (triangle multiplications) on an N × N × 384 tensor. The port compiles one MLX graph per pass, runs it in fp16, folds the pair-bias norm into the projection, fuses projections, and on M5-class GPUs uses a hand-written Metal 4 tensor-op GEMM (bit-identical to MLX's in every test; end-to-end gain roughly 0–10 %: the one paired test measured 1.11×, 95 % CI 1.01–1.22, and the gain is not visible in two longer runs). Features, sampler, heads and decoder stay in PyTorch/MPS and are upstream's unmodified code (the rfd4_mlx.generate wrapper runs that torch side at fp32 by default, not at upstream's bf16-mixed default). The trunk takes its weights from the loaded model (LoRA merged in memory), and a guard routes any unsupported model, training call, gradient call or non-finite output back to the torch trunk with a warning. Details: BENCHMARKS.md; things that did not work are in benchmarks/negative_results.md.

Related work

This repo is an MLX re-implementation of the RFD4-Proteina trunk with published speed and fidelity comparisons against stock PyTorch-MPS and CUDA (searched 27 Oct 2026; other efforts exist for related models). Other Mac efforts we found: javierbq/rfd3-mlx is an MLX port of RFdiffusion3, and Boltz-2 itself supports MPS. Upstream states that macOS is unsupported for RFD4 (see the install notes in the upstream repository).

Limitations

• In silico only. Co-folding predictors have false positives; passing these filters is not evidence of binding or function.
• Validation depth. Statistical comparison with CUDA and the oracle threshold check cover one task (80-residue binders vs GLP-1R ECD). Other features have one-design functional tests, not success rates. Recheck the oracle thresholds on your target. The pipeline also ran end to end on a second target in a clean-room test (1YCR / MDM2, 4 designs: 3 passed the cheap prefilter, 2 of those 3 passed the oracle; the hotspot and --oracle none paths also ran): that shows the code path works, it is not a success rate.
• Boltz-2 is not AF3. Upstream's AF3-calibrated pass criteria (ptm_coi, min_ipAE, binder RMSD) are not reproduced by Boltz-2; read rfd4_mlx.evaluate output as Boltz-2 numbers.
• Hybrid, one machine. PyTorch is still required for the non-trunk parts; timings are from one M5 Pro; other chips are projections.
• Low recycle preset (low/front_loaded, ≈ 1.4× faster) raises binder–target clashes on both MLX and CUDA.

Repo contents

rfd4_mlx/ (hook, layers, generate/binder/prep/evaluate CLIs, oracles) · weights/rfd4_trunk_fp16.safetensors (optional verified fp16 copy of the denoiser trunk, 178.4 M params; the hook builds weights from the loaded model and ignores a file that does not match) · examples/ · tests/ · benchmarks/ · designs/ (binder sequences and CA coordinates of every evaluated set, plus 100 full MLX designs) · paper/ · scripts/, experiments/, modal/ (provenance only, not portable tools; modal/ was used only for the CUDA reference generation and the cloud ESMFold/Boltz-2 evaluation, never for inference of this port) · SHA256SUMS.

Responsible use

This port does not change the model's outputs by design, but it makes running the model cheaper and more accessible; follow the upstream usage norms, the upstream model's terms and your institution's biosafety and biosecurity norms. Computational scores are screening filters, not evidence of function or safety.

License and citation

BSD-3-Clause. Redistributes a precision-converted subset (the denoiser trunk, fp16) of the RFD4-Proteina weights, © 2026 Institute for Protein Design, University of Washington (LICENSE, NOTICE.md). Boltz-2 (MIT), ESMFold and other tools keep their own licenses and are not redistributed. The port is maintained by BIOS Labs; author names and institutions of the original model are not used to endorse it. The upstream model card (nvidia/RFD4-Proteina, LICENSE) states BSD 3-Clause, Copyright (c) 2026 Institute for Protein Design, University of Washington (details in NOTICE.md).

@article{corley2026rfd4proteina,
  title  = {Biomolecular design through latent reasoning with {RFD4-Proteina}},
  author = {Corley, Nathaniel and Didi, Kieran and Kazman, Maxwell W. and Brent, Rafael I. and others},
  year   = {2026}, note = {Preprint}, url = {https://rfd4-proteina.org/assets/paper/RFD4-Proteina.pdf}
}

Built and verified against nvidia/RFD4-Proteina revision a06e2509c7033b010fc43614e9fd840303588502. The current main (48377e45eb25a4fce8d2aff56c1a659b726d2106, checked 28 Oct 2026) has the identical model.safetensors (sha256 5b549dd61adb2f6c6b5cb95bda13f40f5e31556b4e6c755edc0ca6f011c68fb0), config.json, LICENSE and release-checkpoints.json; only inference.yaml (one line: the hf_ckpt_revision pin) and provenance.json (one note line) differ, and the code overrides the pin, so unpinned downloads work. The clean-room install test ran on main. Cite the revision you use.

To cite this port:

@misc{biolabs2026rfd4mlx,
  title        = {{RFD4-Proteina} on Apple silicon ({MLX}/Metal port)},
  author       = {{BIOS Labs}},
  year         = {2026},
  howpublished = {Hugging Face model repository},
  url          = {https://huggingface.co/bios-labs/RFD4-Proteina-MLX-fp16}
}

The accompanying technical report is paper/preprint.pdf: Running RFD4-Proteina on a laptop: an MLX/Metal trunk for latent-reasoning protein design on Apple silicon (unofficial).