[DimaMolod]

Summary

--fast_kernels off|on|auto now covers the AlphaFold 3 backend. It switches on the fork's opt-in fused triangle kernels (KosinskiLab/alphafold3#4). The default stays off, which runs exactly the previous code. AlphaFold 2's behaviour and API are unchanged.

• The AF3 policy has its own module, alphapulldown/prediction/af3_fused_triangles.py:
  • resolve(mode, device) checks the GPU before any weights load: an NVIDIA card with compute capability ≥ 8.0, a known tile table and memory budget, and a smoke compile-and-run of both fused blocks on that GPU.
  • on raises an error with the missing requirement; auto falls back with a log line.
  • Per card: Pallas attention on 8.0, 9.0 and 12.0; tokamax's attention core on 8.6 and 8.9; triangle multiplication on every card from 8.0. sm_100 is not enabled.
  • fast_kernels.py (AF2) is byte-identical to main; the two backends share only normalise_mode.
• Traceable: each AF3 output directory gets inference_kernels.json, keyed seed-N. It records the requested mode, what each operation ran and the reason. Writing it never fails a finished prediction.
• Flags: FAST_KERNEL_FLAGS is shared by AF2 and AF3; AlphaLink still rejects the flag.
• Docs: docs/af3_fused_triangles.md, linked from the README and the AF3 installation section.
• AlphaFold 3 v3.0.4 (separate commit): the submodule moves to the fork's main, which now carries upstream v3.0.4 (Update fork to AlphaFold 3 v3.0.4 alphafold3#3) and the fused kernels (Opt-in fused triangle kernels (Pallas) for inference alphafold3#4). The alphafold3 extra pins jax 0.10.2 / tokamax 0.0.12, matching the submodule's uv.lock, which the image installs. The Dockerfile version stamp and README follow.

Bug fix, separate commit: run_multimer_jobs.py never forwarded --fast_kernels. It takes its flags from run_structure_prediction, but builds its own argument list, which never passed --fast_kernels on for either backend and dropped --jax_compilation_cache_dir for AlphaFold 2. Both are now forwarded to AF2 and AF3; AlphaLink and UniFold still reject them.

Validation

• Unit suite (serial): 726 passed, 10 skipped. That includes the new test_af3_fused_triangles.py and the metadata-helper tests, one of them with a corrupt existing file. The AF2 fast_kernels tests are unchanged.

• AF3 image, CPU smoke: device_settings per card; smoke_test for both attention variants in Pallas interpret mode; resolve fallback and failure; metadata against the real GlobalConfig; importing the backend.

• GPU campaign (AF3 image, DeepMind weights; every job its own SLURM job). Warm speed-up, fused vs stock on the same card:

  H100	A100	A40	L40S	RTX 3090	RTX Pro 6000
  1.15–1.30×	1.18–1.66×	1.14–2.03×	1.11–1.34×	1.08–1.53×	1.12–1.23×

  • The largest fold each card holds is unchanged.
  • --fast_kernels=off is bit-identical to the previous fork on all 6 cards (deterministic XLA); on is deterministic.
  • 12 post-cutoff heterodimers, paired seeds 0–1 (stock also seed 2): Δranking and ΔDockQ are 0.00 on nearly all of them. The two A40 outliers come from stock seeds; a 10-seed follow-up is planned.
  • Protein–RNA, protein–ligand and template inputs give the same ranking with and without the kernels.
  • The campaign ran this feature before a refactor that kept its behaviour: module split, renamed config and seed-key spelling.

• Before any release: the GPU campaign is re-run on the AF3 image that CI builds from this branch (jax 0.10.2). The numbers above were measured on jax 0.9.1, and so far the v3.0.4 stack has been tested on CPU only: the fork's 90 fused tests, the golden config, and this unit suite.

Companion PR: the workflow accepts the flag for AF3 (KosinskiLab/AlphaPulldownSnakemake#SMPR).

🤖 Generated with Claude Code

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.