fast_trimul

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.
• Modular, vendor-agnostic backends with a fallback chain (cuda → torch): the fast CUTLASS kernel when it fits, an always-correct pure-torch path otherwise, so an unsupported shape/dtype degrades gracefully instead of crashing. New hardware (TPU/Intel/…) is a plug-in, not a rewrite.
• 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 — one call, pick the source (names are remapped, and fused a/b projections are split, for you):

module.load_weights(ref.state_dict(), source="openfold")   # or: openfold3 / protenix / boltz / chai

The five named helpers (load_openfold_state_dict, …) still work as thin aliases. 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()

Pick or force a backend (default is "auto" — fastest available, then torch):

from fast_trimul import list_backends
list_backends()                                            # e.g. ['torch', 'cuda']
FastTriangleMultiplication(d_z=128, backend="cuda")        # force the fast path (still falls back)
FastTriangleMultiplication(d_z=128, backend="torch")       # force the portable reference

Opt in explicitly, without any global monkeypatching (safe under strict runtime policies) — @accelerate on a function, or the scoped accelerated():

from fast_trimul import accelerate, accelerated

@accelerate                    # runs this function with the accelerated backend
def infer(z): ...

with accelerated("cuda"):      # scoped: reverts on exit, never leaks
    out = model(z)

One-line correctness check — build your library's TriMul, then:

import fast_trimul
fast_trimul.verify("openfold", my_openfold_trimul)         # -> True if outputs match (fp16 tol)

verify stays a one-liner even when a library changes its API, because you pass the reference module and only the (registered) name-remap is library-specific.

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.

Architecture

The library is a small stable front-end over a registry of interchangeable backends. CUDA is one backend; a pure-torch backend is the universal fallback. Adding new hardware or a new library is a plug-in (one decorated class/function), not a core edit.

  front-end   @accelerate  accelerated()  verify()          # stable, tiny
      │
  guard        contiguous · dtype · align · int64 strides   # NormalizedInput
      │
  dispatch     pick backend, then FALL BACK: cuda -> torch  # never crashes on a bad shape
      │
  backends     cuda_cute (CUTLASS)   torch_ref (portable)   # + future: xla/tpu, xpu/intel

Everything the architecture adds — registries, guard, dispatch, decorators — is O(1) overhead; only the op itself scales with N. The CUDA kernels (_kernels.py) are untouched.

Adding a new backend

from fast_trimul.core.registry import backend

@backend("mybackend", dtypes={torch.float16}, min_align=8)
class MyBackend:
    def __init__(self, caps): self.caps = caps
    def execute(self, inp, params): ...   # inp.tensor is guarded (contiguous, right dtype)

That one file makes backend="mybackend" selectable and slots it into the fallback chain — no changes to the dispatcher or the module.

API

• fast_trimul.FastTriangleMultiplication(d_z, d_c=None, mode="outgoing", residual=True, backend="auto") — the module. forward(z, mask=None), .graphed(z, mask=None), .load_weights(state_dict, source=...) (plus the named aliases .load_openfold_state_dict / .load_openfold3_state_dict / .load_protenix_state_dict / .load_boltz_state_dict).
• fast_trimul.accelerate / fast_trimul.accelerated(backend="auto") — explicit opt-in decorator / scoped context manager.
• fast_trimul.verify(source, reference, ...) — one-line correctness check against a reference module.
• fast_trimul.list_backends() — backends registered on this machine.
• fast_trimul.functional.triangle_multiplication(z, params, mask=None) — low-level functional call.
• fast_trimul.core.registry.{backend, weights_for} — decorators to register a new backend or library weight-map.

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 via load_weights(sd, source=...): openfold (OpenFold/AF2), openfold3 (separate or fused variant), protenix, boltz, and chai (Boltz/Chai/AF3 fuse the a/b projections). Other stacks: supply a @weights_for 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.
• The CUDA kernel is Ampere (sm80) tested; fp16 only. On other hardware, or for bf16/fp32, or non-multiple-of-8 N, the dispatcher falls back to the pure-torch backend (correct, slower). Hopper/Blackwell + fp8 are future work.
• Building a module needs a CUDA GPU + CUTLASS (the fused kernel weights live in a CUTLASS module). import fast_trimul itself is lazy and does not require CUTLASS until you construct FastTriangleMultiplication.

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.