NVIDIA BioNeMo Fold-CP Framework

• NVIDIA BioNeMo Fold-CP is a distributed framework that enables GPU mesh execution to predict large biomolecular structures beyond single-device limits.
• The system shards sequence tokens, MSA features, and atom coordinates to reduce per-device memory to O(N^2/P), facilitating efficient, large-scale co-folding.
• It integrates with AlphaFold-style models using custom distributed primitives to preserve numerical equivalence while accelerating full-context structure prediction.

Searching arXiv for the Fold-CP paper and closely related biomolecular modeling references to ground the article. arxiv_search(query="(Lin et al., 16 Mar 2026)", max_results=5)

I’ll look up the exact arXiv record for Fold-CP and related Boltz references before writing. NVIDIA BioNeMo Fold-CP is a general-purpose framework that lifts the context-length ceiling of AlphaFold-style co-folding models by distributing both inference and training across a mesh of GPUs. It is presented as a context parallelism framework for biomolecular modeling, motivated by the atomic-scale reconstruction of large biomolecular assemblies that are otherwise constrained by the hardware memory requirements of models such as AlphaFold 3. Using the Boltz models as open source reference architectures, Fold-CP shards sequence tokens, multiple-sequence alignments, and atom coordinates so that no single device ever materializes the full $O(N^2)$ pair representation. For an $N$-token input distributed across $P$ GPUs, per-device memory scales as $O(N^2/P)$, enabling structure prediction of assemblies exceeding 30,000 residues on 64 NVIDIA B300 GPUs while preserving exact numerical equivalence to single-device baselines (Lin et al., 16 Mar 2026).

1. Computational setting and motivation

Modern co-folding models such as AlphaFold 3 and Boltz use a dense pairwise representation $z\in\mathbb R^{N\times N\times D}$ whose memory grows as $O(N^2)$ and whose geometric operations incur up to $O(N^3)$ compute. In practice, this imposes a practical ceiling of a few thousand residues that can be processed on a single GPU. Fold-CP is designed to overcome this barrier by distributing a single protein or protein assembly across multiple GPUs rather than constraining execution to a single-device context (Lin et al., 16 Mar 2026).

At the center of the framework is Context Parallelism (CP), a multi-dimensional sharding strategy that partitions three geometric manifolds: sequence tokens, multiple-sequence alignments (MSA), and atom coordinates. The explicit objective is that no single rank ever exceeds $O(N^2/P)$ memory. This design choice is significant because the dominant bottleneck in AlphaFold-style co-folding is not only the number of residues, but the quadratic pair tensor and the cubic-cost geometric kernels that operate on it.

Fold-CP is also framed as a full-pipeline system rather than a narrowly scoped kernel optimization. The framework distributes both inference and training pipelines of co-folding models across multiple GPUs. A plausible implication is that the method targets long-context biomolecular modeling as a systems problem involving memory layout, communication scheduling, autograd behavior, and fused kernel execution, rather than only a single operator-level acceleration.

2. Context Parallelism architecture

Fold-CP arranges the $P$ GPUs into a 2D “context mesh” of size $R\times C$ with $N$0, plus an orthogonal data-parallel dimension. By tiling the pair matrix into $N$1 blocks, each GPU stores a sub-matrix of size $N$2. This reduces per-device memory to $N$3 and enables linear growth of maximum context with $N$4 (Lin et al., 16 Mar 2026).

The distributed state is defined componentwise. Pair representations are 2D-tiled so that each GPU $N$5 holds

$N$6

MSA features $N$7 are sharded along both the depth $N$8 and token $N$9 axes to align with pair shards. Atom coordinates, which map irregularly to tokens, are co-sharded with tokens to preserve local atom–token mappings. For window-batched attention, atom tokens are re-packed into contiguous intervals per shard via small P2P halo exchanges.

The memory-scaling argument is explicit. Let $P$0 be the number of input tokens, defined as residues plus ligand atoms, and let $P$1 be the number of GPUs. The global pair tensor has size $P$2. With $P$3, each shard has dimensions

$P$4

so per-device memory for pair activations is

$P$5

A similar argument holds for MSA and atom features, ensuring that no single rank ever exceeds $P$6 memory. This makes the pair tensor, MSA state, and atom features consistent participants in the same sharding regime rather than independent memory domains.

3. Distributed primitives for Evoformer-style modules

