NVIDIA BioNeMo Inference Runtime Accelerates Proteome-Scale Biomolecular Structure Prediction

NVIDIA BioNeMo Inference Runtime Accelerates Proteome-Scale Biomolecular Structure Prediction NVIDIA BioNeMo Inference Runtime (BioIR) accelerates biomolecular structure-prediction models on NVIDIA GPUs while retaining a standard PyTorch workflow. Designed for high-throughput workloads, it couples kernel and CUDA Graph optimizations with a Ray-based multi-stage execution pipeline to distribute model replicas across GPUs. BioIR supports end-to-end processing across five pipeline stages—Parser, Tokenizer, Feature generator, Folding engine, and Writer—as well as direct PyTorch module integration. Ray multi-GPU scaling deploys full model replicas on individual GPUs within a single node to handle independent inputs concurrently without splitting single forward passes. In an 8xH100 benchmark on 1,000 human dimer targets using Boltz-2, BioIR achieved 58.5K folded residues per GPU-hour versus 20.2K for an open-source baseline, a 2.90× throughput gain. Linear extrapolations to one million comparable targets indicate folding energy consumption drops from 35 MWh to 11 MWh under thermal design power equivalents. Architecture and Pipeline Design BioNeMo Inference Runtime accelerates biomolecular structure models, including Boltz-2, OpenFold2, and OpenFold3, through two primary integration paths: direct PyTorch module adoption and an end-to-end pipeline processor. The processor executes five sequential stages: parsing, tokenization, feature generation, GPU folding, and file writing in PDB or mmCIF formats. Runtime execution does not require nvcc, CMake, or the CUDA toolkit because the package provides precompiled CUBINs. Protein chains require input multiple sequence alignments (MSAs), and users must supply their own structural templates as BioIR does not execute HHsearch or HMMsearch. Scaling Independent Inputs with Ray For batch processing, BioIR offers a serial backend for single-input validation and a Ray backend for scaling throughput across GPUs on a single node. The Ray executor assigns complete model replicas to individual GPUs rather than distributing a single forward pass across multiple devices. CPU stages run concurrently with GPU folding to overlap preparation overhead with model inference. Users must balance worker allocations across stages using configuration parameters such as compute, num_cpus, and num_gpus to avoid bottlenecks where folding engines sit idle waiting for upstream tokenization or feature generation. Layered Optimizations and Benchmark Results BioIR separates latency and throughput optimizations into three distinct layers: custom kernel selection (such as cuEquivariance or PyTorch fallbacks), CUDA Graph module capture, and Ray pipeline scaling. In a controlled benchmark across 1,000 human dimer targets totaling under 2,800 residues on an 8xH100 system, BioIR completed all targets and produced 58.5K folded residues per GPU-hour. By contrast, the public torch-compiled Boltz-2 implementation delivered 20.2K residues per GPU-hour and failed on 29 targets due to out-of-memory errors, demonstrating a 2.90× residue-normalized throughput improvement for BioIR. Energy Extrapolations and Operational Scope Extrapolating the matched 1,000-target benchmark to one million targets on 8xH100 80GB nodes yields estimated folding IT equipment power requirements of 11 MWh for BioIR compared to 35 MWh for the baseline at thermal design power, or 21 MWh versus 64 MWh at full-node maximum power. These calculations exclude data center overhead such as power usage effectiveness. Performance metrics for model-forward latency must remain separate from end-to-end pipeline metrics, as the benchmark measurements exclude MSA generation, CPU preprocessing overhead, storage input/output, and network transfers.

NVIDIA 推出 BioNeMo 推理运行时以加速高通量生物分子结构预测 NVIDIA 推出 BioNeMo Inference Runtime (BioIR)，旨在通过优化算子核、CUDA Graphs 与 Ray 框架的多副本扩展，在保持 PyTorch 工作流的同时加速蛋白质及复合体结构预测模型的高通量推理。 BioIR 提供端到端处理器和直接 PyTorch 模块集成两种模式，预编译 CUBIN 使其无需本地 nvcc 即可运行。 端到端工作流涵盖解析、分词、特征生成、折叠引擎与写出五大阶段，支持串行及单机多 GPU 的 Ray 副本并行。 在针对 1,000 个双聚体目标的 8xH100 基准测试中，BioIR 的残基归一化吞吐量达开源 Boltz-2 的 2.90 倍。 BioIR 的优化分为算子选择、模块优化和流水线扩展三层，模型前向测速与端到端耗时需严格区分评估。 BioIR 架构与工作流设计 BioNeMo Inference Runtime 专为蛋白质组规模的结构预测设计，曾助力 AlphaFold 数据库扩展并完成数千万候选复合体的计算。系统支持直接使用 PyTorch 模块或通过端到端处理器运行。其端到端流水线包含五个依赖阶段：Parser、Tokenizer、Feature generator、Folding engine 与 Writer。运行时提供了串行后端用于验证，以及 Ray 后端用于并发调度。在 Ray 配置下，引擎在单节点各 GPU 上部署完整模型副本以处理独立输入，但不会将单个前向传递拆分跨卡。 三层优化机制与阶段资源调优 BioIR 在三层提供差异化优化：算子层针对硬件与张量形状自动选取定制核、cuEquivariance 或回退实现；模块层使用 optimize() 启用 CUDA Graph 捕获；流水线层则通过 Ray 重叠 CPU 预处理与 GPU 推理。调优流水线时，可通过各阶段配置的 compute、num_cpus 及 num_gpus 匹配资源。若 GPU 闲置等待，需扩展上游特征生成或解析的 CPU 算力；若显存或对象存储承压，则需收缩并发规模以防崩溃。 基准测试表现与能耗推算 在 8 张 H100 GPU 上的 1,000 个复合体目标对照测试中，BioIR 运行 Boltz-2 完整处理了所有样本，并取得每 GPU 小时 58.5K 残基的吞吐表现，而基于 torch.compile 的开源实现在 29 个目标上面临 OOM，吞吐仅为 20.2K 残基，BioIR 实现 2.90 倍残基吞吐加速。线性推算至百万目标时，BioIR 在 8 卡 TDP 下预计仅耗电 11 MWh，较公开实现的 35 MWh 显著降低。

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 (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 (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. Figure 1. BioNeMo Inference Runtime supports an end-to-end processor and direct PyTorch module integration; kernel and eligible module optimizations apply within model execution; the five-stage Ray pipeline applies to the end-to-end processor workflow Prerequisites Python 3.12 or later A compatible NVIDIA GPU and driver, plus a BioIR wheel or supported development environment A staged model checkpoint (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. 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. 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: 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: 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 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: Figure 2. BioNeMo Inference Runtime geometric-mean speedups for GPU-synchronized model-forward latency relative to OSS torch.compile on H100 and H200 GPUs 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 . 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. Figure 3. BioIR delivers 2.90× higher residue-normalized folding throughput than the tested public open-source implementation across a matched 1,000-target human-dimer benchmark. Throughput metrics depend on the composition of the input dataset 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). Figure 4. BioIR reduces the estimated rated-power-equivalent energy to fold one million comparable targets by 23 MWh using nodes with 8x H100 80 GB GPUs at TDP, and by 43 MWh at full-node maximum power 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 To further accelerate drug discovery workflows with agentic orchestration, check out NVIDIA BioNeMo Agent Toolkit (BAT) . For the latest acceleration numbers, consult the API reference and support matrix .