diff --git a/README.md b/README.md index 90c342fe..f7007204 100644 --- a/README.md +++ b/README.md @@ -373,7 +373,7 @@ XlaRuntimeError: UNIMPLEMENTED: ... ptxas too old This cannot be patched from outside the container. jaxlib calls its own bundled `ptxas`, so `XLA_FLAGS=--xla_gpu_cuda_data_dir` and `PATH` have no effect, and bind-mounting a newer `ptxas` still leaves the CUDA runtime and cuDNN too old for the real kernels. From 2.5.0 the images ship a -consistent CUDA >= 12.8 stack (AF3: jax 0.9.1, ptxas 12.9, cuDNN 9.17, Tokamax; AF2: jax 0.5.3, +consistent CUDA >= 12.8 stack (AF3: jax 0.9.1, from AlphaFold 3 v3.0.4 on jax 0.10.2, ptxas 12.9, cuDNN 9.17, Tokamax; AF2: jax 0.5.3, ptxas 12.9, cuDNN 9.2x) and return the same confidence scores as the older cards. All three AF3 attention implementations (`triton`/Tokamax, `cudnn`, `xla`) work, so no `--flash_attention_implementation` override is needed. @@ -940,6 +940,9 @@ structure_inference_arguments: --save_distogram: False --use_ap_style: False # shared with AlphaFold2 ``` + +Experimental fused triangle kernels (`--fast_kernels`, not in released images yet): +[AlphaPulldown docs/af3_fused_triangles.md](https://github.com/KosinskiLab/AlphaPulldown/blob/main/docs/af3_fused_triangles.md). --- diff --git a/alphafold3 b/alphafold3 index 86b9ea3f..7809148b 160000 --- a/alphafold3 +++ b/alphafold3 @@ -1 +1 @@ -Subproject commit 86b9ea3feacc8934e6e2a581c49eb4c37a2a3d20 +Subproject commit 7809148bfefdb42daa3bf56c884e345b2b7a30c6 diff --git a/alphapulldown/folding_backend/alphafold3_backend.py b/alphapulldown/folding_backend/alphafold3_backend.py index 07bca2c1..fa41bd08 100644 --- a/alphapulldown/folding_backend/alphafold3_backend.py +++ b/alphapulldown/folding_backend/alphafold3_backend.py @@ -42,6 +42,7 @@ from alphafold.common import residue_constants from alphafold.common.protein import Protein, to_mmcif from alphapulldown.folding_backend.folding_backend import FoldingBackend +from alphapulldown.prediction import af3_fused_triangles from alphapulldown.prediction.jax_compilation_cache import enable_persistent_compilation_cache from alphapulldown.objects import MultimericObject, MonomericObject, ChoppedObject from alphapulldown.utils.af2_to_af3_msa import ( @@ -122,6 +123,7 @@ class ModelRunner: config: base_config.BaseConfig device: jax.Device model_dir: pathlib.Path + fused_triangles: af3_fused_triangles.FusedTriangleChoice = af3_fused_triangles.OFF @functools.cached_property def model_params(self) -> hk.Params: @@ -566,6 +568,50 @@ def _sequential_residue_ids_per_chain(chain_ids: Sequence[str]) -> list[int]: return residue_ids +def _record_inference_kernels( + output_dir: os.PathLike[str] | str, + seed: int, + model_runner: ModelRunner, + example: features.BatchDict, +) -> None: + """Record in inference_kernels.json which kernels a seed's triangle layers ran. + + Fused and original layers differ slightly in their numerics, so the record lets + ranking and analysis tell them apart, like AlphaFold 2's. Entries are keyed + ``seed-N`` like AF3's other outputs and are read back from the runner's own config, + so they state what actually ran. The record never fails a finished prediction: an + unreadable file is replaced, and any other failure is logged and skipped. + """ + fused_triangles = getattr(model_runner, "fused_triangles", None) + tokens = example.get("token_index", example.get("aatype")) + if fused_triangles is None or tokens is None: # not an AF3 runner (e.g. a test double) + return + key = f"seed-{seed}" + try: + record = af3_fused_triangles.metadata( + model_runner.config.global_config, fused_triangles, num_tokens=int(tokens.shape[0]) + ) + except Exception as exc: # provenance only; the prediction itself succeeded + logging.warning("Not recording the inference kernels of %s: %s", key, exc) + return + path = pathlib.Path(output_dir) / "inference_kernels.json" + try: + records = json.loads(path.read_text()) if path.exists() else {} + if not isinstance(records, dict): + raise ValueError("not a JSON object") + except (OSError, ValueError) as exc: + logging.warning("Replacing unreadable %s: %s", path, exc) + records = {} + records[key] = record + try: + os.makedirs(output_dir, exist_ok=True) + temporary = path.with_suffix(".json.tmp") + temporary.write_text(json.dumps(records, indent=2, sort_keys=True) + "\n") + temporary.replace(path) + except OSError as exc: + logging.warning("Not recording the inference kernels of %s: %s", key, exc) + + def predict_structure( fold_input: folding_input.Input, model_runner: ModelRunner, @@ -637,6 +683,8 @@ def write_from_array(rows: np.ndarray, suffix: str): logging.info( f'Extracting structures for seed {seed} took {time.time() - extract_start:.2f} seconds.' ) + if output_dir is not None: + _record_inference_kernels(output_dir, seed, model_runner, example) # Optional: gather embeddings and distogram embeddings_out = None @@ -763,6 +811,7 @@ def setup( num_recycles: int = 10, return_embeddings: bool = False, return_distogram: bool = False, + fast_kernels: str = "off", **kwargs, ) -> Dict: """Sets up the ModelRunner with the given configurations.""" @@ -778,11 +827,14 @@ def make_model_config( num_recycles: int = 10, return_embeddings: bool = False, return_distogram: bool = False, + global_config_update: Dict[str, Any], ): # The new code approach: config = model_class.Config() if hasattr(config, 'global_config'): config.global_config.flash_attention_implementation = flash_attention_implementation + for key, value in global_config_update.items(): + setattr(config.global_config, key, value) if hasattr(config, 'heads') and hasattr(config.heads, 'diffusion'): config.heads.diffusion.eval.num_samples = num_diffusion_samples # Optional overrides present in upstream AF3 runner @@ -812,6 +864,7 @@ def make_model_config( f' include "{required_flag}".' ) logging.info(f'Found local devices: {gpu_devices}') + fused_triangles = af3_fused_triangles.resolve(fast_kernels, gpu_devices[0]) logging.info('Building model from scratch...') model_runner = ModelRunner( @@ -825,9 +878,11 @@ def make_model_config( num_recycles=num_recycles, return_embeddings=return_embeddings, return_distogram=return_distogram, + global_config_update=fused_triangles.global_config_update(), ), device=gpu_devices[0], model_dir=pathlib.Path(model_dir), + fused_triangles=fused_triangles, ) return {'model_runner': model_runner} diff --git a/alphapulldown/prediction/af3_fused_triangles.py b/alphapulldown/prediction/af3_fused_triangles.py new file mode 100644 index 00000000..3d4b4314 --- /dev/null +++ b/alphapulldown/prediction/af3_fused_triangles.py @@ -0,0 +1,209 @@ +"""Decide whether AlphaFold 3 can use the fork's fused triangle kernels here. + +``--fast_kernels`` with ``--fold_backend=alphafold3`` turns on fused Pallas kernels for +the triangle multiplication and triangle attention layers of AF3's pair stack. They ship +with KosinskiLab's AlphaFold 3 fork (``alphafold3.jax.fused_triangle``), so unlike +AlphaFold 2's ``colabfold-kernels`` they need no extra package. + +The fork decides per layer: a layer outside the measured GPUs, dtype, shapes or size +limits runs the original module body. This module checks the device once, before a model +is built, and runs both kernels on it, so that a GPU, driver or JAX that cannot compile +them fails now rather than in the middle of a prediction: + +* ``off``: the original AF3 layers, exactly as before. +* ``on``: the kernels wherever a layer qualifies, or a ``ValueError`` saying why this + device cannot run them. +* ``auto``: as ``on`` where the device can run them, the original layers (with a log + line) elsewhere. + +JAX and the fork are imported only when a mode other than ``off`` is resolved, or when +the record of a model that requested the kernels is written. +""" + +from __future__ import annotations + +import dataclasses +from typing import TYPE_CHECKING, Any, Mapping + +from absl import logging + +from alphapulldown.prediction.fast_kernels import normalise_mode + +if TYPE_CHECKING: + import jax + +_OPERATIONS = ("triangle_multiplication", "triangle_attention") +# Pair channels of AF3's triangle layers: 128 in the pair stack, 64 in the template stack. +_RECORDED_NUM_CHANNELS = (128, 64) + + +@dataclasses.dataclass(frozen=True) +class FusedTriangleChoice: + """What ``resolve`` decided for AF3's triangle layers, and why.""" + + enabled: bool + reason: str + mode: str = "off" + settings: Mapping[str, Any] = dataclasses.field(default_factory=dict) + + def global_config_update(self) -> dict[str, Any]: + """Fields for AF3's ``GlobalConfig``; empty when the kernels are off.""" + return dict(self.settings) if self.enabled else {} + + +OFF = FusedTriangleChoice(False, "--fast_kernels=off") + + +def device_settings(device: jax.Device) -> dict[str, Any]: + """The ``GlobalConfig`` fields that turn the kernels on for ``device``. + + The fork's policy reads the device's compute capability and its JAX allocator budget, + which set the per-operation size limits. No model parameters are loaded. + + Raises: + RuntimeError: if the installed AF3 has no fused-triangle hooks, or the fork's + policy does not enable the kernels on ``device``. + """ + from alphafold3.model import model_config + + if not hasattr(model_config.GlobalConfig(), "triangle_multiplication_implementation"): + raise RuntimeError( + "this AlphaFold 3 installation has no fused-triangle hooks " + "(alphafold3.jax.fused_triangle); install the KosinskiLab alphafold3 fork" + ) + from alphafold3.jax.fused_triangle import dispatch + + memory_gib = (device.memory_stats() or {}).get("bytes_limit", 0) / 2**30 + policy = dispatch.device_policy(getattr(device, "compute_capability", ""), memory_gib) + if not policy.enabled: + raise RuntimeError(f"{getattr(device, 'device_kind', 'GPU')}: {policy.reason}") + return { + "triangle_multiplication_implementation": "pallas", + "triangle_attention_implementation": "auto", + "fused_triangle_compute_capability": policy.compute_capability, + "fused_triangle_memory_gib": memory_gib, + } + + +def smoke_test(settings: Mapping[str, Any], device: jax.Device) -> None: + """Compile and run both fused triangle layers on ``device``, without model weights. + + One 64-token, 128-channel pair goes through TriangleMultiplication and then + GridSelfAttention, configured with ``settings``. + + Raises: + RuntimeError: if the output is not finite. Whatever compiling or running the + kernels raises is passed on. + """ + import haiku as hk + import jax + import jax.numpy as jnp + from alphafold3.model import model_config + from alphafold3.model.components import utils + from alphafold3.model.network import modules + + global_config = model_config.GlobalConfig(final_init="linear", **settings) + + def forward(act, mask): + with utils.bfloat16_context(): + act = modules.TriangleMultiplication( + modules.TriangleMultiplication.Config(equation="ikc,jkc->ijc"), + global_config, + name="triangle_multiplication", + )(act, mask) + return modules.GridSelfAttention( + modules.GridSelfAttention.Config(), + global_config, + transpose=True, + name="triangle_attention", + )(act, mask) + + transformed = hk.without_apply_rng(hk.transform(forward)) + with jax.default_device(device): + act = jnp.ones((64, 64, 128), jnp.bfloat16) + mask = jnp.ones((64, 64), jnp.bfloat16) + params = transformed.init(jax.random.PRNGKey(0), act, mask) + output = jax.jit(transformed.apply)(params, act, mask) + finite = bool(jnp.all(jnp.isfinite(output))) + if not finite: + raise RuntimeError("the fused triangle layers returned non-finite values") + + +def resolve(mode: Any, device: jax.Device) -> FusedTriangleChoice: + """Turn a ``--fast_kernels`` value into a decision for ``device``, checking it first. + + Raises: + ValueError: for an unknown mode, or for ``on`` when ``device`` cannot run the + kernels. + """ + mode = normalise_mode(mode) + if mode == "off": + return OFF + try: + settings = device_settings(device) + smoke_test(settings, device) + except Exception as exc: # any failure here would recur inside the model + problem = f"AF3 fused triangle kernels are unavailable: {exc}" + if mode == "on": + raise ValueError(f"--fast_kernels=on, but {problem}.") from exc + logging.warning("Fused kernels off (--fast_kernels=auto): %s.", problem) + return FusedTriangleChoice(False, f"--fast_kernels=auto: {problem}", mode) + logging.info( + "AF3 fused triangle kernels on; per-operation size limits apply: %s", settings + ) + return FusedTriangleChoice(True, f"--fast_kernels={mode}", mode, settings) + + +def metadata( + global_config: Any, choice: FusedTriangleChoice, *, num_tokens: int +) -> dict[str, Any]: + """Which implementation AF3's triangle layers dispatch to for one padded bucket. + + Lists the 128-channel pair stack and the 64-channel template stack. Dispatch depends + only on ``global_config`` and the shapes, so the record also holds when a compiled + model comes from the compile cache. It does not claim that templates were supplied. + """ + record = { + "backend": "alphafold3", + "requested_mode": choice.mode, + "padded_tokens": num_tokens, + "fused_kernels": False, + "reason": choice.reason, + } + requested = { + operation: getattr(global_config, f"{operation}_implementation", "default") + for operation in _OPERATIONS + } + if all(implementation == "default" for implementation in requested.values()): + return record + from alphafold3.jax.fused_triangle import dispatch + + policy = dispatch.device_policy( + global_config.fused_triangle_compute_capability, + global_config.fused_triangle_memory_gib, + ) + dtype = "float32" if global_config.bfloat16 == "none" else "bfloat16" + operations = {} + for num_channels in _RECORDED_NUM_CHANNELS: + for operation, implementation in requested.items(): + selected, reason = dispatch.select_implementation( + operation, + implementation, + policy, + (num_tokens, num_tokens, num_channels), + dtype, + (num_tokens, num_tokens), + ) + operations[f"{operation}_c{num_channels}"] = { + "implementation": selected, + "reason": reason, + } + record.update( + fused_kernels=any( + operation["implementation"] != "default" for operation in operations.values() + ), + operations=operations, + policy_version=dispatch.POLICY_VERSION, + device_policy=dataclasses.asdict(policy), + ) + return record diff --git a/alphapulldown/prediction/inference_flags.py b/alphapulldown/prediction/inference_flags.py index e8341d15..ec78565b 100644 --- a/alphapulldown/prediction/inference_flags.py +++ b/alphapulldown/prediction/inference_flags.py @@ -38,8 +38,10 @@ ALPHALINK_EXTRA_FLAGS = frozenset({"crosslinks"}) -# AlphaFold 2 only: AlphaLink runs its own network, and AF3 has its own kernels. -AF2_EXTRA_FLAGS = frozenset({"fast_kernels"}) +# Fused kernels: ColabFold's for AlphaFold-Multimer, the AF3 fork's for AF3's triangle +# layers (see alphapulldown.prediction.fast_kernels and .af3_fused_triangles). AlphaLink +# runs its own network and has none. +FAST_KERNEL_FLAGS = frozenset({"fast_kernels"}) UNIFOLD_UNAVAILABLE_REASON = ( "UniFold is unavailable in this release: the bundled runtime is an AlphaLink " @@ -55,9 +57,9 @@ }) FLAGS_BY_BACKEND: Mapping[str, frozenset] = { - "alphafold2": COMMON_FLAGS | AF2_LIKE_FLAGS | AF2_EXTRA_FLAGS, + "alphafold2": COMMON_FLAGS | AF2_LIKE_FLAGS | FAST_KERNEL_FLAGS, "alphalink": COMMON_FLAGS | AF2_LIKE_FLAGS | ALPHALINK_EXTRA_FLAGS, - "alphafold3": COMMON_FLAGS | AF3_FLAGS, + "alphafold3": COMMON_FLAGS | AF3_FLAGS | FAST_KERNEL_FLAGS, } diff --git a/alphapulldown/scripts/run_multimer_jobs.py b/alphapulldown/scripts/run_multimer_jobs.py index 2b5cb0cc..3a356bbe 100644 --- a/alphapulldown/scripts/run_multimer_jobs.py +++ b/alphapulldown/scripts/run_multimer_jobs.py @@ -132,6 +132,7 @@ def main(argv): "--flash_attention_implementation": getattr(FLAGS, "flash_attention_implementation", None), "--buckets": getattr(FLAGS, "buckets", None), "--jax_compilation_cache_dir": getattr(FLAGS, "jax_compilation_cache_dir", None), + "--fast_kernels": getattr(FLAGS, "fast_kernels", None), "--save_embeddings": getattr(FLAGS, "save_embeddings", None), "--save_distogram": getattr(FLAGS, "save_distogram", None), "--debug_templates": getattr(FLAGS, "debug_templates", None), @@ -169,6 +170,10 @@ def main(argv): "--hb_allowance": FLAGS.hb_allowance, "--plddt_threshold": FLAGS.plddt_threshold, } + if fold_backend == "alphafold2": + # AlphaLink and UniFold reject the JAX compile cache and the fused kernels. + constant_args["--jax_compilation_cache_dir"] = getattr(FLAGS, "jax_compilation_cache_dir", None) + constant_args["--fast_kernels"] = getattr(FLAGS, "fast_kernels", None) command_args = {} for k, v in constant_args.items(): diff --git a/alphapulldown/scripts/run_structure_prediction.py b/alphapulldown/scripts/run_structure_prediction.py index a9688f16..157a360f 100644 --- a/alphapulldown/scripts/run_structure_prediction.py +++ b/alphapulldown/scripts/run_structure_prediction.py @@ -145,11 +145,11 @@ # JAX inference performance tuning. flags.DEFINE_string( 'fast_kernels', 'off', - 'off, on or auto. AlphaFold 2 only: fused Pallas kernels from the colabfold-kernels ' - 'package (pip install "alphapulldown[fast-kernels]"), about 2x faster ' - 'AlphaFold-Multimer inference on NVIDIA GPUs of compute capability 8.0 or newer. ' - '"on" fails if they cannot run here, "auto" falls back to the standard code. ' - 'Monomer models always use the standard code. true/false are read as on/off.', + 'off, on or auto (default off). AF2-Multimer uses colabfold-kernels ' + '(pip install "alphapulldown[fast-kernels]"); AF3 uses the fork\'s bundled ' + 'fused triangle kernels with per-GPU and per-size fallbacks. ' + '"on" fails if the kernel setup cannot run here; "auto" falls back. ' + 'AF2 monomers keep the standard code. true/false are read as on/off.', ) flags.DEFINE_string( 'jax_compilation_cache_dir', diff --git a/docker/alphafold3.dockerfile b/docker/alphafold3.dockerfile index 1b7349ea..c3604967 100644 --- a/docker/alphafold3.dockerfile +++ b/docker/alphafold3.dockerfile @@ -86,7 +86,7 @@ RUN cd /hmmer_build && \ COPY . /app/AlphaPulldown # --------------------------------------------------------------------- -# Install AlphaFold3 from its locked v3.0.3/Tokamax environment +# Install AlphaFold3 from its locked v3.0.4/Tokamax environment # --------------------------------------------------------------------- WORKDIR /app/AlphaPulldown/alphafold3 @@ -96,7 +96,7 @@ ENV PIP_INDEX_URL=${PIP_INDEX_URL} \ PIP_RETRIES=10 \ PIP_NO_CACHE_DIR=1 \ PIP_DISABLE_PIP_VERSION_CHECK=1 \ - SETUPTOOLS_SCM_PRETEND_VERSION_FOR_ALPHAFOLD3=3.0.3 + SETUPTOOLS_SCM_PRETEND_VERSION_FOR_ALPHAFOLD3=3.0.4 ENV CMAKE_BUILD_PARALLEL_LEVEL=1 ENV CFLAGS="-O2 -pipe" diff --git a/docs/af3_fused_triangles.md b/docs/af3_fused_triangles.md new file mode 100644 index 00000000..0b0d07c7 --- /dev/null +++ b/docs/af3_fused_triangles.md @@ -0,0 +1,166 @@ +# Fused triangle kernels for AlphaFold 3 (experimental) + +AlphaFold 3 spends much of its time in the triangle layers of its pair stack: +triangle multiplication and triangle attention. `--fast_kernels` can run these layers +as fused GPU kernels written in Pallas, JAX's language for custom kernels, which +combine several steps into one pass over memory. The kernels are Anthropic's +FlashPairformer triangle kernels, vendored into KosinskiLab's AlphaFold 3 fork. Off by +default; predictions without the flag run exactly as before. + +The same flag turns on ColabFold's fused kernels for AlphaFold-Multimer, a separate +implementation; AlphaFold 3 does not need the `colabfold-kernels` package. + +## Turning them on + +```bash +run_structure_prediction.py --fold_backend=alphafold3 --fast_kernels=auto ... +``` + +| value | effect | +| --- | --- | +| `off` (default) | the original AlphaFold 3 layers | +| `on` | the fused kernels, or an error before any weights load if this GPU cannot run them | +| `auto` | the fused kernels where this GPU can run them, the original layers (with a warning) elsewhere | + +`true`/`false` are read as `on`/`off`, so an unquoted YAML `on` works too. + +Before building the model, AlphaPulldown checks the GPU against the measured +architectures and memory limits below, then compiles and runs both fused layers once +on a small input on that GPU. A GPU, driver or JAX that cannot compile them therefore +fails, or falls back, at that point rather than in the middle of a prediction. + +The kernels need the `alphafold3` package at the submodule revision of this +AlphaPulldown checkout: build the AF3 environment or image from it +([backend installation](backend_installation.md#alphafold3-backend)). AlphaPulldown +2.9.1 and older do not accept `--fast_kernels` for AlphaFold 3. With an older +`alphafold3` package, `on` fails with "no fused-triangle hooks" and `auto` keeps the +original layers. + +## What runs on which GPU + +Each layer is decided on its own, and a layer that does not qualify runs the original +code; the rest of the model is unchanged. A layer qualifies only with bf16 +activations, a square pair representation of 64 or 128 channels whose token count is +a multiple of 64, and (for attention) four heads. These are the shapes of AlphaFold 3's +pair stack (128 channels) and template stack (64 channels). + +- **Triangle multiplication:** fused Pallas steps around XLA's matrix multiplication, on + every measured architecture. +- **Triangle attention** comes in two variants. They differ in the attention core, the + step that computes softmax(QKᵀ + bias)·V without storing the full score matrix + (flash attention): + - `pallas`: the whole layer in Pallas, including Anthropic's own Pallas + flash-attention core. + - `pallas_tokamax_core`: the same Pallas steps before and after the core, around the + flash-attention core of tokamax, the DeepMind kernel library that AlphaFold 3 + already uses for attention (`--flash_attention_implementation` picks its kernel). + Used where Anthropic's core was slower than the original layer. + +| compute capability | GPUs measured | multiplication | attention | +| --- | --- | --- | --- | +| 8.0 | A100 | fused | `pallas` | +| 8.6 | A40, RTX 3090 | fused | `pallas_tokamax_core` | +| 8.9 | L40S | fused | `pallas_tokamax_core` | +| 9.0 | H100 | fused | `pallas` | +| 12.0 | RTX Pro 6000 Blackwell | fused | `pallas` (with the 9.0 tile sizes) | + +Any other architecture keeps the original layers, including compute capability 10.0 +(B200) and anything newer: a higher compute capability is not evidence that the +kernels compile, or are faster, there. + +## Memory limits + +Fused attention processes all rows of the pair representation at once, where the +original layer works through them in chunks, so it needs more memory. Each operation +therefore has a token limit set by the GPU memory JAX may allocate (the allocator's +`bytes_limit`, which `XLA_PYTHON_CLIENT_MEM_FRACTION` and similar settings change), +not by the card's advertised memory. Above its limit a layer runs the original code. + +| JAX memory budget | multiplication up to | attention up to | +| --- | ---: | ---: | +| below 12 GiB | not used | not used | +| 12 to <20 GiB | 1,536 tokens | 768 tokens | +| 20 to <32 GiB | 2,560 tokens | 1,536 tokens | +| 32 to <64 GiB | 3,584 tokens (5,120 on 9.0 and 12.0) | 2,048 tokens | +| 64 GiB and more | 3,584 tokens (5,120 on 9.0 and 12.0) | 3,072 tokens | + +Token counts are AlphaFold 3's padded bucket sizes (`--buckets`). The limits are +conservative experimental settings, not measured capacities of the full model. + +## How it was measured + +The kernels were checked layer by layer ("Gate 0"): each fused layer against the +original layer and a float64 reference, with random and DeepMind weights and padded +masks, on A100, A40, RTX 3090, L40S, H100 and RTX Pro 6000 Blackwell. The full-model +comparison of `off` and `on` predictions (speed, memory, identical-output checks and +accuracy on DeepMind weights) runs on the same cards. Both used JAX 0.9.1 and +tokamax 0.0.11; other versions have not been measured on GPU. + +Layer speed-ups justify an end-to-end pilot, not a claim about production speed. The +option stays opt-in until the full-model comparison is complete. + +The fused kernels are skipped while parameters are initialised, so parameter names, +shapes, initialisers and random-number order are those of the original model, and +DeepMind's weights load unchanged. + +## Output record + +Every AlphaFold 3 output directory gets an `inference_kernels.json` with one entry per +seed, keyed like AlphaFold 3's other outputs: + +```json +{ + "seed-1": { + "backend": "alphafold3", + "requested_mode": "auto", + "reason": "--fast_kernels=auto", + "padded_tokens": 2560, + "fused_kernels": true, + "operations": { + "triangle_attention_c128": {"implementation": "default", "reason": "size_limit"}, + "triangle_multiplication_c128": {"implementation": "pallas", "reason": ""}, + "...": "..." + }, + "policy_version": 1, + "device_policy": {"compute_capability": "8.6", "memory_gib": 44.5, "...": "..."} + } +} +``` + +`operations` states, for the 128-channel pair stack and the 64-channel template stack, +which implementation a layer of that bucket runs (`default` is the original layer) and +why it falls back. It is derived from the configuration and the bucket, so it also +holds for a model loaded from the compile cache. It does not mean that templates were +supplied. When the kernels are off, the entry has no `operations`. Writing the record +never fails a prediction: an unreadable file is replaced, and other problems are +logged. + +## Troubleshooting + +`--fast_kernels=on` stops before loading weights and gives the reason; `auto` logs the +same reason as `Fused kernels off (--fast_kernels=auto): ...`: + +| reason contains | meaning | +| --- | --- | +| `no fused-triangle hooks` | the installed `alphafold3` predates the kernels; rebuild it from the submodule | +| `unvalidated_compute_capability` | the GPU's architecture was not measured (table above) | +| `unknown_memory_budget` | JAX reports no allocator limit for the GPU | +| `memory_budget_below_12_gib` | JAX may allocate less than 12 GiB | +| anything else | compiling or running the small check on this GPU failed | + +During a prediction, a warning such as `triangle_attention_implementation='auto' was +requested, but a layer ... runs the default implementation (size_limit)` is expected +for buckets above the limits. It is logged once per reason. + +Scripts that import `alphapulldown.folding_backend.alphafold3_backend` themselves +should call `jax.devices()` before that import. In the AlphaFold 3 image, importing the +backend first can leave JAX with only its CPU and TPU platforms +(`Backend 'cuda' is not in the list of known backends`), so the GPU check cannot see +the GPU. + +## Source and licence + +The kernel files are vendored unchanged from Anthropic's +[uplifting-biomolecular-modeling](https://github.com/anthropics/uplifting-biomolecular-modeling) +repository (Apache License 2.0). The fork's `src/alphafold3/jax/fused_triangle/SOURCE.md` +names the commit and the files, and `NOTICE` carries the attribution. diff --git a/docs/backend_installation.md b/docs/backend_installation.md index 2e1f9546..b999e861 100644 --- a/docs/backend_installation.md +++ b/docs/backend_installation.md @@ -102,6 +102,7 @@ Notes: - For cluster installs, prefer `python -m pip install ".[alphafold3,test]"` and then build the vendored `alphafold3` package. - The compiled `alphafold3.cpp` extension comes from `python -m pip install --no-deps -e ./alphafold3`, not from the root install. - The vendored AF3 package provides the `build_data` entry point. Use that directly. +- The experimental `--fast_kernels` option (fused triangle kernels) needs the vendored `alphafold3` built from this checkout's submodule revision; see [AF3 fused triangle kernels](af3_fused_triangles.md). - If you are actively developing inside the checkout and want the root package editable as well, add: ```bash diff --git a/pyproject.toml b/pyproject.toml index 24662aeb..d2ee263f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,14 +59,14 @@ fast-kernels = [ ] alphafold3 = [ "dm-haiku==0.0.16", - "jax==0.9.1", - "jax[cuda12]==0.9.1; platform_system == 'Linux' and platform_machine == 'x86_64'", + "jax==0.10.2", + "jax[cuda12]==0.10.2; platform_system == 'Linux' and platform_machine == 'x86_64'", "modelcif>=1.6", "numpy", "openmm>=8.2", "pdbfixer>=1.10", "rdkit==2025.9.4", - "tokamax==0.0.11", + "tokamax==0.0.12", "zstandard", ] test = [ diff --git a/test/unit/test_af3_fused_triangles.py b/test/unit/test_af3_fused_triangles.py new file mode 100644 index 00000000..c424f0a2 --- /dev/null +++ b/test/unit/test_af3_fused_triangles.py @@ -0,0 +1,286 @@ +"""--fast_kernels for AlphaFold 3: checking the device before a model is built.""" + +from __future__ import annotations + +import contextlib +import importlib.util +import sys +import types +from pathlib import Path + +import numpy as np +import pytest + +from alphapulldown.prediction import af3_fused_triangles + +DISPATCH_PATH = ( + Path(__file__).resolve().parents[2] + / "alphafold3/src/alphafold3/jax/fused_triangle/dispatch.py" +) +SETTINGS = { + "triangle_multiplication_implementation": "pallas", + "triangle_attention_implementation": "auto", +} + + +def _fail(*args): + raise AssertionError("off must not inspect the device or import the fork") + + +def _unavailable(*args): + raise RuntimeError("unavailable test device") + + +def _global_config(**overrides) -> types.SimpleNamespace: + fields = dict( + triangle_multiplication_implementation="pallas", + triangle_attention_implementation="auto", + fused_triangle_compute_capability="8.6", + fused_triangle_memory_gib=45, + bfloat16="all", + ) + fields.update(overrides) + return types.SimpleNamespace(**fields) + + +@pytest.fixture +def dispatch(monkeypatch): + """The fork's dispatch module from the submodule, importable without the fork.""" + spec = importlib.util.spec_from_file_location("af3_dispatch_test", DISPATCH_PATH) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + package = types.ModuleType("alphafold3.jax.fused_triangle") + package.dispatch = module + monkeypatch.setitem(sys.modules, package.__name__, package) + return module + + +@pytest.fixture +def model_config(monkeypatch): + """A stub ``alphafold3.model.model_config`` whose GlobalConfig has the hooks.""" + module = types.ModuleType("alphafold3.model.model_config") + module.GlobalConfig = lambda: types.SimpleNamespace( + triangle_multiplication_implementation="default" + ) + parent = types.ModuleType("alphafold3.model") + parent.model_config = module + monkeypatch.setitem(sys.modules, parent.__name__, parent) + monkeypatch.setitem(sys.modules, module.__name__, module) + return module + + +def test_off_checks_nothing(monkeypatch): + monkeypatch.setattr(af3_fused_triangles, "device_settings", _fail) + monkeypatch.setattr(af3_fused_triangles, "smoke_test", _fail) + + choice = af3_fused_triangles.resolve("off", object()) + + assert not choice.enabled + assert choice.mode == "off" + assert choice.global_config_update() == {} + + +def test_unknown_mode_is_rejected(): + with pytest.raises(ValueError, match="one of off, on, auto"): + af3_fused_triangles.resolve("fast", object()) + + +@pytest.mark.parametrize("stage", ["device_settings", "smoke_test"]) +def test_on_fails_loudly_and_auto_falls_back(monkeypatch, stage): + monkeypatch.setattr(af3_fused_triangles, "device_settings", lambda device: SETTINGS) + monkeypatch.setattr(af3_fused_triangles, "smoke_test", lambda *args: None) + monkeypatch.setattr(af3_fused_triangles, stage, _unavailable) + + with pytest.raises(ValueError, match="--fast_kernels=on, but .*unavailable test device"): + af3_fused_triangles.resolve("on", object()) + + choice = af3_fused_triangles.resolve("auto", object()) + assert not choice.enabled + assert choice.mode == "auto" + assert choice.reason.startswith("--fast_kernels=auto: ") + assert "unavailable test device" in choice.reason + assert choice.global_config_update() == {} + + +def test_checks_the_given_device_before_enabling(monkeypatch): + device = object() + seen = [] + monkeypatch.setattr( + af3_fused_triangles, "device_settings", lambda d: seen.append(d) or SETTINGS + ) + monkeypatch.setattr( + af3_fused_triangles, "smoke_test", lambda *args: seen.append(args) + ) + + choice = af3_fused_triangles.resolve("on", device) + + assert seen[0] is device + assert seen[1] == (SETTINGS, device) + assert choice.enabled and choice.mode == "on" + assert choice.global_config_update() == SETTINGS + # The update is a copy: changing it does not change the choice. + choice.global_config_update().clear() + assert choice.global_config_update() == SETTINGS + + +def test_smoke_test_runs_on_the_device_it_checks(monkeypatch): + # Stub JAX, Haiku and the AF3 modules, recording the default device of each step. + device = object() + current = [] + steps = [] + + @contextlib.contextmanager + def default_device(chosen): + current.append(chosen) + try: + yield + finally: + current.pop() + + def step(name): + def run(*args, **kwargs): + steps.append((name, current[-1] if current else None)) + return np.ones(1) + + return run + + jnp = types.SimpleNamespace( + ones=step("ones"), bfloat16="bfloat16", all=np.all, isfinite=np.isfinite + ) + jax = types.SimpleNamespace( + numpy=jnp, + default_device=default_device, + jit=lambda function: step("jit"), + random=types.SimpleNamespace(PRNGKey=step("PRNGKey")), + ) + transformed = types.SimpleNamespace(init=step("init"), apply=None) + haiku = types.SimpleNamespace( + transform=lambda function: transformed, without_apply_rng=lambda t: t + ) + model = types.ModuleType("alphafold3.model") + model.model_config = types.SimpleNamespace(GlobalConfig=lambda **fields: fields) + model.components = types.SimpleNamespace(utils=None) + model.network = types.SimpleNamespace(modules=None) + for name, module in { + "jax": jax, + "jax.numpy": jnp, + "haiku": haiku, + "alphafold3.model": model, + "alphafold3.model.components": model.components, + "alphafold3.model.network": model.network, + }.items(): + monkeypatch.setitem(sys.modules, name, module) + + af3_fused_triangles.smoke_test(SETTINGS, device) + + assert [name for name, _ in steps] == ["ones", "ones", "PRNGKey", "init", "jit"] + assert all(used is device for _, used in steps) + + +def test_device_settings_use_the_allocator_budget(dispatch, model_config): + device = types.SimpleNamespace( + compute_capability="12.0", + device_kind="Blackwell slice", + memory_stats=lambda: {"bytes_limit": 15 * 2**30}, + ) + + settings = af3_fused_triangles.device_settings(device) + + assert settings == { + "triangle_multiplication_implementation": "pallas", + "triangle_attention_implementation": "auto", + "fused_triangle_compute_capability": "12.0", + "fused_triangle_memory_gib": 15, + } + + +@pytest.mark.parametrize( + "compute_capability, memory_stats, reason", + [ + ("12.0", lambda: None, "unknown_memory_budget"), + ("10.0", lambda: {"bytes_limit": 80 * 2**30}, "unvalidated_compute_capability"), + ("8.0", lambda: {"bytes_limit": 8 * 2**30}, "memory_budget_below_12_gib"), + ], +) +def test_device_settings_reject_devices_the_policy_does_not_enable( + dispatch, model_config, compute_capability, memory_stats, reason +): + device = types.SimpleNamespace( + compute_capability=compute_capability, device_kind="GPU", memory_stats=memory_stats + ) + + with pytest.raises(RuntimeError, match=reason): + af3_fused_triangles.device_settings(device) + + +def test_device_settings_reject_an_alphafold3_without_hooks(model_config): + model_config.GlobalConfig = lambda: types.SimpleNamespace() + + with pytest.raises(RuntimeError, match="no fused-triangle hooks"): + af3_fused_triangles.device_settings(types.SimpleNamespace()) + + +def test_metadata_of_the_original_layers_does_not_import_the_fork(monkeypatch): + monkeypatch.setitem(sys.modules, "alphafold3.jax.fused_triangle", None) + + record = af3_fused_triangles.metadata( + types.SimpleNamespace(), af3_fused_triangles.OFF, num_tokens=256 + ) + + assert record == { + "backend": "alphafold3", + "requested_mode": "off", + "padded_tokens": 256, + "fused_kernels": False, + "reason": "--fast_kernels=off", + } + + +def test_metadata_records_per_operation_fallback(dispatch): + choice = af3_fused_triangles.FusedTriangleChoice(True, "--fast_kernels=auto", "auto") + + record = af3_fused_triangles.metadata(_global_config(), choice, num_tokens=2560) + + assert record["fused_kernels"] is True + assert record["requested_mode"] == "auto" + assert record["policy_version"] == dispatch.POLICY_VERSION + assert record["device_policy"]["attention_implementation"] == "pallas_tokamax_core" + operations = record["operations"] + assert set(operations) == { + "triangle_multiplication_c128", + "triangle_attention_c128", + "triangle_multiplication_c64", + "triangle_attention_c64", + } + assert operations["triangle_multiplication_c128"] == { + "implementation": "pallas", + "reason": "", + } + # Attention on a 45 GiB budget stops at 2048 tokens. + assert operations["triangle_attention_c128"] == { + "implementation": "default", + "reason": "size_limit", + } + + +def test_metadata_is_the_same_on_a_compile_cache_hit(dispatch): + choice = af3_fused_triangles.FusedTriangleChoice(True, "--fast_kernels=on", "on") + global_config = _global_config() + + first = af3_fused_triangles.metadata(global_config, choice, num_tokens=1024) + + assert first == af3_fused_triangles.metadata(global_config, choice, num_tokens=1024) + + +def test_metadata_float32_layers_run_the_original_body(dispatch): + choice = af3_fused_triangles.FusedTriangleChoice(True, "--fast_kernels=on", "on") + + record = af3_fused_triangles.metadata( + _global_config(bfloat16="none"), choice, num_tokens=256 + ) + + assert record["fused_kernels"] is False + assert {operation["reason"] for operation in record["operations"].values()} == { + "dtype_not_bfloat16" + } diff --git a/test/unit/test_alphafold3_backend_helpers.py b/test/unit/test_alphafold3_backend_helpers.py index ab2ee3f1..0433cb3f 100644 --- a/test/unit/test_alphafold3_backend_helpers.py +++ b/test/unit/test_alphafold3_backend_helpers.py @@ -842,6 +842,119 @@ def extract_structures(self, batch, result, target_name): assert final_msa.read_text(encoding="utf-8") == ">query\nABC\n" +class FakeKernelRunner: + """A runner with an AF3 fused-triangle choice; logs extraction to ``events``.""" + + def __init__(self, af3_backend_module, events=None): + self.fused_triangles = af3_backend_module.af3_fused_triangles.OFF + self.config = SimpleNamespace(global_config=SimpleNamespace()) + self.events = [] if events is None else events + + def run_inference(self, batch, rng_key): + return {} + + def extract_structures(self, batch, result, target_name): + self.events.append("extract") + return [] + + +OFF_RECORD = { + "backend": "alphafold3", + "requested_mode": "off", + "padded_tokens": 3, + "fused_kernels": False, + "reason": "--fast_kernels=off", +} + + +def test_predict_structure_records_kernels_per_seed_after_extraction( + af3_backend_module, + monkeypatch, + tmp_path, +): + fold_input = af3_backend_module.folding_input.Input( + name="job", chains=("A",), rng_seeds=(7, 8) + ) + example = {"token_index": np.arange(3)} + monkeypatch.setattr( + af3_backend_module.featurisation, + "featurise_input", + lambda **kwargs: [example, example], + ) + events = [] + real_metadata = af3_backend_module.af3_fused_triangles.metadata + monkeypatch.setattr( + af3_backend_module.af3_fused_triangles, + "metadata", + lambda *args, **kwargs: events.append("record") or real_metadata(*args, **kwargs), + ) + + af3_backend_module.predict_structure( + fold_input=fold_input, + model_runner=FakeKernelRunner(af3_backend_module, events), + output_dir=tmp_path, + ) + + # After extraction, so it is not timed as model inference. + assert events == ["extract", "record", "extract", "record"] + record = json.loads((tmp_path / "inference_kernels.json").read_text()) + assert record == {"seed-7": OFF_RECORD, "seed-8": OFF_RECORD} + + +def test_record_inference_kernels_replaces_a_corrupt_file( + af3_backend_module, tmp_path, caplog +): + path = tmp_path / "inference_kernels.json" + path.write_text("{not json") + + af3_backend_module._record_inference_kernels( + tmp_path, 7, FakeKernelRunner(af3_backend_module), {"token_index": np.arange(3)} + ) + + assert json.loads(path.read_text()) == {"seed-7": OFF_RECORD} + assert "Replacing unreadable" in caplog.text + + +def test_record_inference_kernels_keeps_other_entries(af3_backend_module, tmp_path): + path = tmp_path / "inference_kernels.json" + path.write_text(json.dumps({"seed-1": {"fused_kernels": True}})) + + af3_backend_module._record_inference_kernels( + tmp_path, 7, FakeKernelRunner(af3_backend_module), {"token_index": np.arange(3)} + ) + + assert json.loads(path.read_text()) == { + "seed-1": {"fused_kernels": True}, + "seed-7": OFF_RECORD, + } + + +def test_record_inference_kernels_never_fails_a_prediction( + af3_backend_module, monkeypatch, tmp_path, caplog +): + def broken(*args, **kwargs): + raise ImportError("no alphafold3.jax.fused_triangle") + + monkeypatch.setattr(af3_backend_module.af3_fused_triangles, "metadata", broken) + + af3_backend_module._record_inference_kernels( + tmp_path, 7, FakeKernelRunner(af3_backend_module), {"token_index": np.arange(3)} + ) + + assert not (tmp_path / "inference_kernels.json").exists() + assert "no alphafold3.jax.fused_triangle" in caplog.text + + +def test_record_inference_kernels_skips_runners_without_a_choice( + af3_backend_module, tmp_path +): + af3_backend_module._record_inference_kernels( + tmp_path, 7, SimpleNamespace(), {"token_index": np.arange(3)} + ) + + assert not (tmp_path / "inference_kernels.json").exists() + + def test_af3_setup_builds_model_runner_and_validates_gpu_capability( af3_backend_module, monkeypatch, @@ -933,6 +1046,77 @@ class FakeModel: ) +def _stub_af3_setup(af3_backend_module, monkeypatch, device) -> None: + """Stub what AlphaFold3Backend.setup needs besides the fused-triangle decision.""" + + class FakeConfig: + def __init__(self): + self.global_config = SimpleNamespace(flash_attention_implementation=None) + + monkeypatch.setattr( + sys.modules["alphafold3.model.model"], + "Model", + type("FakeModel", (), {"Config": FakeConfig}), + ) + monkeypatch.setattr( + af3_backend_module.jax, "local_devices", lambda backend="gpu": [device] + ) + monkeypatch.setattr( + af3_backend_module, "_initialise_tokamax_trace_context", lambda: None + ) + + +def _af3_setup(af3_backend_module, tmp_path, **kwargs): + return af3_backend_module.AlphaFold3Backend.setup( + num_diffusion_samples=1, + flash_attention_implementation="triton", + buckets=[128], + jax_compilation_cache_dir=None, + model_dir=str(tmp_path / "models"), + **kwargs, + )["model_runner"] + + +def test_af3_setup_keeps_the_original_triangle_layers_by_default( + af3_backend_module, monkeypatch, tmp_path +): + _stub_af3_setup(af3_backend_module, monkeypatch, SimpleNamespace(compute_capability=8.0)) + + runner = _af3_setup(af3_backend_module, tmp_path) + + assert runner.fused_triangles == af3_backend_module.af3_fused_triangles.OFF + assert vars(runner.config.global_config) == {"flash_attention_implementation": "triton"} + + +def test_af3_setup_applies_the_fused_triangle_choice_for_its_gpu( + af3_backend_module, monkeypatch, tmp_path +): + device = SimpleNamespace(compute_capability=8.0) + _stub_af3_setup(af3_backend_module, monkeypatch, device) + settings = { + "triangle_multiplication_implementation": "pallas", + "triangle_attention_implementation": "auto", + "fused_triangle_compute_capability": "8.0", + "fused_triangle_memory_gib": 38.0, + } + choice = af3_backend_module.af3_fused_triangles.FusedTriangleChoice( + True, "--fast_kernels=auto", "auto", settings + ) + decisions = [] + monkeypatch.setattr( + af3_backend_module.af3_fused_triangles, + "resolve", + lambda mode, gpu: decisions.append((mode, gpu)) or choice, + ) + + runner = _af3_setup(af3_backend_module, tmp_path, fast_kernels="auto") + + assert decisions == [("auto", device)] + assert runner.fused_triangles is choice + for key, value in settings.items(): + assert getattr(runner.config.global_config, key) == value + + def test_af3_tokamax_trace_context_is_created_by_its_own_hook( af3_backend_module, monkeypatch ): diff --git a/test/unit/test_inference_flags.py b/test/unit/test_inference_flags.py index a4472275..feadb8b2 100644 --- a/test/unit/test_inference_flags.py +++ b/test/unit/test_inference_flags.py @@ -53,10 +53,10 @@ def test_convert_to_modelcif_is_accepted_by_both_backends(): assert inference_flags.unsupported_flags(backend, ["convert_to_modelcif"]) == [] -def test_fast_kernels_is_an_alphafold2_flag(): - assert inference_flags.unsupported_flags("alphafold2", ["fast_kernels"]) == [] - for backend in ("alphafold3", "alphalink"): - assert inference_flags.unsupported_flags(backend, ["fast_kernels"]) == ["fast_kernels"] +def test_fast_kernels_is_an_alphafold2_and_alphafold3_flag(): + for backend in ("alphafold2", "alphafold3"): + assert inference_flags.unsupported_flags(backend, ["fast_kernels"]) == [] + assert inference_flags.unsupported_flags("alphalink", ["fast_kernels"]) == ["fast_kernels"] def test_unknown_backend_reports_nothing(): diff --git a/test/unit/test_light_import_invariant.py b/test/unit/test_light_import_invariant.py index 7e51ca9e..556f3c13 100644 --- a/test/unit/test_light_import_invariant.py +++ b/test/unit/test_light_import_invariant.py @@ -35,6 +35,8 @@ "alphapulldown.prediction.jax_compilation_cache", # And checks the fused-kernel requirements only when a mode other than off is resolved. "alphapulldown.prediction.fast_kernels", + # The AF3 backend imports it, so it must not import JAX or the fork either. + "alphapulldown.prediction.af3_fused_triangles", # Compatibility paths must retain the same lightweight behavior. "alphapulldown.feature_batch", "alphapulldown.af2_feature_finalizer", diff --git a/test/unit/test_script_entrypoints.py b/test/unit/test_script_entrypoints.py index d2c1be6f..110694fd 100644 --- a/test/unit/test_script_entrypoints.py +++ b/test/unit/test_script_entrypoints.py @@ -359,6 +359,7 @@ def _load_run_multimer_jobs_module(): "flash_attention_implementation": "triton", "buckets": ["64", "128"], "jax_compilation_cache_dir": None, + "fast_kernels": "off", "save_embeddings": False, "save_distogram": False, "debug_templates": False, @@ -1408,6 +1409,71 @@ def test_run_multimer_jobs_forwards_relax_best_score_threshold( assert calls[0][calls[0].index("--relax_best_score_threshold") + 1] == "0.6" +@pytest.mark.parametrize("fold_backend", ["alphafold2", "alphafold3"]) +def test_run_multimer_jobs_forwards_fast_kernels_and_compile_cache( + run_multimer_jobs_module, + monkeypatch, + fold_backend, +): + calls = [] + monkeypatch.setattr( + run_multimer_jobs_module.subprocess, + "run", + lambda command, check, env: calls.append(command), + ) + run_multimer_jobs_module.generate_fold_specifications = ( + lambda input_files, delimiter, exclude_permutations: ["job1"] + ) + + _set_flag(run_multimer_jobs_module.FLAGS, "mode", "custom") + _set_flag(run_multimer_jobs_module.FLAGS, "protein_lists", ["proteins.txt"]) + _set_flag(run_multimer_jobs_module.FLAGS, "dry_run", False) + _set_flag(run_multimer_jobs_module.FLAGS, "fold_backend", fold_backend) + _set_flag(run_multimer_jobs_module.FLAGS, "output_path", "/tmp/output") + _set_flag(run_multimer_jobs_module.FLAGS, "data_dir", "/tmp/models") + _set_flag(run_multimer_jobs_module.FLAGS, "monomer_objects_dir", ["/tmp/features"]) + _set_flag(run_multimer_jobs_module.FLAGS, "fast_kernels", "auto") + _set_flag(run_multimer_jobs_module.FLAGS, "jax_compilation_cache_dir", "/tmp/jax-cache") + + run_multimer_jobs_module.main(["prog"]) + + assert len(calls) == 1 + command = calls[0] + assert command[command.index("--fast_kernels") + 1] == "auto" + assert command[command.index("--jax_compilation_cache_dir") + 1] == "/tmp/jax-cache" + + +def test_run_multimer_jobs_does_not_forward_fast_kernels_to_alphalink( + run_multimer_jobs_module, + monkeypatch, +): + calls = [] + monkeypatch.setattr( + run_multimer_jobs_module.subprocess, + "run", + lambda command, check, env: calls.append(command), + ) + run_multimer_jobs_module.generate_fold_specifications = ( + lambda input_files, delimiter, exclude_permutations: ["job1"] + ) + + _set_flag(run_multimer_jobs_module.FLAGS, "mode", "custom") + _set_flag(run_multimer_jobs_module.FLAGS, "protein_lists", ["proteins.txt"]) + _set_flag(run_multimer_jobs_module.FLAGS, "dry_run", False) + _set_flag(run_multimer_jobs_module.FLAGS, "use_alphalink", True) + _set_flag(run_multimer_jobs_module.FLAGS, "alphalink_weight", "/tmp/alphalink.pt") + _set_flag(run_multimer_jobs_module.FLAGS, "output_path", "/tmp/output") + _set_flag(run_multimer_jobs_module.FLAGS, "monomer_objects_dir", ["/tmp/features"]) + _set_flag(run_multimer_jobs_module.FLAGS, "fast_kernels", "auto") + _set_flag(run_multimer_jobs_module.FLAGS, "jax_compilation_cache_dir", "/tmp/jax-cache") + + run_multimer_jobs_module.main(["prog"]) + + assert len(calls) == 1 + assert "--fast_kernels" not in calls[0] + assert "--jax_compilation_cache_dir" not in calls[0] + + @pytest.mark.parametrize("legacy_switch", [True, False]) def test_run_multimer_jobs_rejects_unifold_before_reading_inputs( run_multimer_jobs_module, monkeypatch, legacy_switch,