fast_trimul

PyPI License: Apache 2.0 Python

Fused Triangle Multiplicative Update (AlphaFold2 / AlphaFold3 family) built on hand-written CUTLASS CuTe DSL kernels — a drop-in nn.Module for the structural-biology stacks (OpenFold, OpenFold-3, Boltz, Chai, Protenix).

  • Numerically matches the stock module (fp16 tolerance) — verified against OpenFold, OpenFold-3, Boltz-1, Protenix, and an AF3/Chai-style reference by loading their weights and comparing outputs.
  • Roughly halves peak memory versus the stock eager module (kernel fusion + CUDA-graph buffer reuse).
  • Fastest at small N, where per-launch overhead dominates and the captured CUDA graph removes it.
  • Drop-in on any shape, no whole-model compilation.

Run the shipped benchmark on your own GPU for numbers — see Benchmark below, and read Limitations for where torch.compile is the better choice.

Install

pip install fast_trimul          # or: uv pip install fast_trimul

Requires a CUDA GPU, torch, nvidia-cutlass-dsl, and cuda-python. Kernels JIT-compile on first use (one-time cost, then cached in-process).

Quick start

On Google Colab (Runtime → Change runtime type → GPU), install first:

!pip install -q uv
!uv pip install fast_trimul

Then use it:

import torch
from fast_trimul import FastTriangleMultiplication

module = FastTriangleMultiplication(d_z=128, d_c=128, mode="outgoing").cuda()
z = torch.randn(1, 256, 256, 128, device="cuda")          # (B, N, N, d_z)
mask = torch.ones(1, 256, 256, device="cuda")             # optional (B, N, N)
out = module(z, mask=mask)                                 # same dtype as z

For fastest inference at a fixed shape, capture a CUDA graph once — this removes the per-launch overhead of the internal kernels, which dominates the runtime at small N:

module.graphed(z, mask)      # capture once at this shape (inference only)
out = module(z, mask=mask)   # subsequent calls replay the graph

Low-level functional API (FlashAttention style):

from fast_trimul import functional
out = functional.triangle_multiplication(z, module._impl, mask=mask)

Load pretrained weights from a target library (parameter names are remapped for you):

module.load_openfold_state_dict(ref.state_dict())    # OpenFold / AF2  (separate a/b projections)
module.load_openfold3_state_dict(ref.state_dict())   # OpenFold-3      (separate OR fused variant)
module.load_protenix_state_dict(ref.state_dict())    # Protenix        (OpenFold-style names, bias-free)
module.load_boltz_state_dict(ref.state_dict())       # Boltz-1 / Chai / AF3 (fused p_in/g_in, split for you)

These target modules apply their residual (+ z) outside the triangle block, so build with residual=False when matching their output exactly:

module = FastTriangleMultiplication(d_z=128, d_c=128, mode="outgoing", residual=False).cuda()

Benchmark

The package ships a benchmark that measures machine ceilings (memory bandwidth, fp16 tensor-core peak, launch floor), a per-iteration median timer, achieved TFLOP/s, peak memory, and a size sweep. It reports fast_trimul both un-graphed and graphed, next to torch.compile and an eager reference, so you can compare on your own hardware:

!pip install -q uv
!uv pip install fast_trimul

from fast_trimul.benchmark import run_benchmark
run_benchmark()              # or: run_benchmark(head_size=384, sweep=(128, 256, 512))

Or from a shell:

python -m fast_trimul.benchmark

It reports these variants:

  • fast no-graph — the kernel, fp16, un-graphed (shows the launch-overhead cost),
  • fast +graph — the same kernel with a captured CUDA graph (.graphed()),
  • compile — torch.compile(mode="reduce-overhead") and default mode,
  • torch eager — the eager reference.

Use CUDA events + synchronize() (as the shipped benchmark does) so timing reflects when the GPU finishes the work, not when the launch is queued. Warm up (or call .graphed()) before timing to exclude the one-time JIT/autotune cost.

Drop-in monkeypatch for the target libraries

Each helper replaces the library's TriMul class with an adapter matching its constructor. Patch before building the model. See Limitations for the pretrained-weight note.

