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

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. 

# 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). 

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.