NVIDIA BioIR Speeds Up Boltz-2 Folding by 2.9x on Eight H100 GPUs

NVIDIA's BioNeMo Inference Runtime (BioIR) accelerates supported biomolecular structure-prediction models on its GPUs while keeping the PyTorch workflow intact. The runtime provides an end-to-end processor that parses input, tokenizes it, generates features, runs GPU inference, and writes PDB or mmCIF output. It also integrates directly as a torch.nn.Module for custom code.

In a benchmark of 1,000 human dimer targets on eight H100 GPUs, BioIR-accelerated Boltz-2 folded 58.5 thousand residues per GPU-hour, versus 20.2 thousand for a torch-compiled open-source implementation. That is a 2.90 times improvement in residue-normalized throughput. The estimated energy to fold one million comparable targets is 11 MWh with BioIR, compared to 35 MWh for the public implementation using eight-GPU TDP equivalents.

How BioIR Optimizes

BioIR applies optimization at three layers. Kernel selection chooses between BioIR custom, cuEquivariance, or PyTorch fallback implementations based on the model, GPU, datatype, and tensor shape. Module optimization uses an optimize() mechanism to capture CUDA Graphs for compatible modules. Pipeline scaling via Ray places a full model replica on each GPU in a single node, overlapping CPU stages with GPU folding. The Ray pipeline follows dependency order: parser, tokenizer, feature generator, folding engine, writer.

For large worklists, the Ray executor sizes CPU stages from torch.cuda.device_count() and places one complete model replica on each visible GPU. An alternative configuration explicitly controls four GPUs, creating four engine actors each reserving one GPU. Stage configuration fields include compute, num_cpus, memory, batch_size, and for engine actors also max_concurrent_batches, accelerator_type, and num_gpus. The capacity rule for the Ray engine stage is that compute times num_gpus must not exceed visible GPUs. Ray scaling does not split a single model forward pass across GPUs; it only runs independent replicas.

A serial backend runs each stage sequentially for one input, useful for validation before using Ray. The end-to-end processor supports ligand structure prediction but not ligand-affinity prediction. Each protein chain input must provide an A3M MSA; paired or unpaired MSAs are accepted for multi-chain inputs.

Real-World Use and Requirements

BioIR was used to expand the AlphaFold Database, processing 31 million candidate complexes across 4,777 proteomes and releasing 1.81 million high-confidence predictions. Supported models include Boltz-2 and OpenFold3; Boltz-2 uses cuEq kernels in the open-source baseline. Early benchmarking showed geometric-mean speedups for GPU-synchronized model-forward latency on H100 and H200 GPUs relative to open-source torch.compile. Measurements used one warm-up run (discarded) and one measurement call across 17 inputs ranging from 29 to 1,734 residues. The open-source baselines for OpenFold3 and Boltz-2 used torch.compile with dynamic=None, fullgraph=False, recompile_limit=128, accumulated_recompile_limit=256, and fail_on_recompile_limit_hit=True. Boltz-2 used cuEq kernels; OpenFold3 used use_cuequivariance=True and use_deepspeed=True.

BioIR requires Python 3.12 or later, a compatible NVIDIA GPU and driver, and a BioIR wheel or development environment. The wheel contains precompiled CUBINs, so runtime use does not need nvcc, CUDA source, CMake, or the CUDA toolkit. Ray scaling effectiveness depends on CPU stage costs, storage, scheduling, and failure handling; users should measure actual throughput. The BioNeMo Agent Toolkit can further accelerate drug-discovery workflows with agentic orchestration. The API reference and support matrix provide the latest acceleration numbers.