NVIDIA Speeds Up Co-Folding Pipeline With GPU MSA, Optimized Kernels, and Context Parallelism

NVIDIA published benchmarks for its BioNeMo Agent Toolkit that accelerate each stage of the biomolecular co-folding pipeline, from multiple sequence alignment (MSA) generation to multi-GPU inference. The company reports up to 177x faster alignment with GPU-based MMseqs2-GPU compared with CPU JackHMMER, up to 3.1x lower forward-pass latency with cuEquivariance kernels on B300 GPUs, and context parallelism that scales per-device memory as O(N²/P), enabling 32,000-token complexes on 64 B300 GPUs with the Boltz-2 model.

GPU MSA Search

MSA generation has traditionally been a CPU-bound step that dominates wall-clock time for co-folding models. NVIDIA's MMseqs2-GPU moves homology search onto GPUs and adds Hopper and Blackwell-specific optimizations, including support for databases larger than GPU memory on Grace systems and improved Blackwell DPX instructions available from CUDA 13.2. These contributions were upstreamed into the main MMseqs2 repository.

The MSA Search NIM uses MMseqs2-GPU, whose Nature Methods paper reports up to 177x faster alignment than CPU JackHMMER on a single L40S. NVIDIA's benchmarking shows the stage scales smoothly past 10,000 tokens on H100 and B300 GPUs. The NIM can be called directly, self-hosted, or wrapped as a tool in an agentic workflow. For large-scale workloads, NVIDIA recommends the self-hosted NIM over the hosted API endpoint on build.nvidia.com, which is not built to handle large volume requests and will time out.

cuEquivariance and OpenFold3 NIM

cuEquivariance is a CUDA-X library of geometric learning primitives for atomistic modeling. It provides accelerated versions of the Triangle Attention, Triangle Multiplication, and Attention Pair Bias kernels that dominate co-folding. On B300, it cuts latency by up to 3.1x. The kernels are integrated directly into open-source models including OpenFold3, OpenFold2, RosettaFold3, Protenix, and Boltz, so researchers get the speedups automatically by running their existing models on an NVIDIA GPU. cuEquivariance also extends maximum sequence length to approximately 5,900 tokens, while PyTorch runs out of memory beyond roughly 1,500 to 2,500 tokens.

The OpenFold3 NIM applies further inference optimizations on top of cuEquivariance, achieving sequence lengths up to approximately 6,400 tokens on a single B300. NVIDIA reports over 3x faster folding on H100 and over 4x faster on B300 compared with the baseline PyTorch implementation.

Fold-CP for Multi-GPU Scale-Out

Single-GPU memory has historically capped co-folding models at a few thousand residues. On B300, larger HBM and Blackwell-generation efficiency combined with the cuEquivariance and inference optimizations push that ceiling higher without model changes. 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. This reaches 32,000 tokens on 64 B300 GPUs with the Boltz-2 model, about a 12x jump over the single-GPU limit.

NVIDIA positions these accelerations as enabling structure-based virtual screening at earlier stages of drug discovery, where co-folding models have historically been too expensive for large compound libraries. For large biomolecular assemblies, the combination of extended single-GPU capacity on B300 and the context-parallel Fold-CP framework makes complexes at the scale of the ribosome or spliceosome tractable for co-folding models. Folding a complex of roughly 10,000 residues, approximately the scale of the bacterial ribosome, would have been prohibitively expensive on a single GPU; on B300 with Fold-CP, such predictions become tractable across a multi-GPU node.