Accelerating End-to-End Co-Folding Performance with NVIDIA BioNeMo Agent Toolkit

NVIDIA announced the BioNeMo Agent Toolkit to accelerate biomolecular structure prediction and co-folding with OpenFold3, using GPU-accelerated MSA generation and multi-GPU scaling to enable virtual screening and large complex prediction. The toolkit integrates MMseqs2-GPU for up to 177× faster alignment and supports agentic workflows for drug discovery.

Biomolecular structure prediction and co-folding with models like OpenFold3 are now mainstream, large-scale workloads powering drug discovery and protein design. Increasingly, they’re driven end-to-end by AI agents. For an agent to run that pipeline well, every step needs to be fast and scalable: Multiple Sequence Alignment MSA generation, co-folding inference, serving, and multi-GPU scale-out. A bottleneck anywhere limits overall throughput. Speed and memory-efficiency are critical for key drug discovery workflows such as virtual screening and prediction of large molecular assemblies. In virtual screening, millions to billions of compounds are screened against one or a few protein targets. While co-folding models often give the best predicted structures, they can be expensive to run, making them impractical for virtual screening applications. That is where NVIDIA acceleration becomes key, making possible the deployment of OpenFold3 and related methods at the scale of large compound libraries. Speed is also important for predicting large molecular assemblies involving multiple proteins and thousands of amino acid residues, as co-folding model runtime scales cubically with the number of residues. An even bigger challenge, however, is memory use, as single GPU memory can be limited, placing a hard ceiling on the size of complexes that can be predicted in one shot. Methods to reduce memory requirements, and that distribute prediction tasks across multiple GPUs, would enable qualitatively new applications that are simply unfeasible today. NVIDIA has built tools to accelerate and improve the efficiency of each step of the structure prediction and co-folding workflow. NVIDIA BioNeMo Agent Toolkit gives agents seamless access to the tools they need to accelerate biology and chemistry workflows. In this post, we break down the accelerations for each stage on NVIDIA B300 and H100 GPUs, then show how those stages can be executed through an agent see Figure 1, below . Remove the MSA bottleneck with GPU MSA For co-folding models, building the MSA has traditionally been a CPU-bound step that can dominate wall-clock time. MMseqs2-GPU moves homology search onto NVIDIA GPUs, reducing this bottleneck while scaling with sequence length on both NVIDIA Hopper and NVIDIA Blackwell architectures. The latest GPU accelerated version https://www.biorxiv.org/content/10.64898/2026.03.11.711233v1 adds Hopper and Blackwell specific optimizations, including efficient support for larger-than-GPU-memory database search on NVIDIA Grace systems and additional speedups from improved Blackwell DPX instructions available from CUDA 13.2 https://docs.nvidia.com/cuda/archive/13.2.1/parallel-thread-execution/index.html ptx-isa-version-9-2 . These GPU contributions have been upstreamed back into the main MMseqs2 repository http://github.com/soedinglab/MMseqs2 so the whole community can benefit from the accelerations. The MSA Search NIM https://build.nvidia.com/colabfold/msa-search uses MMseqs2-GPU, whose Nature Methods paper https://www.nature.com/articles/s41592-025-02819-8 reports up to 177× faster alignment than CPU JackHMMER on a single L40S. In our benchmarking, the stage scales smoothly past 10k tokens on H100 and B300 GPUs see Figure 2, below . The MSA Search NIM can be called directly, self-hosted or wrapped as a tool in an agentic workflow. | Note: We show how to call the NIM API endpoint hosted on | Add the MSA Search NIM skill to your agent browse all: add --list npx skills add NVIDIA-BioNeMo/bionemo-agent-toolkit --skill msa-search-nim --agent claude-code Use hosted API on build.nvidia.com nothing to download You can also use the skill to download the NIM container and provide self-hosted API endpoint. We don't cover that in this tutorial. export NVIDIA API KEY=<key from build.nvidia.com Just prompt the agent: You have to download a sample target.fasta file. You can prompt the agent to download it for you or point to an already existing file. "Build an MSA for the sequence in target.fasta with the MSA Search NIM." Fold at SOTA speed with cuEquivariance and the OpenFold3 NIM cuEquivariance https://github.com/NVIDIA/cuEquivariance is a CUDA-X library of geometric learning primitives for atomistic modeling and it provides accelerated versions of the Triangle Attention, Triangle Multiplication and Attention Pair Bias kernels that dominate co-folding. On B300 it cuts latency up to ~3× see Table 1, below . Sequence length | PyTorch OSS | cuEquivariance | Speedup | | 1,024 H100 | 79.4 s | 41.6 s | 1.9× | | 1,024 B300 | 49.1 s | 27.3 s | 1.8× | | 1,536 B300 | 149.0 s | 56.2 s | 2.7× | | 2,048 B300 | 300.1 s | 97.6 s | 3.1× | Table 1. Comparison of OpenFold3 forward-pass latency with and without cuEquivariance cuEquivariance kernels are integrated directly into the OSS models like OpenFold3 https://github.com/aqlaboratory/openfold-3 provided as an optional dependency , OpenFold2 http://github.com/aqlaboratory/openfold , RosettaFold3 https://github.com/RosettaCommons/foundry/tree/production/models/rf3 , Protenix https://github.com/bytedance/Protenix and Boltz http://github.com/jwohlwend/boltz . Because the accelerations are upstreamed into these OSS models, a researcher gets the speedups automatically simply by running the model they already use on an NVIDIA GPU. CuEquivariance kernels also extends maximum sequence length to ~5.9k tokens while PyTorch runs out of memory beyond ~1.5k–2.5k tokens. On top of cuEquivariance, the OpenFold3 NIM applies further inference optimizations that compound the gain see Figure 4, below , achieving sequence lengths of up to ~ 6,400 on a single B300 . These additional accelerations are delivered through the NIM; for SOTA out of the box, developers can call the NIM endpoint directly or compose it into an agentic workflow. Add the OpenFold3 NIM skill to your agent npx skills add NVIDIA-BioNeMo/bionemo-agent-toolkit --skill openfold3-nim --agent claude-code Hosted API on build.nvidia.com You can also use the skill to download the NIM container and provide self-hosted API endpoint. We don't cover that in this tutorial. export NVIDIA API KEY=<key from build.nvidia.com Prompt the agent: You have to download a sample target.fasta file. You can prompt the agent to download it for you or point to an already existing file. "Fold target.fasta with OpenFold3 using the MSA from the previous step; return the ranked structures with confidence scores." Scale beyond one GPU with Fold-CP Single-GPU memory has historically capped co-folding models at a few thousand residues. On NVIDIA B300 Blackwell Ultra , the larger HBM and Blackwell-generation efficiency combined with the cuEquivariance and advanced inference optimizations above push that ceiling substantially higher with no model changes. In many cases, being single-device bound may not be sufficient. Fold-CP introduces a new parallelization technique such that per-device memory requirement scales as O N²/P where N is token count and P is the number of GPUs , reaching 32,000 tokens on 64 B300 with the Boltz-2 model—about a 12× jump over the single-GPU limit. To try Fold-CP, simply point your agent to the Boltz-CP codebase and ask it to run multi-GPU inference. Context-parallel inference across 4 GPUs with Fold-CP boltz-cp git clone https://github.com/NVIDIA-Digital-Bio/boltz-cp && cd boltz-cp Install dependencies and then run the command below torchrun --nnodes 1 --nproc per node 4 \ src/boltz/distributed/main.py predict /path/to/preprocessed data \ --out dir ./predictions \ --size dp 1 --size cp 4 \ --recycling steps 3 --sampling steps 200 --diffusion samples 5 Accelerating the end-to-end co-folding pipeline Structure prediction performance is now an end-to-end systems problem. For OpenFold3, the practical workflow spans MSA generation, co-folding inference, deployment, and the memory limits that determine how large a biological assembly can be modeled. NVIDIA accelerates each layer of that workflow: MSA Search NIM speeds homology search by 177x, cuEquivariance and OpenFold3 NIM lower inference latency by up to 4x on Blackwell GPUs, and Fold-CP shows how context parallelism can extend co-folding beyond a single GPU to 32,000-token complexes on 64 NVIDIA B300 GPUs. Together, these tools make structure prediction faster, more scalable, and easier to compose into agentic discovery workflows, helping researchers move from model predictions to larger, more useful biological systems. What these accelerations unlock These improvements open up classes of structural biology problems that were previously out of reach. In virtual screening, faster co-folding inference means that structure-based methods, which have historically been reserved for the final stages of a drug discovery campaign, can now be applied at far earlier stages and at much larger library scales, improving the quality and diversity of hits that advance through the pipeline. For large biomolecular assemblies, the combination of extended single-GPU capacity on B300 and the context-parallel Fold-CP framework shifts what is modelable: complexes at the scale of the ribosome, the spliceosome, or large signaling assemblies have been structurally intractable for co-folding models, and these accelerations begin to change that. To make this concrete, folding a complex of ~10,000 residues, roughly the scale of the bacterial ribosome, would have been prohibitively expensive or simply out of reach on a single GPU; on B300 with Fold-CP, such predictions become tractable across a multi-GPU node. These are qualitative shifts in the questions that structural biology can ask computationally, not merely improvements in throughput. Making these tools accessible through open-source integrations and agentic APIs will accelerate the pace at which computational predictions translate into biological insight and, ultimately, into new medicines. Getting started Try the accelerated OpenFold3 workflow with NVIDIA BioNeMo Agent Toolkit, starting with these tools: - MSA Search NIM: https://build.nvidia.com/colabfold/msa-search https://build.nvidia.com/colabfold/msa-search - cuEquivariance: https://github.com/nvidia/cuequivariance https://github.com/nvidia/cuequivariance - OpenFold3 NIM: https://build.nvidia.com/openfold/openfold3 https://build.nvidia.com/openfold/openfold3 - Fold-CP for larger assemblies: https://github.com/NVIDIA-BioNeMo/boltz-cp https://github.com/NVIDIA-BioNeMo/boltz-cp Acknowledgments We’d like to thank our broader NVIDIA team for developing the benchmarks and tools: Franco Pellegrini, Lalit Vaidya, Duc Tran, Tien Pham, Maximilian Stadler, Alejandro Chacon, Quan Vu, Simon Chu, Brian Roland, Dejun Lin, Joseph Chang, Hoa La, Jonathan Mitchell, Vishanth Iyer, Timur Rvachov, Christian Dallago, Christian Hundt. We’d also like to thank the broader OpenFold and OMSF teams for our collaboration and their contributions.