Accelerating End-to-End Co-Folding Performance with NVIDIA BioNeMo Agent Toolkit
Mirrored from NVIDIA Developer Blog for archival readability. Support the source by reading on the original site.
Accelerating End-to-End Co-Folding Performance with NVIDIA BioNeMo Agent Toolkit
AI-Generated Summary
- NVIDIA has delivered end-to-end acceleration for biomolecular structure prediction pipelines, including MSA generation with MMseqs2-GPU, co-folding inference via cuEquivariance and OpenFold3 NIM, and multi-GPU scaling through Fold-CP, all integrated within the BioNeMo Agent Toolkit.
- MMseqs2-GPU on NVIDIA Hopper and Blackwell architectures achieves up to 177x faster MSA generation compared to CPU JackHMMER, while cuEquivariance cuts OpenFold3 forward-pass runtime by up to 3x and extends sequence length limits to ~5.9k tokens, with OpenFold3 NIM achieving up to 6,400 tokens on a single B300.
- Fold-CP enables context-parallel inference, reducing per-GPU memory requirements to O(N/P) and allowing large assemblies (up to 32,000 tokens on 64 B300 GPUs) to be modeled, unlocking previously intractable structural biology problems such as ribosome-scale complexes.
AI-generated content may summarize information incompletely. Verify important information. Learn more
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 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. These GPU contributions have been upstreamed back into the main MMseqs2 repository so the whole community can benefit from the accelerations.
The MSA Search NIM uses MMseqs2-GPU, whose Nature Methods paper 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 build.nvidia.com, but for large-scale workloads, the self-hosted NIM is the preferred recommended endpoint. build.nvidia.com endpoints are not built to handle large volume requests and will time-out. |
# 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 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× |
cuEquivariance kernels are integrated directly into the OSS models like OpenFold3 (provided as an optional dependency), OpenFold2, RosettaFold3, Protenix and 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
- cuEquivariance: https://github.com/nvidia/cuequivariance
- OpenFold3 NIM: https://build.nvidia.com/openfold/openfold3
- Fold-CP for larger assemblies: 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.
Tags
About the Authors
Kyle Tretina is a product marketing leader at NVIDIA, focused on advancing AI for digital biology and drug discovery. He drives the strategy and storytelling behind BioNeMo and our work with BioPharma, shaping how next-generation foundation models and GPU-accelerated microservices transform molecular and protein design. With a PhD in molecular microbiology and immunology, Kyle bridges science and strategy, translating breakthroughs in AI, chemistry, and biology into platforms that accelerate discovery for researchers, startups, and pharmaceutical companies worldwide.
Roy is a senior alliance manager for BioNeMo, NVIDIA’s digital biology platform. His academic background is at the intersection of biology and machine learning with an MSc from UCL, MS from École Normale Supérieure and BSc from Imperial College. Roy has 7 years of industry experience in AI-driven drug discovery.
Zoey Zhang is the product manager for AI training in Digital Biology at NVIDIA. Her background spans software engineering and machine learning research roles, with a degree in Biomedical Engineering from the University of Waterloo specializing in Medical AI and Computing. Zoey is passionate about accelerating scientific discovery and the development of life-saving treatments through AI and accelerated computing.
Emine Kucukbenli is R&D manager with the NVIDIA BioNeMo team. In the last decade, they have been working on first principles atomistic modeling with machine learning for material science and drug discovery applications.
Jared Wilber is a senior software engineer for digital biology on the BioNeMo team at NVIDIA, where he trains and interprets biological foundation models and builds tools to visualize the data behind them. Before NVIDIA, he founded a Y Combinator–backed startup, worked at AWS, and lectured at UC San Diego. He holds a degree in statistics and computer science from UC Berkeley.
Mohammed AlQuraishi is an assistant professor of systems biology (with a joint appointment in computer science) at Columbia University, where he leads the AlQuraishi Laboratory and is a member of the Program for Mathematical Genomics. His research sits at the intersection of machine learning, biophysics, and systems biology, operating across two scales: at the molecular level, developing ML models for protein structure/function prediction, protein-ligand interactions, and learned protein representations; at the systems level, applying those models proteome-wide to study signal transduction network organization, population-level variation, and cancer dysregulation.
Comments
Comments are closed.
More from NVIDIA Developer Blog
-
Serve Qwen3.8-2.4T-A95B, a 2.4T-Parameter Model, with Configurable Reasoning on NVIDIA GB300 NVL72
Aug 12
-
How to Choose Full-Stack Observability for NVIDIA AI Factories
Aug 12
-
NVIDIA JetPack 7.2.1 Adds Agentic Video Skills and T3000 Emulation
Aug 11
-
NVIDIA Nemotron 3.5 Lightning Delivers Fast, Accurate Specialized Task Execution for Long-Running Agents
Aug 11
Discussion (0)
Sign in to join the discussion. Free account, 30 seconds — email code or GitHub.
Sign in →No comments yet. Sign in and be the first to say something.