OpenFold

import fast_trimul.integrations as fti
fti.patch_openfold()          # patches Outgoing + Incoming
# ... now build your OpenFold model as usual ...

Equivalent manual form:

import openfold.model.triangular_multiplicative_update as of_tri
from fast_trimul.integrations import adapter
of_tri.TriangleMultiplicationOutgoing = adapter("outgoing")
of_tri.TriangleMultiplicationIncoming = adapter("incoming")

Boltz-1 / BoltzDesign

import fast_trimul.integrations as fti
fti.patch_boltz()

Manual form:

import boltz.model.layers.triangular_mult as b_tri
from fast_trimul.integrations import adapter
b_tri.TriangleMultiplicationOutgoing = adapter("outgoing")
b_tri.TriangleMultiplicationIncoming = adapter("incoming")

Protenix

import fast_trimul.integrations as fti
fti.patch_protenix()

Manual form:

import protenix.model.modules.pairformer as p_tri
from fast_trimul.integrations import adapter
p_tri.TriangleMultiplication = adapter("outgoing")

Chai-1

Chai's module path is version-dependent, so patch the attribute explicitly (replace the import path with the one in your installed version):

from fast_trimul.integrations import adapter
import chai_lab.model.<...>.triangle_mult as c_tri   # <- verify path for your version
c_tri.TriangleMultiplicationOutgoing = adapter("outgoing")
c_tri.TriangleMultiplicationIncoming = adapter("incoming")

API

  • fast_trimul.nn.FastTriangleMultiplication(d_z, d_c=None, mode="outgoing", residual=True) — high-level module, forward(z, mask=None), .graphed(z, mask=None); weight loaders .load_openfold_state_dict / .load_openfold3_state_dict / .load_protenix_state_dict / .load_boltz_state_dict.
  • fast_trimul.functional.triangle_multiplication(z, params, mask=None) — low-level functional call.
  • fast_trimul.integrations.{patch_openfold, patch_boltz, patch_protenix, adapter} — monkeypatch helpers.

Limitations (read before relying on it)

  • torch.compile(mode="reduce-overhead") is competitive and often faster above small N. On an A100 it is frequently faster per call in the mid-range and, on several stacks, uses similar peak memory. These kernels are not yet epilogue-fused (future work), so the reasons to prefer this are drop-in-ness and robustness, not raw latency: reduce-overhead needs static shapes and recompiles per sequence length (awkward for variable-length inputs) and can break on some models, whereas this is a plain nn.Module that works on any shape with no compilation step. Benchmark both on your workload.
  • First call is slow: JIT compile + GEMM autotune. On the first forward at a new shape, the GEMM configs are auto-tuned (one-time, cached). Disable with the env var FAST_TRIMUL_AUTOTUNE=0. Warm up (or call .graphed()) before timing.
  • fp16 only. bf16/fp32 inputs are cast to fp16 and back; keep the module in fp16 (do not call .float()/.bfloat16() on it).
  • Pretrained weights need name remapping. Each library names its projections/norms differently, so a strict checkpoint load will not line up. Automated for the common stacks: load_openfold_state_dict (OpenFold/AF2), load_openfold3_state_dict (OpenFold-3, separate or fused variant), load_protenix_state_dict (Protenix), and load_boltz_state_dict (Boltz-1 / Chai / AF3, which fuse the a/b projections). Other stacks: patch-then-train, or supply a parameter remap.
  • Mask semantics are approximate. The mask is applied to the pair tensor in and out; validate against each library's exact masking before production use.
  • Backward is correct but not fast (torch recompute), so it helps inference more than training throughput.
  • Ampere (sm80) tested. Hopper/Blackwell + fp8 are future work.
  • import fast_trimul needs a CUDA GPU (device properties are read at import).

License

Apache License 2.0 (this project) — see LICENSE. The GEMM core is derived from NVIDIA CUTLASS and is licensed under BSD 3-Clause — see NOTICE.

Project links

Key dates

PyPI data
Data sourced directly from PyPI's database.
  • Released:
Newer release available (3.0.4)