Fold-CP introduces custom multidimensional primitives for each Evoformer-style module. The framework interleaves local computation with asynchronous peer-to-peer transfers, implemented with torch.distributed send/recv, so that forward and backward passes proceed without ever gathering the full $P$7 tensor on a single device (Lin et al., 16 Mar 2026).

Triangle Attention computes row- or column-wise multi-head attention on $P$8, adding a triangular bias $P$9 orthogonal to the mask. Queries remain local, while keys, values, bias, and masks are first re-aligned via a 2D transpose and then circulated through a $O(N^2/P)$0-step ring. At step $O(N^2/P)$1, each rank computes partial logits

$O(N^2/P)$2

and merges softmax outputs using a numerically stable tiled softmax with running log-sum-exp. Double buffering overlaps P2P transfer of the next block with the current compute. The reported compute per rank is $O(N^2/P)$3 and the memory per rank is $O(N^2/P)$4.

Triangle Multiplication updates

$O(N^2/P)$5

which is structurally a batched matmul with element-wise gating. Fold-CP transposes one operand across the 2D mesh, then performs a $O(N^2/P)$6-step Cannon-style 2D ring. In each step, each rank multiplies its local blocks of $O(N^2/P)$7 and $O(N^2/P)$8 while shifting $O(N^2/P)$9 along rows and $z\in\mathbb R^{N\times N\times D}$0 along columns, accumulating into the local output. The compute per rank is $z\in\mathbb R^{N\times N\times D}$1 and the memory per rank is $z\in\mathbb R^{N\times N\times D}$2.

Pair Weighted Averaging forms attention weights from pair features and applies them to MSA-projected values:

$z\in\mathbb R^{N\times N\times D}$3

Both $z\in\mathbb R^{N\times N\times D}$4 and the MSA-projected $z\in\mathbb R^{N\times N\times D}$5 are sharded. Pair weights are transposed and then ring-shifted across the 2D mesh. Each step computes partial softmax over a local $z\in\mathbb R^{N\times N\times D}$6 block and merges via tiled softmax. The compute per rank is $z\in\mathbb R^{N\times N\times D}$7 and the memory per rank is $z\in\mathbb R^{N\times N\times D}$8.

Outer Product Mean projects MSA features into new pair features through

$z\in\mathbb R^{N\times N\times D}$9

Fold-CP transposes $O(N^2)$0 for row alignment, holds $O(N^2)$1 in columns, and executes a $O(N^2)$2 Cannon ring where each rank shifts $O(N^2)$3 along rows, $O(N^2)$4 along columns, and accumulates partial outer-products. A final all-reduce completes the mean. The compute per rank is $O(N^2)$5 and the memory per rank is $O(N^2)$6.

Attention Pair Bias is implemented in two variants. In the “ring” variant, dense full-sequence attention biases $O(N^2)$7 by $O(N^2)$8; keys, values, bias, and mask are transposed and then ring-shifted over $O(N^2)$9 steps, merging through tiled softmax. In the “shardwise” variant, local atom-level windows are handled independently once upstream window batching supplies the correct key halos. The ring variant scales as $O(N^3)$0 per rank, while the shardwise variant is $O(N^3)$1 with zero further communication.

Window batching and distributed gather target atom-level attention. Atom-level attention groups spatially neighboring atoms into windows of size $O(N^3)$2 with default $O(N^3)$3 and halo $O(N^3)$4. Instead of materializing an $O(N^3)$5 sparse index matrix, Fold-CP exploits the block-Toeplitz structure of window assignments to generate a zero-memory virtual index and exchanges only small halo regions of approximately 48 atoms with adjacent shards. Token-to-atom mapping is performed via bounding-interval P2P gather, where each shard fetches only the token embeddings overlapping its atom range, with $O(N^3)$6 per window. The resulting space complexity is $O(N^3)$7 and the time complexity is $O(N^3)$8 per rank, with no global sync.

Kernel fusion is delegated to cuEquivariance and FlexAttention. Triangle Attention uses cuEquivariance’s BF16-enabled fused kernels, halving memory versus FP32, and Attention Pair Bias calls FlexAttention’s score_mod kernel to fuse pair-bias addition into attention. Each local block, whether ring or shardwise, is fed into these kernels.

4. Execution model, autograd strategy, and model integration

The implementation is organized around a device mesh, communication schedule, and explicit autograd behavior. The device mesh is a 2D CP grid, $O(N^3)$9, plus a DP axis. The CP mesh must be square to simplify symmetric row/column rings and transposes. This square-topology constraint is later identified as one of the key challenges that remains (Lin et al., 16 Mar 2026).

