AF3: --fast_kernels with fused triangle kernels; AlphaFold 3 v3.0.4

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).
 </details>
 
 ---

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}
 

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

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,
 }
 
 

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():