High-Throughput Structure Prediction with BioNeMo Inference Runtime

NVIDIA's BioNeMo Inference Runtime (BioIR) was used in the expansion of the AlphaFold Database (AFDB), accelerating protein-complex structure generation across 4,777 proteomes — about 31 million candidate complexes, of which 1.81 million were released as high-confidence predictions. BioIR runs supported biomolecular structure-prediction models such as Boltz-2 on NVIDIA GPUs using optimized kernels and CUDA Graphs while preserving a PyTorch workflow, and can run a full model replica per GPU via Ray for large batches of independent inputs. The runtime ships as a wheel with precompiled CUBINs, so runtime use requires no nvcc, CUDA source, CMake, or CUDA toolkit, and requires Python 3.12 or later.

Biomolecular structure prediction is now often run at proteome scale, where the goal is to move an entire worklist through the pipeline efficiently. NVIDIA BioNeMo Inference Runtime https://docs.nvidia.com/bionemo/inference-runtime/overview/ BioIR helps accelerate supported biomolecular structure-prediction models on NVIDIA GPUs while keeping the familiar PyTorch workflow. It uses optimized kernels and, where applicable, CUDA Graphs to speed model execution. For large batches of independent inputs, Ray can run a complete model replica on each GPU in a single node to increase overall throughput. BioIR has also been used in real proteome-scale work, including the recent expansion of the AlphaFold Database https://doi.org/10.64898/2026.03.27.714458 AFDB , accelerating the generation of protein-complex structures across 4,777 proteomes, about 31 million candidate complexes in total, with 1.81 million released as high-confidence predictions. You can use it in two ways see Figure 1, below : - The end-to-end processor moves an InputRequest through parsing, tokenization, feature generation, GPU inference, and PDB or mmCIF writing. - Direct PyTorch integration lets you construct a supported model torch.nn.Module or reuse selected modules in custom code. This tutorial walks through BioIR’s end-to-end processor, from input preparation to GPU inference and PDB or mmCIF output, and shows how to track structures per hour and resource efficiency. Prerequisites - Python 3.12 or later - A compatible NVIDIA GPU and driver, plus a BioIR wheel or supported development environment - A staged model checkpoint https://www.google.com/url?q=https://docs.nvidia.com/bionemo/inference-runtime/references/dev/&sa=D&source=docs&ust=1788374353897146&usg=AOvVaw14PU3b1hx9Ph1E05tQBeXy in the below example, Boltz-2 and required chemical metadata - Each protein chain requires an A3M MSA. For inputs with multiple non-identical protein chains, paired or unpaired MSA are accepted - For Ray throughput scaling, use several visible GPUs on the same node and more independent records than replicas The wheel contains precompiled CUBINs, so runtime use does not require nvcc , CUDA source, CMake, or the CUDA toolkit. Step 1. Choose a supported structure-prediction workflow In the following, we demonstrate the end-to-end workflow for Boltz2 in BioIR. Use model source="boltz-2" . Protein chains require an MSA; paired or unpaired MSAs are accepted for inputs with multiple non-identical protein chains. You can optionally supply templates yourself because BioIR does not run HHsearch or HMMsearch. The end-to-end processor supports ligand structure prediction, but not ligand-affinity prediction. python from bionemo ir.data.schemas import InputRequest, MSARecord, Polymer request = InputRequest input id="demo", polymers= Polymer polymer type="protein", chain id= "A" , sequence="GSHMSL...", msas= MSARecord path="msa.a3m", format="a3m" , paired msas= , templates=None, , rows = {"record": request, " record id": request "input id" } Replace the truncated sequence and MSA path with valid values. For Ray tests, build rows from a real worklist with more records than replicas; do not repeat one row as evidence of useful scaling. Step 2. Validate one prediction with the serial processor BioIR has two executor backends for the end-to-end processor workflow: serial and Ray. The serial backend runs each stage in sequence for one input, completing the full workflow before moving to the next input. This makes it useful for checking your setup before using the Ray backend to process independent inputs concurrently in Step 3. python import json from bionemo ir.pipeline.processor.engine proc import EngineProcessorConfig, build processor, from bionemo ir.pipeline.stages.configs import FeatureGeneratorStageConfig, WriterStageConfig, serial config = EngineProcessorConfig model source="boltz-2", executor backend=None, None mean the serial processor runtime args={ "recycling steps": 3, "num sampling steps": 50, "diffusion samples": 1, }, feature generator stage=FeatureGeneratorStageConfig init context={"random seed": 42}, , writer stage=WriterStageConfig output path="output/serial", format="cif", , engine kwargs={"profile inference": True}, serial processor = build processor serial config serial outputs = serial processor rows for row in serial outputs: scores = json.loads row "scores" print row "output path" print row "model inference time" print scores.get "confidence score" model inference time is BioIR’s CUDA-synchronized folding-model forward measurement. It excludes parsing, tokenization, feature generation, postprocessing, and writing. Set the seed through the feature generator’s init context . The scores field must be decoded, because it is a JSON string. Step 3. Scale independent inputs with Ray replicas The Ray backend can be selected as follows, for the default replica layout: python from bionemo ir.pipeline.processor.engine proc import EngineProcessorConfig ray config = EngineProcessorConfig.create default replica mode config model source="boltz-2", output dir="output/ray", output format="cif", This configuration places one complete model replica on each visible GPU on the current node, and sizes CPU stages from torch.cuda.device count . Below is an alternative configuration that controls four GPUs explicitly: python import ray from bionemo ir.pipeline.processor.engine proc import EngineProcessorConfig, build processor, from bionemo ir.pipeline.stages.configs import EngineStageConfig, FeatureGeneratorStageConfig, ParallelismMode, ParserStageConfig, TokenizerStageConfig, WriterStageConfig, ray config = EngineProcessorConfig model source="boltz-2", executor backend="ray", selects the ray backend parser stage=ParserStageConfig compute=4 , tokenizer stage=TokenizerStageConfig compute=4, num cpus=2 , feature generator stage=FeatureGeneratorStageConfig compute=8, num cpus=4, init context={"random seed": 42}, , engine stage=EngineStageConfig parallelism mode=ParallelismMode.REPLICA, compute=4, num gpus=1.0, num cpus=4, , writer stage=WriterStageConfig compute=4, output path="output/ray", format="cif", , ray processor = build processor ray config dataset = ray.data.from items rows ray outputs = list ray processor dataset .materialize .iter rows The capacity rule is engine stage.compute × engine stage.num gpus ≤ visible GPUs . This four-replica example is a single-node configuration that assumes four visible GPUs. This tutorial does not cover multi-node deployment. Here Ray creates four engine actors and reserves one GPU for each. Every actor loads the full model. build processor initializes Ray if needed. Actual throughput depends on input distribution, stage balance, storage, scheduling, and failures, so measure it. Step 4. Balance the five processor stages In the ray end-to-end processor, the five processor stages consume inputs in this dependency order: Parser → Tokenizer → Feature generator → Folding engine → Writer. Configure each stage with its matching StageConfig; the Step 3 example shows the relevant fields. The enabled field is not a public skip control. Each stage exposes compute ; relevant stages also expose num cpus , memor y, and batch size . The Ray engine adds max concurrent batches , accelerator type , and num gpus . To tune the Ray pipeline, start with EngineProcessorConfig.create default replica mode config ... . Increase a stage’s compute to add workers; use num cpus , memory , and, for engine actors, num gpus to set resource reservations. Add parser, tokenizer, or feature workers if engines wait for inputs. Reduce concurrency or separate large inputs when GPU or object-store memory causes failures. Ray is designed to overlap CPU stages with inference, but whether this improves the target workload depends on the run time cost of the parsing, feature generation, and output writing stages for a given input on a given hardware configuration. Refer to Figure 4 at ScaleFold https://arxiv.org/pdf/2404.11068 to see the diversity in pre-processing times for OpenFold. Step 5. Separate per-replica acceleration from pipeline scaling BioIR provides optimization at three distinct layers: - Kernel selection: Supported operations select compatible BioIR custom, cuEquivariance, or PyTorch fallback implementations based on the model configuration, GPU, data type, and tensor shape. - Module optimization: Where supported, the separate optimize mechanism enables CUDA Graph capture for compatible modules. - Pipeline scaling: The Ray executor places complete model replicas on GPUs and distributes independent inputs among them. These layers target different bottlenecks. Kernel and module optimizations reduce model-forward time within a replica. Ray can increase worklist throughput by overlapping CPU stages with GPU folding and by running full-model replicas on separate GPUs for independent inputs. Ray does not split a single model forward pass across GPUs. Figure 1, above, distinguishes the processor and direct-integration paths; Ray scaling applies only to the processor path. Our early benchmarking with BioNeMo Inference Runtime estimated the following model-forward accelerations: Speedups measured using 1 warmup run discarded and 1 measurement call across 17 inputs spanning 29–1,734 residues. OpenFold3 and Boltz2 OSS baselines used torch.compile with dynamic=None , fullgraph=False , recompile limit=128 , accumulated recompile limit=256 , fail on recompile limit hit=True . Boltz2 OSS used cuEq; OpenFold3 OSS used use cuequivariance=True and use deepspeed=True . These results quantify acceleration within one model replica. They do not measure parsing, feature generation, output writing, Ray scheduling, multi-GPU throughput, or complete-worklist wall time. Figure 3, below, shows why model-forward and end-to-end measurements must remain separate. To determine what additional GPUs unlock for a real deployment, measure the same representative worklist with one, two, and four Ray replicas on a single node. Step 6 defines the required metrics and comparison method. Step 6. Benchmark folding-stage efficiency and end-to-end delivery: AFDB – A Case Study To make these measurements concrete, we ran a matched benchmark on 1,000 human dimer targets with combined sequence lengths below 2,800 residues, representing a large collection of independent biomolecular structure-prediction tasks similar to the recently added dataset in the AlphaFold Database https://doi.org/10.64898/2026.03.27.714458 . This representative folding-stage benchmark compares BioIR-accelerated Boltz-2 with a torch compiled open-source Boltz-2 implementation, on 8xH100 GPUs. Both implementations used the same targets, staged MSAs, inference recipe, and GPU configuration; throughput metrics and other results are specific to this configuration and should not be generalized to all BioIR-supported models, datasets, or hardware. The workflow used three recycles, 200 sampling steps, and five diffusion samples per target. BioIR completed all 1,000 targets and delivered 58.5K successfully folded residues per allocated GPU-hour, compared with 20.2K for the public implementation—a 2.90× improvement in residue-normalized throughput; the open-source implementation ran out of memory on 29 targets. The left panel of Figure 3, above, compares the model-forward times of the BioIR and torch-compiled open-source implementations and directly shows the lower model-forward time delivered by BioIR. The right panel of Figure 3 compares the throughput delivered by BioIR with the open-source implementation, where throughput is the total number of residues in the predicted structures normalized by allocated GPU-hours. The right panel of Figure 3 shows the speedup delivered by the kernel-level, module-level, and pipeline-level optimizations in BioIR. The left panel of Figure 3 shows the speedup delivered by the kernel-level and module-level implementations. Similar to Boltz2 accelerations, BIR also enables faster inference for other biomolecular cofolding models, such as OpenFold2 and OpenFold3. An early version of BIR contributed accelerated modules to an NVIDIA-internal version of OpenFold2-MM, which enabled protein structure predictions at scale for the AFDB with 31 million protein complex structures. We linearly extrapolated the 1000-target matched benchmark from Figure 3 to one million comparable targets using rated-power equivalents for an 8 x H100 80GB HBM3 node see Figure 4, below . BioIR is estimated to require 11 MWh versus 35 MWh for the public implementation using 8-GPU TDP Thermal Design Power equivalents, and 21 MWh versus 64 MWh using full-node maximum-power equivalents. These are folding-only estimates for IT equipment, not metered energy measurements, and exclude data center overhead, such as power usage effectiveness PUE . The controlled comparison measures folding throughput with the same inputs and MSAs for each implementation; it excludes MSA generation, preprocessing CPU allocations, storage, data transfer, retries, and engineering overhead. Report end-to-end pipeline performance metrics as distinct from model-forward metrics, including completed structures per hour, GPU and CPU utilization, peak GPU memory, completion rate, failures, and retries. Troubleshooting - Only one GPU is active: Confirm Ray, REPLICA , more than one replica, several visible GPUs, and enough independent records. - Processor construction raises ValueError : Check that compute num gpus doesn’t exceed visible GPUs. - Protein input fails: Check the required unpaired A3M and worker-visible paths. - GPUs wait: Inspect CPU stages, CPU reservations, queueing, and Ray object-store capacity before adding replicas. - Rows wait after inference: Inspect writer concurrency and destination throughput. - Ray cannot place actors: Check CPU, GPU, memory, and accelerator type labels. - One record stops the job: The fail-fast default raises FoldingPredictionError . Set should continue on error=True only for intended row-level continuation, then inspect inference error . Get started Explore BioNeMo Inference Runtime BioIR and integrate it into your structure prediction workflows at scale: http://github.com/NVIDIA-BioNeMo/BioNeMo-Inference-Runtime http://github.com/NVIDIA-BioNeMo/BioNeMo-Inference-Runtime To further accelerate drug discovery workflows with agentic orchestration, check out NVIDIA BioNeMo Agent Toolkit BAT https://github.com/NVIDIA-BioNeMo/bionemo-agent-toolkit . For the latest acceleration numbers, consult the API reference https://docs.nvidia.com/bionemo/inference-runtime/references/api/ and support matrix https://docs.nvidia.com/bionemo/inference-runtime/references/support-matrix/ .