Fold-CP: A Context Parallelism Framework for Biomolecular Modeling

Dejun Lin; Simon Chu; Vishanth Iyer; Youhan Lee; John St John; Kevin Boyd; Brian Roland; Xiaowei Ren

q-bio.QMMar 16, 2026

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

Organizations: NVIDIA · Rezo Therapeutics · Proxima · Earendil Labs

Abstract

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(N^2/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	$BHN^{3}D$	$BHN^{2}D$	$BHN^{3}D/P$	$BHN^{2}D/P$	2D ring + transpose
Triangle Multiplication	$BN^{3}D$	$BN^{2}D$	$BN^{3}D/P$	$BN^{2}D/P$	2D ring (Cannon-style)
Pair Weighted Avg.	$BHSN^{2}D$	$BHSND$	$BHSN^{2}D/P$	$BHSND/P$	2D ring + tiled softmax
Outer Product Mean	$BSN^{2}D^{2}$	$BN^{2}D^{2}$	$BSN^{2}D^{2}/P$	$BN^{2}D^{2}/P$	2D ring + all-reduce
Attn Pair Bias (Ring)	$BHN^{2}D$	$BHND$	$BHN^{2}D/\sqrt{P}$	$BH\frac{N}{\sqrt{P}}D$	1D ring + transpose
Attn Pair Bias (Shard)	$BAHD$	$BAHD$	$B\frac{A}{\sqrt{P}}HD$	$B\frac{A}{\sqrt{P}}HD$	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(\cdot)

; atoms are partitioned into fixed-size windows with overlapping key neighborhoods. All ring patterns use

\sqrt{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, $\sqrt{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/ $\sqrt{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
# GPUs	Tokens	Walltime (mins)	Tokens	Walltime (mins)
$1\times 1$	2,700	13	2,400	12
$2\times 2$	4,000	13	3,800	16
$4\times 4$	6,500	37	6,000	33
$8\times 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 ( $\epsilon$ =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 $\text{CP}=2\times 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.

Panel of Figure 10 — 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).

Context Parallelism

See all 10 papers

Design-CP: Context Parallelism for Design of Protein Nanoparticles
Jul 3, 2026Lorenzo Tarricone, Helen E. Eisenach, Aiko Muraishi +1GPU Acceleration High-Performance Computing
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
Folding, Reasoning, and Scaling with Open-source Drug Discovery Engine
Jul 4, 2026Aureka AI OpenDDE projectProtein Structure Prediction Molecular Foundation Models
OpenDDE Project, Aureka AI Research
Few-step Cofolding with All-Atom Flow Maps
Jun 7, 2026Gianluca Scarpellini, Ron Shprints, Peter Holderrieth +7Protein Structure Prediction Diffusion Distillation
1Genesis Molecular AI · Massachusetts Institute of Technology · Carnegie Mellon University +2

Jul 3, 2026J/K move · Enter open · S save

Design-CP: Context Parallelism for Design of Protein Nanoparticles

Lorenzo Tarricone, Helen E. Eisenach, Aiko Muraishi, Charlotte M. Deane

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

GPU Acceleration High-Performance Computing Protein Design Distributed Inference Context Parallelism

Open paper AskPDF arXiv

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.

Abstract

Context Parallelism

Explore similar work

Design-CP: Context Parallelism for Design of Protein Nanoparticles

Folding, Reasoning, and Scaling with Open-source Drug Discovery Engine

Few-step Cofolding with All-Atom Flow Maps