Fold-CP: A Context Parallelism Framework for Biomolecular Modeling
Authors: Dejun Lin, Simon Chu, Vishanth Iyer, Youhan Lee, John St John, Kevin Boyd, Brian Roland, Xiaowei Ren, Guoqing Zhou, Zhonglin Cao, Polina Binder, Yuliya Zhautouskaya, Jakub Zakrzewski, Maximilian Stadler, Kyle Gion, Yuxing Peng, Xi Chen, Tianjing Zhang, Philipp Junk, Michelle Dimon, Paweł Gniewek, Fabian Ortega, McKinley Polen, Ivan Grubisic, Ali Bashir, Graham Holt, Danny Kovtun, Matthias Grass, Luca Naef, Rui Wang, Jian Peng, Anthony Costa, Saee Paliwal, Eddie Calleja, Timur Rvachov, Neha Tadimeti, Roy Tal, Emine Kucukbenli
Understanding cellular machinery requires atomic-scale reconstruction of large biomolecular assemblies. However, predicting the structures of these systems has been constrained by hardware memory requirements of models like AlphaFold 3, imposing a practical ceiling of a few thousand residues that can be processed on a single GPU. Here we present NVIDIA BioNeMo Fold-CP, a context parallelism framework that overcomes this barrier by distributing the inference and training pipelines of co-folding models across multiple GPUs. We use the Boltz models as open source reference architectures and implement custom multidimensional primitives that efficiently parallelize both the dense triangular updates and the irregular, data-dependent pattern of window-batched local attention. Our approach achieves efficient memory scaling; for an N-token input distributed across P GPUs, per-device memory scales as O(N2/P), enabling the structure prediction of assemblies exceeding 30,000 residues on 64 NVIDIA B300 GPUs. We demonstrate the scientific utility of this approach through successful developer use cases: Fold-CP enabled the scoring of over 90% of Comprehensive Resource of Mammalian protein complexes (CORUM) database, as well as folding of disease-relevant PI4KA lipid kinase complex bound to an intrinsically disordered region without cropping. By providing a scalable pathway for modeling massive systems with full global context, Fold-CP represents a significant step toward the realization of a virtual cell.
Figures & tables
Figure 1: Software Architecture of the Fold-CP Framework. Fold-CP design takes a bottom-up approach starting with foundational torch.distributed APIs, making module-level operators CP-aware, and incrementally building up to model-specific workflows. This design allows developers to port specific distributed modules into diverse model organizations without managing low-level synchronization logic or altering the model flow.
Component
Serial Time
Serial Space
CP Time (per rank)
CP Space (per rank)
Comm. Pattern
Triangle Attention
BHN3D
BHN2D
BHN3D/P
BHN2D/P
2D ring + transpose
Triangle Multiplication
BN3D
BN2D
BN3D/P
BN2D/P
2D ring (Cannon-style)
Pair Weighted Avg.
BHSN2D
BHSND
BHSN2D/P
BHSND/P
2D ring + tiled softmax
Outer Product Mean
BSN2D2
BN2D2
BSN2D2/P
BN2D2/P
2D ring + all-reduce
Attn Pair Bias (Ring)
BHN2D
BHND
BHN2D/P
BHPND
1D ring + transpose
Attn Pair Bias (Shard)
BAHD
BAHD
BPAHD
BPAHD
None (upstream gather)
Table 1: Complexity of distributed modules and the scaling factors. N : number of tokens; P : GPU count; B : batch size; S : number of MSA sequences; H : number of attention heads; D : feature dimension (per-head dimension for attention-based modules); A : number of atoms. Shard-wise complexity absorbs constant window and key-span sizes into O(⋅) ; atoms are partitioned into fixed-size windows with overlapping key neighborhoods. All ring patterns use P steps with double-buffered overlap. For Window Batching, the serial column reflects the pre-Toeplitz baseline; the CP column includes both the Toeplitz optimization and distribution (see Section 3.7 ).
Figure 2: Distributed forward-pass flow for the three core ring-communication modules. (a) Triangle Attention: cross-axis transpose of triangular bias followed by a 1D ring rotation of keys, values, and bias with tiled-softmax merge. (b) Triangle Multiplication: Cannon-style 2D ring where projections shift along orthogonal axes with additive accumulation. (c) Ring Attention Pair Bias: 1D ring rotation of keys, values, and pair bias with tiled-softmax merge. Pair Weighted Averaging (Section 3.4 ) shares the softmax-merge archetype of (a)/(c) over a 2D ring; Outer Product Mean (Section 3.5 ) shares the Cannon accumulation archetype of (b), replacing matrix multiplication with outer products.
Figure 3: Boltz-2 CP scalability on NVIDIA B300 GPUs via maximum context length. For a given GPU count P , maximum context length reached is plotted. Linear trend lines are used to inform with respect to number of GPUs in a CP rank, P , due to the square topology requirement of Fold-CP. (Left) Inference scaling performance: The maximum context length scales linearly at a rate of approximately 4,000 tokens/ P , with runtime annotated accordingly. (Right) Training scaling performance: The training follows a similar scaling with a approximately half the slope of that of inference.
# GPUs
Confidence prediction off
Confidence prediction on
Tokens
Walltime (mins)
Tokens
Walltime (mins)
1×1
2,700
13
2,400
12
2×2
4,000
13
3,800
16
4×4
6,500
37
6,000
33
8×8
11,500
47
11,000
55
Table 2: Boltz-1 CP inference scaling on NVIDIA H100 80GB GPUs. Confidence prediction requires caching activation from structure prediction, leading to about 10% reduction in maximum context length. The reported token count is the output of tokenization process, e.g. a sample with PDB ID: 7NYC consists of 6,279 residues and gets tokenized to 6,500 tokens.
Figure 4: Inference consistency validation on Boltz-1 test set. Numerical equivalence between the DP baseline and Fold-CP implementation verified via a Two One-Sided Tests (TOST) procedure on a Wilcoxon signed-rank test with a margin of error ( ϵ =0.01). Error bar estimated from standard error of mean (SEM) from 5 diffusion samples for the respective x- (DP, horizontal bar) and y-axis (CP, vertical bar).
Figure 5: Training consistency validation. Validation lDDT trajectories for five DP baseline with random seeding and one CP=2×2 configuration with a model trained on 256-token crop size. The Fold-CP validation lDDT curve largely mirrors the DP references, with variation stemming from stochastic nature of the model. The model used in this validation experiment is a truncated Boltz-1 architecture, hence it is not expected to match the published reference. See text for details.
Figure 6: Boltz-2 prediction of the TTC7A/PI4KA/FAM126A/EFR3A(700–823) complex (3,605 residues) . PI4KA (green), TTC7A (cyan), FAM126A (magenta), and the intrinsically disordered EFR3A C-terminus (yellow). The prediction was generated in a single inference pass using Fold-CP across 4 NVIDIA H100 80GB GPUs.
Figure 7: Enrichment of different methods for identifying protein complexes, using Boltz-1 CP for structural validation. From left to right: Rezo , Complexes predicted by Rezo’s Network AI; CORUM , CORUM protein complexes; STRING Clique , High-confidence STRING interactions that form cliques; and STRING Baseline , High-confidence STRING interactions that form connected components but not cliques. To ensure comparability, all evaluated sets were restricted to complexes containing three or four proteins.
Figure 8: Structural comparison of the ELP123-ELP456 complex. (a) A side-by-side comparison of the complete complex. On the left is the yeast structure (PDB 8ASV [ 27 ] ), and on the right is the Boltz-1 CP predicted structure of the human Elongator complex. Outlined sections represent XL-MS contact sites [ 28 ] that were either captured (contact sites 1 and 3, CS1+3) or missed (contact site 2, CS2) by the model. (b) A magnified view of the Boltz-1 CP predicted structure, overlaid on 8ASV, appears to recapitulate previously identified interface contact sites (1 and 3) [ 27 ] that were successfully mapped from yeast to human. The remaining contact site 2 resides in a region of low sequence and structural homology between the human (dashed rectangle) and yeast proteins, and likely represents an error in the predicted structure, as detailed in the text.
Figure 9: Structure prediction of 5YZO homo 4-mer S9 peptidase mutant (S514A) from Deinococcus radiodurans R1. Yellow is the ground truth while blue is the predicted version from Earendil’s model. The prediction was generated in a single inference pass using Fold-CP across 64 NVIDIA H100 80GB GPUs.
Figure 10: Structural artifacts in long-context homomeric assemblies. Boltz-1 CP predicted structures exhibit chain overlap and clumping. Experiment structures are in gray. Confidence scale (pLDDT): 90–100 (dark blue), 70–90 (light blue), 50–70 (yellow), 0–50 (orange).
Many all-atom generative protein models can in principle design large multimeric complexes by jointly modelling all chains, but their quadratic token- and atom-pair representations quickly exceed single-GPU memory as the number of chains and residues modelled grows. We introduce Design-CP, two context-parallel (CP) inference strategies for RFdiffusion 3 (1D row-sharding and 2D grid sharding with ring attention) that distribute the quadratic activations across a multi-GPU mesh while preserving pretrained weights. We characterise their scaling when sampling icosahedral assemblies, showing that the maximum feasible asymmetric subunit (ASU) size grows with the expected square-root trend in GPU count and that 2D sharding achieves better wall-clock scaling. Moreover, we show how strong point-group symmetry constraints make CP usable out of the box for end-to-end, all-atom design of icosahedral nanoparticles, yielding favourable in silico structural and interface metrics. Finally, we demonstrate octahedral nanoparticle design on a small cluster of workstation-grade 16GB GPUs, illustrating how Design-CP can be a practical path towards democratising large-assembly protein design.
Lorenzo Tarricone, Helen E. Eisenach, Aiko Muraishi +1
Department of Statistics, University of Oxford, Oxford, UK · Ellison Institute of Technology Oxford, Oxford, UK · Institute for Protein Design, University of Washington, Seattle, USA +2
Accurately modeling biomolecular interactions is a central bottleneck in biology and therapeutic discovery. Here, we introduce Open Drug Discovery Engine (OpenDDE), an open-source, all-atom biomolecular foundation model that uses co-folding as the entry point to a scalable AI-driven drug discovery engine. Rather than treating structure prediction as an isolated endpoint, OpenDDE is designed as a shared structural reasoning layer for modeling sequence-structure-function relationships across biomolecular complexes, enabling complex structure prediction today while providing a foundation for de novo design, affinity estimation, structure-conditioned optimization, and more. OpenDDE integrates advances in all-atom architecture, atomic latent reasoning, inference optimization, and large-scale data processing to achieve IsoDDE-level co-folding accuracy within a reproducible and openly accessible framework. We also identify two scaling-law directions for co-folding models, revealing practical routes for continued improvement through data, model, inference, and training scaling. By releasing training code, inference pipelines, checkpoints, and benchmarks, OpenDDE aims to democratize access to frontier biomolecular intelligence, accelerate global collaboration, and lay an open foundation for next-generation drug discovery systems that can move from predicting molecular structures toward designing, scoring, and optimizing therapeutic candidates for human health.
All-atom generative modeling of 3D biomolecular complexes has emerged as the dominant paradigm for predicting the structure of proteins and protein-ligand systems. Generating structures at the atomic level of fidelity, however, typically requires expensive iterative diffusion rollouts, making both conventional deployment and inference-time search techniques computationally costly. In this paper, we introduce the Denoiser Cofolding All-Atom Flowmap (DeCAF) framework for distilling state-of-the-art all-atom cofolding models into all-atom flow maps that produce high-quality samples in only a few inference steps. We build DeCAF on a denoiser-based formulation of flow maps with endpoint losses that naturally support SE(3) rigid alignment, which we show is critical for training accurate models. We further derive a simple change of variables that lets DeCAF operate in the σ-space noise schedule of EDM-style architectures, enabling direct distillation from pretrained cofolding diffusion models. Equipped with DeCAF's flowmap lookahead, we introduce a purpose-built inference-time framework that improves sampling through reward-guided search. Empirically, DeCAF-Boltz statistically improves over Boltz-1x in both accuracy (RMSD) and physical validity scores of protein-ligand poses at strict NFE budgets on the challenging Runs N' Poses, while also showing a more optimal Pareto frontier across all inference compute budgets on PoseBusters. Distilling the state-of-the-art Pearl cofolding model, DeCAF-Pearl outperforms diffusion-based cofolding models and matches its teacher on success rate while using 5x fewer NFEs. We release our code at https://github.com/genesistherapeutics/decaf.
Gianluca Scarpellini, Ron Shprints, Peter Holderrieth +7
1Genesis Molecular AI · Massachusetts Institute of Technology · Carnegie Mellon University +2