Loading…
Accelerating End-to-End Co-Folding Performance with NVIDIA BioNeMo Agent Toolkit
NVIDIA Developer BlogElizabeth Goodman
Summary
Biomolecular structure prediction and co-folding pipelines encounter latency and memory ceilings during sequence alignment and multi-residue inference. NVIDIA addresses these constraints through the BioNeMo Agent Toolkit, which connects AI agents to accelerated biological computing components. The pipeline offloads multiple sequence alignment to MMseqs2-GPU, while cuEquivariance and the OpenFold3 NIM accelerate core geometric attention kernels to cut inference latency up to fourfold. To handle large molecular assemblies beyond single-GPU memory limits, the Fold-CP framework distributes context across multiple GPUs with O(N²/P) memory scaling. Together, these tools enable the prediction of complexes up to 32,000 tokens on 64 NVIDIA B300 GPUs, making large-scale virtual screening and ribosome-scale modeling computationally feasible.
Context
Biomolecular structure prediction and co-folding workloads like OpenFold3 face severe throughput and memory bottlenecks in agentic drug discovery and protein design pipelines. Multiple Sequence Alignment generation has traditionally been CPU-bound, while co-folding runtime scales cubically with residue count. Single-GPU memory limits also cap the size of complexes that can be predicted in a single pass.
Approach / What changed
NVIDIA created an end-to-end acceleration workflow within the BioNeMo Agent Toolkit. Homology search is accelerated via MMseqs2-GPU in the MSA Search NIM. Inference relies on cuEquivariance kernels and the OpenFold3 NIM for optimized execution. For massive assemblies, the Fold-CP framework distributes memory requirements across multiple GPUs using context parallelism.
Takeaways
- MMseqs2-GPU moves Multiple Sequence Alignment to GPUs, achieving up to 177x faster alignment than CPU JackHMMER on an L40S and scaling past 10k tokens on H100 and B300 GPUs.
- The cuEquivariance library accelerates Triangle Attention, Triangle Multiplication, and Attention Pair Bias kernels, cutting OpenFold3 latency by up to 3.1x on B300 compared to PyTorch.
- Fold-CP distributes per-device memory at O(N²/P), enabling the Boltz-2 model to scale to 32,000 tokens across 64 NVIDIA B300 GPUs.
Related reading
Build an AI Scientist for Life Science Discovery with NVIDIA BioNeMo Agent Toolkit
General-purpose AI agents often struggle with biomolecular discovery because scientific workflows require precise model selection, strict input formatting, and artifact interpretation. The NVIDIA BioNeMo Agent Toolkit bridges this gap by providing BioNeMo Skills and Model Context Protocol wrappers on top of accelerated NVIDIA NIM microservices. These skills instruct agents on tool capabilities, required inputs, expected file formats like CIF and SMILES, and specific failure modes. Developers can route requests through hosted endpoints for testing or deploy local NIM microservices on GPU nodes to lower warm per-call latency during iterative candidate generation loops. Internal benchmarks using Codex CLI demonstrate that equipping agents with BioNeMo Skills improves task completion rates from 57.1% to 100% while doubling token efficiency across passing assertions.
Kyle Tretina