BioIR Support Matrix

What BioNeMo Inference Runtime (BioIR) runs, on which GPUs, and which fused kernels it uses. How to construct a model or call build_processor is in Python API.

Model keys are the strings in bionemo_ir.hubs.FoldingSupportMatrix. Pass them as EngineProcessorConfig.model_source or as model_name= on the constructor.

Models and Data Pipeline

Every model below runs on the optimized PyTorch backend.

Distributed execution in build_processor is replica mode only: each engine worker holds a full model copy on a single GPU and folds independent inputs (ParallelismMode.REPLICA; api.md Ray replicas). Splitting one forward pass across GPUs (model / context parallelism, as in Fold-CP) is planned, not available yet.

protenix-v2 and boltz-2-affinity have an nn.Module but no tokenizer, feature factory, or post-processor. Do not pass them to build_processor. Ligand structure prediction on Boltz-1/2 and OpenFold3 is supported; affinity prediction is not.

Input coverage for the bundled data pipeline (InputRequest → parser → writer). “Supported” means the pipeline accepts the input and produces a structure. ”—” means the mode does not apply to that family.

Notes:

• Nucleic acids and ligands are Boltz-1/2 and OpenFold3 only. OpenFold2 / AlphaFold2 fold protein chains exclusively.
• Templates are allowed only on polymer_type="protein". BioIR does not run HHsearch / HMMsearch; pass CIF (or PDB) hits you already have.
• Nucleic-acid and ligand chains carry no MSA (msas / paired_msas are empty).
• AF2 monomer takes a single unpaired a3m per chain. AF2 multimer accepts paired a3m (some sample builders require it on every chain). Boltz-1/2 and OpenFold3 treat paired MSAs as optional.

Polymer.polymer_type is "protein", "rna", "dna", "ccd_ligand" (a CCD code or _-joined list), or "smiles_ligand" (a SMILES string in sequence). Schema details: api.md input requests.

GPUs

This table is backend compatibility, not release qualification. It says which fused kernels apply on each architecture; every row runs. The devices this release is qualified on are H200, H100, A100, L40S, GB200, and GB300, with measured results for each in Benchmarks. The rest of the table works and is not part of that set.

Optimized CuTeDSL kernels cover Ampere through Hopper, plus SM100 / SM103 for pair-weighted averaging, outer-product mean and AdaLN; the kernel table below lists the exact SMs per kernel. On SM100, SM103, SM120, and SM121, triangle attention and dual GEMM fall back to cuEquivariance. The PyTorch backend still runs on any of these SKUs; rows below describe which fused kernels apply.

Check the device:

$
python -c 'import torch; print(torch.cuda.get_device_capability())'

(8, 0) is A100 / Ampere, (8, 9) is L40 / Ada, (9, 0) is H100 / Hopper, (10, 0) is B200 / SM100, (10, 3) is B300 / SM103, (12, 0) is SM120, (12, 1) is DGX Spark / SM121.

On Boltz-1, Boltz-2, OpenFold3, and Protenix (protenix-v2) the diffusion module (including the token transformer) can be captured as a CUDA graph and replayed across sampling steps (largest win on short sequences). OpenFold2 / AlphaFold2 have no CUDA-graph module. How to enable it: api.md CUDA graphs.

Fused Kernels

get_pretrained_config selects CuTeDSL triangle and pairwise attention on SM80 / SM86 / SM89 / SM90 for fp16 / bf16. On SM100, SM103, SM120, and SM121, triangle attention and dual GEMM use cuEquivariance (no CuTeDSL CUBIN). Other SKUs or fp32 fall back to cuEquivariance (triangle attention if installed) or PyTorch SDPA.

Call through the dispatchers below (bionemo_ir._torch.attention_backend and bionemo_ir._torch.custom_ops). Layers wrap the same ops: TriangleAttention / AttentionPairBias, TriangleMultiplicationNode, PairWeightedAveraging, OuterProductMean, AdaLN, LNProjMoveaxisPad, Transition.

Override backends on a config if you need a reference path, for example config.trunk.set_triangle_attention_backend("SDPA").

CuTeDSL source is not open-sourced, and there is no plan to publish it. The public tree and wheels ship the Python callables plus precompiled CUBIN payloads for every CuTeDSL kernel in the table above. They do not ship the CuTeDSL kernel implementations used to generate those CUBINs. Loading a CUBIN skips CuTeDSL JIT (cute.compile), so those kernels run at full speed on the first iterations — no kernel-JIT warmup is required to hide compile latency. If a CUBIN is missing for the current SM / dtype, those ops fall back to PyTorch (SDPA or a vanilla reference) rather than JIT-compiling CuTeDSL. Triton fused ops in bionemo_ir.dsl_kernels.triton remain ordinary Python source and still JIT on first use.