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

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.

Takeaways

1. 1

   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.

2. 2

   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.

3. 3

   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.