Metadata-Version: 2.5
Name: fast_trimul
Version: 0.0.26
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: MIT
License-File: LICENSE
License-File: NOTICE
Keywords: alphafold,cute,cutlass,gpu-kernel,openfold,triangle-multiplication
Requires-Python: >=3.10
Requires-Dist: cuda-python
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: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
[![Python](https://img.shields.io/pypi/pyversions/fast_trimul)](https://pypi.org/project/fast_trimul/)

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

**At matched fp16 precision (A100, N=512), `fast_trimul` is 1.2–1.75× faster and
~2.3× lower peak memory than the stock TriMul in OpenFold, Boltz-1, Protenix, and
an AF3/Chai-style reference** (see Results). It is *not* faster than a fully-tuned
`torch.compile(mode="reduce-overhead")` above small N (see *Limitations*), but it
needs no whole-model compilation and matches PyTorch fp16 numerically.

## Results

Measured on **NVIDIA A100-SXM4-40GB**, `B=1, N=512, d_z=d_c=128`, all at **matched
fp16** precision (baselines via autocast), 30 timed runs:

| Implementation (fp16)                | Latency (ms) ↓ | Peak VRAM (GB) ↓ | vs fast_trimul       |
|:-------------------------------------|---------------:|-----------------:|:---------------------|
| **fast_trimul** (fp16 + CUDA graph)  |       **6.35** |         **0.64** | —                    |
| AF3 reference (Chai-style)           |           7.46 |             1.45 | 1.17× slower · 2.3× RAM |
| Boltz-1                              |           7.59 |             1.45 | 1.20× slower · 2.3× RAM |
| Protenix                             |          10.36 |             1.39 | 1.63× slower · 2.2× RAM |
| OpenFold                             |          11.10 |             1.39 | 1.75× slower · 2.2× RAM |

Against the **stock fp32** modules (what many stacks actually run), the speed gap is
larger (1.8–2.5×). The **memory win is dtype-independent** (fusion + CUDA-graph
buffer reuse); part of the speed win over these eager baselines is the CUDA graph.
Reproduce with `from fast_trimul.benchmark import run_benchmark` (see *Full benchmark*).

## 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).

## 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 Python overhead of the internal kernels, which dominates
the runtime at small/medium N:

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

Benchmark it against torch and `torch.compile` in one line (see the full report below):

```python
from fast_trimul.benchmark import run_benchmark
run_benchmark()
```

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 (parameter names are remapped for you):

```python
module.load_openfold_state_dict(ref.state_dict())    # OpenFold / AF2 (separate a/b projections, biased)
module.load_protenix_state_dict(ref.state_dict())    # Protenix (OpenFold-style names, bias-free linears)
module.load_boltz_state_dict(ref.state_dict())       # Boltz-1 / Chai / AF3 (FUSED p_in/g_in, split for you)
```

These target modules have **no internal residual** (the residual lives in the enclosing
block), so construct the module with `residual=False` when matching their output exactly.

## Colab / Jupyter quickstart (with an event-based timer)

Install:

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

Run it and time it. The timer uses **CUDA events + `synchronize()`**, so it measures
when the GPU actually *finishes the work* — not when the launch is queued:

```python
import time, torch
from fast_trimul import FastTriangleMultiplication

assert torch.cuda.is_available(), "Need a CUDA GPU (Colab: Runtime -> Change runtime type -> GPU)."
print("GPU:", torch.cuda.get_device_name(0))

B, N, d_z, d_c = 1, 256, 128, 128
module = FastTriangleMultiplication(d_z=d_z, d_c=d_c, mode="outgoing").cuda()
z    = torch.randn(B, N, N, d_z, device="cuda")     # (B, N, N, d_z)
mask = torch.ones(B, N, N, device="cuda")           # optional (B, N, N)
print(f"input : {tuple(z.shape)}  {z.dtype}")

# first call: one-time CuTe JIT compile + GEMM autotune (wall clock is fine here)
t0 = time.perf_counter()
with torch.no_grad():
    out = module(z, mask=mask)
torch.cuda.synchronize()
print(f"first call (JIT compile + autotune): {time.perf_counter()-t0:5.2f} s")
print(f"output: {tuple(out.shape)}  {out.dtype}   mean={out.mean():.4f}  std={out.std():.4f}")

module.graphed(z, mask)      # capture a CUDA graph -> the fast steady-state path

def bench(fn, iters=50, warmup=10):
    for _ in range(warmup):                 # warmup: compiled + caches hot
        fn()
    torch.cuda.synchronize()
    start = torch.cuda.Event(enable_timing=True)
    end   = torch.cuda.Event(enable_timing=True)
    start.record()
    for _ in range(iters):
        fn()
    end.record()
    torch.cuda.synchronize()                # read a COMPLETED timestamp, not a queued one
    return start.elapsed_time(end) / iters  # ms per call (GPU timeline)

with torch.no_grad():
    ms = bench(lambda: module(z, mask=mask))
elems = z.numel()
print(f"\nsteady-state (CUDA events):")
print(f"  {ms*1e3:8.1f} us / call")
print(f"  {elems/1e6:6.1f}M elements  ->  {elems/(ms/1e3)/1e9:6.2f} Gelem/s")
```

## Scope vs `torch.compile`

At **matched fp16**, `torch.compile(mode="reduce-overhead")` can match or beat these
kernels on **latency** above small N (see *Limitations*) — `fast_trimul` does not
target that comparison, for two concrete reasons:

- **Static shapes.** `reduce-overhead` relies on CUDA graphs, which require a fixed
  input shape. Protein inputs are variable-length, which is why the target stacks
  ship hand-written kernels rather than whole-model `compile`.
- **Memory.** A CUDA graph over the eager op still materializes the full intermediate
  tensors, so `torch.compile` does **not** deliver the ~2.2× peak-memory reduction —
  that comes from kernel *fusion*, and it holds **even where compile is faster**.

So the win here is **memory + larger-N reach + drop-in with no whole-model
compilation**, not raw fp16 latency above small N.

## Full benchmark (machine ceilings + roofline)

To **compare correctly**, the package ships a rigorous benchmark — measured
machine ceilings (memory bandwidth, fp16 tensor-core peak, launch floor), a
per-iteration **median** timer (median / min / p95 / CV, not a mean), roofline
placement (% of peak, × above roofline, × launch floor), effective GB/s, achieved
TFLOP/s, and a size sweep. It reports fast_trimul **both un-graphed and graphed**,
so you can see what the CUDA graph buys, next to the fair baseline:

* **`fast no-graph`** — this kernel, fp16, un-graphed (host/launch bound),
* **`fast +graph`** — the same kernel with a captured **CUDA graph** (`.graphed()`),
* **`compile16`** — `torch.compile(mode="reduce-overhead")` in fp16 (also CUDA graphs) — the fair fight,
* **`compile32`** — `torch.compile()` in fp32 (default mode) — reference,
* **`torch eager`** — naive fp32 reference.

**On Google Colab** (Runtime → Change runtime type → **GPU**), just two cells:

```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 prints something like:

```text
GPU: NVIDIA A100-SXM4-40GB
  measured mem bandwidth peak :     1490 GB/s
  measured fp16 matmul peak   :      270 TFLOP/s
  launch-overhead floor       :      4.6 us

Head-to-head  N=256, d_z=128, d_c=128   (17.2 GFLOP/call, fp16 err vs torch = 3.8e-03)
  metric              fast no-graph   fast +graph   torch eager     compile32     compile16
  -----------------------------------------------------------------------------------------
  median (us)                    ...
  p95 (us)                       ...
  TFLOP/s                        ...
  x above roofline               ...
  speedup vs compile16           ...

Size sweep (median us/call). fast_ng = un-graphed, fast_g = fp16+CUDA graph,
compile32 = fp32 default, compile16 = fp16 reduce-overhead:
      N   fast_ng    fast_g     eager   compile32   compile16  fast TFLOP/s
     64      ...
```

(Numbers are illustrative — run it on your GPU. `fast_ng` = un-graphed (shows the
launch-overhead cost), `fast_g` = CUDA graph, `compile16` = fair fp16 baseline,
`compile32`/`eager` = fp32 references.)

## Drop-in monkeypatch for the 4 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 caveat.

### OpenFold
```python
import fast_trimul.integrations as fti
fti.patch_openfold()          # patches Outgoing + Incoming
# ... now build your OpenFold model as usual ...
```
Equivalent manual form:
```python
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
```python
import fast_trimul.integrations as fti
fti.patch_boltz()
```
Manual form:
```python
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
```python
import fast_trimul.integrations as fti
fti.patch_protenix()
```
Manual form:
```python
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):
```python
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")` — high-level module, `forward(z, mask=None)`; `.load_openfold_state_dict(sd)` / `.load_protenix_state_dict(sd)` / `.load_boltz_state_dict(sd)` load pretrained weights from those stacks with name remapping.
- `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(reduce-overhead)` is a strong baseline — and it matches this on
  both speed and memory at N≤~2560.** Measured on A100: at matched settings it is
  ~10-15% *faster* per call and uses the *same* peak memory (Inductor fuses and its
  CUDA-graph pool reuses buffers). The honest reasons to use this instead are
  **robustness and drop-in-ness**, not raw latency: `reduce-overhead` needs *static*
  shapes and recompiles per sequence length (painful for variable-length inputs) and
  breaks on some models, whereas this is a plain `nn.Module` that works on any shape
  with no compilation step. The kernels are not yet epilogue-fused (future work), so
  beating compile on latency is not yet a goal.
- **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 four common stacks: `load_openfold_state_dict` (OpenFold/AF2),
  `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

MIT (this project). The GEMM core is derived from NVIDIA CUTLASS and is licensed
under BSD 3-Clause — see [NOTICE](NOTICE).
