[DimaMolod]

Summary

Opt-in fused triangle kernels for inference. They are Anthropic's FlashPairformer Pallas kernels for TriangleMultiplication and GridSelfAttention, from uplifting-biomolecular-modeling @ f4f62fa6 (Apache-2.0). Off by default; with the defaults the model runs exactly the original code.

• Vendored unchanged (commit 9aedcb1; trimul_pallas.py, triattn_pallas.py and NOTICE are byte-identical to the kit):
  • in src/alphafold3/jax/fused_triangle/, the place AF3 kept its own Pallas kernels before tokamax;
  • provenance in SOURCE.md.
• Our code:
  • jax/fused_triangle/dispatch.py: per-device tile tables and size and memory limits. The tile rows are copied from the kit and credited.
  • model/network/fused_triangle.py: Haiku adapters that read the stock parameters, so the weights load unchanged. They use DeepMind's ending-node bias convention, not OpenFold's.
  • Two hooks in modules.py. init always runs the stock module.
• Config, named like flash_attention_implementation:
  • triangle_multiplication_implementation: 'default' or 'pallas';
  • triangle_attention_implementation: 'default', 'pallas', 'pallas_tokamax_core' or 'auto';
  • plus fused_triangle_compute_capability and fused_triangle_memory_gib, which the caller sets.
  • Any shape, card or size the kernels don't serve runs the stock path. An explicit request that falls back logs a warning.
• Device policy:
  • Pallas attention on compute capability 8.0, 9.0 and 12.0 (12.0 on the 9.0 tiles);
  • the kit's prologue and epilogue around tokamax's attention core on 8.6 and 8.9, where the kit's generic tiles are slower;
  • triangle multiplication on every card from 8.0;
  • sm_100 is not enabled, because nothing was measured there.
• test_data/model_config.json gains the new fields.

Validation

This branch carries the same commits as #2, replayed onto the v3.0.4 update (#3); the fused code is byte-identical between the two. The GPU numbers below were measured on the v3.0.3 / jax 0.9.1 stack. The v3.0.4 / jax 0.10.2 stack is validated on CPU here, and gets the full GPU campaign on the rebuilt AlphaPulldown image before any release.

• fused_triangle_test.py (absltest, added to CI): 90 tests, passing on jax 0.9.1 and on 0.10.2 / tokamax 0.0.12 on CPU in Pallas interpret mode. They cover:

  • parity with the stock layers and an f64 reference;
  • both trimul equations and both attention orientations;
  • padded and asymmetric masks, and C = 128 and C = 64;
  • determinism and finiteness;
  • that the kernels actually run;
  • the fallbacks and their warnings;
  • the OpenFold orientation failing.

• Layer level on GPU, 6 cards:

  • Accuracy: every fused cell passes the gate against an f64 reference (error at most 1.2× stock's), with leak, determinism and finiteness checks.
  • Speed per layer: trimul 1.24–1.71×; attention 1.11–1.54× on 8.0/9.0/12.0, and 1.03–1.36× with the tokamax core on 8.6/8.9.

• End to end through AlphaPulldown, DeepMind weights, warm calls, fused vs stock on the same card:

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

  • The largest complex each card can fold is unchanged.
  • With the flag off, outputs are bit-identical to the original fork on every card (deterministic XLA).
  • 12 heterodimers, paired seeds: Δranking and ΔDockQ are 0.00 on nearly all of them. Protein–RNA, protein–ligand and template inputs match too.

🤖 Generated with Claude Code