Metadata-Version: 2.5
Name: fast_trimul
Version: 2.1.2
Summary: Fused Triangle Multiplicative Update (AlphaFold/OpenFold) on CUTLASS CuTe DSL kernels.
Project-URL: Homepage, https://github.com/tiagomonteiro0715/fast_trimul
Project-URL: Repository, https://github.com/tiagomonteiro0715/fast_trimul
Project-URL: Issues, https://github.com/tiagomonteiro0715/fast_trimul/issues
Author: Tiago Monteiro
License: Apache-2.0
License-File: LICENSE
License-File: NOTICE
Keywords: alphafold,cute,cutlass,gpu-kernel,openfold,triangle-multiplication
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Requires-Dist: cuda-python<13
Requires-Dist: nvidia-cutlass-dsl
Requires-Dist: torch>=2.2
Provides-Extra: test
Requires-Dist: pytest; extra == 'test'
Description-Content-Type: text/markdown

# fast_trimul

[![PyPI](https://img.shields.io/pypi/v/fast_trimul)](https://pypi.org/project/fast_trimul/)
[![License: Apache 2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](LICENSE)
[![Python](https://img.shields.io/pypi/pyversions/fast_trimul)](https://pypi.org/project/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.**

## Quickstart: accelerate OpenFold-3 in one line

```python
import fast_trimul
fast_trimul.patch_openfold3()      # OpenFold-3 now runs on the fast kernel

# ...build and run OpenFold-3 exactly as usual — it's now fast_trimul-powered...
```

That's the whole integration. `patch_*` swaps the library's TriMul for a
fast_trimul adapter (call it **before** building the model). Same one-liner for the
other stacks: `patch_openfold()`, `patch_boltz()`, `patch_protenix()`.

**Verified end-to-end on the OpenFold-3 Pairformer trunk (A100).** Ready-to-run
copies with the install line are in [`quickstart/`](quickstart/) — one for **Google
Colab**, one for **Lightning AI**.

Prefer no global patch? Use the module directly (`FastTriangleMultiplication`, see
*Quick start*), and use `@accelerate` / `with accelerated("cuda"):` to pin which
**backend** runs your own code (they control cuda-vs-torch selection, not the
model swap).

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

```bash
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).

**Driver note:** `cuda-python` must match your CUDA **driver**. If `nvidia-smi`
shows CUDA 12.x (e.g. driver 570), pin the CUDA-12 line — otherwise you get
`cudaErrorInsufficientDriver (35)`:
```bash
pip install "cuda-python<13"
```
`fast_trimul` pins `cuda-python<13` by default (most drivers are still CUDA 12.x);
override it if your driver is CUDA 13+.

## Quick start

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

```python
!pip install -q uv
!uv pip install fast_trimul
```

Then use it:

```python
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:

```python
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):

```python
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):

```python
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:

```python
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):

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

```python
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:

```python
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.

## Swap into a real model (whole-trunk example)

Drop `fast_trimul` into a library's trunk by replacing its TriMul class with a
thin adapter, then capturing a CUDA graph per layer. **Verified end-to-end on the
OpenFold-3 Pairformer trunk (A100).**

```python
import torch
import openfold3.core.model.latent.base_blocks as blocks
from openfold3.core.model.latent.pairformer import PairFormerStack
from fast_trimul import FastTriangleMultiplication

# 1) a thin adapter: match the library's constructor, swallow its extra forward kwargs,
#    and use residual=False (the block adds the residual itself).
class FastTriMul(FastTriangleMultiplication):
    def __init__(self, c_z, c_hidden=None, *args, **kw):
        super().__init__(d_z=c_z, d_c=c_hidden or c_z, mode="outgoing", residual=False)
    def forward(self, z, mask=None, **kw):
        return super().forward(z, mask=mask)

# 2) patch the library's TriMul classes BEFORE building the model
blocks.TriangleMultiplicationOutgoing = FastTriMul
blocks.TriangleMultiplicationIncoming = FastTriMul
blocks.FusedTriangleMultiplicationOutgoing = FastTriMul
blocks.FusedTriangleMultiplicationIncoming = FastTriMul

model = PairFormerStack(c_s=384, c_z=128, no_blocks=8, ...).cuda().eval()

# 3) capture a CUDA graph for each swapped-in layer (fixed shape, inference only)
dummy_z = torch.randn(1, N, N, 128, device="cuda")
dummy_mask = torch.ones(1, N, N, device="cuda")
for m in model.modules():
    if isinstance(m, FastTriangleMultiplication):
        m.graphed(dummy_z, dummy_mask)

# ... now run model(s, z, single_mask, pair_mask) as usual ...
```

Install for this example (note the CUDA-12 driver pin — see *Install*):
```bash
uv pip install --system openfold3 "fast_trimul>=2.0.0" "cuda-python<13"
```

## Gradients (training)

The forward runs the fused fp16 kernel; the backward is a **correct torch recompute**
(correct, not yet fast), so gradients flow and you can train / fine-tune with it. Do
**not** call `.graphed()` for training — graphs are inference only. **Verified on A100:**

```python
import torch
from fast_trimul import FastTriangleMultiplication

trimul = FastTriangleMultiplication(d_z=128, d_c=128, mode="outgoing").cuda()
z = torch.randn(1, 64, 64, 128, device="cuda", requires_grad=True)   # N a multiple of 8
out = trimul(z, mask=torch.ones(1, 64, 64, device="cuda"))
out.sum().backward()                                                 # gradients recomputed through the kernel
assert z.grad is not None and torch.isfinite(z.grad).all()           # input grads flow
assert all(p.grad is not None for p in trimul.parameters())          # parameter grads flow
```

## 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:

```python
!pip install -q uv
!uv pip install fast_trimul
```
```python
from fast_trimul.benchmark import run_benchmark
run_benchmark()              # or: run_benchmark(head_size=384, sweep=(128, 256, 512))
```

Or from a shell:
```bash
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
```

| Module | Role |
|---|---|
| `core/registry.py` | `@backend` / `@weights_for` decorators + O(1) lookup tables |
| `core/context.py` | the input **guard** → `NormalizedInput` (contiguous, dtype, alignment, int64 strides) |
| `core/dispatch.py` | picks a backend and walks the **fallback chain** (`cuda → torch`) |
| `core/decorators.py` | `@accelerate` + `accelerated()` — explicit, scoped, no global patching |
| `backends/torch_ref.py` | universal pure-torch backend (any device torch supports) |
| `backends/cuda_cute.py` | thin wrapper over the existing CUTLASS kernels (unchanged) |
| `ops/triangle.py` | `FastTriangleMultiplication` (dispatch + graph capture + `load_weights`) |
| `integrations/loaders.py` | the five per-library weight maps (`@weights_for`) |
| `integrations/checks.py` | `verify()` |

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

```python
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](LICENSE). The GEMM core is
derived from NVIDIA CUTLASS and is licensed under BSD 3-Clause — see [NOTICE](NOTICE).
