Spaces:

---

kernels-community

/

README

Like 8

Running

Hello everyone!!!

I created fast_trimul, an open-source, hardware-agnostic Python library providing Fused Triangle Multiplicative Updates for AlphaFold3 family models (including OpenFold-3, Boltz-1, Protenix, etc.).

The kernel was written in the Python CuTe DSL under Apache-2.0.

Key Highlights

• Performance: Runs 4.5–6.8× faster on short sequences compared to standard implementations (and up to 6.1× faster for small sequence lenghts).
• VRAM Savings: Cuts peak GPU memory by 2.2–2.4× (N=2048–2560), enabling ~1.4× longer sequences before OOM.
• Zero Recompilation Overhead: Handles dynamic sequence lengths without re-triggering CUDA graph/compilation overheads (unlike torch.compile).
• Safety & Fallback: Exact output parity (difference ~0.0006%) with automatic fallback to PyTorch on failure to prevent runtime crashes.
• Extensible Architecture: Written using Python CuTe DSL with modular plugins, making it easy to adapt for new hardware targets (e.g., H100, B200) or frameworks.

In End-to-End OpenFold-3 Performance, it provides an overall 1.15–1.33× full-stack model speedup when integrated into OpenFold-3 for small sequence lenghts.

I would like to submit the kernel here of the library, but I need access to do so.

Benchmark & Code

• Repository: https://github.com/tiagomonteiro0715/fast_trimul
• PyPI: https://pypi.org/project/fast_trimul/

All these statistics are explained in detail in the GitHub repo.