1 maintainer

PyPI data
Data sourced directly from PyPI's database.
Avatar for tiagomonteiro0715 from gravatar.com tiagomonteiro0715

Credits

Author: Tiago Monteiro

License

Apache Software License (Apache-2.0)

Requires

Python >=3.10

Provides Extra

test

Tags

alphafold cute cutlass gpu-kernel openfold triangle-multiplication

Classifiers

Report project as malware

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

fast_trimul-1.0.0.tar.gz (89.2 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

fast_trimul-1.0.0-py3-none-any.whl (39.7 kB view details)

Uploaded Python 3

File details

Details for the file fast_trimul-1.0.0.tar.gz.

File metadata

  • Download URL: fast_trimul-1.0.0.tar.gz
  • Upload date:
  • Size: 89.2 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.9.18 {"installer":{"name":"uv","version":"0.9.18","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":null,"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}

File hashes

Hashes for fast_trimul-1.0.0.tar.gz
Algorithm Hash digest
SHA256 ab02b75f2688553ede7dfb959c69eab56bd3a3aa870f1df1d6e614c65320b060
MD5 84d555451e524281be86c88e88ad6bc6
BLAKE2b-256 d608b372ed93ff47bbba11101132439b234e4181e6399beb627ddad138b0b933

See more details on using hashes here.

File details

Details for the file fast_trimul-1.0.0-py3-none-any.whl.

File metadata

  • Download URL: fast_trimul-1.0.0-py3-none-any.whl
  • Upload date:
  • Size: 39.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: uv/0.9.18 {"installer":{"name":"uv","version":"0.9.18","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":null,"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":null}

File hashes

Hashes for fast_trimul-1.0.0-py3-none-any.whl
Algorithm Hash digest
SHA256 e48f130004fa7b66dcd2af818ebc5d999e8e6dd101f2e365ab4d0997326d56ac
MD5 d85816883d2c1dba8054612c12b28706
BLAKE2b-256 4847b370e9f32ac4675458facf4417e99f05639ca9e821b00dad82d927930bf7

See more details on using hashes here.

Release history Release notifications | RSS feed

3.0.4

2 files

3.0.3

2 files

3.0.2

2 files

3.0.1

2 files

3.0.0

2 files

2.4.34

2 files

2.4.33

2 files

2.4.32

2 files

2.4.31

2 files

2.4.30

2 files

2.4.29

2 files

2.4.28

2 files

2.4.27

2 files

2.4.26

2 files

2.4.25

2 files

2.4.24

2 files

2.4.23

2 files

2.4.22

2 files

2.4.21

2 files

2.4.20

2 files

2.4.19

2 files

2.4.18

2 files

2.4.17

2 files

2.4.16

2 files

2.4.15

2 files

2.4.14

2 files

2.4.13

2 files

2.4.12

2 files

2.4.11

2 files

2.4.10

2 files

2.4.9

2 files

2.4.8

2 files

2.4.7

2 files

2.4.6

2 files

2.4.5

2 files

2.4.4

2 files

2.4.3

2 files

2.4.2

2 files

2.4.1

2 files

2.4.0

2 files

2.3.3

2 files

2.3.2

2 files

2.3.1

2 files

2.3.0

2 files

2.2.2

2 files

2.2.1

2 files

2.2.0

2 files

2.1.4

2 files

2.1.3

2 files

2.1.2

2 files

2.1.1

2 files

2.0.1

2 files

2.0.0

2 files

This release

1.0.0 This release

2 files

0.0.30

2 files

0.0.29

2 files

0.0.28

2 files

0.0.27

2 files

0.0.26

2 files

0.0.25

2 files

0.0.24

2 files

0.0.22

2 files

0.0.21

2 files

0.0.20

2 files

0.0.16

2 files

0.0.15

2 files

0.0.14

2 files

0.0.13

2 files

0.0.12

2 files

0.0.11

2 files

0.0.10

2 files

0.0.9

2 files

0.0.8

2 files

0.0.7

2 files

0.0.6

2 files

0.0.5

2 files

0.0.4

2 files

0.0.1

2 files

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page