DeepFoldProtein/openfold3-cuequiv-trt

TensorRT BF16 engines for OpenFold3 (pairformer + token transformer), built with Bio-TRT.

Triangle attention backend: CUEQUIV

Note: This repo contains CUEQUIV (cuEquivariance) triangle attention engines. Requires Bio-TRT commit with the FP32 triangle_bias cast fix in thirdparty/tensorrt_bionemo/_trt/layers/attention.py — earlier Bio-TRT versions produce numerically incorrect results due to a cuEquivariance SM100f fallback bug.

TRT engines are GPU-architecture specific. Download the directory matching your GPU's compute capability.

Available Engines

Directory	GPU	SM	TRT Version	CUDA	Built
B200_SM10.0_TRT10.15/	unknown	SM unknown	unknown	unknown	unknown

Contents (per GPU directory)

<gpu_tag>/
├── pairformer_engine/trt/
│   ├── rank0.engine
│   └── config.json
└── token_transformer_engine/trt/
    ├── rank0.engine
    └── config.json

Usage

# Download engines for your GPU
pip install huggingface-hub
huggingface-cli download DeepFoldProtein/openfold3-cuequiv-trt --local-dir engines \
    --include "B200_SM10.0_TRT10.15/**"

# Run inference
bash openfold3/run_inference.sh \
    --query examples/openfold3/ubiquitin.json \
    --checkpoint /path/to/of3-p2-155k.pt \
    --engines-dir engines/<gpu_tag> \
    --output-dir results/ubiquitin