Communication is asynchronous and point-to-point. Ring patterns use $O(N^2/P)$0 point-to-point shifts, while Cannon-style rings shift along rows and columns simultaneously. Computation is overlapped with communication through double buffering. The system-level point is not only that communication occurs, but that it is orchestrated so that the global tensors remain distributed throughout execution.

A central implementation detail is the use of a custom torch.autograd.Function scaffold in each module. Each distributed module defines explicit forward and backward communication patterns, including ring or Cannon-style algorithms, thereby avoiding PyTorch’s implicit All-Gather in the backward pass and eliminating OOM failures at scale. This is one of the mechanisms by which Fold-CP preserves exact numerical equivalence to single-device baselines while remaining executable at long context.

Model integration is conservative with respect to architecture. Fold-CP wraps Boltz-1 and Boltz-2 architectures, replacing all pairwise, MSA, and atom attention modules with their CP-aware counterparts. The overall module flow remains unchanged, preserving functional equivalence. This makes Fold-CP a framework for distributed execution of AlphaFold-style co-folding models rather than a new co-folding model family.

5. Scaling behavior and benchmark results

For inference with Boltz-2 CP on NVIDIA B300 in BF16, Fold-CP scales maximum context linearly with $O(N^2/P)$1. The reported progression is approximately 4 k tokens at $O(N^2/P)$2 on 1 GPU, approximately 8 k tokens at $O(N^2/P)$3 on 4 GPUs, and approximately 32 k tokens at $O(N^2/P)$4 on 64 GPUs. The paper summarizes this as approximately 4,000 tokens per $O(N^2/P)$5, with per-sample runtimes inversely varying from approximately 20 s to approximately 90 s depending on $O(N^2/P)$6. Training scaling is reported as similar but at approximately 1,900 tokens per $O(N^2/P)$7 (Lin et al., 16 Mar 2026).

On NVIDIA H100 in FP32 with Boltz-1 CP, the maximum tokens reached are reported as 2,700 tokens in 13 min for a $O(N^2/P)$8 mesh, 4,000 tokens in 13 min for a $O(N^2/P)$9 mesh, 6,500 tokens in 37 min for a $P$0 mesh, and 11,500 tokens in 47 min for an $P$1 mesh. These values are compared to a 2,048-token cap on a single A100 80 GB, demonstrating more than a fivefold context extension with Fold-CP.

The framework therefore shifts the feasible regime for AlphaFold-style co-folding from a few thousand residues on a single GPU to assemblies exceeding 30,000 residues on 64 NVIDIA B300 GPUs. This suggests that the principal benefit is not raw latency minimization at fixed context, but access to a new context regime in which full global context can be retained for very large biomolecular systems.

6. Scientific use cases, limitations, and future directions

The scientific utility of Fold-CP is demonstrated through several developer use cases. Using Boltz-1 CP, Rezo Therapeutics scored structures for 93% of the CORUM database, up from less than 30% with single-GPU Boltz-1, enabling systematic structural validation of mammalian assemblies (Lin et al., 16 Mar 2026).

A second use case involves the PI4KA lipid kinase complex. A 3,605-residue heterotrimer plus an intrinsically disordered region was folded uncropped on 4 H100 80 GB GPUs in under five minutes for five samples, approximately 54 s per sample. Against a cryo-EM reference, PDB 9BAX, the model achieved TM-score 0.83, backbone lDDT 0.72, and DockQ 0.56, while revealing a novel EFR3A–TTC7A–PI4KA interface unreachable by cropping. The paper presents this example as evidence that preserving full global context can expose interfaces that are inaccessible to cropped workflows.

Other collaborations extend the reported range further. Proxima, through the NeoLink platform, and Earendil Labs integrated Fold-CP into their foundation models to fold complexes up to 150 k atoms, described as ribosome scale, and 24 k tokens, respectively, improving throughput by up to 27× in the ribosome case and eliminating the chunking bottleneck.

The paper also identifies several remaining challenges. These include relaxing the square-topology constraint, adapting to heterogeneous interconnects, and resolving long-context training artifacts such as chain “clumping.” It further notes that data scarcity of uncropped, large complexes will drive the creation of synthetic long-context datasets for next-generation CP-native model training. In the paper’s framing, sharding the dominant $P$2 memory and $P$3 compute kernels across a CP mesh is an essential step toward virtual-cell simulations, and Fold-CP is presented as providing a scalable pathway for modeling massive systems with full global context.