From f76b0720c3451e81493579d4ba2958112f80a1c8 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 6 Oct 2026 15:09:43 +0200 Subject: [PATCH 01/11] Enable optional AF3 fused triangles with device checks and output provenance --- alphafold3 | 2 +- .../folding_backend/alphafold3_backend.py | 26 ++++ alphapulldown/prediction/fast_kernels.py | 114 +++++++++++++++++- alphapulldown/prediction/inference_flags.py | 4 +- .../scripts/run_structure_prediction.py | 10 +- docs/af3_fused_triangles.md | 41 +++++++ docs/backend_installation.md | 3 + test/unit/test_alphafold3_backend_helpers.py | 26 ++++ test/unit/test_fast_kernels.py | 87 +++++++++++++ test/unit/test_inference_flags.py | 5 +- 10 files changed, 305 insertions(+), 13 deletions(-) create mode 100644 docs/af3_fused_triangles.md diff --git a/alphafold3 b/alphafold3 index 86b9ea3f..8772b6fc 160000 --- a/alphafold3 +++ b/alphafold3 @@ -1 +1 @@ -Subproject commit 86b9ea3feacc8934e6e2a581c49eb4c37a2a3d20 +Subproject commit 8772b6fcfdac4774c0b135b2ba0863cb4118f2e3 diff --git a/alphapulldown/folding_backend/alphafold3_backend.py b/alphapulldown/folding_backend/alphafold3_backend.py index 07bca2c1..d9b71dbc 100644 --- a/alphapulldown/folding_backend/alphafold3_backend.py +++ b/alphapulldown/folding_backend/alphafold3_backend.py @@ -43,6 +43,7 @@ from alphafold.common.protein import Protein, to_mmcif from alphapulldown.folding_backend.folding_backend import FoldingBackend from alphapulldown.prediction.jax_compilation_cache import enable_persistent_compilation_cache +from alphapulldown.prediction import fast_kernels as fast_kernels_policy from alphapulldown.objects import MultimericObject, MonomericObject, ChoppedObject from alphapulldown.utils.af2_to_af3_msa import ( Af2ToAf3TranslationResult, @@ -122,6 +123,8 @@ class ModelRunner: config: base_config.BaseConfig device: jax.Device model_dir: pathlib.Path + kernel_mode: str = 'off' + kernel_reason: str = '' @functools.cached_property def model_params(self) -> hk.Params: @@ -625,6 +628,22 @@ def write_from_array(rows: np.ndarray, suffix: str): rng_key = jax.random.PRNGKey(seed) result = model_runner.run_inference(example, rng_key) + if output_dir is not None and hasattr(model_runner, 'kernel_mode'): + token_features = example.get('token_index', example.get('aatype')) + if token_features is not None: + record = fast_kernels_policy.af3_metadata( + model_runner.config.global_config, + requested_mode=model_runner.kernel_mode, + tokens=int(token_features.shape[0]), + reason=model_runner.kernel_reason, + ) + os.makedirs(output_dir, exist_ok=True) + metadata_path = pathlib.Path(output_dir) / 'inference_kernels.json' + previous = json.loads(metadata_path.read_text()) if metadata_path.exists() else {} + previous[f'seed_{seed}'] = record + temporary = metadata_path.with_suffix('.json.tmp') + temporary.write_text(json.dumps(previous, indent=2, sort_keys=True) + '\n') + temporary.replace(metadata_path) logging.info( f'Model inference for seed {seed} took {time.time() - inference_start_time:.2f} seconds.' ) @@ -763,6 +782,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.""" @@ -783,6 +803,8 @@ def make_model_config( config = model_class.Config() if hasattr(config, 'global_config'): config.global_config.flash_attention_implementation = flash_attention_implementation + for key, value in kernels.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 +834,8 @@ def make_model_config( f' include "{required_flag}".' ) logging.info(f'Found local devices: {gpu_devices}') + kernels = fast_kernels_policy.resolve( + fast_kernels, backend='alphafold3', device=gpu_devices[0]) logging.info('Building model from scratch...') model_runner = ModelRunner( @@ -828,6 +852,8 @@ def make_model_config( ), device=gpu_devices[0], model_dir=pathlib.Path(model_dir), + kernel_mode=fast_kernels_policy.normalise_mode(fast_kernels), + kernel_reason=kernels.reason, ) return {'model_runner': model_runner} diff --git a/alphapulldown/prediction/fast_kernels.py b/alphapulldown/prediction/fast_kernels.py index 379046c5..9eea0a95 100644 --- a/alphapulldown/prediction/fast_kernels.py +++ b/alphapulldown/prediction/fast_kernels.py @@ -1,4 +1,4 @@ -"""Decide whether AlphaFold 2 can use ColabFold's fused Pallas kernels here. +"""Resolve optional fused kernels before constructing AF2 or AF3 models. ``--fast_kernels`` turns on the optional fused kernels of the ``colabfold-kernels`` package (attention, LayerNorm, triangle multiplication) in AlphaFold-Multimer, through @@ -13,7 +13,8 @@ * ``on``: the kernels, or a ``ValueError`` saying which requirement is missing. * ``auto``: the kernels where they work, the stock code (with a log line) elsewhere. -Monomer models run in fp32 and the kernels need bf16, so only multimer models use them. +AF2 monomer models run in fp32 and the kernels need bf16, so only multimer models use them. +AF3 uses the fork's vendored triangle kernels with independent per-operation limits. JAX is imported only when a mode other than ``off`` is resolved. """ @@ -144,8 +145,12 @@ def _requirement_problem() -> tuple[Optional[str], Optional[int]]: return None, capability -def resolve(mode: str) -> KernelChoice: +def resolve(mode: str, *, backend="alphafold2", device=None) -> KernelChoice | AF3KernelChoice: """Turn a ``--fast_kernels`` value into a decision, checking requirements first.""" + if backend == "alphafold3": + return _resolve_af3(mode, device=device) + if backend != "alphafold2": + raise ValueError(f"Fused kernels are not supported by {backend}") mode = normalise_mode(mode) if mode == "off": return KernelChoice(False, "--fast_kernels=off") @@ -166,3 +171,106 @@ def resolve(mode: str) -> KernelChoice: raise ValueError(f"--fast_kernels=on, but {problem}.") logging.warning("Fused kernels off (--fast_kernels=auto): %s.", problem) return KernelChoice(False, f"--fast_kernels=auto: {problem}", compute_capability=capability) + + +@dataclasses.dataclass(frozen=True) +class AF3KernelChoice: + """Validated AF3 device policy, separate from AF2's optional dependency.""" + + enabled: bool + reason: str + settings: dict = dataclasses.field(default_factory=dict) + + def global_config_update(self) -> dict: + return dict(self.settings) if self.enabled else {} + + +def _af3_smoke_test(settings): + """Compile and execute both complete fused blocks before loading weights.""" + 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 + + config = model_config.GlobalConfig(final_init='linear', **settings) + def fn(x, mask): + with utils.bfloat16_context(): + x = modules.TriangleMultiplication( + modules.TriangleMultiplication.Config(equation='ikc,jkc->ijc'), + config, name='trimul')(x, mask) + return modules.GridSelfAttention( + modules.GridSelfAttention.Config(), config, transpose=True, + name='attention')(x, mask) + transformed = hk.without_apply_rng(hk.transform(fn)) + x = jnp.ones((64, 64, 128), jnp.bfloat16) + mask = jnp.ones((64, 64), jnp.bfloat16) + params = transformed.init(jax.random.PRNGKey(0), x, mask) + output = jax.jit(transformed.apply)(params, x, mask) + output.block_until_ready() + if not bool(jnp.all(jnp.isfinite(output))): + raise RuntimeError('AF3 fused smoke returned non-finite values') + + +def _af3_settings(device): + """Inspect the fork and allocator budget; do not load model parameters.""" + from alphafold3.model import model_config + from alphafold3.model.network.fused_triangle import fpf_pallas_serve as dispatch + + if not hasattr(model_config.GlobalConfig(), 'fused_triangle_multiplication'): + raise RuntimeError('the AlphaFold 3 fork has no fused-triangle hooks') + if device is None: + import jax + device = jax.local_devices(backend='gpu')[0] + memory = (device.memory_stats() or {}).get('bytes_limit', 0) / 2**30 + policy = dispatch.card_policy(getattr(device, 'compute_capability', ''), memory) + if not policy.enabled: + raise RuntimeError(f'{getattr(device, "device_kind", "GPU")}: {policy.reason}') + return dict(fused_triangle_multiplication=True, fused_triangle_attention='auto', + fused_triangle_compute_capability=policy.compute_capability, + fused_triangle_memory_gib=memory) + + +def _resolve_af3(mode, *, device=None): + mode = normalise_mode(mode) + if mode == 'off': + return AF3KernelChoice(False, '--fast_kernels=off') + try: + settings = _af3_settings(device) + _af3_smoke_test(settings) + except Exception as exc: + 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 AF3KernelChoice(False, problem) + logging.info('AF3 fused triangle kernels enabled; per-operation size limits apply: %s', settings) + return AF3KernelChoice(True, f'--fast_kernels={mode}', settings) + + +def af3_metadata(config, *, requested_mode, tokens, reason=''): + """Resolve provenance for this padded bucket, including cache-hit runs. + + Lists dispatch for AF3's standard C=128 pair and C=64 template modules. + It is not a trace counter or a claim that populated templates exist. + """ + enabled = (getattr(config, 'fused_triangle_multiplication', False) + or getattr(config, 'fused_triangle_attention', 'off') != 'off') + result = dict(backend='alphafold3', requested_mode=requested_mode, + padded_tokens=tokens, fused_kernels=False, reason=reason) + if not enabled: + return result + from alphafold3.model.network.fused_triangle import fpf_pallas_serve as dispatch + dtype = 'float32' if config.bfloat16 == 'none' else 'bfloat16' + selections = {} + for channels in (128, 64): + for kind in ('trimul', 'attention'): + backend, fallback = dispatch.select( + config, kind, (tokens, tokens, channels), dtype, (tokens, tokens)) + selections[f'{kind}_c{channels}'] = dict(backend=backend, reason=fallback) + result.update(fused_kernels=any(v['backend'] != 'stock' for v in selections.values()), + operations=selections, source_commit=dispatch.SOURCE_COMMIT, + policy_version=dispatch.POLICY_VERSION, + device_policy=dataclasses.asdict(dispatch.policy_from_config(config))) + return result diff --git a/alphapulldown/prediction/inference_flags.py b/alphapulldown/prediction/inference_flags.py index e8341d15..f694b92b 100644 --- a/alphapulldown/prediction/inference_flags.py +++ b/alphapulldown/prediction/inference_flags.py @@ -38,7 +38,7 @@ ALPHALINK_EXTRA_FLAGS = frozenset({"crosslinks"}) -# AlphaFold 2 only: AlphaLink runs its own network, and AF3 has its own kernels. +# AlphaLink runs its own network. AF3 accepts the same flag below. AF2_EXTRA_FLAGS = frozenset({"fast_kernels"}) UNIFOLD_UNAVAILABLE_REASON = ( @@ -51,7 +51,7 @@ "jax_compilation_cache_dir", "buckets", "flash_attention_implementation", "num_diffusion_samples", "num_seeds", "debug_templates", "debug_msas", "num_recycles", "save_embeddings", "save_distogram", "use_ap_style", - "convert_to_modelcif", + "convert_to_modelcif", "fast_kernels", }) FLAGS_BY_BACKEND: Mapping[str, frozenset] = { 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/docs/af3_fused_triangles.md b/docs/af3_fused_triangles.md new file mode 100644 index 00000000..8febb1f5 --- /dev/null +++ b/docs/af3_fused_triangles.md @@ -0,0 +1,41 @@ +# Experimental AF3 fused triangles + +Use `--fold_backend=alphafold3 --fast_kernels=on` to enable the bundled triangle +kernels. `off` is the default. `auto` enables them after a device check and a +small compile/run check, and otherwise logs the reason and uses stock AF3. +`on` reports an error if that check fails. Both modes fall back independently +for layers outside the supported dtype, shape or size limits. + +Install the matching `alphafold3` submodule revision in the AF3 environment, or +build an AF3 image from this branch. Existing released images do not contain +these hooks. AF3 does not need the AF2 `colabfold-kernels` optional package. + +The measured architectures are sm_80, sm_86, sm_89, sm_90 and sm_120. Multiplication +uses fused Pallas steps around XLA matrix multiplication. Attention uses the +kit's core on sm_80/90/120 and tokamax's core on sm_86/89. Only bf16 square pair +layers with 64 or 128 channels and four attention heads are enabled. Unmeasured +architectures fall back; compute capability alone is not evidence for new cards. + +Attention removes row chunking and needs more memory. Initial token limits use +the JAX allocator budget, rather than the card's advertised memory: + +| Allocator budget | Multiplication limit | Attention limit | +| --- | ---: | ---: | +| 12–<20 GiB | 1,536 | 768 | +| 20–<32 GiB | 2,560 | 1,536 | +| 32–<64 GiB | 3,584 (5,120 on sm_90/120) | 2,048 | +| ≥64 GiB | 3,584 (5,120 on sm_90/120) | 3,072 | + +These are conservative experimental limits, not guarantees of full-model +capacity. The kernels are off during parameter initialisation, preserving names, +shapes, initialisers and random-number ordering for existing DeepMind weights. + +Each output directory contains `inference_kernels.json`, keyed by seed, recording +the padded bucket, device policy, kernel source revision and selections for the +128-channel pair and 64-channel template layers. It records dispatch eligibility, +including cache hits; it does not claim that populated templates were supplied. + +Layer measurements justify an end-to-end pilot, not a production speed claim. +Keep this feature opt-in until the paired speed, memory and DeepMind-weight +accuracy validation is complete. Kernel source and attribution are recorded in +the fork's `fused_triangle/SOURCE.md` and `NOTICE`. diff --git a/docs/backend_installation.md b/docs/backend_installation.md index 2e1f9546..eb26f4c6 100644 --- a/docs/backend_installation.md +++ b/docs/backend_installation.md @@ -2,6 +2,9 @@ This guide is for direct AlphaPulldown use without Snakemake. +The experimental AF3 triangle-kernel option and its required fork revision are +described in [AF3 fused triangles](af3_fused_triangles.md). + Two points matter in practice: 1. Run the install commands from the AlphaPulldown repo root. diff --git a/test/unit/test_alphafold3_backend_helpers.py b/test/unit/test_alphafold3_backend_helpers.py index ab2ee3f1..155b7828 100644 --- a/test/unit/test_alphafold3_backend_helpers.py +++ b/test/unit/test_alphafold3_backend_helpers.py @@ -796,6 +796,7 @@ def test_predict_structure_writes_final_msa_and_collects_optional_outputs( example = { "msa": np.asarray([[1, 2, 3], [4, 5, 6]], dtype=np.int32), "num_alignments": 1, + "token_index": np.arange(3), } inference_result = af3_backend_module.model.InferenceResult( predicted_structure=None, @@ -814,6 +815,10 @@ def test_predict_structure_writes_final_msa_and_collects_optional_outputs( ) class FakeRunner: + kernel_mode = 'off' + kernel_reason = '--fast_kernels=off' + config = SimpleNamespace(global_config=SimpleNamespace()) + def run_inference(self, batch, rng_key): return { "single_embeddings": np.asarray([[1.0], [2.0], [3.0]], dtype=np.float32), @@ -840,6 +845,9 @@ def extract_structures(self, batch, result, target_name): assert results[0].distogram.shape == (2, 2) final_msa = tmp_path / "job_name_seed-7_final_complex_msa.a3m" assert final_msa.read_text(encoding="utf-8") == ">query\nABC\n" + metadata = json.loads((tmp_path / 'inference_kernels.json').read_text()) + assert metadata['seed_7']['padded_tokens'] == 3 + assert metadata['seed_7']['fused_kernels'] is False def test_af3_setup_builds_model_runner_and_validates_gpu_capability( @@ -918,6 +926,24 @@ class FakeModel: # Before the first trace, or the second prediction recompiles the model. assert tokamax_context == [True] + decisions = [] + settings = {'fused_triangle_multiplication': True, + 'fused_triangle_attention': 'auto', + 'fused_triangle_compute_capability': '8.0', + 'fused_triangle_memory_gib': 38.0} + def resolve(mode, *, backend, device): + decisions.append((mode, backend, device.compute_capability)) + return af3_backend_module.fast_kernels_policy.AF3KernelChoice(True, 'tested', settings) + monkeypatch.setattr(af3_backend_module.fast_kernels_policy, 'resolve', resolve) + fast_runner = 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'), fast_kernels='auto')['model_runner'] + assert decisions == [('auto', 'alphafold3', 8.0)] + for key, value in settings.items(): + assert getattr(fast_runner.config.global_config, key) == value + assert (fast_runner.kernel_mode, fast_runner.kernel_reason) == ('auto', 'tested') + monkeypatch.setattr( af3_backend_module.jax, "local_devices", diff --git a/test/unit/test_fast_kernels.py b/test/unit/test_fast_kernels.py index 923e098d..b5269db7 100644 --- a/test/unit/test_fast_kernels.py +++ b/test/unit/test_fast_kernels.py @@ -172,3 +172,90 @@ def missing(name): monkeypatch.setitem(sys.modules, "colabfold_kernels", types.SimpleNamespace(__version__="0.4.0")) assert fast_kernels._package_version() == "0.4.0" + + +def test_af3_off_does_not_inspect_or_import_the_fork(monkeypatch): + def fail(*args): + raise AssertionError('off must not inspect devices or import fused kernels') + monkeypatch.setattr(fast_kernels, '_af3_settings', fail) + monkeypatch.setattr(fast_kernels, '_af3_smoke_test', fail) + choice = fast_kernels.resolve('off', backend='alphafold3') + assert not choice.enabled + assert choice.global_config_update() == {} + assert fast_kernels.af3_metadata(types.SimpleNamespace(), requested_mode='off', tokens=256)['fused_kernels'] is False + + +@pytest.mark.parametrize('stage', ['settings', 'smoke']) +def test_af3_on_fails_and_auto_falls_back(monkeypatch, stage): + def fail(*args): + raise RuntimeError('unavailable test device') + monkeypatch.setattr(fast_kernels, '_af3_settings', lambda device: {'sentinel': True}) + monkeypatch.setattr(fast_kernels, '_af3_smoke_test', lambda settings: None) + monkeypatch.setattr(fast_kernels, '_af3_' + ('settings' if stage == 'settings' else 'smoke_test'), fail) + with pytest.raises(ValueError, match='unavailable test device'): + fast_kernels.resolve('on', backend='alphafold3') + choice = fast_kernels.resolve('auto', backend='alphafold3') + assert not choice.enabled and choice.global_config_update() == {} + + +def test_af3_uses_requested_device_and_smokes_before_enabling(monkeypatch): + device = object() + seen = [] + settings = {'fused_triangle_multiplication': True, 'fused_triangle_attention': 'auto'} + monkeypatch.setattr(fast_kernels, '_af3_settings', lambda d: seen.append(d) or settings) + monkeypatch.setattr(fast_kernels, '_af3_smoke_test', lambda s: seen.append(dict(s))) + choice = fast_kernels.resolve('on', backend='alphafold3', device=device) + assert seen == [device, settings] + assert choice.enabled and choice.global_config_update() == settings + choice.global_config_update().clear() + assert choice.global_config_update() == settings + + +def test_unknown_backend_cannot_enable_kernels(): + with pytest.raises(ValueError, match='not supported'): + fast_kernels.resolve('auto', backend='alphalink') + + +@pytest.fixture +def af3_dispatch(monkeypatch): + from pathlib import Path + source = Path(__file__).resolve().parents[2] / 'alphafold3/src/alphafold3/model/network/fused_triangle/fpf_pallas_serve.py' + spec = importlib.util.spec_from_file_location('af3_dispatch_test', source) + module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, spec.name, module) + spec.loader.exec_module(module) + package = types.ModuleType('alphafold3.model.network.fused_triangle') + package.fpf_pallas_serve = module + monkeypatch.setitem(sys.modules, package.__name__, package) + return module + + +def test_af3_metadata_records_partial_fallback_on_cache_hits(af3_dispatch): + config = types.SimpleNamespace( + fused_triangle_multiplication=True, fused_triangle_attention='auto', + fused_triangle_compute_capability='8.6', fused_triangle_memory_gib=45, + bfloat16='all') + first = fast_kernels.af3_metadata(config, requested_mode='auto', tokens=2560) + assert first == fast_kernels.af3_metadata(config, requested_mode='auto', tokens=2560) + assert first['fused_kernels'] + assert first['operations']['trimul_c128']['backend'] == 'pallas' + assert first['operations']['attention_c128'] == dict(backend='stock', reason='size_limit') + config.bfloat16 = 'none' + assert not fast_kernels.af3_metadata(config, requested_mode='auto', tokens=256)['fused_kernels'] + + +def test_af3_device_policy_uses_allocator_budget(af3_dispatch, monkeypatch): + config_module = types.ModuleType('alphafold3.model.model_config') + config_module.GlobalConfig = lambda: types.SimpleNamespace(fused_triangle_multiplication=False) + parent = types.ModuleType('alphafold3.model') + parent.model_config = config_module + monkeypatch.setitem(sys.modules, parent.__name__, parent) + monkeypatch.setitem(sys.modules, config_module.__name__, config_module) + device = types.SimpleNamespace(compute_capability='12.0', device_kind='Blackwell slice', + memory_stats=lambda: {'bytes_limit': 15 * 2**30}) + result = fast_kernels._af3_settings(device) + assert result['fused_triangle_memory_gib'] == 15 + assert result['fused_triangle_compute_capability'] == '12.0' + device.memory_stats = lambda: None + with pytest.raises(RuntimeError, match='unknown_memory_budget'): + fast_kernels._af3_settings(device) diff --git a/test/unit/test_inference_flags.py b/test/unit/test_inference_flags.py index a4472275..583113f5 100644 --- a/test/unit/test_inference_flags.py +++ b/test/unit/test_inference_flags.py @@ -53,9 +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(): +def test_fast_kernels_is_an_af2_and_af3_flag(): assert inference_flags.unsupported_flags("alphafold2", ["fast_kernels"]) == [] - for backend in ("alphafold3", "alphalink"): + assert inference_flags.unsupported_flags("alphafold3", ["fast_kernels"]) == [] + for backend in ("alphalink",): assert inference_flags.unsupported_flags(backend, ["fast_kernels"]) == ["fast_kernels"] From a2d47f16031946141cc3a979567143d8fdef976c Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 6 Oct 2026 18:17:33 +0200 Subject: [PATCH 02/11] chore: alphafold3 submodule at the fork's renamed fused-triangle API Move the submodule to KosinskiLab/alphafold3 73f260e (branch exp/af3-fused-triangle), where the fused triangle kernels moved under alphafold3/jax/fused_triangle and the config fields are named like flash_attention_implementation: fused_triangle_multiplication=True -> triangle_multiplication_implementation='pallas' fused_triangle_attention='auto' -> triangle_attention_implementation='auto' network.fused_triangle.fpf_pallas_serve.card_policy/select -> jax.fused_triangle.dispatch.device_policy/select_implementation The per-device choices are unchanged. Dispatch now names the original body 'default' (was 'stock', reason 'not_requested' instead of 'off'), so the per-seed record keys its operations triangle_multiplication_cN / triangle_attention_cN with an 'implementation' field. The fork no longer exposes the kit's source commit; its SOURCE.md records it, and the record keeps the dispatch policy version. Co-Authored-By: Claude Opus 5.5 --- alphafold3 | 2 +- alphapulldown/prediction/fast_kernels.py | 40 ++++++++++++-------- test/unit/test_alphafold3_backend_helpers.py | 4 +- test/unit/test_fast_kernels.py | 20 ++++++---- 4 files changed, 39 insertions(+), 27 deletions(-) diff --git a/alphafold3 b/alphafold3 index 8772b6fc..73f260ee 160000 --- a/alphafold3 +++ b/alphafold3 @@ -1 +1 @@ -Subproject commit 8772b6fcfdac4774c0b135b2ba0863cb4118f2e3 +Subproject commit 73f260eea1c307bd6daa28f1b8df01943f0b9d42 diff --git a/alphapulldown/prediction/fast_kernels.py b/alphapulldown/prediction/fast_kernels.py index 9eea0a95..25603b07 100644 --- a/alphapulldown/prediction/fast_kernels.py +++ b/alphapulldown/prediction/fast_kernels.py @@ -216,18 +216,20 @@ def fn(x, mask): def _af3_settings(device): """Inspect the fork and allocator budget; do not load model parameters.""" from alphafold3.model import model_config - from alphafold3.model.network.fused_triangle import fpf_pallas_serve as dispatch - if not hasattr(model_config.GlobalConfig(), 'fused_triangle_multiplication'): + if not hasattr(model_config.GlobalConfig(), 'triangle_multiplication_implementation'): raise RuntimeError('the AlphaFold 3 fork has no fused-triangle hooks') + from alphafold3.jax.fused_triangle import dispatch + if device is None: import jax device = jax.local_devices(backend='gpu')[0] memory = (device.memory_stats() or {}).get('bytes_limit', 0) / 2**30 - policy = dispatch.card_policy(getattr(device, 'compute_capability', ''), memory) + policy = dispatch.device_policy(getattr(device, 'compute_capability', ''), memory) if not policy.enabled: raise RuntimeError(f'{getattr(device, "device_kind", "GPU")}: {policy.reason}') - return dict(fused_triangle_multiplication=True, fused_triangle_attention='auto', + return dict(triangle_multiplication_implementation='pallas', + triangle_attention_implementation='auto', fused_triangle_compute_capability=policy.compute_capability, fused_triangle_memory_gib=memory) @@ -255,22 +257,28 @@ def af3_metadata(config, *, requested_mode, tokens, reason=''): Lists dispatch for AF3's standard C=128 pair and C=64 template modules. It is not a trace counter or a claim that populated templates exist. """ - enabled = (getattr(config, 'fused_triangle_multiplication', False) - or getattr(config, 'fused_triangle_attention', 'off') != 'off') + requested = { + operation: getattr(config, f'{operation}_implementation', 'default') + for operation in ('triangle_multiplication', 'triangle_attention') + } result = dict(backend='alphafold3', requested_mode=requested_mode, padded_tokens=tokens, fused_kernels=False, reason=reason) - if not enabled: + if all(value == 'default' for value in requested.values()): return result - from alphafold3.model.network.fused_triangle import fpf_pallas_serve as dispatch + from alphafold3.jax.fused_triangle import dispatch + policy = dispatch.device_policy(config.fused_triangle_compute_capability, + config.fused_triangle_memory_gib) dtype = 'float32' if config.bfloat16 == 'none' else 'bfloat16' selections = {} for channels in (128, 64): - for kind in ('trimul', 'attention'): - backend, fallback = dispatch.select( - config, kind, (tokens, tokens, channels), dtype, (tokens, tokens)) - selections[f'{kind}_c{channels}'] = dict(backend=backend, reason=fallback) - result.update(fused_kernels=any(v['backend'] != 'stock' for v in selections.values()), - operations=selections, source_commit=dispatch.SOURCE_COMMIT, - policy_version=dispatch.POLICY_VERSION, - device_policy=dataclasses.asdict(dispatch.policy_from_config(config))) + for operation, implementation in requested.items(): + selected, fallback = dispatch.select_implementation( + operation, implementation, policy, (tokens, tokens, channels), dtype, + (tokens, tokens)) + selections[f'{operation}_c{channels}'] = dict( + implementation=selected, reason=fallback) + result.update( + fused_kernels=any(v['implementation'] != 'default' for v in selections.values()), + operations=selections, policy_version=dispatch.POLICY_VERSION, + device_policy=dataclasses.asdict(policy)) return result diff --git a/test/unit/test_alphafold3_backend_helpers.py b/test/unit/test_alphafold3_backend_helpers.py index 155b7828..bb2ece84 100644 --- a/test/unit/test_alphafold3_backend_helpers.py +++ b/test/unit/test_alphafold3_backend_helpers.py @@ -927,8 +927,8 @@ class FakeModel: assert tokamax_context == [True] decisions = [] - settings = {'fused_triangle_multiplication': True, - 'fused_triangle_attention': 'auto', + settings = {'triangle_multiplication_implementation': 'pallas', + 'triangle_attention_implementation': 'auto', 'fused_triangle_compute_capability': '8.0', 'fused_triangle_memory_gib': 38.0} def resolve(mode, *, backend, device): diff --git a/test/unit/test_fast_kernels.py b/test/unit/test_fast_kernels.py index b5269db7..a766a9b6 100644 --- a/test/unit/test_fast_kernels.py +++ b/test/unit/test_fast_kernels.py @@ -201,7 +201,8 @@ def fail(*args): def test_af3_uses_requested_device_and_smokes_before_enabling(monkeypatch): device = object() seen = [] - settings = {'fused_triangle_multiplication': True, 'fused_triangle_attention': 'auto'} + settings = {'triangle_multiplication_implementation': 'pallas', + 'triangle_attention_implementation': 'auto'} monkeypatch.setattr(fast_kernels, '_af3_settings', lambda d: seen.append(d) or settings) monkeypatch.setattr(fast_kernels, '_af3_smoke_test', lambda s: seen.append(dict(s))) choice = fast_kernels.resolve('on', backend='alphafold3', device=device) @@ -219,34 +220,37 @@ def test_unknown_backend_cannot_enable_kernels(): @pytest.fixture def af3_dispatch(monkeypatch): from pathlib import Path - source = Path(__file__).resolve().parents[2] / 'alphafold3/src/alphafold3/model/network/fused_triangle/fpf_pallas_serve.py' + source = Path(__file__).resolve().parents[2] / 'alphafold3/src/alphafold3/jax/fused_triangle/dispatch.py' spec = importlib.util.spec_from_file_location('af3_dispatch_test', source) module = importlib.util.module_from_spec(spec) monkeypatch.setitem(sys.modules, spec.name, module) spec.loader.exec_module(module) - package = types.ModuleType('alphafold3.model.network.fused_triangle') - package.fpf_pallas_serve = module + package = types.ModuleType('alphafold3.jax.fused_triangle') + package.dispatch = module monkeypatch.setitem(sys.modules, package.__name__, package) return module def test_af3_metadata_records_partial_fallback_on_cache_hits(af3_dispatch): config = types.SimpleNamespace( - fused_triangle_multiplication=True, fused_triangle_attention='auto', + triangle_multiplication_implementation='pallas', + triangle_attention_implementation='auto', fused_triangle_compute_capability='8.6', fused_triangle_memory_gib=45, bfloat16='all') first = fast_kernels.af3_metadata(config, requested_mode='auto', tokens=2560) assert first == fast_kernels.af3_metadata(config, requested_mode='auto', tokens=2560) assert first['fused_kernels'] - assert first['operations']['trimul_c128']['backend'] == 'pallas' - assert first['operations']['attention_c128'] == dict(backend='stock', reason='size_limit') + assert first['operations']['triangle_multiplication_c128']['implementation'] == 'pallas' + assert first['operations']['triangle_attention_c128'] == dict( + implementation='default', reason='size_limit') config.bfloat16 = 'none' assert not fast_kernels.af3_metadata(config, requested_mode='auto', tokens=256)['fused_kernels'] def test_af3_device_policy_uses_allocator_budget(af3_dispatch, monkeypatch): config_module = types.ModuleType('alphafold3.model.model_config') - config_module.GlobalConfig = lambda: types.SimpleNamespace(fused_triangle_multiplication=False) + config_module.GlobalConfig = lambda: types.SimpleNamespace( + triangle_multiplication_implementation='default') parent = types.ModuleType('alphafold3.model') parent.model_config = config_module monkeypatch.setitem(sys.modules, parent.__name__, parent) From e100030c95e574d4b8e5025ee9fb8cf11f2326f0 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 6 Oct 2026 19:52:17 +0200 Subject: [PATCH 03/11] refactor: AF3 fused-triangle policy in its own module Move the AlphaFold 3 decision out of prediction/fast_kernels.py into prediction/af3_fused_triangles.py, so fast_kernels.py is AF2-only again, byte-identical to main, and its resolve(mode) API is unchanged. The new module reuses fast_kernels.normalise_mode and exposes: resolve(mode, device) -> FusedTriangleChoice (enabled, reason, mode, settings; global_config_update()) OFF the choice for --fast_kernels=off device_settings(device) GlobalConfig fields for the device (was _af3_settings) smoke_test(settings) (was _af3_smoke_test) metadata(global_config, choice, *, num_tokens) (was af3_metadata) ModelRunner keeps the resolved FusedTriangleChoice (fused_triangles) instead of a re-normalised mode string and a reason. Decisions, the GlobalConfig fields set and the record's fields are unchanged. The AF3 tests move from test_fast_kernels.py (back to main's) to test_af3_fused_triangles.py, and the backend's own checks are separate tests instead of additions to the MSA and setup tests. The light-import test covers the new module: the AF3 backend imports it, and it imports neither JAX nor the fork until a kernel mode is resolved. Co-Authored-By: Claude Opus 5.5 --- .../folding_backend/alphafold3_backend.py | 22 +- .../prediction/af3_fused_triangles.py | 207 ++++++++++++++++ alphapulldown/prediction/fast_kernels.py | 122 +--------- test/unit/test_af3_fused_triangles.py | 229 ++++++++++++++++++ test/unit/test_alphafold3_backend_helpers.py | 136 +++++++++-- test/unit/test_fast_kernels.py | 91 ------- test/unit/test_light_import_invariant.py | 2 + 7 files changed, 560 insertions(+), 249 deletions(-) create mode 100644 alphapulldown/prediction/af3_fused_triangles.py create mode 100644 test/unit/test_af3_fused_triangles.py diff --git a/alphapulldown/folding_backend/alphafold3_backend.py b/alphapulldown/folding_backend/alphafold3_backend.py index d9b71dbc..e39f7657 100644 --- a/alphapulldown/folding_backend/alphafold3_backend.py +++ b/alphapulldown/folding_backend/alphafold3_backend.py @@ -42,8 +42,8 @@ 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.prediction import fast_kernels as fast_kernels_policy from alphapulldown.objects import MultimericObject, MonomericObject, ChoppedObject from alphapulldown.utils.af2_to_af3_msa import ( Af2ToAf3TranslationResult, @@ -123,8 +123,7 @@ class ModelRunner: config: base_config.BaseConfig device: jax.Device model_dir: pathlib.Path - kernel_mode: str = 'off' - kernel_reason: str = '' + fused_triangles: af3_fused_triangles.FusedTriangleChoice = af3_fused_triangles.OFF @functools.cached_property def model_params(self) -> hk.Params: @@ -628,14 +627,13 @@ def write_from_array(rows: np.ndarray, suffix: str): rng_key = jax.random.PRNGKey(seed) result = model_runner.run_inference(example, rng_key) - if output_dir is not None and hasattr(model_runner, 'kernel_mode'): + if output_dir is not None and hasattr(model_runner, 'fused_triangles'): token_features = example.get('token_index', example.get('aatype')) if token_features is not None: - record = fast_kernels_policy.af3_metadata( + record = af3_fused_triangles.metadata( model_runner.config.global_config, - requested_mode=model_runner.kernel_mode, - tokens=int(token_features.shape[0]), - reason=model_runner.kernel_reason, + model_runner.fused_triangles, + num_tokens=int(token_features.shape[0]), ) os.makedirs(output_dir, exist_ok=True) metadata_path = pathlib.Path(output_dir) / 'inference_kernels.json' @@ -803,7 +801,7 @@ def make_model_config( config = model_class.Config() if hasattr(config, 'global_config'): config.global_config.flash_attention_implementation = flash_attention_implementation - for key, value in kernels.global_config_update().items(): + for key, value in fused_triangles.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 @@ -834,8 +832,7 @@ def make_model_config( f' include "{required_flag}".' ) logging.info(f'Found local devices: {gpu_devices}') - kernels = fast_kernels_policy.resolve( - fast_kernels, backend='alphafold3', device=gpu_devices[0]) + fused_triangles = af3_fused_triangles.resolve(fast_kernels, gpu_devices[0]) logging.info('Building model from scratch...') model_runner = ModelRunner( @@ -852,8 +849,7 @@ def make_model_config( ), device=gpu_devices[0], model_dir=pathlib.Path(model_dir), - kernel_mode=fast_kernels_policy.normalise_mode(fast_kernels), - kernel_reason=kernels.reason, + 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..7d2ce882 --- /dev/null +++ b/alphapulldown/prediction/af3_fused_triangles.py @@ -0,0 +1,207 @@ +"""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]) -> None: + """Compile and run both fused triangle layers, 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)) + 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) + if not bool(jnp.all(jnp.isfinite(output))): + 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) + 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, 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/fast_kernels.py b/alphapulldown/prediction/fast_kernels.py index 25603b07..379046c5 100644 --- a/alphapulldown/prediction/fast_kernels.py +++ b/alphapulldown/prediction/fast_kernels.py @@ -1,4 +1,4 @@ -"""Resolve optional fused kernels before constructing AF2 or AF3 models. +"""Decide whether AlphaFold 2 can use ColabFold's fused Pallas kernels here. ``--fast_kernels`` turns on the optional fused kernels of the ``colabfold-kernels`` package (attention, LayerNorm, triangle multiplication) in AlphaFold-Multimer, through @@ -13,8 +13,7 @@ * ``on``: the kernels, or a ``ValueError`` saying which requirement is missing. * ``auto``: the kernels where they work, the stock code (with a log line) elsewhere. -AF2 monomer models run in fp32 and the kernels need bf16, so only multimer models use them. -AF3 uses the fork's vendored triangle kernels with independent per-operation limits. +Monomer models run in fp32 and the kernels need bf16, so only multimer models use them. JAX is imported only when a mode other than ``off`` is resolved. """ @@ -145,12 +144,8 @@ def _requirement_problem() -> tuple[Optional[str], Optional[int]]: return None, capability -def resolve(mode: str, *, backend="alphafold2", device=None) -> KernelChoice | AF3KernelChoice: +def resolve(mode: str) -> KernelChoice: """Turn a ``--fast_kernels`` value into a decision, checking requirements first.""" - if backend == "alphafold3": - return _resolve_af3(mode, device=device) - if backend != "alphafold2": - raise ValueError(f"Fused kernels are not supported by {backend}") mode = normalise_mode(mode) if mode == "off": return KernelChoice(False, "--fast_kernels=off") @@ -171,114 +166,3 @@ def resolve(mode: str, *, backend="alphafold2", device=None) -> KernelChoice | A raise ValueError(f"--fast_kernels=on, but {problem}.") logging.warning("Fused kernels off (--fast_kernels=auto): %s.", problem) return KernelChoice(False, f"--fast_kernels=auto: {problem}", compute_capability=capability) - - -@dataclasses.dataclass(frozen=True) -class AF3KernelChoice: - """Validated AF3 device policy, separate from AF2's optional dependency.""" - - enabled: bool - reason: str - settings: dict = dataclasses.field(default_factory=dict) - - def global_config_update(self) -> dict: - return dict(self.settings) if self.enabled else {} - - -def _af3_smoke_test(settings): - """Compile and execute both complete fused blocks before loading weights.""" - 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 - - config = model_config.GlobalConfig(final_init='linear', **settings) - def fn(x, mask): - with utils.bfloat16_context(): - x = modules.TriangleMultiplication( - modules.TriangleMultiplication.Config(equation='ikc,jkc->ijc'), - config, name='trimul')(x, mask) - return modules.GridSelfAttention( - modules.GridSelfAttention.Config(), config, transpose=True, - name='attention')(x, mask) - transformed = hk.without_apply_rng(hk.transform(fn)) - x = jnp.ones((64, 64, 128), jnp.bfloat16) - mask = jnp.ones((64, 64), jnp.bfloat16) - params = transformed.init(jax.random.PRNGKey(0), x, mask) - output = jax.jit(transformed.apply)(params, x, mask) - output.block_until_ready() - if not bool(jnp.all(jnp.isfinite(output))): - raise RuntimeError('AF3 fused smoke returned non-finite values') - - -def _af3_settings(device): - """Inspect the fork and allocator budget; do not load model parameters.""" - from alphafold3.model import model_config - - if not hasattr(model_config.GlobalConfig(), 'triangle_multiplication_implementation'): - raise RuntimeError('the AlphaFold 3 fork has no fused-triangle hooks') - from alphafold3.jax.fused_triangle import dispatch - - if device is None: - import jax - device = jax.local_devices(backend='gpu')[0] - memory = (device.memory_stats() or {}).get('bytes_limit', 0) / 2**30 - policy = dispatch.device_policy(getattr(device, 'compute_capability', ''), memory) - if not policy.enabled: - raise RuntimeError(f'{getattr(device, "device_kind", "GPU")}: {policy.reason}') - return dict(triangle_multiplication_implementation='pallas', - triangle_attention_implementation='auto', - fused_triangle_compute_capability=policy.compute_capability, - fused_triangle_memory_gib=memory) - - -def _resolve_af3(mode, *, device=None): - mode = normalise_mode(mode) - if mode == 'off': - return AF3KernelChoice(False, '--fast_kernels=off') - try: - settings = _af3_settings(device) - _af3_smoke_test(settings) - except Exception as exc: - 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 AF3KernelChoice(False, problem) - logging.info('AF3 fused triangle kernels enabled; per-operation size limits apply: %s', settings) - return AF3KernelChoice(True, f'--fast_kernels={mode}', settings) - - -def af3_metadata(config, *, requested_mode, tokens, reason=''): - """Resolve provenance for this padded bucket, including cache-hit runs. - - Lists dispatch for AF3's standard C=128 pair and C=64 template modules. - It is not a trace counter or a claim that populated templates exist. - """ - requested = { - operation: getattr(config, f'{operation}_implementation', 'default') - for operation in ('triangle_multiplication', 'triangle_attention') - } - result = dict(backend='alphafold3', requested_mode=requested_mode, - padded_tokens=tokens, fused_kernels=False, reason=reason) - if all(value == 'default' for value in requested.values()): - return result - from alphafold3.jax.fused_triangle import dispatch - policy = dispatch.device_policy(config.fused_triangle_compute_capability, - config.fused_triangle_memory_gib) - dtype = 'float32' if config.bfloat16 == 'none' else 'bfloat16' - selections = {} - for channels in (128, 64): - for operation, implementation in requested.items(): - selected, fallback = dispatch.select_implementation( - operation, implementation, policy, (tokens, tokens, channels), dtype, - (tokens, tokens)) - selections[f'{operation}_c{channels}'] = dict( - implementation=selected, reason=fallback) - result.update( - fused_kernels=any(v['implementation'] != 'default' for v in selections.values()), - operations=selections, policy_version=dispatch.POLICY_VERSION, - device_policy=dataclasses.asdict(policy)) - return result diff --git a/test/unit/test_af3_fused_triangles.py b/test/unit/test_af3_fused_triangles.py new file mode 100644 index 00000000..4d7a51aa --- /dev/null +++ b/test/unit/test_af3_fused_triangles.py @@ -0,0 +1,229 @@ +"""--fast_kernels for AlphaFold 3: checking the device before a model is built.""" + +from __future__ import annotations + +import importlib.util +import sys +import types +from pathlib import Path + +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 "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][0] == SETTINGS + 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_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 bb2ece84..9db9e0fd 100644 --- a/test/unit/test_alphafold3_backend_helpers.py +++ b/test/unit/test_alphafold3_backend_helpers.py @@ -796,7 +796,6 @@ def test_predict_structure_writes_final_msa_and_collects_optional_outputs( example = { "msa": np.asarray([[1, 2, 3], [4, 5, 6]], dtype=np.int32), "num_alignments": 1, - "token_index": np.arange(3), } inference_result = af3_backend_module.model.InferenceResult( predicted_structure=None, @@ -815,10 +814,6 @@ def test_predict_structure_writes_final_msa_and_collects_optional_outputs( ) class FakeRunner: - kernel_mode = 'off' - kernel_reason = '--fast_kernels=off' - config = SimpleNamespace(global_config=SimpleNamespace()) - def run_inference(self, batch, rng_key): return { "single_embeddings": np.asarray([[1.0], [2.0], [3.0]], dtype=np.float32), @@ -845,9 +840,45 @@ def extract_structures(self, batch, result, target_name): assert results[0].distogram.shape == (2, 2) final_msa = tmp_path / "job_name_seed-7_final_complex_msa.a3m" assert final_msa.read_text(encoding="utf-8") == ">query\nABC\n" - metadata = json.loads((tmp_path / 'inference_kernels.json').read_text()) - assert metadata['seed_7']['padded_tokens'] == 3 - assert metadata['seed_7']['fused_kernels'] is False + + +def test_predict_structure_records_fused_triangle_kernels_per_seed( + af3_backend_module, + monkeypatch, + tmp_path, +): + fold_input = af3_backend_module.folding_input.Input( + name="job", chains=("A",), rng_seeds=(7,) + ) + example = {"token_index": np.arange(3)} + monkeypatch.setattr( + af3_backend_module.featurisation, "featurise_input", lambda **kwargs: [example] + ) + + class FakeRunner: + fused_triangles = af3_backend_module.af3_fused_triangles.OFF + config = SimpleNamespace(global_config=SimpleNamespace()) + + def run_inference(self, batch, rng_key): + return {} + + def extract_structures(self, batch, result, target_name): + return [] + + af3_backend_module.predict_structure( + fold_input=fold_input, model_runner=FakeRunner(), output_dir=tmp_path + ) + + record = json.loads((tmp_path / "inference_kernels.json").read_text()) + assert record == { + "seed_7": { + "backend": "alphafold3", + "requested_mode": "off", + "padded_tokens": 3, + "fused_kernels": False, + "reason": "--fast_kernels=off", + } + } def test_af3_setup_builds_model_runner_and_validates_gpu_capability( @@ -926,24 +957,6 @@ class FakeModel: # Before the first trace, or the second prediction recompiles the model. assert tokamax_context == [True] - decisions = [] - settings = {'triangle_multiplication_implementation': 'pallas', - 'triangle_attention_implementation': 'auto', - 'fused_triangle_compute_capability': '8.0', - 'fused_triangle_memory_gib': 38.0} - def resolve(mode, *, backend, device): - decisions.append((mode, backend, device.compute_capability)) - return af3_backend_module.fast_kernels_policy.AF3KernelChoice(True, 'tested', settings) - monkeypatch.setattr(af3_backend_module.fast_kernels_policy, 'resolve', resolve) - fast_runner = 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'), fast_kernels='auto')['model_runner'] - assert decisions == [('auto', 'alphafold3', 8.0)] - for key, value in settings.items(): - assert getattr(fast_runner.config.global_config, key) == value - assert (fast_runner.kernel_mode, fast_runner.kernel_reason) == ('auto', 'tested') - monkeypatch.setattr( af3_backend_module.jax, "local_devices", @@ -959,6 +972,77 @@ def resolve(mode, *, backend, device): ) +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_fast_kernels.py b/test/unit/test_fast_kernels.py index a766a9b6..923e098d 100644 --- a/test/unit/test_fast_kernels.py +++ b/test/unit/test_fast_kernels.py @@ -172,94 +172,3 @@ def missing(name): monkeypatch.setitem(sys.modules, "colabfold_kernels", types.SimpleNamespace(__version__="0.4.0")) assert fast_kernels._package_version() == "0.4.0" - - -def test_af3_off_does_not_inspect_or_import_the_fork(monkeypatch): - def fail(*args): - raise AssertionError('off must not inspect devices or import fused kernels') - monkeypatch.setattr(fast_kernels, '_af3_settings', fail) - monkeypatch.setattr(fast_kernels, '_af3_smoke_test', fail) - choice = fast_kernels.resolve('off', backend='alphafold3') - assert not choice.enabled - assert choice.global_config_update() == {} - assert fast_kernels.af3_metadata(types.SimpleNamespace(), requested_mode='off', tokens=256)['fused_kernels'] is False - - -@pytest.mark.parametrize('stage', ['settings', 'smoke']) -def test_af3_on_fails_and_auto_falls_back(monkeypatch, stage): - def fail(*args): - raise RuntimeError('unavailable test device') - monkeypatch.setattr(fast_kernels, '_af3_settings', lambda device: {'sentinel': True}) - monkeypatch.setattr(fast_kernels, '_af3_smoke_test', lambda settings: None) - monkeypatch.setattr(fast_kernels, '_af3_' + ('settings' if stage == 'settings' else 'smoke_test'), fail) - with pytest.raises(ValueError, match='unavailable test device'): - fast_kernels.resolve('on', backend='alphafold3') - choice = fast_kernels.resolve('auto', backend='alphafold3') - assert not choice.enabled and choice.global_config_update() == {} - - -def test_af3_uses_requested_device_and_smokes_before_enabling(monkeypatch): - device = object() - seen = [] - settings = {'triangle_multiplication_implementation': 'pallas', - 'triangle_attention_implementation': 'auto'} - monkeypatch.setattr(fast_kernels, '_af3_settings', lambda d: seen.append(d) or settings) - monkeypatch.setattr(fast_kernels, '_af3_smoke_test', lambda s: seen.append(dict(s))) - choice = fast_kernels.resolve('on', backend='alphafold3', device=device) - assert seen == [device, settings] - assert choice.enabled and choice.global_config_update() == settings - choice.global_config_update().clear() - assert choice.global_config_update() == settings - - -def test_unknown_backend_cannot_enable_kernels(): - with pytest.raises(ValueError, match='not supported'): - fast_kernels.resolve('auto', backend='alphalink') - - -@pytest.fixture -def af3_dispatch(monkeypatch): - from pathlib import Path - source = Path(__file__).resolve().parents[2] / 'alphafold3/src/alphafold3/jax/fused_triangle/dispatch.py' - spec = importlib.util.spec_from_file_location('af3_dispatch_test', source) - 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 - - -def test_af3_metadata_records_partial_fallback_on_cache_hits(af3_dispatch): - config = types.SimpleNamespace( - triangle_multiplication_implementation='pallas', - triangle_attention_implementation='auto', - fused_triangle_compute_capability='8.6', fused_triangle_memory_gib=45, - bfloat16='all') - first = fast_kernels.af3_metadata(config, requested_mode='auto', tokens=2560) - assert first == fast_kernels.af3_metadata(config, requested_mode='auto', tokens=2560) - assert first['fused_kernels'] - assert first['operations']['triangle_multiplication_c128']['implementation'] == 'pallas' - assert first['operations']['triangle_attention_c128'] == dict( - implementation='default', reason='size_limit') - config.bfloat16 = 'none' - assert not fast_kernels.af3_metadata(config, requested_mode='auto', tokens=256)['fused_kernels'] - - -def test_af3_device_policy_uses_allocator_budget(af3_dispatch, monkeypatch): - config_module = types.ModuleType('alphafold3.model.model_config') - config_module.GlobalConfig = lambda: types.SimpleNamespace( - triangle_multiplication_implementation='default') - parent = types.ModuleType('alphafold3.model') - parent.model_config = config_module - monkeypatch.setitem(sys.modules, parent.__name__, parent) - monkeypatch.setitem(sys.modules, config_module.__name__, config_module) - device = types.SimpleNamespace(compute_capability='12.0', device_kind='Blackwell slice', - memory_stats=lambda: {'bytes_limit': 15 * 2**30}) - result = fast_kernels._af3_settings(device) - assert result['fused_triangle_memory_gib'] == 15 - assert result['fused_triangle_compute_capability'] == '12.0' - device.memory_stats = lambda: None - with pytest.raises(RuntimeError, match='unknown_memory_budget'): - fast_kernels._af3_settings(device) 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", From 4bed2a4d40a5109649f6973250aa55db22ad2ea0 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 6 Oct 2026 19:52:43 +0200 Subject: [PATCH 04/11] fix: AF3 kernel check runs on the GPU it checks; reason names the mode smoke_test compiled and ran the two fused triangle layers on JAX's default device, while device_settings had read the limits of the device the model runs on. It now takes that device and runs inside jax.default_device(device), so the check and the model see the same GPU. An --fast_kernels=auto fallback now records its reason as "--fast_kernels=auto: ...", like AlphaFold 2's, instead of the bare problem, so inference_kernels.json says which mode fell back. Co-Authored-By: Claude Opus 5.5 --- .../prediction/af3_fused_triangles.py | 20 ++++--- test/unit/test_af3_fused_triangles.py | 59 ++++++++++++++++++- 2 files changed, 69 insertions(+), 10 deletions(-) diff --git a/alphapulldown/prediction/af3_fused_triangles.py b/alphapulldown/prediction/af3_fused_triangles.py index 7d2ce882..3d4b4314 100644 --- a/alphapulldown/prediction/af3_fused_triangles.py +++ b/alphapulldown/prediction/af3_fused_triangles.py @@ -85,8 +85,8 @@ def device_settings(device: jax.Device) -> dict[str, Any]: } -def smoke_test(settings: Mapping[str, Any]) -> None: - """Compile and run both fused triangle layers, without model weights. +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``. @@ -119,11 +119,13 @@ def forward(act, mask): )(act, mask) transformed = hk.without_apply_rng(hk.transform(forward)) - 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) - if not bool(jnp.all(jnp.isfinite(output))): + 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") @@ -139,13 +141,13 @@ def resolve(mode: Any, device: jax.Device) -> FusedTriangleChoice: return OFF try: settings = device_settings(device) - smoke_test(settings) + 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, problem, mode) + return FusedTriangleChoice(False, f"--fast_kernels=auto: {problem}", mode) logging.info( "AF3 fused triangle kernels on; per-operation size limits apply: %s", settings ) diff --git a/test/unit/test_af3_fused_triangles.py b/test/unit/test_af3_fused_triangles.py index 4d7a51aa..c424f0a2 100644 --- a/test/unit/test_af3_fused_triangles.py +++ b/test/unit/test_af3_fused_triangles.py @@ -2,11 +2,13 @@ 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 @@ -96,6 +98,7 @@ def test_on_fails_loudly_and_auto_falls_back(monkeypatch, stage): 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() == {} @@ -113,7 +116,7 @@ def test_checks_the_given_device_before_enabling(monkeypatch): choice = af3_fused_triangles.resolve("on", device) assert seen[0] is device - assert seen[1][0] == SETTINGS + 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. @@ -121,6 +124,60 @@ def test_checks_the_given_device_before_enabling(monkeypatch): 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", From 796fc70274ba868a56679e732369afc18725d219 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 6 Oct 2026 19:54:09 +0200 Subject: [PATCH 05/11] fix: AF3 kernel record never fails a prediction, keyed seed-N The per-seed inference_kernels.json update was inline in the seed loop: a corrupt existing file, or a fork without the dispatch module, raised out of predict_structure after the model had run, and its time was counted in "Model inference for seed N took". It is now _record_inference_kernels, mirroring AlphaFold 2's: it runs after structure extraction, replaces an unreadable file with a warning, and logs and skips any other failure. Entries are keyed seed-N, AF3's own spelling in output names, instead of seed_N. make_model_config read the fused-triangle decision from the enclosing scope, assigned only later in setup(); the GlobalConfig update is now passed in explicitly. Co-Authored-By: Claude Opus 5.5 --- .../folding_backend/alphafold3_backend.py | 65 +++++++--- test/unit/test_alphafold3_backend_helpers.py | 116 ++++++++++++++---- 2 files changed, 144 insertions(+), 37 deletions(-) diff --git a/alphapulldown/folding_backend/alphafold3_backend.py b/alphapulldown/folding_backend/alphafold3_backend.py index e39f7657..328e9ace 100644 --- a/alphapulldown/folding_backend/alphafold3_backend.py +++ b/alphapulldown/folding_backend/alphafold3_backend.py @@ -568,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, @@ -627,21 +671,6 @@ def write_from_array(rows: np.ndarray, suffix: str): rng_key = jax.random.PRNGKey(seed) result = model_runner.run_inference(example, rng_key) - if output_dir is not None and hasattr(model_runner, 'fused_triangles'): - token_features = example.get('token_index', example.get('aatype')) - if token_features is not None: - record = af3_fused_triangles.metadata( - model_runner.config.global_config, - model_runner.fused_triangles, - num_tokens=int(token_features.shape[0]), - ) - os.makedirs(output_dir, exist_ok=True) - metadata_path = pathlib.Path(output_dir) / 'inference_kernels.json' - previous = json.loads(metadata_path.read_text()) if metadata_path.exists() else {} - previous[f'seed_{seed}'] = record - temporary = metadata_path.with_suffix('.json.tmp') - temporary.write_text(json.dumps(previous, indent=2, sort_keys=True) + '\n') - temporary.replace(metadata_path) logging.info( f'Model inference for seed {seed} took {time.time() - inference_start_time:.2f} seconds.' ) @@ -654,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 @@ -796,12 +827,13 @@ 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 fused_triangles.global_config_update().items(): + 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 @@ -846,6 +878,7 @@ 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), diff --git a/test/unit/test_alphafold3_backend_helpers.py b/test/unit/test_alphafold3_backend_helpers.py index 9db9e0fd..0433cb3f 100644 --- a/test/unit/test_alphafold3_backend_helpers.py +++ b/test/unit/test_alphafold3_backend_helpers.py @@ -842,45 +842,119 @@ def extract_structures(self, batch, result, target_name): assert final_msa.read_text(encoding="utf-8") == ">query\nABC\n" -def test_predict_structure_records_fused_triangle_kernels_per_seed( +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,) + 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] + 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), ) - class FakeRunner: - fused_triangles = af3_backend_module.af3_fused_triangles.OFF - config = SimpleNamespace(global_config=SimpleNamespace()) + af3_backend_module.predict_structure( + fold_input=fold_input, + model_runner=FakeKernelRunner(af3_backend_module, events), + output_dir=tmp_path, + ) - def run_inference(self, batch, rng_key): - return {} + # 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 extract_structures(self, batch, result, target_name): - return [] - af3_backend_module.predict_structure( - fold_input=fold_input, model_runner=FakeRunner(), output_dir=tmp_path +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)} ) - record = json.loads((tmp_path / "inference_kernels.json").read_text()) - assert record == { - "seed_7": { - "backend": "alphafold3", - "requested_mode": "off", - "padded_tokens": 3, - "fused_kernels": False, - "reason": "--fast_kernels=off", - } + 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, From 2748c80742b61eb277902ebb66bcd12e4d61d651 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 6 Oct 2026 19:54:34 +0200 Subject: [PATCH 06/11] refactor: one FAST_KERNEL_FLAGS set for AlphaFold 2 and 3 --fast_kernels was listed twice, as AF2_EXTRA_FLAGS and inside AF3_FLAGS, under a comment that still said AlphaFold 2 only. Both backends now take it from one FAST_KERNEL_FLAGS set, whose comment says which kernels each backend uses and why AlphaLink has none. The flags each backend accepts are unchanged. Co-Authored-By: Claude Opus 5.5 --- alphapulldown/prediction/inference_flags.py | 12 +++++++----- test/unit/test_inference_flags.py | 9 ++++----- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/alphapulldown/prediction/inference_flags.py b/alphapulldown/prediction/inference_flags.py index f694b92b..ec78565b 100644 --- a/alphapulldown/prediction/inference_flags.py +++ b/alphapulldown/prediction/inference_flags.py @@ -38,8 +38,10 @@ ALPHALINK_EXTRA_FLAGS = frozenset({"crosslinks"}) -# AlphaLink runs its own network. AF3 accepts the same flag below. -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 " @@ -51,13 +53,13 @@ "jax_compilation_cache_dir", "buckets", "flash_attention_implementation", "num_diffusion_samples", "num_seeds", "debug_templates", "debug_msas", "num_recycles", "save_embeddings", "save_distogram", "use_ap_style", - "convert_to_modelcif", "fast_kernels", + "convert_to_modelcif", }) 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/test/unit/test_inference_flags.py b/test/unit/test_inference_flags.py index 583113f5..feadb8b2 100644 --- a/test/unit/test_inference_flags.py +++ b/test/unit/test_inference_flags.py @@ -53,11 +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_af2_and_af3_flag(): - assert inference_flags.unsupported_flags("alphafold2", ["fast_kernels"]) == [] - assert inference_flags.unsupported_flags("alphafold3", ["fast_kernels"]) == [] - for backend in ("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(): From c235cd8562cfc374c24e5e3257fa9fea0ab1ef0f Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 6 Oct 2026 19:56:46 +0200 Subject: [PATCH 07/11] docs: AF3 fused triangle kernels in sections, with terms and measurements Rewrite docs/af3_fused_triangles.md in sections like docs/mmseqs2_rna.md: turning the kernels on, what runs on which GPU, memory limits, how they were measured, the output record, troubleshooting, and source and licence. Define the terms the old text assumed (Pallas, the attention core, Anthropic's Pallas flash-attention core vs tokamax's), drop "this branch", and state the measurement: Gate 0 and the full-model campaign on A100, A40, RTX 3090, L40S, H100 and RTX Pro 6000 Blackwell with JAX 0.9.1 / tokamax 0.0.11; Pallas attention on compute capability 8.0/9.0/12.0, the tokamax core on 8.6/8.9, multiplication on all of them, nothing on 10.0 (B200) or newer. The record section follows the renamed fields and seed-N keys. Troubleshooting says that scripts importing the AF3 backend should call jax.devices() first. Link the page from the README's AlphaFold3 flags, and move the link in backend_installation.md from the top into the AlphaFold3 backend notes. Co-Authored-By: Claude Opus 5.5 --- README.md | 3 + docs/af3_fused_triangles.md | 195 ++++++++++++++++++++++++++++------- docs/backend_installation.md | 4 +- 3 files changed, 164 insertions(+), 38 deletions(-) diff --git a/README.md b/README.md index 90c342fe..7eeb0aa8 100644 --- a/README.md +++ b/README.md @@ -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/docs/af3_fused_triangles.md b/docs/af3_fused_triangles.md index 8febb1f5..0b0d07c7 100644 --- a/docs/af3_fused_triangles.md +++ b/docs/af3_fused_triangles.md @@ -1,41 +1,166 @@ -# Experimental AF3 fused triangles +# Fused triangle kernels for AlphaFold 3 (experimental) -Use `--fold_backend=alphafold3 --fast_kernels=on` to enable the bundled triangle -kernels. `off` is the default. `auto` enables them after a device check and a -small compile/run check, and otherwise logs the reason and uses stock AF3. -`on` reports an error if that check fails. Both modes fall back independently -for layers outside the supported dtype, shape or size limits. +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. -Install the matching `alphafold3` submodule revision in the AF3 environment, or -build an AF3 image from this branch. Existing released images do not contain -these hooks. AF3 does not need the AF2 `colabfold-kernels` optional package. +The same flag turns on ColabFold's fused kernels for AlphaFold-Multimer, a separate +implementation; AlphaFold 3 does not need the `colabfold-kernels` package. -The measured architectures are sm_80, sm_86, sm_89, sm_90 and sm_120. Multiplication -uses fused Pallas steps around XLA matrix multiplication. Attention uses the -kit's core on sm_80/90/120 and tokamax's core on sm_86/89. Only bf16 square pair -layers with 64 or 128 channels and four attention heads are enabled. Unmeasured -architectures fall back; compute capability alone is not evidence for new cards. +## Turning them on -Attention removes row chunking and needs more memory. Initial token limits use -the JAX allocator budget, rather than the card's advertised memory: +```bash +run_structure_prediction.py --fold_backend=alphafold3 --fast_kernels=auto ... +``` -| Allocator budget | Multiplication limit | Attention limit | +| 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 | | --- | ---: | ---: | -| 12–<20 GiB | 1,536 | 768 | -| 20–<32 GiB | 2,560 | 1,536 | -| 32–<64 GiB | 3,584 (5,120 on sm_90/120) | 2,048 | -| ≥64 GiB | 3,584 (5,120 on sm_90/120) | 3,072 | - -These are conservative experimental limits, not guarantees of full-model -capacity. The kernels are off during parameter initialisation, preserving names, -shapes, initialisers and random-number ordering for existing DeepMind weights. - -Each output directory contains `inference_kernels.json`, keyed by seed, recording -the padded bucket, device policy, kernel source revision and selections for the -128-channel pair and 64-channel template layers. It records dispatch eligibility, -including cache hits; it does not claim that populated templates were supplied. - -Layer measurements justify an end-to-end pilot, not a production speed claim. -Keep this feature opt-in until the paired speed, memory and DeepMind-weight -accuracy validation is complete. Kernel source and attribution are recorded in -the fork's `fused_triangle/SOURCE.md` and `NOTICE`. +| 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 eb26f4c6..b999e861 100644 --- a/docs/backend_installation.md +++ b/docs/backend_installation.md @@ -2,9 +2,6 @@ This guide is for direct AlphaPulldown use without Snakemake. -The experimental AF3 triangle-kernel option and its required fork revision are -described in [AF3 fused triangles](af3_fused_triangles.md). - Two points matter in practice: 1. Run the install commands from the AlphaPulldown repo root. @@ -105,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 From 8f6c55e6bdc3c435c0262806e0e37a36874d0e74 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 6 Oct 2026 19:57:10 +0200 Subject: [PATCH 08/11] chore: double-quoted fast_kernels default in AF3 setup Co-Authored-By: Claude Opus 5.5 --- alphapulldown/folding_backend/alphafold3_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/alphapulldown/folding_backend/alphafold3_backend.py b/alphapulldown/folding_backend/alphafold3_backend.py index 328e9ace..fa41bd08 100644 --- a/alphapulldown/folding_backend/alphafold3_backend.py +++ b/alphapulldown/folding_backend/alphafold3_backend.py @@ -811,7 +811,7 @@ def setup( num_recycles: int = 10, return_embeddings: bool = False, return_distogram: bool = False, - fast_kernels: str = 'off', + fast_kernels: str = "off", **kwargs, ) -> Dict: """Sets up the ModelRunner with the given configurations.""" From d8b898c524960728b3ee274d7b06a0aea3f4b066 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 6 Oct 2026 16:27:41 +0200 Subject: [PATCH 09/11] AF3: move the alphafold3 submodule to the fork's v3.0.4 update The submodule points at KosinskiLab/alphafold3 exp/af3-fused-triangle-v3.0.4 (ab98e22): upstream v3.0.4 merged into the fork without conflicts (de1577e), with the fused-triangle commits replayed on top. The fork's fused-triangle test passes 90/90 on it, on jax 0.9.1 and on jax 0.10.2 / tokamax 0.0.12 (CPU, Pallas interpret mode), and make_model_config() matches the golden model_config.json. v3.0.4 moves AF3 to jax 0.10.2 and tokamax 0.0.12, so the alphafold3 extra, the image's version stamp and the README follow. The image installs AF3 from the submodule's uv.lock, which v3.0.4 already updates. Experimental: needs a rebuilt AF3 image and the GPU campaign rerun on it before any PR. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- alphafold3 | 2 +- docker/alphafold3.dockerfile | 4 ++-- pyproject.toml | 6 +++--- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 7eeb0aa8..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. diff --git a/alphafold3 b/alphafold3 index 73f260ee..ab98e22c 160000 --- a/alphafold3 +++ b/alphafold3 @@ -1 +1 @@ -Subproject commit 73f260eea1c307bd6daa28f1b8df01943f0b9d42 +Subproject commit ab98e22cb66d523799f37d86e4b7bd75898ee90d 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/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 = [ From 18e5fcc5451dba0352668cf2f0267d1dc860181c Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Tue, 6 Oct 2026 20:31:11 +0200 Subject: [PATCH 10/11] fix: run_multimer_jobs forwards --fast_kernels and the compile cache run_multimer_jobs takes its flags from run_structure_prediction but built its own argument list and never passed --fast_kernels on, for either backend, and dropped --jax_compilation_cache_dir on the AlphaFold 2 path. Both are now forwarded to AlphaFold 2 and AlphaFold 3. AlphaLink and UniFold reject them, so they stay out of those commands. Co-Authored-By: Claude Opus 5.5 --- alphapulldown/scripts/run_multimer_jobs.py | 5 ++ test/unit/test_script_entrypoints.py | 66 ++++++++++++++++++++++ 2 files changed, 71 insertions(+) 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/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, From ef46bc278eff9870084c4513b42952464e1056f5 Mon Sep 17 00:00:00 2001 From: Dima Molodenskiy Date: Wed, 7 Oct 2026 07:17:33 +0200 Subject: [PATCH 11/11] chore: alphafold3 submodule at the merged fork (KosinskiLab/alphafold3#3, #4) Fork main 7809148: upstream v3.0.4 (#3) plus the fused triangle kernels (#4). Same tree as the tested ab98e22 apart from upstream's data-test fix (43525e6). Co-Authored-By: Claude Opus 5.5 --- alphafold3 | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/alphafold3 b/alphafold3 index ab98e22c..7809148b 160000 --- a/alphafold3 +++ b/alphafold3 @@ -1 +1 @@ -Subproject commit ab98e22cb66d523799f37d86e4b7bd75898ee90d +Subproject commit 7809148bfefdb42daa3bf56c884e345b2b7a30c6