fast_trimul

概要

• AlphaFold3ファミリーのモデルに、Fused Triangle Multiplicative Updatesを簡単に組み込めるPythonライブラリです。
• OpenFold3やProtenixで利用できる、実運用への導入しやすさを重視したモジュール構成を採用しています。

主な特徴

• GPUメモリ使用量をおよそ半分に抑え、配列長にかかわらずアクティベーションデータをほぼ使わずに処理できます。
• 短い配列で特に高い性能を発揮し、CUDAカーネルにはPython版CuTe DSLを採用してコードを小さく保っています。
• ハードウェアに依存しにくい設計で、H100やB200など別のGPUへも適応しやすい構成です。

こんな人におすすめ

• AlphaFold3系のオープンソースモデルを運用し、推論時のGPUメモリ消費や処理速度を改善したい研究者・開発者。
• OpenFold3やProtenixに高速なカスタムCUDA処理を導入し、既存コードへの変更を最小限に抑えたいチーム。