Ab Initio Accuracy Beyond the Cutoff: Distributed Message-Passing Equivariant Attention at Tens-of-Millions-of-Atom Scale

Ab Initio Accuracy Beyond the Cutoff: Distributed Message-Passing Equivariant Attention at Tens-of-Millions-of-Atom Scale Allex Kan, Ji Zhang, Hao Yin, Hao Li, Owen Fu, Xunxing Chen, Lin Huang, Zion Wang, Lucas Li, Jiajun Cheng, Arthur Jiang* , Junping Zhao* , and Jia Zhang* IQuest Research, UBio Team and HPC Team, China Molecular dynamics predicts how molecules move and interact by computing the force on every atom and integrating its motion.

arXiv:2611.06122v1 [cs.DC] 4 Nov 2026

Density functional theory supplies those forces from first principles, but its 𝑂 (𝑁 3 ) cost stops at thousands of atoms; empirical force fields reach millions and miss bond breaking and charge transfer. Message-passing equivariant Transformers regress the ab initio surface at linear cost, and their stacked layers carry information past any single cutoff, compounding four layers into a receptive field of up to 23 Å. That is where the accuracy frontier now sits. Getting there at scale is a communication problem: every layer exchanges boundary features across devices, the conservative force repeats that exchange in the backward pass, and the neighbor graph is re-derived from moving atoms, so tensor shapes keep changing under a runtime that would rather capture them once. We present Tessera, a distributed runtime that holds the workload’s shape fixed between rebuilds. It sorts atoms along a space-filling curve and cuts the sorted order into equal-count shards: every rank holds the same number of atoms, every tensor has one shape, and each rank’s peer set grows sublinearly, roughly as 𝐾 0.5 rather than 𝐾 over the device counts we measure. A conservative neighbor skin keeps that topology valid across many coordinate updates, so the forward and backward passes of an entire rebuild interval replay as one captured CUDA Graph. Fused kernels, a three-level neighbor search, and an overlapped All-to-All complete it. On 512 GPUs Tessera sustains 30.7M atoms at 89.7% weak-scaling efficiency (77–90% across three per-GPU loads), with per-step time rising 1.11–1.29× as machine and problem grow 16×. Graph reuse cuts full-step time 1.87×, and neighbor-search cost grows 1.7× over an 8× device increase where a Morton-order baseline grows 3.0×. On a live 10.24M-atom salt solution, energy drift stays at 0.42 𝜇eV/atom/ps over 10 ps of NVE, and structure matches a serial reference. A 10.06M-atom Norovirus capsid then holds a 100 ps NVT trajectory at ≈ 8 ps/day on 256 A100 GPUs, with no broken bonds, no anomalously close pairs, and only three compiled variants across 200,000 steps—no graph break, no fallback to eager execution.

Keywords: molecular dynamics; equivariant graph neural networks; distributed systems; CUDA Graphs; space-filling curves; neighbor search 1

INTRODUCTION

Some computations are hard to distribute for a reason that arithmetic alone does not explain: the data they need is not where they compute. Predicting the force on every atom is one of them, and it is where molecular simulation begins. Atoms are not independent. An atom’s force is set by its whole environment: electrostatic interactions decay slowly, and the effects that decide chemical behaviour—polarization, charge transfer, the making and breaking of bonds—are neither pairwise nor confined to a few neighbours. The energy is therefore a function of the entire configuration, and it admits no division into independent per-atom problems that keeps those effects. For any machine that computes it, two consequences follow. The work is global: each atom’s evaluation gathers information from around it, so pulling the system apart across 𝐾 machines does not pull the dependency apart with it. Data has to cross whatever boundary the machine imposes, and the further the physics reaches, the more of it has to move. The pattern is also not fixed. Positions change every step, so which atoms depend on which changes with them, and a plan built for one configuration cannot be assumed valid for the next. The systems where this matters most are the solvated, charged ones: salt solutions, electrolyte interfaces, ion channels, the binding of a ligand to a protein. There, polarization and charge transfer reach past any single cutoff, and a model that stops at one cutoff cannot represent them. Those systems also have to be large, because the interface, the pore and the binding site are properties of a whole assembly rather than of a few hundred atoms. * Corresponding authors: arthursjiang@iquestlab.com, zhaojp@gmail.com, jialrs.z@iquestlab.com. The authors above are listed in random order.

1

a Spatial decomposition

b Equivariant compute core irreps 𝑙 ≤ 𝑙 max

1D map

Rank 0

Level 1: inter-device AABB filter

Rank 2

Rank 3

equal-count cut ⇒ 𝑁 /𝐾 atoms per rank

Clebsch–Gordan ⊗

𝑙=0 (scalar)

Level 2: balanced query chunks

equiv. attention

×𝐿 layers

3D-near ⇒ index-near

Rank 1

exact 𝐸 (3) symmetry 𝑓 (𝑅x) = 𝐷 (𝑅) 𝑓 (x)

neighbor gather residual

𝑁 ∼ 107 atoms 1D index

c Multi-level neighbor search

𝑙=1 (vector)

Level 3: exact FAISS search

gating + norm

𝑙=2 (tensor)

+ residual → 𝑉𝑙′

mixed TF32/FP32

no true neighbor pruned unless 𝑘 saturates

𝑄 loc𝐶 cand →

static shapes ⇒ CUDA-Graph capture

≤ 𝑘 spatial peers, 𝑘 ≪ 𝐾 async launch comm compute

fused custom Triton kernels 23 Å receptive field per atom

> 10× lower neighbor cost at 512 GPUs

e Decentralized dynamics

d Multi-stream overlap remote data arrives

GPU 0

remote

remote exchange

extrapolates demonstrated

r → r′

owned + halo

owned + halo

forces F𝑖 = −𝜕𝐸/𝜕r𝑖 (autograd) overlap window local compute

f Scaling set by locality

GPU 1

r → r′

remote exchange local branch

𝑖 𝑄 𝑖 𝐶𝑖 , 𝐶𝑖 ≪ 𝐶 cand

cache: 9/10 steps skip rebuild

throughput

peers grow as roughly 𝐾 0.5 , not 𝐾

Í

ide

77–90% efficiency

al

512 GPUs

Velocity–Verlet integration

fixed-size All-Reduce: energy, 𝑇 , virial remote All-to-All hidden under local compute

global re-code every 100 steps

pre/post: same overlap fwd & bwd

each rank owns & integrates its atoms

launch order decouples scheduling from readiness

no per-step scatter/reduce

64

128

256

512

GPUs (log) no ceiling observed through 512 GPUs

Fig. 1. Overview of the algorithm–system co-design. a, Equal-count cut of a 1D space-filling-curve order (Hilbert shown), giving invariant tensor shapes and a peer set growing as roughly 𝐾 0.5 rather than 𝐾. b, Per-GPU E2Former V2 stack: sparse neighbor gather, Clebsch–Gordan tensor products, equivariant attention over irreps up to 𝑙 max , stacked to a 23 Å receptive field. c, Three-level neighbor filter—AABB, balanced chunking, exact FAISS search, conservative throughout—with cross-step cache. d, Communication and compute on separate streams, the overlapped All-to-All hidden under local work in both passes. e, Per-rank on-GPU integration (Velocity–Verlet, autograd forces F𝑖 = −𝜕𝐸/𝜕r𝑖 ), synchronizing a fixed handful of scalars, re-partitioning every 100 steps. f, Measured weak scaling against the ideal linear reference (Fig. 12).

Empirical force fields have long been the practical answer. Their force law is cheap: fixed partial charges, bonded springs and torsions, a fitted van der Waals term. That cheapness carries the method far—a single GPU holds millions of atoms, and a distributed run reaches a hundred million [1, 41]—which is why classical MD became the default rather than a compromise. What the force law cannot do is change. Charges do not polarize, bonds do not break, and charge transfer never happens, so the model is wrong wherever the chemistry is interesting. Nor is the size of that error known in advance for a new system, which matters because the quantities that decide a biological outcome—a conformational free-energy difference, an ion-selectivity ratio—are often smaller than the error itself [14]. Ab initio methods, density functional theory among them [21], supply the missing physics at a cost that grows as 𝑂 (𝑁 3 ): in practice they stop at hundreds to a few thousand atoms, too small to hold a solvated interface and too slow to follow one for long. Machine-learning force fields are the route that has worked [5, 35, 43, 4, 3]. They regress the ab initio energy surface and evaluate it at a cost linear in atom count, and the strictly local ones have taken that to a hundred million atoms: Deep Potential MD’s Gordon-Bell-winning Summit run [17, 50] confines each atom’s energy to its own cutoff ball, 6 and 8 Å for water and copper, and Allegro uses cutoffs as short as 4 Å in some benchmarks, forgoing message passing by design [25]. The choice buys scale at the cost of the property the physics needs: information never travels the molecular graph, so a learned feature never propagates past a single cutoff [11, 49, 20]. Message-passing equivariant Transformers restore the propagation [24, 39, 46]. Stacked finite-radius attention compounds radii, so a few layers reach 20 Å and beyond, 23 Å in the stack used here, and this class now sets the accuracy frontier in molecular simulation. It is also the class the argument above bites hardest on: the further a model reaches into the system, the more of it each evaluation needs, and the more communication the machine must carry. To our knowledge, no published system has taken the class past tens of GPUs; the closest, SevenNet [29], reports 32. 2

This paper asks a single question. How can we propagate graph attention across a live molecular graph at tens-ofmillions-of-atom scale without letting communication and variable topologies dominate the step? Three difficulties stand in the way, one per layer of the stack. Acceleration-hostile equivariant operators. Accuracy hinges on 𝐸 (3)-equivariant Transformers, which keep forces rotating with the system. Their components—Clebsch–Gordan tensor products, radial-basis encodings, equivariant gating—match neither the dense kernels FlashAttention [9] targets nor the BF16 precision it assumes. Single-GPU efficiency therefore has to be rebuilt from custom kernels. A decomposition that carries its own cost. Ten million atoms exceed any single node, so the graph must be cut, yet cutting it does not make the per-layer exchange small. The graph is also the physical configuration itself. Nobody partitions it offline, once, the way a training pipeline partitions a dataset. Topology, neighbor lists, and plan are therefore re-derived from live coordinates mid-trajectory. Recapture on the critical path. Every re-derivation changes edge counts, which is what forbids replay. Conservative forces make the recapture expensive, because the full forward–backward execution is the step. There is no smaller unit to fall back on. Adjacent systems each address a subset. LLM training scales to thousands of GPUs [27], but over dense, statically shaped collectives whose routing is fixed for the run. Classical domain decomposition [32, 41] assumes shape-flexible updates that a CUDA Graph cannot tolerate. Distributed GNN training [51, 18] partitions a fixed graph once, where ours is re-derived from moving coordinates every step. The equivariant models that reach biomolecular size do so by forgoing propagation [26]. Universal MLFFs such as UMA are released for single-GPU inference [47], with no distributed runtime reported. No single optimization suffices. The three difficulties sit at different layers of the stack, and a fix at one pays off only if the others cooperate. We therefore treat the model and the distributed runtime as one object of design, an algorithm–system co-design resting on a common foundation: how atoms are placed on devices. Standard rules force a choice there. Equal-volume cells bound each rank to 26 neighbouring cells, 27 with its own, but counts drift with density, and every tensor shape with them. Equal atom shares give exact counts but ignore geometry, so they destroy locality and leave per-edge work unbalanced. One construction avoids the choice. Quantize atomic coordinates onto a space-filling curve [37], sort, and cut the sorted list into equal-count contiguous segments, one per rank. The cut fixes per-rank counts, and with them every node-indexed tensor shape; edge-indexed tensors and the exchange plan are pinned separately by the rebuild epoch. The curve then confines each segment’s neighbors to a set of adjacent subdomains that grows far more slowly than the machine. At 25k atoms per GPU the heaviest rank requests candidates from 6, 11, 14, and 19 source ranks at 64, 128, 256, and 512 GPUs, an exponent near 0.5 rather than 1 (Section 3.2). The confinement is tight for Hilbert, where Z-order shards fragment. Two further leverage points come from the physics. At 300 K, hydrogen moves about 0.1 Å between one rebuild and the next, well inside the 0.25 Å displacement trigger that the 0.5 Å skin allows, and about 1 Å over the 50 fs between re-partitions, small against the ∼ 67 Å subdomain. A molecular graph’s topology therefore outlives its geometry on two clocks: nine steps in ten reuse the edge set and the communication plan, and the partition is re-cut only every hundred steps. Neither reuse stands alone. A rebuild turns over the edge shapes, so every rank must rebuild on the same step. An all_reduce(MAX) over local triggers agrees that step, which keeps the exchange plans mutually consistent and the captured graph valid. Partition, neighbor search, and graph capture hold or fail together. Propagation is not a design preference here: the distributed cost follows from the model class, not from one implementation choice (Section 2.1). We call this runtime Tessera, after the mosaic tile.1 Figure 1 traces one step through it, from the equal-count partition to the decentralized integrator. The four contributions below build it from the bottom up, kernels to trajectory. High-performance equivariant kernels. A unified MCH layout, dual CSR/CSC indexing, and deep operator fusion in custom Triton kernels, together with torch.compile, reduce step latency 2.6× and peak memory 4.4× on 1 A tessera is a small block, cut to the same size as its neighbours and laid beside them so that many make a whole; it has fixed faces and a definite boundary; and in Rome it was also a token a messenger carried, telling the next person what to do. The design has one of each: the equal-count shards are the blocks, the fixed tensor shapes are their faces, and the epoch every rank agrees on before a rebuild is the token.

3

a 10k-atom single-GPU ablation; adding TF32 and recomputation brings these to 3.2× and 9.2×, at the accuracy cost Section 4 characterizes. Live-graph decomposition with topology reuse. An equal-count space-filling-curve cut fixes every rank’s atom count and keeps its shard local, supplying static shapes and a peer set that grows as roughly 𝐾 0.5 rather than 𝐾. A Í conservative AABB/chunking/FAISS [19] hierarchy then rebuilds the graph from live coordinates at 𝑖 𝑄𝑖 𝐶𝑖 rather than 𝑄 loc𝐶 cand cost. Separating reusable topology from per-step geometry cuts neighbor cost by an order of magnitude at 512 GPUs, and its growth over an 8× device increase falls to 1.7×, against 3.0× for a Morton-order all-gather baseline. Sparse, overlapped, decentralized communication. Plan reuse, sender-side SKV projection, and a pre-registered buffer pool strip non-payload overhead. Paired pre/post autograd operators keep the exchange hidden in both passes, worth 84.7% against 72.9% efficiency at 512 GPUs. An on-GPU decentralized engine replaces the rank-0 scatter/reduce loop with a constant-payload barrier. Demonstration and validation at scale. Tessera sustains the full force-evaluation step (forward, exchange, autograd backward) at 30.7M atoms on 512 GPUs at 89.7% weak-scaling efficiency. Per-step time rises 1.11–1.29× across a 16× growth in machine and problem. On a live 10.24M-atom trajectory it holds energy drift to 0.42 𝜇eV/atom/ps and matches the serial reference’s structure. A matched pair against a 10.06M-atom Norovirus capsid (protein-dense against water) then locates the cost of an atom-count cut: 13–33% more per step across trajectory phase, 14.5% in the reported window, in edges per atom. To our knowledge, no prior system has run a message-passing equivariant Transformer at tens of millions of atoms on hundreds of GPUs. The forces are taken as the analytic gradient, over a neighbor graph re-derived from live coordinates, with information propagating along that graph rather than stopping at a single cutoff. The exchange each GPU performs is set by local geometry and grows only slowly with device count, so the exchange cost extends in principle toward billion-atom systems, though the 𝐾- and 𝑁 -dependent terms of Section 6 still grow. Roadmap. Section 2 supplies the background the rest assumes. It covers equivariant message passing and the receptive-field arithmetic, the conservative-force path and why it doubles the exchange, why static graph execution is fragile, and the three requirements a partition must satisfy. Section 3 develops the co-design against them. Section 4 evaluates it, Section 5 positions it, and Section 6 states what the measurements do and do not establish. The derivations the argument turns on—the overlap bound, the local-neighbor fraction, the payload and plan-reuse arithmetic, and the cache exactness invariant—are developed in place, in the sections that use them. 2

BACKGROUND AND PROBLEM FORMULATION

This section supplies the two bodies of background the paper assumes—equivariant message passing and static graph execution—and states, as requirements, what a partition must deliver for the second to survive the first. The receptive field and the force path are derived here because the design choices of Section 3 are only legible against them. 2.1

Equivariant Message Passing

A molecular configuration is a set of 𝑁 atoms with positions r𝑖 ∈ R3 and species 𝑧𝑖 , under periodic boundary conditions. An equivariant message-passing potential represents each atom not by a flat vector but by a collection of irreducible representations of the rotation group: for each degree 𝑙 up to a maximum 𝑙 max , a block of 2𝑙 + 1 components that mix among themselves under a rotation as the spherical harmonics 𝑌𝑙𝑚 do. A block carries 𝐶 channels, so the feature at ) atom 𝑖 after layer 𝑠 is a family h𝑠,(𝑙 ∈ R (2𝑙+1) ×𝐶 for 𝑙 = 0, . . . , 𝑙 max . The deployed model uses 𝑙 max = 2 and 𝐶 = 256 𝑖

(Appendix A), so one atom’s feature holds (𝑙 max + 1) 2𝐶 = 9 × 256 = 2,304 scalars—the count that reappears as the equivariant Value width in Section 3.4. One layer performs three operations. First, edge construction: an edge (𝑖, 𝑗) exists when the minimum-image displacement r𝑖 𝑗 = r 𝑗 − r𝑖 satisfies |r𝑖 𝑗 | ≤ 𝑟𝑐 , described by a radial basis encoding of the distance (256 Gaussian-smear functions here) and the spherical harmonics of its direction. Second, Clebsch–Gordan tensor products [40, 10]: features of degrees

4

𝑙 1 and 𝑙 2 combine into a feature of degree 𝐿out by 𝑙1 ∑︁
𝐿 x𝑙1 ⊗ y𝑙2 𝑀out =

𝑙2 ∑︁

𝐺𝑙𝐿𝑀 𝑥 𝑙 1 𝑦𝑙 2 , 1𝑚 1 𝑙 2𝑚 2 𝑚 1 𝑚 2

𝐿out = |𝑙 1 − 𝑙 2 |, . . . , 𝑙 1 + 𝑙 2,

(1)

𝑚 1 =−𝑙 1 𝑚 2 =−𝑙 2

with 𝐺𝑙𝐿𝑀 the Clebsch–Gordan coefficients. The product is what makes the representation equivariant: rotating 1𝑚 1 𝑙 2𝑚 2

both inputs rotates the output by the same rotation, so a network built only from such products, scalar contractions, and gating maps a rotated configuration to a rotated prediction. It is also not a dense matrix multiply: the index space is a direct sum over degree triples, and the number of surviving paths grows with 𝑙 max , which is why the kernels of Section 3.1 are custom. Third, attention [34, 24]: for each edge the layer forms a scalar score 𝑠𝑖 𝑗 as an invariant contraction of a query feature on 𝑖 with a key feature on 𝑗; with 𝐻 = 128 heads of scalar width 𝑑𝑘 = 8 the score is a per-head scalar, and a softmax over 𝑖’s neighbors turns the scores into weights 𝛼𝑖 𝑗 . The message is an equivariant Value—the degree-𝑙 block of atom 𝑗, Í possibly tensor-produced with the edge’s spherical harmonics—aggregated as 𝑗 𝛼𝑖 𝑗 V 𝑗 and combined with the residual stream. The score and the message are per-edge objects, so the attention cost scales with the number of edges rather than the number of atoms, while the per-atom projections and gating still scale with 𝑁 . That is what Section 4.6 confronts when an atom-count cut meets a protein-dense region. Receptive field. Stacking layers compounds their cutoffs. Let layer 𝑠 have cutoff 𝑟 (𝑠 ) . After one layer, h𝑖1 depends on Í ′ atoms within 𝑟 (1) of 𝑖. If h𝑠𝑗 depends on atoms within 𝑠 ′ ≤𝑠 𝑟 (𝑠 ) of 𝑗, then h𝑠+1 depends on every atom within 𝑟 (𝑠+1) of 𝑖 Í ′) (𝑠 some such 𝑗, hence on atoms within 𝑠 ′ ≤𝑠+1 𝑟 of 𝑖: the bound is additive by induction, because each layer spends its cutoff to reach one more hop. Our stack has three short-range layers at 5 Å and one long-range layer at 8 Å, so the per-atom receptive field is bounded by 3 × 5 Å + 1 × 8 Å = 23 Å,

(2)

the radius we quote throughout as the nominal receptive field. A strictly local potential has a receptive field equal to its single cutoff—4–8 Å in the systems Section 5 compares—so its per-atom computation touches only its own cutoff ball. Here the energy of one atom can depend on atoms up to 23 Å away, through intermediate atoms whose features were themselves computed from remote neighbors. 2.2

Conservative Forces

The potential is trained to reproduce an ab initio energy surface, and the dynamics must conserve energy. A network that emitted forces as a second output head would not do: the force field would not be the gradient of the energy it reports, and a trajectory integrated with it would drift for reasons unrelated to the physics. We therefore take the force on atom 𝑖 as the analytic gradient F𝑖 = −

𝜕𝐸 , 𝜕r𝑖

𝐸=

∑︁

𝐸𝑖 {r 𝑗 } 𝑗 ∈ R (𝑖 ) ,

(3)

𝑖

where R (𝑖) is the receptive field of Section 2.1—the R distinguishing it from the one-hop neighbour set N (𝑖) of a single layer, used in Section 3.1. Two consequences follow. First, a step costs a forward and a backward pass. The forward pass evaluates 𝐸; the forces come from reverse-mode differentiation through the same graph, and because every atom’s energy depends on its neighbors’ features, the gradient with respect to one coordinate collects terms from every edge on which that coordinate appears—as a direct displacement and through the features of intermediate atoms. There is no cheaper local formula: the whole propagated graph participates. The backward pass is a first-class half of the step, comparable in cost to the forward and, as Section 3.4 shows, separately schedulable. Second, the backward pass must exchange the same boundary features again, in reverse. In the forward pass a remote atom 𝑗’s Key, source embedding, and Value travel from 𝑗’s owner to 𝑖’s owner. Differentiating with respect to r𝑖 and r 𝑗 requires the upstream gradient of the score, 𝑔𝑖 𝑗 = 𝜕𝐸/𝜕𝑠𝑖 𝑗 , to travel back along the same edge, and requires the Value and edge geometry on 𝑗’s owner to accumulate 𝑗’s contribution. A forward-only exchange leaves the backward exchange

5

behind the forward compute, a serial tail no stream juggling repairs—which is why Section 3.4 splits the interface into a pre operator that communicates in the forward pass and a post operator that communicates in the backward. Both consequences constrain correctness. Because the force is the exact gradient of the energy actually evaluated, the backward pass must use the neighbor set the forward pass used, with the same active pairs; an approximate query that drops an in-cutoff pair changes the force, not merely the speed. The cache of Section 3.3 is built to satisfy this: it reuses topology only under a displacement bound that guarantees the reused superset still contains every pair that can be active, and it re-applies the exact cutoff test on every hit. 2.3

Why Static Graph Execution Is Fragile

The per-step work is thousands of small, irregular kernels—gathers, scatters, gated multiplications, short tensor-product contractions—rather than a few large matrix multiplies. At that granularity the CPU-side cost of dispatching a kernel is comparable to its GPU-side cost, so the runtime is worth capturing. Two mechanisms do the capturing, and they fail differently. torch.compile and guards. The compiler traces the model and specializes it [2], emitting guards: predicates on shapes, strides, and dtypes that must hold for the compiled code to be valid. Every call checks them, and a guard that fails forces a re-trace and often a recompilation; the work is not wrong, but it is paid on the step where the shape moved. Compiling for dynamic shapes—symbolic extents rather than constants—removes the re-trace for extents that vary within the compiled range. CUDA Graph capture and replay [28]. A captured graph is a different object. Capture records the sequence of kernel launches, their arguments, and their grid and block geometry into a DAG, together with the device addresses the kernels read and write; replay submits that DAG with a single launch, bypassing the dispatcher, the allocator, and the Python layer. That is where the saving comes from, and it is why the mechanism is rigid: replay consults no per-launch shape, because no per-launch host code is left. The addresses are fixed, so the allocator must hand out stable buffers; the launch geometry is fixed, so every tensor extent that determines it must be fixed. The two compose in one direction only. Compiling for dynamic shapes removes the re-trace, but not the capture: the captured graph still replays the allocations and launch parameters recorded at capture time, so a step whose tensors change extent cannot be replayed and must be captured again. Section 5 states this once; here is the consequence. The neighbor graph is re-derived from live coordinates, and its edge count moves whenever atoms cross the cutoff; if the node count also moved, no captured graph would survive more than a step, and the runtime would pay capture—not merely dispatch—every step. A design that wants replay must pin shapes by construction. 2.4

Problem Formulation

Let 𝐾 be the number of ranks (GPUs), 𝑁 the number of atoms, and r = (r1, . . . , r𝑁 ) the positions in a periodic cell. The stack carries two cutoffs—5 Å for its three short-range layers and 8 Å for its long-range one—and at step 𝑡 the live edge set is  E𝑡 = (𝑖, 𝑗) : |r𝑖 𝑗 | ≤ 𝑟𝑐(𝑠 ) for the layer 𝑠 that owns the edge ,

(4)

with 𝑟𝑐 standing for the cutoff of whichever layer is being built. The cache of Section 3.3 does not store E𝑡 : it stores a rebuild-time superset built at 𝑟𝑐 + 𝛿, and a hit re-applies the test in (4) to recover the active pairs. A partition is a map 𝜋𝑡 : {1, . . . , 𝑁 } → {1, . . . , 𝐾 } assigning each atom to a rank; write 𝑃𝑝 for the set of ranks from which rank 𝑝 requests candidates. A step advances the positions, rebuilds E𝑡 as needed, and evaluates the force by the forward–backward path of Section 2.2. The runtime can replay that step only if the partition delivers three properties together. (R1) Exact per-rank atom count. Every rank owns ⌊𝑁 /𝐾⌋ or ⌈𝑁 /𝐾⌉ atoms, so the shape of every node-indexed tensor is fixed. The cut is a deterministic function of the coordinates at each re-cut, so all ranks derive the same ownership map without a consensus round, and the map is frozen between re-cuts. (R2) A peer set that grows sublinearly in 𝐾. The number of ranks a rank must talk to should be set mainly by local density and cutoff geometry, and grow only sublinearly with machine size. Neither classical rule delivers that

6

alone. An equal-volume grid bounds its peers at 27—the 26 face-, edge-, and corner-adjacent cells plus the cell itself—but only by letting counts float. An equal-count cut with no geometry keeps the counts exact but degenerates toward ≈ 𝐾 peers. The design requires both properties at once, and Section 3.2 shows that one construction supplies them. There the measured peer degree is 6, 11, 14, and 19 source ranks at 64, 128, 256, and 512 GPUs (25k atoms/GPU, Hilbert with rank-box): 3.2× over an 8× increase in devices, an exponent near 0.5 rather than 1. This is sublinear growth, not a constant. The peer count still rises with the machine, and Section 6 lists it among the terms that do. (R3) Tensor shapes constant within a cache epoch. Between rebuild events every tensor in both passes keeps a fixed shape, so one capture serves the whole interval. Node tensors are pinned by (R1); edge-indexed tensors follow the cached candidate superset of Section 3.3, and both they and the communication plan change only on a rebuild, which is therefore an epoch boundary and must be unanimous: if one rank rebuilds and another does not, the exchange plans disagree and the captured graph is invalid. These are requirements, not results; Section 3 shows that one placement mechanism supplies all three. Notation. 𝑁 atoms, 𝐾 ranks; r𝑖 , v𝑖 , F𝑖 position, velocity, and force of atom 𝑖; 𝑚 atomic mass; Δ𝑡 timestep; 𝑟𝑐 cutoff, with 5 Å and 8 Å the two values in our stack; 𝛿 displacement skin (0.5 Å); 𝐼 rebuild interval (10 steps); E𝑡 the active edge set at step 𝑡 and Π the communication plan, with the rebuild-time candidate superset and the plan cached between rebuilds; 𝜎 the sort permutation produced by a re-partition (Algorithm 1); 𝑄 loc and 𝐶 cand the local query and candidate counts of the neighbor search (Section 3.3), with 𝑄𝑖 and 𝐶𝑖 the per-chunk values; 𝐶 alone is the channel count of Section 2.1; 𝐻 attention heads and 𝑑𝑘 scalar head width; 𝑙 max maximum degree; 𝑅𝑟 the number of Transformer blocks sharing cutoff 𝑟 ; 𝑇𝑐 , 𝑇ℓ , 𝑇𝑟 the exchange, local-branch, and remote-branch times of the overlap analysis; 𝜌 number density. 3

ALGORITHM-SYSTEM CO-DESIGN

We develop the co-design in the order its dependencies run. Section 3.1 restores single-GPU efficiency for operators the mainstream stack cannot serve. Everything after it rests on how atoms are placed on devices: the equal-count space-filling-curve cut of Section 3.2 supplies invariant tensor shapes and a peer set that grows sublinearly with device count, and on those two guarantees rest the multi-level neighbor search (Section 3.3), the communication path that shrinks, fuses, and overlaps the remaining exchange (Section 3.4), and the decentralized integrator that removes the per-step 𝑂 (𝑁 ) rank-0 gather/scatter path, leaving that loop’s collectives at constant payload (Section 3.5). Two threads recur: a graph’s topology outlives its geometry, letting most steps reuse both the graph and the partition; and a rebuild must be unanimous to preserve the shape invariant of Section 2.4, so these mechanisms hold or fail together rather than composing independently. 3.1

Single-GPU Optimization

E2Former V2’s sparse equivariant attention, Clebsch–Gordan tensor products, and radial-basis encodings neither reduce to the dense matmuls FlashAttention-style kernels target nor tolerate the reduced-precision fast paths those kernels assume, so both throughput and HBM budget must be rebuilt from custom kernels upward. Triton-based Kernel Optimization. Custom Triton [42] kernels, registered as opaque torch.compile [2] operators so hand-written and compiler fusion compose cleanly, are optimized along four axes; two are standard practice— hardware-aware tile-size autotuning and deep operator fusion—and are deferred to Appendix B. The other two follow from the two-pass execution of Section 2.2. Dual sparse indexing. The forward pass and the 𝑑k backward pass traverse the same sparse structure in transposed orders. Forward reads query-grouped, accumulating over each query’s neighbors, o𝑖 =

∑︁

𝛼𝑖 𝑗 V 𝑗 ,

(5)

𝑗 ∈ N (𝑖 )

while the key gradient 𝑑k 𝑗 =

𝑖 𝑔𝑖 𝑗 q𝑖 —with 𝑔𝑖 𝑗 = 𝜕𝐸/𝜕𝑠𝑖 𝑗 carrying the upstream gradient through the softmax Jacobian—

Í

is a sparse matrix–vector product that reads neighbor-grouped, accumulating over each neighbor’s queries. The two traversals want opposite index orders: a row-major view is coalesced for one and strided for the other. A single row7

major format therefore leaves the backward pass on strided loads, and since the backward pass runs every step, the penalty would be paid every step. We keep a hybrid CSR/CSC representation instead—CSR for the forward, CSC for 𝑑k—so each pass reads its natural order. Both are built once per neighbor-list rebuild and cached; because the cache lives as long as the graph does (Section 3.3), that construction is charged to rebuild steps only, and a cache hit pays nothing for it. A layout that survives both passes. Equivariant attention indexes atoms, channels, and heads in different orders at different stages: attention iterates over edges and heads, the tensor product contracts over channels and degree, and the residual stream is atom-major. A naive implementation pays repeated permute traffic between these orders, and the traffic lands in exactly the small irregular kernels the GPU is worst at. We adopt a unified MCH (𝑀 × 𝐶 × 𝐻 : rows × channels × heads) layout end-to-end and block our kernels over it: a C-blocked kernel walks a 3D grid Í (𝑀, ⌈𝐶/𝐵𝐶⌉, ⌈𝐻 /𝐵𝐻 ⌉), loads contiguous [𝐵𝐶, 𝐵𝐻 ] tiles, and accumulates out = 𝛼 · 𝑉 and its gradient entirely in registers. Because each tile is reused across the rows of a CSR row, both passes are held to a single HBM round-trip per tensor; the same layout that removes the transposes also removes the intermediate materialization that would otherwise dominate the peak allocation. Compiler, Precision, and Memory Budget. On top, torch.compile fuses pointwise ops and cuts launch overhead, valuable since runtime is dominated by many small gather, scatter, and gating kernels rather than large matmuls. Dense projections and tensor products additionally run in TF32, trading latency for non-trivial energy/force drift. Two further mechanisms buy HBM headroom instead: a compiler-driven activation-recomputation budget lets torch.compile rather than the author place checkpoints, roughly halving the activation footprint at moderate runtime cost; a two-tier allocator policy then keeps fragmentation from eroding that headroom over long runs (Appendix B). The TF32 trade is not free, which is why the trajectory validation of Section 4.6, where drift is the measurement, runs with it off. 3.2

Spatial Partitioning via Space-Filling Curves

A partition over hundreds of GPUs must meet the three goals Section 2.4 stated as requirements: (i) exact atom-count balance, (ii) local, sparse communication, and (iii) static tensor shapes the per-rank computation can be captured under. Neither classical choice reaches all three. An equal-count cut ignoring geometry balances perfectly but lets a rank’s neighbors live anywhere, degenerating toward all-to-all (≈ 𝐾 peers); an equal-volume grid bounds peers at 27—the 26 adjacent cells plus its own—but lets counts follow local density and pins no shape without padding to fixed capacity. The count drift is structural, not a nuisance: under uniform density a cell’s population is Binomial(𝑁 , 1/𝐾) ≈ Poisson(𝑁 /𝐾), √︁ so the extreme-value spread across 𝐾 cells is 𝑛 max − 𝑛 min ≈ 2 2𝜆 log 𝐾 with 𝜆 = 𝑁 /𝐾—at 𝑁 = 1.024 × 107 , 𝐾 = 512 this is ≈ 19,500–20,500 atoms/rank—an estimate that assumes uniform density and therefore degrades further on an inhomogeneous system. One mechanism reaches all three, with no capacity to guess: quantize the 3D coordinates onto an integer lattice, map them onto a 1D space-filling curve (SFC) that keeps spatially-close atoms close in the linear order, radix-sort by SFC code, and cut the sorted list into equal-count contiguous segments, one per rank. The cut gives exact balance: ⌊𝑁 /𝐾⌋ or ⌈𝑁 /𝐾⌉ atoms per rank, independent of density. Because work is per-edge rather than per-atom, exact counts balance the load only up to the local edge-density variation Section 4.6 measures. The curve then confines each segment’s neighbors to a small set of adjacent subdomains. That set grows sublinearly with the machine rather than staying fixed: at 25k atoms/GPU the heaviest rank requests from 6, 11, 14, and 19 source ranks at 64, 128, 256, and 512 GPUs (Section 4.3; counted per Appendix C), an 8× increase in devices costing 3.2× in peers, which is an exponent near 0.5 rather than 1. This is the sparse peer set Section 3.4 relies on, and it is small enough to keep the exchange off the critical path at our scale. Goal (i) is the means to goal (iii): only because counts are fixed exactly can shapes be pinned. Static shapes make the step capturable. Section 2.3 explains why capture fails when shapes move; here is how the cut removes the cause. Once per-device counts move, torch.compile re-triggers guard failures and any captured CUDA Graph is discarded, leaving CPU-side scheduling latency to idle the Streaming Multiprocessors (SMs). The equal-count cut removes that cause for the node tensors, and is also deterministic—a pure function of the quantized coordinates with fixed offsets, so identical coordinates always yield the identical ownership map. Two timescales then follow. Per-rank atom counts—and with them every node-indexed tensor—change only when ownership is re-cut; 8

(a) Morton (Z-order)

r1 AABB of rank 2

(b) Hilbert

r2 r3 foreign cell

r4 long jump

Fig. 2. One partitioning rule, two curves: equal atom counts either way, different shard geometry. An 8 × 8 grid cut into four shards of equal atom count (cell counts differ, since density does). (a) Morton’s power-of-two jumps attach a remote sliver to a shard (orange arrow): rank 2 fragments into two pieces whose bounding box—what the Level-1 filter tests—encloses 23 cells it does not own. (b) Hilbert’s consecutive indices stay adjacent, keeping shards connected and boxes tight (mean AABB/owned-cell ratio: 1.42× Morton vs. 1.06× Hilbert). Shown in 2D for legibility (the real partition is 3D); consequences at scale in Fig. 11.

edge-indexed tensors and the communication plan follow the neighbor graph, rebuilt on the schedule of Section 3.3 and constant between rebuilds. Within such a cache epoch every tensor in both passes has a fixed shape, so we trace the computation once and capture it as a static CUDA Graph, replayed on each intervening step rather than re-dispatched kernel by kernel. Section 4.6 reports three compiled variants, no graph break, and no fallback over a 200,000-step run. Which curve: locality decides the exchange cost. Balance, determinism, and static shapes belong to the cut, not the curve, so the choice between the two we implement turns entirely on goal (ii). Morton (Z-order) interleaves the bits of the quantized 𝑥, 𝑦, 𝑧—cheap, but it jumps across the domain at every power-of-two boundary, so a contiguous code interval can span several disconnected fragments. Hilbert’s recursive rotation keeps consecutive indices adjacent, mapping an interval to one connected, near-cubic region. Load balance cannot see the difference; the communication plan can. Figure 2 shows why: the Level-1 filter summarizes each rank by a bounding box, and a fragmented Morton shard’s box also covers the gaps between its fragments—enclosing 23 of 64 cells it does not own—so the rank requests from far more peers than geometry requires, and the count keeps climbing with the machine where Hilbert’s stays sublinear: over 64–512 GPUs Morton rises from 5 to 27 and Hilbert from 6 to 19 at 25k atoms/GPU, both measured (Fig. 11c). Neither series saturates, and the slopes differ. Neighbor-construction latency then grows 2.8–2.9× under Morton from 64 to 512 GPUs at 25k atoms/GPU, against 1.6–1.7× under Hilbert (Section 4.3). We therefore adopt Hilbert equal-count partitioning as the default, keeping Morton as the ablation baseline; its higher encoding cost is paid once per repartition, off the per-step path. One caveat: our two implementations normalize the quantization axes differently, so on strongly anisotropic cells part of the gap reflects resolution, not curve topology—though it persists at 64 GPUs, where a cubic cell makes both coincide (Appendix C).

9

Inter-device AABB filtering device 𝑗 D1 device 𝑖

D0

D2

b

𝑄2

SFC-sparse active peers

Chunk-local range search candidates 𝐶𝑖

exact ≤ 𝑟𝑐

𝑄4

sparse peer set

D3

c

Recursive query chunking

per-chunk candidates

chunk 𝑄𝑖

a

𝑟𝑐 𝑄3

𝑄1

tight per-chunk AABB: 𝐶 →𝐶𝑖

exact, at

Í

𝑖 𝑄 𝑖 𝐶𝑖 ≪ 𝑄 loc𝐶 cand

Fig. 3. Multi-level neighbor search. (a) Per-device AABBs and a 𝐾 ×𝐾 overlap test bound each device’s active peers—the sparse peer set the communication layer reuses. (b) Recursive bisection into compact chunks 𝑄𝑖 , each screened by its own tight box, cuts Í candidates from 𝐶 to 𝐶𝑖 ≪ 𝐶 (detail: Fig. 4). (c) The exact cutoff search then scores only block-diagonal 𝑄𝑖 –𝐶𝑖 pairs, at 𝑖 𝑄𝑖 𝐶𝑖 rather than 𝑄 loc𝐶 cand .

Maintaining the partition under MD dynamics. Since positions evolve, a partition fixed at 𝑡 = 0 would lose both balance and locality, so we re-establish it every 100 steps. This re-coding is the only event that changes which atoms a rank owns: counts stay fixed, so node shapes never move, only the identities behind them. It also invalidates the neighbor graph built against the old ownership, whose edge shapes turn over on the faster clock of Section 3.3. The two cadences compose rather than compete: a repartition is a rebuild that also reassigns atoms, so its cost is the partition-and-broadcast plus one graph re-capture, measured at 407–418 ms in Section 4 and amortized over 100 steps. It is also the one event that invalidates the topology cache outright, since the cached edge sets, chunk assignments, and plans are all indexed by ownership; Section 3.3 states the invariant, and Algorithm 1 makes the invalidation unconditional rather than inferred. 3.3

Multi-Level Neighbor Search

Neighbor construction is the stage that scales worst: each device owns a subset of atoms but must find their neighbors among nearby remote regions. Figure 3 shows our answer—three filters pruning the candidate set progressively, across devices, then within a device, then exactly. Exactness survives the pruning: every filter is conservative, discarding only atoms provably beyond the cutoff, and Level 3 is an exact brute-force cutoff-radius search via FAISS [19], not an approximate nearest-neighbor query. Nor do we impose a fixed-𝑘 truncation of our own: where the backend’s 𝑘-selection is unsaturated (Section 6), the edge set is the complete radius graph, at a fraction of the cost of an all-pairs search. At Level 1, each of the 𝐾 devices summarizes its atoms by a coarse axis-aligned bounding box (AABB), and a 𝐾 × 𝐾 test marks (𝑖, 𝑗) active whenever atoms on device 𝑗 may fall within the cutoff of device 𝑖’s queries. Nominally 𝑂 (𝐾 2 ), it runs on coarse boxes from one small metadata all-gather, and this is where the partition pays off: SFC locality confines each device’s neighbors to a set of adjacent subdomains that grows sublinearly, keeping every row sparse as 𝐾 grows. Candidates then move as a selective-subgroup variable-length All-to-All, whose sparse peer set the communication plan of Section 3.4 reuses. Under periodic boundaries a shard wraps the cell, so each device publishes one AABB per occupied image: at most 27 including the identity image, with the corresponding counts, 8 + 27 × 8 doubles and under 2 KiB per rank. Because that descriptor is small and fixed in size, one metadata all-gather lets every rank reconstruct the complete rank×rank×image request matrix locally, with no second exchange of request masks. Level-1 granularity. The three realizations of that filter differ in what they publish, and the comparison is instructive because the scheme that prunes hardest loses. All-gather performs no geometric filtering: every device broadcasts its full PBC-expanded candidate set to every other device. It builds no plan and issues one regular collective, which makes it the correctness/performance baseline and competitive at small 𝐾, but each device’s candidate population then grows with the global atom count—the memory and scalability wall the design exists to remove. Rank-box (ours) publishes one query AABB per device plus up to 27 periodic-image candidate AABBs per source: fixed small metadata, no request-mask exchange. Spatial-bin replaces the per-device boxes with a regular grid sharing a global origin and publishes the AABB and count of every occupied query and candidate bin, capped at 8192 bins per side and auto-coarsened if the cap is hit; surviving bins double as query chunks for Level 3. Spatial-bin therefore prunes the 10

b

D0

D1

D2

D3

AABB candidate filtering

c

FAISS range search candidates 𝐶

𝑟𝑐

queries owned by rank 𝑘

one AABB stretched over outliers

all 𝑄 × 𝐶 pairs scored

compact core + a few outliers

large candidate set 𝐶

cost ∝ 𝑄 𝐶

e

Recursive bisection D0

D1

D2

D3

Per-chunk AABB filtering

FAISS range search

f

local 𝐶𝑖 chunks 𝑄𝑖

With chunking

d

SFC rank partition

queries 𝑄

Without chunking

a

longest-axis split into 𝑄 1 . . .𝑄 4

tight per-chunk AABBs isolate outliers

same atoms, chunked

small 𝐶𝑖 ≪ 𝐶

block-diagonal pairs only

cost ∝

Í

𝑖 𝑄 𝑖 𝐶𝑖

𝑄𝐶 query

candidate

outlier

AABB

cutoff 𝑟𝑐

Í

𝑖 𝑄 𝑖 𝐶𝑖

𝐶𝑖 ≪𝐶

Fig. 4. Level 2 in detail: why chunking shrinks the candidate set. The mechanism behind card b of Fig. 3. (a–c) Without chunking: a single device-level AABB is stretched by a few spatial outliers, so the retained candidate set 𝐶 is large and the exact search scores all 𝑄 × 𝐶 pairs (∝ 𝑄 loc𝐶 cand ). (d–f) With chunking: recursive bisection along the longest axis produces tight per-chunk AABBs that Í isolate those outliers, so only the block-diagonal 𝑄𝑖 –𝐶𝑖 pairs are scored and the cost drops to 𝑖 𝑄𝑖 𝐶𝑖 with 𝐶𝑖 ≪ 𝐶 cand . Same atoms in both rows.

candidate side hardest—about 0.5M entries per rank, essentially independent of 𝐾, against 1.9M for Hilbert rank-box at 512 GPUs—and is nonetheless slower at every scale we measured, because its metadata is ≈ 1 MiB per rank rather than under 2 KiB, it needs a second all-gather to distribute the request mask, and it hands Level 3 roughly four times as many query groups to schedule. Two further readings support attributing the gap to the scheme rather than to the partition it runs on. Spatial-bin’s latency varies by under 5% between the Hilbert and Morton curves, where rank-box’s varies by 2.0×: replacing perdevice boxes with a shared grid does make shard geometry irrelevant, exactly as designed, and makes the scheme pay for the grid instead. And where the per-rank box is genuinely inflated the grid does help: on Morton, spatial-bin brings the peer degree down to 19, matching Hilbert rank-box, and at 50k atoms/GPU it is ≈ 10% faster than Morton rank-box. That is the only configuration in which it wins: at 25k atoms/GPU on Morton the two are level, and against Hilbert rank-box it is behind throughout: the latency the smaller box buys back never covers the cost of the plan that produced it. A filter’s cost includes publishing what it filtered with. Because the sub-stage timers are MAX-reduced independently (Section D), we report this as consistent with the telemetry rather than as a measured stage decomposition. Boundary atoms and outliers still inflate a device’s AABB and the candidates it retains, so Level 2 bisects local queries at the median of the longest axis into compact chunks and screens each against its own dilated box (Fig. 4)—conservative, since any atom within the cutoff of a chunk’s query lies inside that box. The exact search then runs per chunk over only its survivors, as one batched GPU call. With 𝑄 loc local queries and 𝐶 cand candidates after Level 1, an unchunked Í Í search costs ∝ 𝑄 loc𝐶 cand against 𝑖 𝑄𝑖 𝐶𝑖 chunked, where 𝑄𝑖 and 𝐶𝑖 are the per-chunk counts; since 𝑖 𝑄𝑖 = 𝑄 loc , the gain is entirely 𝐶𝑖 ≪ 𝐶 cand , not fewer queries (Appendix C). The chunking itself adds an 𝑂 (𝑄 log 𝑄) median bisection and one AABB test per chunk, both small against the pairwise work they remove. Cache exactness. A step’s graph splits into a topology—edge sets, periodic-image mapping, and the exchange/ communication plans, all indexed by current ownership—and a geometry, the positions and displacements that feed it. Only geometry must change every step. Over one cache epoch of ten steps (5 fs at our timestep) a 300 K hydrogen

11

moves ≈ 0.1 Å rms, inside the 𝛿/2 = 0.25 Å trigger that the 0.5 Å skin allows, so the topology survives. Over the hundred steps between re-partitions it moves ≈ 1 Å rms, far inside the ≈ 67 Å subdomain at ∼ 30k atoms/rank (𝜌 = 0.1 atoms Å−3 ), so ownership survives too. Reuse is only legitimate if it is exact, so the scheme is stated as an invariant. Two search radii are in play, and which one runs depends on whether the cache is enabled. With the cache off—the interval-1 configuration used as the reference throughout—every filter runs at 𝑟𝑐 and each step’s edge set is built from scratch. With the cache on, the rebuild instead constructs an untruncated candidate superset at 𝑟𝑐 + 𝛿 (𝛿 = 0.5 Å by default), and every intervening hit step recomputes minimum-image geometry over that superset and activates only pairs with 𝑟 ≤ 𝑟𝑐 . Activation is a mask over the fixed superset: we never compact the edge list within an epoch, so edge-indexed tensors keep the superset’s extent. The two therefore produce the same active edge set, not merely a similar one: the interval-1 and interval-10 runs agree on the active edge list bit for bit, while their cached supersets and plans differ in length by construction. The model carries two cutoffs, 5 Å for the three short-range layers and 8 Å for the long-range one, and each is searched in its own stage under the same rule: the stage for cutoff 𝑟𝑐 builds its superset at 𝑟𝑐 + 𝛿 and activates 𝑟 ≤ 𝑟𝑐 against it. The 𝛿 is what the displacement trigger is sized against. A rebuild fires once any rank’s own maximum displacement since the last build exceeds 𝛿/2, so two atoms moving directly toward each other close their separation by less than 𝛿 between builds, and no pair beyond 𝑟𝑐 + 𝛿 at build time can enter 𝑟𝑐 unnoticed—one trigger therefore covers both stages, since the smaller 𝛿/2 margin, not the cutoff, bounds the crossing. The decision must be unanimous: every rank agrees via all_reduce(MAX) over local triggers, so all ranks re-enter the same cache epoch together and the exchange plans stay mutually consistent. On a hit we re-expand positions under PBC, recompute edge geometry, and re-apply the 𝑟 ≤ 𝑟𝑐 test over the superset, reusing its indices and plan; the indices are detached from autograd while the recomputed geometry stays in it, so forces are not taken through stale positions. A repartition invalidates the cache outright, since ownership, chunk assignment, and the plan all change with it (Algorithm 1). Left as a separate caveat is the 𝑘-selection clamp of Section 6: the skin bounds which pairs the candidate set must contain, not how many neighbors the backend will return for a single query. This is reuse, not approximation—the skin is sized so no pair crosses undetected between rebuilds—and interval-1 and interval-10 runs give bitwise-identical topology hashes and forces agreeing to 7.2 × 10−7 eV/Å. Even so, every step must still exchange remote Key, source-embedding, and Value features, and at scale this—not floating-point throughput—is the dominant obstacle. 3.4

Communication Optimization

A remote exchange rarely costs just the payload’s bandwidth: per-call handshakes, full-width equivariant transfers, separate launches, and on-path allocation can dominate instead, and irregularly. Three mechanisms strip that overhead; two more then hide what remains. Removing What Is Not Payload. Plan reuse (Fig. 5): blocks sharing a cutoff share a peer set, so the routing plan—peers, send/receive lengths, offsets, index maps, buffer addresses—is built once per cutoff and reused by every later all2allv there. With 𝑅𝑟 blocks at cutoff 𝑟 , building the plan once turns 𝑅𝑟 handshake-plus-payload exchanges into one handshake and 𝑅𝑟 payload-only transfers: 𝑅𝑟 × (handshake + payload) −→ handshake + 𝑅𝑟 × payload,

(6)

cutting per-block handshakes to one per cutoff—two per step, at 5 and 8 Å, rather than four—and the saving grows with the number of blocks sharing a cutoff. It is realized again on every cache-hit step, since each plan is indexed by topology rather than by step number. Sender-side projection (Fig. 6): of the three per-atom quantities attention consumes, only the Value is aggregated as an equivariant feature and must travel at full width (𝑙 max + 1) 2𝐶; the source embedding is an invariant scalar of width 𝐶, and the Key is used solely inside scalar query–key dot products, so it need not be shipped equivariant. Projecting the Key on the sender to 𝐻 per-head vectors of width 𝑑𝑘 (𝐻𝑑𝑘 scalars) leaves the per-atom payload at 𝐶 + 𝐻𝑑𝑘 + (𝑙 max + 1) 2𝐶, against 𝐶 + 2(𝑙 max + 1) 2𝐶 if the Key were sent in its equivariant form. With the deployed 𝑙 max = 2, 𝐶 = 256, 𝐻 = 128,

12

Communication-plan reuse eliminates per-round handshake overhead R0

a Naive: per-round handshake

R1

h.s.

all2allv payload

R2

h.s.

all2allv payload

R3

h.s.

all2allv payload

R3

variable-length All-to-All among ≤ 𝑘 spatially adjacent peers — same topology within a cutoff

R4

c Steady state: plan reuse

b Build one plan per cutoff

negotiate rank 𝑗 rank 𝑖 lengths · offsets · buffer addrs

R1

R2

negotiate once / cutoff rank 𝑗 rank 𝑖 build routing structure

rank 𝑖

Cached communication plan peer 0→1 0→2 0→5 0→6

. . . handshake recurs for all 𝑅 blocks

recv

send 128 64 96 72

96 80 48 72

offset 0 128 208 256

payload only rank 𝑗 routed by cached plan

R1

all2allv payload

R2

all2allv payload

R3

all2allv payload

no handshake

Per step: 𝑅 Transformer blocks ⇒ 𝑅 All-to-All (all2allv) rounds

. . . one plan reused by that cutoff’s blocks

+ index maps · pre-registered buffers

control-plane cost ∝ 𝑅

peers, lengths, offsets, index maps

pure P2P transfer, no negotiation

latency-bound, bandwidth idle

built once per cutoff, cached

bandwidth-bound payload

handshake / negotiation

naive

P2P payload transfer

ours time saved

Fig. 5. Communication-plan reuse eliminates per-round handshake overhead. Every simulation step issues 𝑅 variable-length All-to-All (all2allv) rounds—one per Transformer block—among the same sparse set of ≤ 𝑘 spatially adjacent peers (top). a, A naive implementation renegotiates transfer lengths, offsets, and buffer addresses (handshake, orange) before every round, so control-plane activity recurs each round. b, Because the blocks at one cutoff share a neighbor topology, we pay the handshake once per cutoff to build a cached communication plan recording the participating peers, send/receive lengths, offsets, index mappings, and pre-registered buffer addresses. c, All subsequent rounds reuse the cached plan and perform pure P2P payload transfer (blue) with no further negotiation. Across the 𝑅𝑟 blocks at cutoff 𝑟 the exposed cost drops from 𝑅𝑟 × (handshake + payload) to handshake + 𝑅𝑟 × payload, pushing the communication path into a bandwidth-bound regime; the production stack has two cutoffs and so pays two handshakes per step, not one and not 𝑅 (Section 3.4).

𝑑𝑘 = 8 (Appendix A), we have (𝑙 max + 1) 2𝐶 = 2,304 and 𝐻𝑑𝑘 = 1,024, so the payload falls from 4,864 to 3,584 scalars per atom—×4 bytes in TF32/FP32, i.e. 19,456 → 14,336 bytes—a 26.3% reduction realized entirely on the Key. The saving grows with model order: at fixed 𝐶 the equivariant Key width scales as (𝑙 max + 1) 2 while its projected width 𝐻𝑑𝑘 is constant, so higher-𝑙 max models—where communication is most acute—benefit more. Fusion and pooling: the three tensors then travel as one physical all2allv from a pre-registered buffer (Figs. 6–7), so launches per exchange fall from three to one, and the received SKV data is consumed directly from the fused buffer with the operators’ autograd dependencies rewritten so the backward pass traverses the fused path in a well-defined way. Send and receive buffers are carved from a pool whose Memory-Region registration happens in advance during plan construction, so steady-state transfers repeat neither buffer creation nor dynamic registration. What remains exposed is close to the payload itself, which can only be hidden. Overlapping the Exchange. Attention splits into a local branch, whose source atoms already reside on this GPU, and a remote branch, fetched from peers; running them concurrently hides the smaller of the two times. The bound follows directly. Let 𝑇𝑐 be the time to exchange remote features, 𝑇ℓ the local branch’s execution time, and 𝑇𝑟 the remote branch’s time once the data has arrived. Ignoring terms common to both schedules—the initial QKV projection and the final merge—the sequential schedule costs 𝑇seq = 𝑇𝑐 + 𝑇ℓ + 𝑇𝑟 ,

(7)

whereas launching communication and local computation concurrently exposes only 𝑇ovlp = max(𝑇𝑐 ,𝑇ℓ ) + 𝑇𝑟 . 13

(8)

Sender-side projection with fused SKV communication a Conventional: send the full feature, three separate launches Rank 0 · sender

three separate collective launches launch 1 → Key

Rank 1 · receiver projection repeated on receiver (per received copy)

ℎ𝑖 (full)

launch 2 → Source 𝐶

launch 3 → Value

ℎ𝑖 (full)

key_linear

K

type-embed

S

proj_value

V

K, S, V re-derived

recv. ℎ𝑖

full equivariant feature ℎ𝑖

Key, Value full; Source always compact

(𝑙 max +1) 2 ×𝐶 (large)

b Ours: project on the sender, pack, single fused launch Rank 0 · sender

ℎ𝑖

Rank 1 · receiver

key_linear

K

𝐻 ×𝑑𝑘

type-embed

S

𝑑 embed

proj_value

V (𝑙 max +1) 2 ×𝐶

full equivariant feature (large)

single fused all2allv K

S

K

S

V

V

pack: one contiguous buffer

compact K | S | V (fused)

fused receive buffer

compact payload

structured views by offset

K, S, V consumed directly

conventional ours ≈ 26% less traffic, 3 → 1 launch

Fig. 6. Sender-side projection with fused SKV communication. a, In the conventional scheme the Key, source-embedding, and Value projections are repeated on the receiver, so three separate collective launches fire; two of them (Key, Value) must still carry 2 the full equivariant intermediate ℎ𝑖 ∈ R (𝑙max +1) ×𝐶 , since only Value is aggregated equivariant and the receiver cannot recover a sender-side scalar Key without it, while the source embedding is already a 𝐶-wide scalar and was never large. b, In our scheme, the Key (𝐻 ×𝑑𝑘 ), source-embedding (𝑑 embed ), and Value ((𝑙 max +1) 2 ×𝐶) tensors are projected locally on the sender and packed into a single peer-specific contiguous buffer, exchanged in one fused all2allv; the receiver reconstructs structured tensor views from deterministic offsets and consumes them directly. This lowers the transmitted volume by ≈ 26%, realized entirely on the Key (Section 3.4), and reduces the number of collective launches from three to one.

The saving is therefore Δ𝑇 = 𝑇seq − 𝑇ovlp = min(𝑇𝑐 ,𝑇ℓ ),

𝜂 hide =

  Δ𝑇 𝑇ℓ = min 1, , 𝑇𝑐 𝑇𝑐

(9)

so the exchange is fully hidden exactly when the local branch is at least as long as the transfer, and otherwise hidden in proportion. Note that 𝑇ovlp is a conservative upper bound: in practice the remote branch can begin as soon as the transfer finishes, rather than only after 𝑇ℓ completes, so a second level of overlap is available whenever communication lands early. The split pays only if 𝑇ℓ is a large share of the total, and that follows from geometry alone, not from tuning. For a cubic subdomain of side 𝐿 and cutoff 𝑟𝑐 , the expected fraction of an atom’s in-cutoff neighbors that are locally owned is 𝑓local = 1 −

9𝑟𝑐 + O (𝑟𝑐 /𝐿) 2 , 8𝐿

(10)

via the box covariogram of a uniform-density cube (verified against Monte Carlo sampling). The coefficient follows from the cube’s geometry. For a cube of side 𝐿, the width along a unit direction n is 𝐿(|n𝑥 | + |n𝑦 | + |n𝑧 |), so a point uniformly placed in the cube leaves it under a displacement 𝑡n with probability 𝑡 (|n𝑥 | + |n𝑦 | + |n𝑧 |)/𝐿; averaging over the sphere gives 32 𝑡/𝐿, and averaging 𝑡 over a ball of radius 𝑟𝑐 gives 34 𝑟𝑐 , so the two averages multiply to 9𝑟𝑐 /(8𝐿). At

the weak-scaling loads we study (10k–60k atoms/GPU), 𝐿 ≈ 46–84 Å at 𝜌 = 0.1 atoms Å−3 , so 𝑓local ≥ 80% even for the 8 Å long-range layer and ≥ 87% for the 5 Å short-range layers; at the 8k-atom/GPU strong-scaling points 𝐿 falls

14

Pooled buffer management with one-time registration a Naive: per-round setup

c Fast-path reuse

b One-time pool initialization

pre-registered pool

R1

alloc MR reg

transfer

free

R2

alloc MR reg

transfer

free

R3

alloc MR reg

transfer

free

1

allocate contiguous pool (once)

one contiguous device-memory pool 2

. . . repeated every round

register pinned MR (once)

R1 borrow

transfer

return

R2 borrow

transfer

return

R3 borrow

transfer

return

pre-registered pinned MR + symmetric memory

. . . no alloc / registration

alloc & MR registration on the critical path

done outside the steady-state path

borrow and return buffers from the pool

repeated setup, fragmented

pooled, pre-registered, reusable

no registration latency

short-lived buffers ⇒ fragmentation

alloc + MR registration (setup)

payload transfer

borrow / return

naive ours setup off critical path

Fig. 7. Pooled buffer management with one-time registration. a, Allocating buffers and performing Memory-Region (MR) registration on the critical path of every communication round incurs repeated setup cost and leaves scattered, short-lived buffers that fragment device memory. b, We instead allocate a single contiguous device-memory pool once and pre-register it as a pinned MR (together with symmetric-memory registration), entirely outside the steady-state path. c, During communication, buffers are borrowed from and returned to the pool with no further allocation or registration, removing dynamic-registration latency from every transfer and leaving only the useful payload on the critical path.

to ≈ 40 Å and the 8 Å bound to ≈ 78%. The local branch therefore carries the larger share of the neighbor work at the loads that dominate the weak-scaling study, which is what the overlap exploits, and the fraction depends on the ratio 𝑟𝑐 /𝐿: at fixed density a rank holding more atoms has a larger 𝐿, so the bound improves with per-GPU load rather than degrading with device count. There is one implementation trap. A naive forward implementation issues the remote exchange and the local branch on two asynchronous streams and lets the scheduler interleave them. On SIMT hardware this is not enough: once a compute-heavy local branch is enqueued it occupies the scheduling queue and can delay the communication stream’s own launch, so the realized overlap window shrinks below the logical min(𝑇𝑐 ,𝑇ℓ ) of the bound above. We instead submit the exchange first, before any local compute is enqueued, and record a CUDA event when the transfer lands; only the remote branch is gated on that event. Separating launch order from data readiness is what makes the logical bound attainable. Placing the Launch in Both Passes. Figure 1(d) sketches the schedule; Figs. 8 and 9 give it in full. Backward does not follow the same recipe. Autograd traverses in reverse, so a forward-only scheme launches the exchange late and leaves a serial tail that no event placement can repair—the consequence, in scheduling terms, of the second force-pass requirement of Section 2.2. We therefore split the interface into two asymmetric operators: pre launches communication in the forward pass and is empty in backward, while post is a forward no-op that launches it in backward. Because post sits later in the forward graph, its backward runs first, issuing the reverse exchange early, with the remote-backward branch waiting on the resulting event. The local/remote split follows communication dependency, not a strict partition of the source atoms: the remote branch still consumes local source atoms, so the two branches share part of the same local features. Implementing them as independent operators with separately materialized inputs would duplicate those features in the autograd graph, inflating activation residency and complicating gradient flow. We instead let both branches consume a single unified representation of the local source data, differing only in the communication-dependent indices they access, and rewrite

15

Forward-pass communication–computation overlap a Ideal mechanism: communication on a separate stream launch async

remote data arrives

exposed cost

remote feature exchange (in flight)

comm stream compute stream

𝑇seq = 𝑇comm + 𝑇local + 𝑇remote 𝑇ovlp = max(𝑇comm ,𝑇local ) + 𝑇remote

local branch (local compute)

remote branch

Δ𝑇 = min(𝑇comm ,𝑇local ) 𝜂 hide = min(1, 𝑇local /𝑇comm )

wait overlap window time

c Correct order: async exchange first

b Naive order: local kernels first

overlap window

comm stream

comm queued (SMs busy)

compute stream

heavy local kernel

in flight

compute stream

remote

saturated by local

SM use

remote branch

local branch

remote

compute + copy-engine

SM use

comm starved ⇒ overlap collapses

local compute

remote exchange (in flight)

comm stream

comm hidden behind local compute

communication in flight

async launch

queued / wait

SM saturated

Fig. 8. Forward-pass communication–computation overlap. a, The attention operator is split into a local branch (source atoms already resident on the GPU) and a remote branch; the remote feature exchange is launched asynchronously on the communication stream while the local branch executes on the compute stream, and the remote branch waits on the completion event only immediately before it consumes the received data; the panel’s cost model gives the resulting exposed cost. b, Launch order is decisive: if the heavy local kernels are enqueued first, the exchange is issued late behind them and the overlap window collapses. c, Issuing the asynchronous exchange first gets the communication kernels resident before the local branch fills the device, preserving a large overlap window.

the backward pass to accumulate through the shared dependency exactly once, preserving the overlap schedule with no additional activation memory. 3.5

A Decentralized On-GPU MD Engine

One global collective per step survives everything above—not in the attention operator but in the integration loop wrapping it. Conventionally rank 0 orchestrates every step (Fig. 10a): it scatters positions and reduces forces through its single network endpoint, which 𝐾 −1 ranks contend for every step. Both per-step collectives funnel through one endpoint: the force reduction is a many-to-one incast and the position scatter its one-to-many mirror, so 𝐾 − 1 ranks contend for a single network interface—congestion rather than aggregate bandwidth—and the volume crossing it grows with total system size instead of staying local. The deeper cost sits between those two collectives: the loop re-encodes the partition on one device every step, quantizing and radix-sorting the full coordinate set, and that sort is most of the partition-and-broadcast bundle it sits in (≈ 70.7 ms of the bundle’s ≈ 87 ms at 10.24M atoms). Re-encoding re-derives ownership, and since the edge sets and exchange plan of Section 3.3 are indexed by ownership, it invalidates the topology cache on every step—erasing exactly what the cache exists to preserve. A centralized driver therefore cannot hold a rebuild off for even the ten steps the cache needs, let alone reach the hundred steps between re-partitions, however fast its collectives run. Integrating on-device additionally removes the host–device coordinate and force transfers a host-driven integrator performs each step. Neither centralization is needed every step: the subdomain size bounds ownership, and the skin bounds which edges can cross, on the two timescales of Section 3.3. We therefore integrate per rank, on the GPU (Fig. 10b; one full step is Algorithm 1). Rank 0 encodes and scatters once at initialization; from then on each rank advances the atoms it owns from its own positions and the globally computed forces, emitting only a fixed handful of scalars—so the per-step force reduction and the centralized integrator collapse into a constant-payload barrier. A global gather and re-partition

16

Autograd-compatible overlap via paired pre/post operators operator chain (forward order) ···

equivariant attention (local ∥ remote)

pre

fwd: launch comm · bwd: empty

post

backward traverses in reverse, so post.backward fires first

···

fwd: no-op · bwd: launch comm

a Naive forward-only overlap forward pass

comm compute

backward pass

pre launch exchange local compute

exchange remote remote-grad

remote

overlap window

local-grad compute

serial backward tail

b Paired pre/post: overlap in both directions forward pass — pre launches

comm

backward pass — post.backward launches

pre launch exchange

post launch exchange

wait

compute

local compute

local-grad compute

remote

overlap window

async launch

remote exchange

scatter-sum

overlap window

local compute

remote branch

exposed serial tail

Fig. 9. Autograd-compatible overlap via paired pre/post operators. a, A naive forward-only overlap works in the forward pass, where pre launches the remote exchange before local computation. In backward, however, communication is not relocated early enough: the remote exchange, remote-gradient computation, and local-gradient computation appear as a serial tail. b, With paired operators, pre launches communication in the forward pass, while post is a forward no-op whose backward method launches the reverse remote exchange. Because post.backward is reached early under reverse traversal, the exchange goes out as soon as its payload—the gradient with respect to the received remote features—exists, and overlaps with unblocked local-gradient computation. Writing the forward halo gather as xhalo = 𝐺 xowned makes the bar after the wait explicit: the owner applies 𝐺⊤ , accumulating ∇xowned += 𝐺⊤ ∇xhalo over every rank and periodic image that requested the atom.

fire only every 100 steps (≈ 50 fs at our timestep)—the shape-changing event of Section 3.2—leaving the steady-state trajectory to run as replayed static graphs. Section 4.6 measures that bundle: ≈ 87–88 ms for the partition-and-broadcast half (≈ 16–18 ms of it cross-rank communication) plus ≈ 320–330 ms for the rebuild it triggers, 407–418 ms in total, amortized to roughly 4 ms of per-step cost against an ≈ 4.55 s full step. Algorithm 1 spells this out in full, executed identically (SPMD) by every rank 𝑝: the half-kick/drift, the unanimous rebuild-or-reuse decision of Section 3.3, the fixed-schedule re-partition of Section 3.2 that unconditionally forces a rebuild on top of it, the overlapped forward–backward pass of Section 3.4, and the closing half-kick. Lines 12–15 are the only branch that can change tensor shape—forced by the unanimous rebuild decision at line 5, taken from the trigger at line 4, or unconditionally by a repartition event (lines 6–11)—while lines 1–5 and 16–19 run identically whether or not the step rebuilds, including the mask re-application of line 16, which sees the same (E, Π) shape on a hit as a rebuild just produced. The shape-changing branch is explicitly fenced; everything outside it is data-independent in extent, which is what makes the captured graph reusable. 4 4.1

EXPERIMENTAL EVALUATION Experimental Setup

Experiments run on 64 compute nodes, each with 8 NVIDIA H100 GPUs (80 GB HBM3) on NVLink 4.0 and linked by a 400 Gbps InfiniBand NDR fabric; the largest run uses all 512. The stack is CUDA 13.0, PyTorch 2.11, NCCL 2.28, Triton 3.7 and FAISS 1.14.2. A second cluster of 256 A100 GPUs (40 GB HBM2e, PCIe, no NVLink) carries the strong-scaling sweep and the two MD trajectory studies of Section 4.6; it was swapped in for the H100/NVLink/InfiniBand fabric for

17

a Conventional MD engine

rank 0: SFC encode + atom partition

b Decentralized on-GPU engine

𝑂 (𝑁 ) / step

rank 0: SFC encode + partition (once per re-partition)

serial hub

scatter once / re-partition 1 → 𝐾

scatter positions 1 → 𝐾

next step

fixed size / step

gather every 100 steps

ranks 0 . . . 𝐾 −1: distributed inference ⇒ forces

ranks 0 . . . 𝐾 −1: distributed inference ⇒ forces ×100 steps (≈50 fs): local update + fixed-size sync

reduce forces 𝐾 → 1

rank 0: integrate ⇒ update positions

Inter-rank comm / step Rank-0 work / step CPU ↔ GPU transfer / step

each rank: local integrate on GPU

serial hub

parallel

Conventional

Ours

𝑂 (𝑁 ) : 1 →𝐾 scatter + 𝐾 → 1 reduce

✓ fixed size: sync signal

integrates every atom (serial hub)

✓ idle between re-partitions

every step

✓ none — fully on-GPU

Fig. 10. Conventional versus decentralized molecular-dynamics engine. a, A standard distributed loop is rank-0 centric: each step scatters atom positions from rank 0 to all ranks (1 →𝐾), runs distributed inference, reduces all forces back to rank 0 (𝐾 → 1), and lets rank 0 integrate and re-encode. The two collective hot spots through rank 0 form a communication bottleneck, and the serial curve encoding and integration on rank 0 form a compute bottleneck—both worsening with system size. b, Our engine keeps every operation on the GPU. After the initial scatter, each rank integrates its own atoms locally using the globally computed forces and emits only a fixed-size synchronization signal, eliminating the per-step 𝐾 → 1 reduction and the central integration. A global gather and re-partition are performed only every 100 steps (≈ 50 fs at the 0.5 fs step), which is safe because atoms move far less than the neighbor-search margin over that interval (Section 3.2). The scorecard contrasts the two engines per step: the decentralized design removes the 𝑂 (𝑁 ) collective bottlenecks, the rank-0 serial-compute bottleneck, and the per-step CPU–GPU transfers entirely.

Algorithm 1 Decentralized MD step (rank 𝑝; identical on every rank) Require: owned positions r, velocities v, forces F from the previous step, build-time reference positions rref , cached topology (E, Π), step 𝑡, timestep Δ𝑡, rebuild interval 𝐼 , displacement skin 𝛿 1 1: v ← v + 2 Δ𝑡 F/𝑚 ⊲ half-kick, force from previous step 2: r ← r + Δ𝑡 v ⊲ drift 3: 𝑑 ← max displacement of an owned atom since rref 4: trigger ← [ 𝑡 mod 𝐼 = 0 ] ∨ [ 𝑑 > 𝛿/2 ] 5: rebuild ← all_reduce(MAX, trigger) ⊲ unanimous, Sec. 3.3 6: if 𝑡 mod 100 = 0 then ⊲ fixed schedule, not data-dependent 7: rall ← Gather(r, rank 0) 8: 𝜎 ← HilbertSort(rall ), cut into 𝐾 equal-count contiguous segments 9: Broadcast(𝜎); migrate owned positions and velocities to match this rank’s new segment 10: rebuild ← true ⊲ which atoms this rank owns changed: cached topology and plan invalid, regardless of line 5 11: end if 12: if rebuild then 13: (E, Π) ← ThreeLevelSearch(r, 𝑟𝑐 +𝛿) ⊲ candidate superset; AABB → chunk → exact FAISS; also resets rref 14: re-capture the static CUDA Graph for this (E, Π) shape 15: end if 16: re-expand r under PBC; recompute edge geometry from E and activate the pairs with 𝑟 ≤ 𝑟𝑐 ⊲ both paths; indices detached, geometry stays in the autograd graph 17: F ← ReplayGraph(r, E, Π) ⊲ fwd + autograd bwd; All-to-All overlapped both passes, Sec. 3.4 1 18: v ← v + 2 Δ𝑡 F/𝑚 ⊲ half-kick, force from this step 19: all_reduce(SUM, 𝐸,𝑇 , virial) ⊲ constant payload, independent of 𝑁 availability, and Section 6 scopes what a comparison across the two interconnects can mean. Every timing reported below is MAX-reduced across ranks, so it measures the distributed critical path rather than the rank-0 view. 18

Table 1. Run configurations. “n/s” marks a value the experimental text does not state. TF32 is the mixed-precision mode of Section 3.1; recomputation is activation recomputation; the cache interval is the rebuild period in steps. The H100 cluster is the 64-node, 8-GPUper-node NVLink/InfiniBand system above, the A100 cluster the 256-GPU PCIe system; the two are never mixed within a single comparison. Configuration

GPUs

Atoms/GPU

TF32

Recomp.

Cache

Cluster

Single-GPU ablation Neighbor, 25k Neighbor, 50k Overlap A/B Weak scaling, 10k Weak scaling, 25k Weak scaling, 60k Full-step cache A/B Strong scaling KCl trajectory Capsid pair

1 64–512 64–512 32–512 32–512 32–512 32–512 512 32–256 256 256

10k 25k 50k 25k 10k 25k 60k 60k 40k→5k 40k 39.3k

off/on on on on on on on on off off off

off/on n/s n/s off off off on on on n/s n/s

n/s 1/10 1/10 10 10 10 10 1/10 n/s 10 10

H100 H100 H100 H100 H100 H100 H100 H100 A100 A100 A100

All runs use a variant of the open-source UBio-MolFM potential [14], itself built on E2Former V2 [23, 13]—three short-range layers at 5 Å then one long-range layer at 8 Å, a 23 Å receptive field per atom, 23.5M parameters (full configuration in Appendix A). This paper’s contribution is the model’s distributed execution, not its chemistry: its DFT-level accuracy on biomolecular systems is established separately, and Table 2 validates that our execution path preserves it. Module benchmarks use a 10k-atom toy system and bulk water to the single-node limit; the scaling study grows bulk water with the device count, at 10k, 25k and 60k atoms per GPU, to 30.7M atoms on 512 GPUs; the MD studies simulate a 10.24M-atom aqueous KCl solution at 40k atoms per GPU on 256 devices and a 10.06M-atom Norovirus capsid in solvent at 39.3k. Sections 4.3–4.5 use TF32 mixed precision, with the accuracy trade-off characterized below; Section 4.6’s trajectory validation runs with TF32 off, and the TF32 A/B of Appendix D runs on this section’s 1.28M-atom strong-scaling system, where its off arm reproduces the 64-GPU step time to 0.9%. Table 1 collects the configurations the section reports. The cluster column is uniform within every comparison but not across the section: the weak-scaling, overlap, neighbor-scaling and single-GPU results are H100 numbers, while the strong-scaling sweep and the trajectory studies are A100 numbers with no NVLink. The recomputation column marks the runs that carry it—the 60k weak-scaling load and the strong-scaling sweep, whose 32/64-GPU points exceed its cluster’s 40 GB budget otherwise. Timing protocol. Every timing in this section is a steady-state number: step 0 is excluded from every series, since it carries warm-up, torch.compile tracing, autotuning and first-touch allocation. The single-GPU ablation averages the last 50 of 100 repetitions; the end-to-end curves average steps 300–399 of a 500-step run, a window that spans exactly ten interval-10 rebuild periods so the mean is the amortized per-step cost rather than a sample biased by where rebuilds land; the strong-scaling points average the last five logged windows of their run. Sub-stage timers are each MAX-reduced independently, and their maxima may originate from different ranks, so we do not decompose a stage additively—an attribution resting on those timers is consistent with the telemetry, not a measured stage decomposition. One official run per point, so no error bars are drawn and differences of a few percent are not treated as separable. Appendix D gives the protocol in full. 4.2

Single-GPU Acceleration

We first quantify the single-GPU stack of Section 3 with a waterfall ablation on a fixed 10k-atom toy system (Table 2). The largest gains come from restructuring the sparse execution path itself: Triton and torch.compile together take latency 363.4 → 141.1 ms and peak memory 66.3 → 15.0 GiB (Table 2), consistent with fusion removing both execution overhead and intermediate materialization. TF32 and recomputation are complementary but non-equivalent optional modes: TF32 buys 1.47× throughput at a ∼ 3 × 104 × larger energy error (3.05 × 10−5 → 8.93 × 10−1 eV) and 6.5× the NVE drift rate—but not the temperature, and not the water or solvation structure, whose first peaks hold to 0.5% while

19

Table 2. Ablation on a 10k-atom toy system on H100. Config abbreviates T=Triton, C=torch.compile, F=TF32, R=activation recomputation, applied cumulatively left to right. Precision metrics are relative to the base (no T/C/F/R). Config

Latency(ms)

Mem(MiB)

|Δ𝐸 |(eV)

max |Δ𝐹 |

363.4 234.9 141.1 95.7 165.0 114.4

67918 21078 15370 15370 7368 7368

0 3.05 × 10 −5 3.05 × 10 −5 8.93 × 10 −1 3.05 × 10 −5 8.96 × 10 −1

0 2.38 × 10 −7 1.01 × 10 −6 1.43 × 10 −3 1.19 × 10 −6 1.37 × 10 −3

base +T +T+C +T+C+F +T+C+R +T+C+F+R

the ion–ion pairs move on the scale of the scatter between two halves of one run (Appendix D); recomputation buys memory headroom at moderate latency cost and leaves that error untouched. The attention kernel itself we isolate against a vanilla sparse baseline and dense SDPA in Appendix B. What that comparison can mean is fixed by the workload: molecular attention’s MFU ceiling is set by sparsity, not the GPU’s dense peak, so the informative comparison is within sparse implementations. 4.3

Neighbor Search at Scale

Setup. This stage is the design’s primary scaling risk. We measure curve, exchange granularity, and cache, separately and composed, from 64 to 512 GPUs at fixed atoms per GPU (12.8M atoms at 25k/GPU, 25.6M at 50k), 𝜌 = 0.1 atoms Å−3 , a common 10k query-chunk limit. Reported neighbor-pipeline time excludes the partition and the forward/backward pass, so these reductions do not transfer one-to-one to a full step (Section 4.5 measures that; protocol in Appendix D). The curve sets the slope. At identical per-rank atom counts and rank-box granularity, neighbor cost grows 2.8–2.9× under Morton against 1.6–1.7× under Hilbert over the range (Fig. 11a,b)—the slope, not the constant, separates a stage that scales from one that does not. Panels (c,d) give the mechanism: a fragmented Z-order shard’s box encloses the gaps between its pieces, so Morton’s peer degree and candidate pool keep climbing fast—5 to 27 source ranks over 64–512 GPUs at 25k atoms/GPU—where Hilbert’s connected shard rises from 6 only to 19 (Section 3.2). Both grow; the exponents differ, near 0.5 for Hilbert against a steeper Morton. Some of the penalty lands downstream of the exchange: handed identical all-gather pools, Morton is still 2.0× slower in the exact search itself. Pruning hardest loses. Spatial-bin holds the smallest candidate pool (panel d) yet runs 1.9–2.4× slower than rank-box on Hilbert at every scale, because pruning must itself be published. All-gather fails at the other extreme: latency indistinguishable from rank-box’s, but a pool growing with the global atom count and no sparse peer set for Section 3.4’s plan to reuse. Appendix C accounts for both. We adopt Hilbert + rank-box, which at this scale costs 3% against all-gather rather than saving; the two differ by about 9% at most across the Hilbert points, a margin we do not treat as separable given the single run per point, and rank-box is the default for its small, fixed-size metadata and reusable peer set, not for latency. The cache, and the three composed. A hit costs an order of magnitude less than a rebuild (Fig. 12a, interval 1 against 10), and is worth more at larger scale: it skips plan construction, metadata collectives, and candidate transport, exactly the part that grows with 𝐾, while the geometry refresh it retains reuses the fixed halo plan and recomputes distances locally. The reuse is exact: the interval-1 and interval-10 runs give bitwise-identical topology hashes and forces agreeing to 7.2 × 10−7 eV/Å (Section 3.3). Composed at 512 GPUs and 25.6M atoms (Fig. 12b), a Morton+all-gather baseline of 393.6 ms falls to 170.8 ms on the curve, holds at 176.3 ms on granularity (kept for footprint, not latency), and reaches 25.1 ms with the cache, 15.7× overall—that last figure composed from separately measured rebuild and hit steps rather than timed on one interval-10 run (Appendix D). Again the slope matters as much: the baseline grows 3.0× over the range, capping weak-scaling efficiency at 33% from this stage alone, against 1.7× and 60% for the full configuration, at an absolute cost small enough that the stage no longer limits the system. That 1.7× is not flat, and the telemetry names the next target: the metadata all-gather, still rising with 𝐾. Table 3 gives the per-mechanism numbers behind the composition, including each row’s own 64-to-512-GPU growth—a scaling behavior Fig. 12(b) does not show, since it plots only the 512-GPU endpoint. 20

Neighbor pipeline (ms)

(a) Rebuild, 25k/GPU

(b) Rebuild, 50k/GPU

300

400

150

200

0

0 64

256

64

512

512

(d) Candidate pool

(c) Active source ranks 512

5

Candidates (M)

Source ranks

256

64

8

1

0.2 64

256

64

512

GPUs Exchange:

allgather

256

512

Morton

Hilbert

GPUs rank-box

spatial-bin

Partition:

Fig. 11. Weak-scaling ablation of the distributed neighbor search on H100, rebuilding every step; color is the Level-1 exchange scheme, line style the partition curve. (a,b) Neighbor-pipeline latency at 25k/50k atoms per GPU—Morton bends upward (2.8–2.9× across the two panels), Hilbert rises gently (1.6–1.7×); rank-box (default) is bold, allgather (within 11% of it throughout) is lightened for clutter, spatial-bin for being the slowest scheme on Hilbert instead, as detailed next. (c) Source ranks requested by the heaviest rank, all schemes at equal weight (all-gather = 𝐾 by construction, drawn once for (c,d)). (d) Candidate population: spatial-bin prunes hardest—2.6–11.3× below all-gather, 2.2–3.6× below rank-box—yet is slowest on Hilbert at every scale: volume alone does not predict cost.

Table 3. Composing the three neighbor-search mechanisms (a-g=all-gather, r-b=rank-box) at 512 GPUs, 50k atoms/GPU; each row adds one mechanism to the row above. “Growth” is the 512-/64-GPU cost ratio, so 1.0 is ideal weak scaling. Cached rows are a composed interval-10 amortized cost, not a timed interval-10 run; the composition and its known bias are stated below. Neighbor pipeline only. Rank-box costs 3% against all-gather at this scale rather than saving: the two differ by under 9% at every Hilbert point in the observed runs, a margin we do not treat as separable given the single run per point (Section 4.3), and rank-box is the default for its small, fixed-size metadata and reusable peer set, not for latency. At 25k atoms/GPU the baseline grows 3.2× over the same range. Partition

Exchange

ms

Reduc.

Growth

Morton Hilbert Hilbert Hilbert

a-g (base) a-g r-b r-b+cache

393.6 170.8 176.3 25.1

1.0× 2.3× 2.2× 15.7×

3.0× 1.6× 1.6× 1.7×

How much weight the composed figure carries. The cached row is (𝑡 rebuild + 9 𝑡 hit )/10, where 𝑡 hit comes from an interval-10 run and 𝑡 rebuild from a separate interval-1 run of the same configuration. They are not interchangeable: the interval-1 arm searches at each layer’s bare cutoff, while an interval-10 rebuild builds the 𝑟𝑐 +𝛿 superset the following nine steps consume—1.33× the candidate volume at the 5 Å layers and 1.20× at the 8 Å one—so 𝑡 rebuild is measured on 21

(c) Overlap on vs. off

(b) Cumulative, 512 GPUs 400

ms/step

300

30

Full step (s)

Pipeline (ms/step)

(a) Graph cache, 50k/GPU

200

0.7

−14.5%

15.7× 0

3 64

256

0.5

512

Mor. a-g

Hil. r-b

32

+c.

128

512

(f) Weak-scaling efficiency

(e) End-to-end step time

(d) Reduction by step type 20

2

10

0

Efficiency (%)

100

Full step (s)

Reduction (%)

Hil. a-g

1

70

0.2 32

64

128

256

32

512

128

512

32

128

512

(g) Aggregate throughput

M atoms/s

(a,b) 10

Mort.+a.g. +cache

Hil.+r.b. hit-only

off hit

on rebuild

10k

25k

(c,d)

(e-g)

2

32

128

512

60k∗ ∗ recompute on (HBM limit)

Fig. 12. H100 results across varying GPU counts. (a,b) Neighbor-graph cache: (a) per-step cost at 50k atoms/GPU—cache hits (orange) an order of magnitude cheaper than rebuilds; (b) cumulative reduction composing curve + exchange + cache at 512 GPUs, the cache bar an upper estimate (Appendix D). (c,d) Overlap on the full step at 25k atoms/GPU: (c) per-step time on vs. off, gap widening with scale; (d) reduction by step type, larger on cache hits. (e-g) End-to-end weak scaling, 32–512 GPUs, three per-GPU loads: (e) step time, rising 1.11–1.29× over 16×; (f) parallel efficiency vs. 32 GPUs; (g) aggregate throughput, each load against its own 32-GPU-anchored linear reference (dashed)—a curve’s gap to its line is (f)’s efficiency. Heaviest load: 30.7M atoms on 512 GPUs. ∗ At 60k atoms/GPU (e-g), recomputation is on since the load exceeds the 80 GB HBM budget otherwise; 10k/25k run without it.

a strictly cheaper search than the one it stands in for, and these runs contain no steady-state interval-10 rebuild to substitute. Where both quantities do exist in one place—Table 4, whose rebuilds fall on steps 10 and 20—a steady-state interval-10 rebuild runs 0.98–1.62× its interval-1 counterpart. The composition therefore understates the amortized cost by up to 1.4×, and the composed 15.7× is an upper estimate rather than a measured speedup. It does not affect the cache-hit latencies of Fig. 12(a), measured directly, or the full-step 1.87× of Section 4.5, an A/B on the complete pipeline. Appendix D carries the full accounting. Measured directly on the full pipeline at 20k atoms/GPU instead of composed, the cache leaves the same 13–16% of the rebuild-every-step cost, row by row in Table 4. Each row’s exchange scheme is tuned for its own curve, so the two rows are not a controlled exchange-scheme comparison. 4.4

Communication Optimization at Scale

What is the communication optimization of Section 3.4 worth once embedded in the complete step? We run a controlled A/B on a single flag, under the 25k-atoms-per-GPU configuration of Section 4.5: the two runs differ only in whether the remote exchange runs on a side stream awaited late or synchronously before the forward pass. Because one flag separates the arms, the difference is attributable to the schedule, and the A/B inherits the end-to-end protocol of Section 4.5. Figure 12(c) shows a result whose shape matters more than its magnitude: at 32 GPUs the flag changes step time by under 1%, and beyond that the time it recovers grows with the device count—the signature of a cost proportional to inter-node traffic, not a fixed overhead. Because overlap acts on exactly the term that grows, its effect on scaling exceeds its per-point size: it removes 52% of the scale-induced increase over the measured range, lifting

22

Table 4. Neighbor-stage cost measured during full-inference-mode runs at 20k atoms per GPU (the timer is the neighbor pipeline, not end-to-end step wall time), the largest per-GPU configuration in this cache-validation sweep that runs without activation recomputation. These runs execute the complete model—forward, autograd forces, torch.compile, TF32, and remote-communication overlap—over a 20-step window after step 0 holding exactly 2 interval-10 rebuilds and 18 hits, so the amortized column is measured, not composed—though that 2:18 split makes it equal to (Rebuild + 9 × Cache hit)/10. Each row’s candidate-exchange scheme is tuned for its own curve, so the two rows are not a controlled exchange-scheme comparison; that remains the neighbor-only matrix of Fig. 11. The load-bearing column is 𝑡 hit /𝑡 rebuild , which licenses composing the amortized cost for Table 3; the last column is the resulting 12.6–16.0% of the rebuild-every-step cost. Configuration

GPUs

Rebuild (ms)

Cache hit (ms)

𝑡 hit /𝑡 rebuild

Amortized (ms)

Amortized/rebuild

Morton + spatial-bin

64 128 256 512

162.8 184.2 204.0 316.2

4.8 5.6 8.0 14.4

0.029 0.031 0.039 0.046

20.6 23.5 27.6 44.6

0.126 0.128 0.135 0.141

Hilbert + multi-box

64 128 256 512

91.5 122.4 182.8 230.0

4.4 5.1 8.3 15.3

0.048 0.041 0.045 0.067

13.1 16.8 25.7 36.8

0.143 0.137 0.141 0.160

512-GPU efficiency from 72.9% to the 84.7% of Section 4.5. Neighbor-construction time stays close between the two arms at every scale, confirming the flag only reschedules the exchange. Panel (d) resolves this by step type: the reduction is roughly 3× larger on cache-hit steps—almost entirely model forward plus the exchange it depends on—than on rebuild steps, which also carry graph construction the forward-pass overlap cannot hide; cache and overlap are therefore complementary, not redundant, and rebuild remains the costlier step type even with overlap on. Nsight timelines corroborate the mechanism qualitatively, showing the layer-wise exchange largely masked behind the local branch, but we do not count them as a second measurement: they are an aggregate, run-level view on a different denominator than the per-step timers the A/B compares. 4.5

End-to-End Weak Scaling

Integrating all modules, we measure the full inference step from 32 to 512 GPUs—PBC expansion, neighbor construction, remote-atom exchange, the E2Former forward pass, autograd forces—with torch.compile, graph capture, TF32 and overlap enabled and the cache at interval 10. We fix the per-GPU atom count and grow the system with the device count at three loads spanning the usable memory range: 10k, 25k, and 60k atoms per GPU; the 60k load exceeds the 80 GB HBM budget under the otherwise-identical settings, so activation recomputation (Section 3.1) is enabled for it alone, trading the moderate latency cost characterized there for the memory headroom this load needs. Every point is the mean over steps 300–399 of a 500-step run: the window excludes the torch.compile first step and spans exactly ten rebuild periods at interval 10, which makes the mean the amortized per-step cost rather than a sample biased by where rebuilds happen to land, and one official run per point means no error bars (protocol in Appendix D). Figure 12(e,f) shows that the integrated step is close to flat under weak scaling, though not flat. Growing 16× in device count and in total system size, the per-step time rises 1.29× at 10k atoms per GPU, 1.18× at 25k, and 1.11× at 60k, for parallel efficiencies at 512 GPUs of 77.3%, 84.7%, and 89.7% respectively. The ordering is the informative part: efficiency improves with the per-GPU load, which is the signature of a fixed residual overhead being amortized against a larger local compute window rather than of a cost that grows with the device count. One caveat on the 60k point specifically: it is the only load carrying activation recomputation, which enlarges the local compute window for a reason unrelated to atom count, so it is not a like-for-like continuation of the 10k→25k trend. The mechanism is cleanest between 10k and 25k, which differ only in load. Aggregate throughput tracks each load’s own linear reference (Fig. 12(g)). The largest configuration is 60k atoms per GPU on 512 GPUs—30.7M atoms at 1.684 s per step, or 18.2M atoms per second. Two mechanisms can be attributed directly at this level, both measured as controlled A/B pairs of the complete pipeline. Overlap’s own contribution is already quantified in Section 4.4; the second is the neighbor cache. Rebuilding the graph every step instead of every tenth costs 1.87× on the full step at 512 GPUs and 60k atoms/GPU over steps 300– 23

(b) Cl–O

(a) K–O

(c) O–O

3

4

2

𝑔 (𝑟 )

2 2

1

1

0

0

0 2

4

6

Distance 𝑟 (Å)

8

2

4

6

Distance 𝑟 (Å) single-GPU (14k)

8

2

4

6

8

Distance 𝑟 (Å)

256 GPUs (10.24M)

Fig. 13. Structural validation of the distributed trajectory. K–O, Cl–O, and O–O radial distribution functions, 256-GPU/10.24Matom run (solid) vs. 14k-atom single-GPU serial reference (dashed, Section 4.6’s protocol at single-GPU scale), same concentration and temperature. First-shell peaks match the 0.08 Å bin width in all three; peak heights agree to about 2% (Cl–O, the widest gap) and 0.4% for O–O, the bulk solvent structure.

399—smaller than the neighbor-only ratio of Section 4.3, as expected: the forward/backward compute that dominates a full step is identical whether or not that step rebuilds, so embedding the stage-only gap in the full step can only shrink it, never widen it. The same A/B at 25k atoms/GPU costs 1.49× over steps 300–399: a lighter step spends proportionally less of itself in the neighbour stage, so the same rebuild cost buys a smaller saving. Both match the design argument: the equal-count cut removes load-balancing stalls, the peer set grows sublinearly in device count, and the overlap hides the residual latency—so the compute-to-communication ratio degrades gracefully rather than collapsing into the communication-bound regime that limits a naive distribution of long-range Transformers. Reaching ns/day—against the ≈ 8 ps/day of Section 4.6’s A100 run—is left to future work: it depends on per-GPU memory headroom and fabric bandwidth at the target size, neither of which this study varies. 4.6

MD Trajectory Validation

Every result above measures throughput and scaling on repeated, independent steps, not whether a live trajectory behaves as designed once atoms move and ownership churns. We close that gap on a separate 256-GPU A100 cluster (40 GB HBM2e, PCIe, no NVLink, swapped in for Section 4.1’s H100/NVLink/InfiniBand fabric for availability): a 150 mmol/L aqueous KCl solution, 10.24M atoms (40k/GPU average), 0.5 fs steps, 20,000 steps of NVT equilibration at 300 K followed by 20,000 of NVE at the same temperature to expose drift a thermostat would mask—10 ps per segment, TF32 off throughout, with cache, sync and skin at Sections 3.3 and 3.5’s defaults (Appendix D). The trajectory rebuilds the neighbor graph every 10 steps and re-partitions every 100, with the 0.5 Å displacement skin triggering an early rebuild past half that margin (0.25 Å). Energy conservation and structure. Over the full 10 ps of NVE the per-atom drift rate holds steady rather than growing, at ≈ 0.42 𝜇eV/atom/ps (6 × 10−9 of total energy, a gauge-dependent ratio), and the final frame shows no broken water molecules and no anomalously close atom pairs. K–O, Cl–O, and O–O radial distribution functions agree with the 14k-atom single-GPU serial reference within the bin width in peak position and to about 2% in peak height (Fig. 13), confirming the distributed path reproduces a correct serial integrator’s physics. Model accuracy is established separately (Section 4.1); these are stability and structure checks on partitioning, live rebuilding, and cross-rank force reduction, not a same-input equivalence test. Cache and driver cost under real displacement. Section 4.5’s nine-in-ten hit rate comes from static repetition—an upper bound, since a real trajectory can also trigger a rebuild by displacement, and a global-sync event forces one outright. The measured hit rate is ≈ 79% against that 90% bound, and the 11-point gap is itself the accounting: a schedule alone rebuilds on one step in ten, so displacement roughly doubles the rebuild rate—the two causes are comparable, not lopsided. Only 1–2% of rank-step records are a rank crossing its own skin (Section 3.3), but that counts which rank initiated a rebuild: the all_reduce(MAX) makes the decision unanimous, so one early crossing pulls every rank into the

24

(a) Strong-scaling step time

(b) Efficiency (ref. 64 GPUs) Efficiency (%)

Step (ms)

4000

1000

100 75 50

500 32

64

128

64

256

128

200

under-utilized

𝜇 s / atom

ms / event

256

(d) Fwd-labeled per atom

(c) Comm. & rebuild cost

100

near-OOM

100 80 60

0 32

64

128

GPUs (log)

256 (c)

repartition

32

64

128

rebuild

GPUs (log)

256

Fig. 14. Strong scaling on a fixed 1.28M-atom KCl-in-water system, NVE, 2,000 steps. (a) Step time, activation recomputation ·64 on (forward-pass detail and recomputation off are in Table 5, Appendix D). (b) Efficiency(𝐾 ) = 𝑇𝑇 (64) (𝐾 ) ·𝐾 × 100%, referenced to 64 GPUs, not 32: that point sits near recomputation’s own memory ceiling, allocator garbage collection inflating its step time. (c) Repartition (rank-0 sort + broadcast) and rebuild cost per event. Amortized, repartition stays under 0.1% of step time; neighbor-cache management overall (rebuild + hit handling) runs 2.2 → 6.8% from 64–256 GPUs—4.6 points against (b)’s 36-point drop. (d) Per-atom forward time (utilization proxy, not profiled)—bottoms at 64 GPUs’ 20k atoms/GPU, rising both ways, consistent with allocator pressure at 32 and too small a batch at 256.

same rebuild, and those induced records carry most of the cost (Appendix D). A global-sync event is a separate, larger bundle: partition, broadcast, and the rebuild it unconditionally forces cost ≈ 87–88 ms for the partition-and-broadcast half (≈ 70.7 ms rank 0 serial quantize-and-sort, ≈ 16–18 ms cross-rank communication) plus ≈ 320–330 ms for the rebuild, 407–418 ms in total—what our prior rank-0 implementation paid every step. The decentralized engine pays it once every 100 steps, roughly 4 ms per step against the ≈ 4.55 s full step (Section 6 scopes this measurement). Strong scaling. Where every scaling result above grew the system with the device count, a separate sweep holds the 1.28M-atom system of Fig. 14 fixed and grows only the count, from 32 to 256 GPUs, so the per-GPU load falls from 40k to 5k atoms. Activation recomputation (Section 3.1) is not an optional latency/memory trade at this end of the range: without it, 32 and 64 GPUs, at 40k and 20k atoms each, exceed this cluster’s 40 GB budget and do not run at all, so every point plotted has it on. Recomputation off is therefore undefined at those two points; where both fit it is 4.7–18.6% faster, flattening at the same knee, and the margin narrows with device count (18.6% at 96 GPUs, 4.7% at 256), so an efficiency figure is only comparable against a shared reference point. Referenced to 64 GPUs rather than 32, for the reason Fig. 14b gives, efficiency declines to 64% at 256 GPUs (61.7% at 224), non-monotonically, the fall steepening past 160 GPUs, where each GPU holds only 8k atoms. Around that knee the returns stall rather than plateau: 192 GPUs is within 1% of 160’s step time, and the step time resumes falling at 224 and 256. Section 4.7 records this as a finding. Only one of the two candidate mechanisms is large enough to matter. Neighbor-cache management grows from 2.2% of wall time at 64 GPUs to 6.8% at 256 and repartition stays under 0.1% amortized, which is 4.6 points against a 36-point drop (Fig. 14c). The deviation sits instead in the forward-labeled phase, which bottoms out per atom at 64 GPUs’ 20k atoms/GPU and rises on both sides of it, following the per-GPU batch size (Fig. 14d). That is consistent with small-batch 25

Table 5. Strong-scaling step time with recomputation on/off, and the forward-pass hit/rebuild split (recomputation on). GPUs

Step, on

Step, off

Fwd, hit

Fwd, rebuild

32 64 96 128 160 192 224 256

3915.1 1404.5 1092.0 829.3 699.4 706.2 650.4 546.3

OOM OOM 888.6 728.5 615.1 598.2 567.6 520.5

3766.7 1369.7 1072.2 827.9 680.0 674.8 591.2 521.9

3776.0 1366.4 1070.9 823.6 711.1 655.8 626.0 532.2

Table 6. Per-step cost under matched hyperparameters on the same A100 cluster, near-uniform KCl vs. heterogeneous Norovirus capsid, over steps 3000–4000. The capsid holds fewer atoms per GPU, so the penalty is not a load artifact. Rows compose as Step ≈ hit rate×(fwd,hit+hit ovh.) + (1−hit rate)×(fwd,rebuild+rebuild ovh.) + residual. KCl (10.24M) Atoms per GPU Step time forward, cache hit forward, rebuild rebuild overhead hit overhead residual (integrate, I/O) Repartition (per event) Cache hit rate

Virus (10.06M)

Diff.

40.0k

39.3k

−1.8%

4552.3 ms 3650.8 ms 3369.0 ms 334.9 ms 33.8 ms 862.7 ms 92.7 ms 79.3%

5212.6 ms 4258.0 ms 4030.5 ms 367.4 ms 33.7 ms 898.2 ms 87.5 ms 79.5%

+14.5% +16.6% +19.6% +9.7% −0.3% +4.1% −5.6% –

under-utilization, though the same timer spans the remote branch and its waits, so it does not isolate a cause. The forward-pass hit/rebuild split of Table 5 shows no consistent ordering there (−35 to +19 ms, with no consistent sign), unlike the hit-slower-than-rebuild gap the next paragraph reports at the much larger per-step scale of the trajectory studies. Density heterogeneity. Every system above is near-uniform in density, bulk water or this KCl solution, which is the easiest case for a cut that equalizes atom counts. We therefore run a matched pair on 256 GPUs, identical in hyperparameters (𝑑𝑡=0.5 fs, rebuild/10 steps, sync/100, 0.5 Å skin) and differing only in the system: the 10.24M-atom KCl solution, and a 10.06M-atom Norovirus capsid in solvent whose protein-dense regions sit against a nearly protein-free water bath (a separate NVT pair, not comparable to the NVE run above). Table 6 splits a steady-state window: the capsid costs 14.5% more per step, almost entirely from forward compute (+17–20% on both paths), while rebuild/hit overhead, hit rate, and the per-event cost of repartition—the one mechanism here that touches every rank—are essentially unchanged. At 40.0k vs. 39.3k atoms per GPU the gap is not a load-size confound. The margin moves with trajectory phase (13–33% across three windows), so read it as a range; every window puts forward compute dominant. Its capsid arm is a 100 ps NVT trajectory, ≈ 8 ps/day on these 256 A100s, that ran with no broken bond, no anomalously close pair, no energy excursion, and only three compiled variants across its 200,000 steps, with no graph break and no fallback to eager (Appendix D). One oddity: cache-hit forward costs 5.6–8.4% more than rebuild forward in both columns, despite skipping the search a rebuild pays for—not a compute effect but cross-rank wait-time drift (Appendix D), and taken up as a finding in Section 4.7. One mechanism is consistent with the capsid’s margin, and scopes Section 3.2’s balance claim: the cut fixes atom count exactly, which is compute balance only insofar as work is per-atom rather than per-edge—and a protein-dense region carries more edges per atom than the water beside it, so equal counts no longer mean equal work, in the term where the cost lands. Node shapes survive untouched, fixed by counts alone; edge shapes and the peer set follow geometry, holding only within a cache epoch. What degrades is the work behind the invariants, not the invariants. An 26

edge-weighted cut would rebalance it, at the cost of the exact-count invariant the captured graph rests on; whether the gap widens with device count is untested (Section 6). 4.7

Negative and Unverified Results

Several of this section’s results are negative or unresolved, and they bound what the positive ones establish. We collect them here rather than leaving them scattered as caveats. Strong scaling stops paying at the knee. At 192 GPUs the step time is 706.2 ms against 160 GPUs’ 699.4 ms—within 1%, and inside the run’s own window-to-window spread—so a 20% device increase buys nothing this run can resolve, at 8k atoms per GPU and beyond. The knee appears in both recomputation modes, and it is the axis on which the timescale gap would have to be closed (Appendix D). A cache hit is slower than a rebuild, in the forward pass. Cache-hit forward costs 5.6–8.4% more than rebuild forward in both columns of Table 6, despite skipping the search a rebuild pays for. The effect is not compute: in the separate 256-GPU profiling run of Appendix D, the model-compute kernel is if anything 1.6% cheaper on hits (1804.3 vs. 1833.4 ms), while the cross-rank NCCL exchange is 724 ms cheaper on rebuild (1235.4 vs. 1959.6 ms), and per-step communication on hit steps spans 735–2451 ms (a 3.3× spread) while compute holds within 1%. The gap is wait time at the attention All-to-All: a rebuild’s ghost-atom exchange re-aligns ranks, and consecutive hit steps carry no equivalent, so timing differences build up until the next exchange waits on whichever rank is slowest. It is a property of the loaded regime, not a universal one—it does not appear at the Table 5 points, where hit is the faster path at four of the eight points (Appendix D). The heterogeneous penalty is a range, not a number. The capsid costs 14.5% more per step in the window of Table 6, but the margin moves with trajectory phase, 13–33% across three windows. The single-window figure is one sample of that range. The mechanism we offer— equal atom counts are not equal edge counts—is consistent with the margin rather than isolated by it, and whether the gap widens with device count is untested. TF32 and ion–ion structure are unresolved. In the 1.28M-atom, ≈ 201 K TF32 A/B of Appendix D, TF32 leaves water and solvation structure alone, with first-peak heights agreeing to 0.5% (K–O), 0.2% (Cl–O) and 0.0% (O–O), and coordination numbers to 0.2%. The three ion–ion first peaks differ by −3.2% (K–Cl), +23.8% (K–K) and +11.6% (Cl–Cl), but splitting the TF32-off run alone into halves moves the same peaks by 10.3%, 14.1% and 12.5%, so only K–K clears the noise floor, and only by 1.7×. The effective sample behind a K–K peak is tens of pairs, not the 103 frames suggest, because ion diffusion decorrelates far more slowly than the 5 ps window. Resolving whether TF32 touches ion–ion structure needs a trajectory long enough to decorrelate the ions, not a longer averaging window over the same one (Appendix D). The 60k weak-scaling point is not like-for-like. It is the only load carrying activation recomputation, which enlarges the local compute window for a reason unrelated to atom count, so its 89.7% efficiency is not a continuation of the 10k→25k trend; the mechanism is cleanest between those two loads, which differ only in load. There is no external baseline. Every comparison in this section is against another configuration of our own system. The closest message-passing implementation stops at 32 GPUs (Section 5), so it cannot supply an end-to-end comparator at these scales, and the kernel comparison of Appendix B is against a vanilla sparse baseline and dense SDPA rather than another production MD engine. The speedups we report are ablation and composition numbers, not a claim of superiority over an external system. Read together, these results mark where the evidence stops: the cache’s composed 15.7× is an upper estimate, one official run per point leaves no error bars, and differences of a few percent—rank-box’s 3% cost against all-gather, the ±35 ms forward split of Table 5—are not separable by these measurements. 5

RELATED WORK

Prior work partitions a large graph-like state and hides the communication. We inherit more from them than we replace: space-filling-curve decomposition is long-standing practice [37], skin-plus-displacement neighbor reuse dates to Verlet [44], and distributing an equivariant potential has direct precedent. SevenNet [29] spatially decomposes a message-passing potential for MD across GPUs, Allegro’s scale-out [26] runs a deep equivariant model at biomolecular 27

Table 7. Where the design sits against systems cited in this paper, with the cells confined to what those citations and our own results establish. “Model class” separates potentials whose energy depends on one cutoff ball from those that propagate features across layers. “Per-layer exch.” asks whether a distributed step synchronizes remote features once per layer, as propagation requires, rather than once per step as a ghost-atom halo does. Exact forces means forces are the gradient of the evaluated energy rather than a separately regressed output. Live-graph reuse means the neighbor topology, or the exchange plan built from it, survives across steps of a moving trajectory. Static-shape capture means a step is replayed under a captured graph whose allocations and launch geometry are fixed. An em dash marks an item this paper could not establish from the cited work; no cell is a ranking, the systems differ in task, model, and hardware, so the columns are not additive.

Model class

Per-layer exch.

Max GPUs

Max atoms

Exact forces

Live-graph reuse

Static-shape capture

DPMD [17, 50] LAMMPS [32, 41] Allegro scale-out [26] SevenNet [29] Xia et al. [48] NequIP rebuild [38] MACE [3] UMA [47]

strictly local strictly local strictly local message passing message passing message passing message passing message passing

no no no yes yes — — —

— — — 32 8 — — 1

100M — — — 8 10 (CPU) — — —

yes yes yes yes — — yes —

— yes — — — — — —

— no — — — — — —

This work

message passing

yes

512

30.7M

yes

yes

yes

System

scale, and a NequIP-framework rebuild [38] pairs compiler integration with custom Triton and vendor equivariant kernels—the same two levers as our Section 3.1, aimed at data-parallel training rather than at splitting one system across ranks. MACE takes the algorithmic route instead, using two high-body-order message-passing layers, which also bounds its receptive field [3]; we keep the depth and pay in the runtime. None of these cited systems combines propagation with scale: strictly-local models decompose without propagation; the message-passing ones stop well short of hundreds of GPUs—32 GPUs at over 80% weak-scaling efficiency on silica in the closest scheme [29], 8 in a recent MPI design whose 108 -atom run is on CPUs [48]—where holding shapes, plans, and captured execution invariant as the graph moves becomes binding. Table 7 makes the shape of the gap concrete without scoring anyone. Two columns do most of the work. Per-layer boundary exchange is a property of the model class, and it is what turns a distributed step from one halo update into a synchronization per Transformer block; the systems marked “no” are distributed but either forgo propagation or replicate the whole system on every rank, and both choices avoid the problem rather than solve it, while a dash marks a system that is not distributed in this sense at all. Live-graph reuse and static-shape capture are properties of the runtime, and they are the two that a moving molecular graph actively attacks: a training system partitions a fixed graph once, and a strictly local MD code can rebuild a halo every step because the halo is a handful of scalars per ghost atom rather than an equivariant payload. The combination in the last row—propagation, exact gradients, reuse, and capture, at hundreds of GPUs—is the one no cited system reports, which is the scoped form of our claim. Spatial domain decomposition in classical molecular dynamics. Spatial domain decomposition with ghostatom communication has been standard in classical MD for three decades [32, 41, 1, 31], routing only boundary updates between adjacent domains, and neutral-territory schemes choose where each interaction is evaluated to minimize imported volume [7, 6], a decision Anton commits to silicon [36]. Our equal-count space-filling-curve partition (Section 3.2) descends from that line, since all bound communication by geometry, not device count, but optimizes for a per-rank shape that does not move between steps, not the smallest import set. Classical short-range force fields are cheap per atom and tolerant of a domain’s count fluctuations; nothing there needs shapes bit-identical across steps, as torch.compile guards and a captured CUDA Graph do—and compiling for dynamic shapes [2] removes the re-trace, not the capture, since a static graph replays fixed allocations and launch parameters. That is why we insist on equal-count, and why holding it under MD dynamics is part of the contribution. Those codes also exchange a handful of scalars per ghost atom, where ours carries a projected equivariant SKV payload (Section 3.4). The difference is worth stating in the currency of the two costs. A ghost-atom halo is sized by the cutoff and the domain surface, and its per-step volume is small because the imported object is a position and a species tag; a per-layer feature exchange is sized by the receptive field of one layer and the feature width, and it recurs once per layer in each pass. The former can be recomputed from scratch every step at negligible cost, so classical codes can rebuild their halo 28

every step and need no separate topology/geometry split, even when they do reuse topology. The latter cannot be rebuilt cheaply, which is why the split of Section 3.3 is load-bearing rather than an optimization—and why its displacement bound must make staleness impossible rather than small, where parameter training can absorb a stale neighborhood as noise but a force evaluation integrates it into the trajectory. Distributed GNN training. DistDGL [51] and ROC [18] partition large graphs across machines—METIS-style and through an offline learned cost model, respectively—then pipeline cross-partition communication for mini-batch GNN training. Ours also splits a graph across devices with boundary message passing, but three properties differ. Atoms carry a 3D embedding, so a geometry-aware curve suffices where a generic sparse graph needs METIS-style cuts. GNN training partitions a fixed graph once and sub-samples from it; ours is re-derived every step from current coordinates, with the entire simulation state as the “batch.” That is why a plan built once and reused for a partition’s lifetime (Section 3.4) is worth the engineering, where a mini-batch system’s routing changes every batch. And both backpropagate, but over different objects: GNN training differentiates a loss with respect to parameters, where stale neighborhoods are an accepted trade—PipeGCN [45] buys its overlap by pipelining one-epoch-stale boundary features— whereas conservative forces are the analytic gradient with respect to coordinates, taken over the exact neighborhood every step, so the exchange recurs in the backward pass and needs the pre/post pair to stay overlapped under reverse traversal (Section 3.4). Communication–computation overlap in large-scale transformer training. Overlapping collectives with compute is standard in LLM training [27, 33], and systems work pushes that overlap below the operator boundary, into kernel fusion, tile-level triggering, and hierarchical collective scheduling [16, 30, 8]. We follow the same principle (Section 3.4) but hide a different collective: those methods overlap dense collectives of rank-uniform size fixed for training’s duration, where ours is sparse, variable-length, per-edge, and re-derived from the evolving neighbor graph. Mixture-of-Experts routing is the closest analogue, variable-length at runtime [12, 15], but it is learned and so negotiable— tokens rebalanced or dropped under a capacity factor—where ours is fixed by geometry and cannot drop an edge without changing the force. That is why a reusable plan, sender-side SKV fusion, and the asymmetric pre/post operators earn their place here, where a fixed-size all-reduce needs none of them. The bound itself is the same arithmetic in all three settings—the smaller concurrent phase is what disappears—but its value differs: here the exchange is sized by a subdomain’s boundary and the local branch by its interior, so the hidden fraction improves with per-rank load rather than degrading with the machine. 6

DISCUSSION AND OUTLOOK

Scaling beyond the present deployment. The per-step path is not tied to our 512-GPU scale: per-GPU cost follows the local problem size—fixed atom count, constant-size scalar reductions (Sections 3.2, 3.5)—rather than device count. Peer degree is the exception in kind, though not in magnitude: it rises from 6 to 19 over 64–512 GPUs at 25k atoms/GPU, near 𝐾 0.5 , so it grows with the machine far more slowly than the 𝐾 an all-gather contacts. No weak-scaling ceiling is reached within the measured range: efficiency stays above 77% through 512 GPUs, though it declines with 𝐾—the 𝐾 ×𝐾 request matrix and the world reductions’ latency still grow with 𝐾, the rank-0 sort with 𝑁 , and the metadata all-gather remains the next target the telemetry names (Section 4.3). Saying the design extends in principle toward billion-atom systems is a statement about the exchange, not about every term in the step. 6.1

Threats to Validity

Internal. Every scaling point is one official run, so no error bars are drawn; where a gap is smaller than the run-torun spread we do not treat it as separable. The neighbor pipeline’s sub-stage timers are MAX-reduced across ranks independently, and different sub-stages’ maxima may come from different ranks, so the pipeline is not decomposed additively—why the attribution of spatial-bin’s cost to plan overhead is described as consistent with the telemetry rather than as a measured stage decomposition (Section 3.3), and why the forward-pass split is not read as a stage breakdown. The composed neighbor-search figure carries a further caveat: its cached rows combine a rebuild measured on a separate interval-1 run with a hit measured on an interval-10 run, and the interval-1 arm searches at each layer’s bare cutoff while an interval-10 rebuild builds the wider 𝑟𝑐 + 𝛿 superset, so the composed 15.7× is an upper estimate 29

rather than a measured speedup (Section 4). The figures the argument leans on—the directly measured hit latencies and the full-step 1.87× A/B—do not depend on it. External. Every system measured is near-uniform in density, bulk water or the KCl solution, the easiest case for a cut that equalizes atom counts; the one heterogeneous system, the Norovirus capsid, is a single device-count pair, and we probe neither sharper density contrasts nor anisotropic cells, where AABBs could inflate past equal-count’s unfixed peer bound. One model is measured throughout—a UBio-MolFM variant built on E2Former V2, at 23.5M parameters (Section 4.1)—so the claim is about this class of equivariant Transformer. And throughput and weak scaling come from an H100/InfiniBand cluster, the trajectory and strong-scaling results from an A100/PCIe one, so numbers from the two are not interchangeable. Construct. The central invariant is atom-count balance, and atom count is not work: work is per-edge, so a proteindense region carries more edges per atom than the water beside it. The capsid pair is consistent with this account, its 13–33% per-step margin pointing to edges per atom rather than atoms (Section 4.6); node shapes survive untouched, fixed by counts alone, while edge shapes and the peer set follow geometry and hold only within a cache epoch, so what degrades is the work behind the invariants. The exact-search backend also clamps to FAISS’s 𝑘-selection limit of 2048 wherever the per-query candidate set is large enough to saturate it—a function of local density and cutoff rather than of device count—not necessarily truncation, but the range’s two ends may not run an identical search, so the edge set is the complete radius graph only where that selection is unsaturated. Limitations that remain. Precision: keeping only precision-sensitive contractions in higher precision could recover TF32’s throughput without its drift; we have not built that split. Cadence: two fixed clocks set the schedule—re-partition every 100 steps, neighbor rebuild every 10—where a displacement-driven trigger would act only when locality degrades. Faster dynamics push the rebuild clock toward every step, an A/B measured at 1.87× on the full step (Section 4.5), not collapse. The re-partition clock is a separate cost: its partition-and-broadcast half runs ≈ 87–88 ms per event, amortized over 100 steps here. No external baseline: we compare only our own configurations, since the closest message-passing implementation stops at 32 GPUs (Section 5). Timescale: the measured ≈ 8 ps/day covers the checks reported here, not the nanosecond windows a binding study needs; closing it is strong scaling at fixed atom count, the axis on which returns go flat past 160 GPUs on the 1.28M-atom strong-scaling system. The drift and structure check runs on a 10.24M-atom trajectory, 3× under the 30.7M scaling run, and one trajectory is not a replicate; the ps/day figure comes from the 10.06M-atom capsid run. 6.2

Complexity Accounting

Table 8 lists what still grows per step. Four terms grow with the machine—peer count, the 𝐾 ×𝐾 request matrix, the metadata all-gather, and world-collective latency—one with the problem, the rank-0 sort, and one with the local density, per-GPU edges, which weak scaling holds fixed at constant atoms per GPU. None is removed by the design. Their sum grows slowly enough that measured weak-scaling efficiency stays between 77% and 90% through 512 GPUs—77.3% and 84.7% at the two loads that run without activation recomputation, and 89.7% at the 60k load, which carries it and so is not a like-for-like continuation—because each term is small against the local compute window at these sizes. Whether that holds at a billion atoms, about 33× the 30.7M measured here, and the correspondingly larger 𝐾, is an extrapolation, which is why we say “extends in principle” rather than predicting an efficiency. The 𝐾-dependent terms are the ones the telemetry already names. The request matrix is 𝐾 2 only nominally: SFC locality keeps each row sparse, but the test is still a loop over all 𝐾 descriptors, and the metadata all-gather feeding it is the term still rising with 𝐾 at the top of our range (Section 4.3). The world collectives are constant in payload but not in latency—a barrier over 𝐾 ranks has a floor that does not shrink with per-GPU work. The rank-0 sort is 𝑂 (𝑁 log 𝑁 ) in the full coordinate set, measured at ≈ 70.7 ms at 10.24M atoms and amortized over 100 steps; a distributed sort would remove it rather than amortize it, and we have not built one. 6.3

What Generalizes, and What Does Not

The placement mechanism is not specific to this model: its three guarantees follow from the placement argument, and apply to any potential that exchanges boundary features per layer over a live graph. The requirements are stated in 30

Table 8. What still grows per step, and in what. Weak scaling holds per-GPU atom count fixed, so the question is which terms grow while the local forward–backward work does not. The last column gives what the design does about each term; none is removed outright. The table scopes the extension argument rather than supporting it. Term

Grows in

What the design does

Per-GPU edges Peer count 𝐾 × 𝐾 request matrix Metadata all-gather World-collective latency Rank-0 𝑂 (𝑁 log 𝑁 ) sort

local density geometry, 6 → 19 over 64–512 GPUs 𝐾 2 nominally 𝐾 log 𝐾 𝑁

𝑂 (1) in 𝐾 at fixed atoms/GPU sublinear in 𝐾, roughly 𝐾 0.5 sparse rows, coarse AABBs small fixed descriptor (< 2 KiB) constant payload per step amortized over 100 steps

Section 2.4 without reference to E2Former, and nothing in the cut, the cache invariant, or the plan-reuse accounting depends on the particular tensor product or attention form. What is bound to our stack is worth naming. The kernel work of Section 3.1 assumes sparse equivariant attention with 𝑙 max = 2 and a scalar-per-head score, though the CSR/CSC duality argument applies to any two-pass sparse operator. The payload arithmetic of Section 3.4 uses the fact that only the Value must travel equivariant, which holds for attention-style aggregation but not for an architecture that ships full features between layers. And the empirical support for the generalization is itself stack-bound, and splits into measurement and derivation: the 6 → 19 peer degree over 64–512 GPUs and the 1.87× cache A/B are measured, while the 80% local-fraction figure is a bound derived from the box covariogram, falling to ≈ 78% at the 8k-atom/GPU strong-scaling points; all are properties of 𝜌 = 0.1 atoms Å−3 , this cutoff stack, and this range of device counts, and the formulas behind them are not. 7

CONCLUSION

This paper built a runtime for message-passing equivariant molecular dynamics at tens-of-millions-of-atom scale, treating the model and the distribution as one object of design. Its foundation is a placement mechanism: quantize coordinates onto a space-filling curve, sort, and cut the order into equal-count contiguous shards. That single decision supplies exact per-rank atom counts and with them invariant tensor shapes, a peer set that grows as roughly 𝐾 0.5 rather than 𝐾, and a deterministic ownership map, so each cache epoch’s whole forward–backward pass replays as a static CUDA Graph. On it rest custom Triton kernels with dual CSR/CSC indexing and a unified MCH layout, a conservative three-level neighbor search with exact cross-step topology reuse, an All-to-All reduced to its payload and overlapped in both passes, and a decentralized on-GPU MD engine. The system sustains 30.7M atoms on 512 GPUs at 89.7% weak-scaling efficiency (77–90% across three per-GPU loads), with graph reuse cutting full-step time 1.87× and neighbor-search cost growth over an 8× device-count increase held to 1.7× rather than 3.0×. On a live 10.24M-atom trajectory the same code path holds energy drift to 0.42 𝜇eV/atom/ps and reproduces the serial reference’s structure; a 10.06M-atom Norovirus capsid locates the cut’s limit in edges per atom. What remains is scale and scope. The exchange grows only slowly with the machine, but it does grow, and the request matrix, the metadata all-gather, and world-collective latency still scale with it, while the rank-0 sort scales with the problem; a distributed sort and a sparser request structure are the next terms to remove. Beyond that, the timescale a binding study needs is a strong-scaling problem at fixed atom count, and the placement mechanism is stated without reference to this model, so any potential exchanging boundary features per layer over a live graph is the natural next object to carry it. CONTRIBUTORS • Neural Network Architecture Optimization: Lin Huang, Jia Zhang, Arthur Jiang • Core Operator Optimization: Jiajun Cheng, Lin Huang, Zion Wang, Junping Zhao • Communication Optimization: Hao Yin, Hao Li, Zion Wang, Ji Zhang, Junping Zhao • MD Engine Development: Jia Zhang 31

• Experimental Design and Analysis: Zion Wang, Hao Yin, Hao Li, Lin Huang, Jia Zhang • Computing Cluster Support: Ji Zhang, Lucas Li, Allex Kan • Project Management and Support: Owen Fu, Xunxing Chen • Technical Report Writing: All authors REFERENCES [1] Mark James Abraham, Teemu Murtola, Roland Schulz, Szilárd Páll, Jeremy C Smith, Berk Hess, and Erik Lindahl. GROMACS: High performance molecular simulations through multi-level parallelism from laptops to supercomputers. SoftwareX, 1:19–25, 2015. [2] Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky, Bin Bao, Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary DeVito, Elias Ellison, Will Feng, Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalambarkar, Laurent Kirsch, Michael Lazos, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, C. K. Luk, Bert Maher, Yunjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Michael Suo, Phil Tillet, Eikan Wang, Xiaodong Wang, William Wen, Shunting Zhang, Xu Zhao, Keren Zhou, Richard Zou, Ajit Mathews, Gregory Chanan, Peng Wu, and Soumith Chintala. PyTorch 2: Faster machine learning through dynamic Python bytecode transformation and graph compilation. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2024. [3] Ilyes Batatia, David P Kovacs, Gregor Simm, Christoph Ortner, and Gabor Csanyi. Mace: Higher order equivariant message passing neural networks for fast and accurate force fields. Advances in Neural Information Processing Systems, 35:11423–11436, 2022. [4] Simon Batzner, Albert Musaelian, Lixin Sun, Mario Mario, Anders Johansson, and Boris Kozinsky. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022. [5] Jörg Behler and Michele Parrinello. Generalized neural-network representation of high-dimensional potential-energy surfaces. Physical review letters, 98(14):146401, 2007. [6] Kevin J. Bowers, Edmond Chow, Huafeng Xu, Ron O. Dror, Michael P. Eastwood, Brent A. Gregersen, John L. Klepeis, Istvan Kolossvary, Mark A. Moraes, Federico D. Sacerdoti, John K. Salmon, Yibing Shan, and David E. Shaw. Scalable algorithms for molecular dynamics simulations on commodity clusters. In Proceedings of the ACM/IEEE Conference on Supercomputing (SC), 2006. [7] Kevin J. Bowers, Ron O. Dror, and David E. Shaw. The midpoint method for parallelization of particle simulations. The Journal of Chemical Physics, 124(18):184109, 2006. [8] Chang Chen, Xiuhong Li, Qianchao Zhu, Jiangfei Duan, Peng Sun, Xingcheng Zhang, and Chao Yang. Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2024. [9] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-efficient exact attention with IO-awareness. In Advances in Neural Information Processing Systems (NeurIPS), 2022. [10] Mario Geiger and Tess Smidt. e3nn: Euclidean neural networks. arXiv preprint arXiv:2207.09453, 2022. [11] Andrea Grisafi and Michele Ceriotti. Incorporating long-range physics in atomic-scale machine learning. The Journal of Chemical Physics, 151(20):204105, 2019. [12] Jiaao He, Jidong Zhai, Tiago Antunes, Haojie Wang, Fuwen Luo, Shangfeng Shi, and Qin Li. FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming (PPoPP), 2022. [13] Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, and Jia Zhang. E2former-v2: On-the-fly equivariant attention with linear activation memory. arXiv preprint arXiv:2601.16622, 2026. [14] Lin Huang, Frank Peng, JiaJun Cheng, Zion Wang, Hao Yin, Hao Li, Ji Zhang, Jack Jia, Junping Zhao, Arthur Jiang, and Jia Zhang. Ubio-molfm: Enabling biomolecular dynamics at dft accuracy and 105 atoms with one untuned potential. arXiv preprint arXiv:2608.18623, 2026. [15] Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, Joe Chau, Peng Cheng, Fan Yang, Mao Yang, and Yongqiang Xiong. Tutel: Adaptive mixture-of-experts at scale. In Proceedings of Machine Learning and Systems (MLSys), 2023. [16] Abhinav Jangda, Jun Huang, Guodong Liu, Amir Hossein Nodehi Sabet, Saeed Maleki, Youshan Miao, Madanlal Musuvathi, Todd Mytkowicz, and Olli Saarikivi. Breaking the computation and communication abstraction barrier in distributed machine learning workloads. In Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), pages 402–416, 2022. [17] Weile Jia, Han Wang, Mohan Chen, Denghui Lu, Lin Lin, Roberto Car, E Weinan, and Linfeng Zhang. Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning. In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–14. IEEE, 2020. [18] Zhihao Jia, Sina Lin, Mingyu Gao, Matei Zaharia, and Alex Aiken. Improving the accuracy, scalability, and performance of graph neural networks with roc. In Proceedings of Machine Learning and Systems (MLSys), volume 2, pages 187–198, 2020. [19] Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data, 7(3):535–547, 2019. [20] Tsz Wai Ko, Jonas A Finkler, Stefan Goedecker, and Jörg Behler. A fourth-generation high-dimensional neural network potential with accurate electrostatics including non-local charge transfer. Nature communications, 12(1):398, 2021. [21] Walter Kohn and Lu Jeu Sham. Self-consistent equations including exchange and correlation effects. Physical review, 140(4A):A1133, 1965. [22] Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro. Reducing activation recomputation in large transformer models. In Proceedings of Machine Learning and Systems (MLSys), volume 5, pages 341–353, 2023. [23] Yunyang Li, Lin Huang, Zhihao Ding, Xinran Wei, Chu Wang, Han Yang, Zun Wang, Chang Liu, Yu Shi, Peiran Jin, Tao Qin, Mark Gerstein, and Jia Zhang. E2former: An efficient and equivariant transformer with linear-scaling tensor products. In Advances in Neural Information Processing Systems, 2025. [24] Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. arXiv preprint arXiv:2206.11990, 2022.

32

[25] Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Anders S Christensen, and Boris Kozinsky. Learning local equivariant representations for large-scale atomistic dynamics. Nature Communications, 14(1):579, 2023. [26] Albert Musaelian, Anders Johansson, Simon Batzner, and Boris Kozinsky. Scaling the leading accuracy of deep equivariant models to biomolecular simulations of realistic size. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), 2023. [27] Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on GPU clusters using Megatron-LM. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), pages 1–15, 2021. [28] NVIDIA. Cuda graphs. https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html, 2023. [29] Yutack Park, Jaesun Kim, Seungwoo Hwang, and Seungwu Han. Scalable parallel algorithm for graph neural network interatomic potentials in molecular dynamics simulations. Journal of Chemical Theory and Computation, 20(11):4857–4868, 2024. [30] Suchita Pati, Shaizeen Aga, Mahzabeen Islam, Nuwan Jayasena, and Matthew D. Sinclair. T3: Transparent tracking and triggering for fine-grained overlap of compute and collectives. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems (ASPLOS), 2024. [31] James C. Phillips, Rosemary Braun, Wei Wang, James Gumbart, Emad Tajkhorshid, Elizabeth Villa, Christophe Chipot, Robert D. Skeel, Laxmikant Kalé, and Klaus Schulten. Scalable molecular dynamics with NAMD. Journal of Computational Chemistry, 26(16):1781–1802, 2005. [32] Steve Plimpton. Fast parallel algorithms for short-range molecular dynamics. Journal of Computational Physics, 117(1):1–19, 1995. [33] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. ZeRO: Memory optimizations toward training trillion parameter models. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC), pages 1–16. IEEE, 2020. [34] Kristof Schütt, Oliver Unke, and Michael Gastegger. Equivariant message passing for the prediction of tensorial properties and molecular spectra. International Conference on Machine Learning, pages 9377–9388, 2021. [35] Kristof T Schütt, Pieter-Jan Kindermans, Felix E Sauceda, Stefan Chmiela, Alexandre Tkatchenko, and Klaus-Robert Müller. Schnet: A continuousfilter convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30, 2017. [36] David E. Shaw, Martin M. Deneroff, Ron O. Dror, Jeffrey S. Kuskin, Richard H. Larson, John K. Salmon, Cliff Young, Brannon Batson, Kevin J. Bowers, Jack C. Chao, Michael P. Eastwood, Joseph Gagliardo, J. P. Grossman, C. Richard Ho, Douglas J. Ierardi, István Kolossváry, John L. Klepeis, Timothy Layman, Christine McLeavey, Mark A. Moraes, Rolf Mueller, Edward C. Priest, Yibing Shan, Jochen Spengler, Michael Theobald, Brian Towles, and Stanley C. Wang. Anton, a special-purpose machine for molecular dynamics simulation. In Proceedings of the 34th Annual International Symposium on Computer Architecture (ISCA), pages 1–12, 2007. [37] Volker Springel. The cosmological simulation code GADGET-2. Monthly Notices of the Royal Astronomical Society, 364(4):1105–1134, 2005. [38] Chuin Wei Tan, Marc L. Descoteaux, Mit Kotak, Gabriel de Miranda Nascimento, Seán R. Kavanagh, Laura Zichi, Menghang Wang, Aadit Saluja, Yizhong R. Hu, Tess Smidt, Anders Johansson, William C. Witt, Boris Kozinsky, and Albert Musaelian. High-performance training and inference for deep equivariant interatomic potentials. Digital Discovery, 5(4):1558–1567, 2026. [39] Philipp Thölke and Gianni De Fabritiis. Equivariant transformers for neural network based molecular potentials. In International Conference on Learning Representations (ICLR), 2022. [40] Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds. arXiv preprint arXiv:1802.08219, 2018. [41] Aidan P Thompson, H Metin Aktulga, Richard Berger, Dan S Bolintineanu, W Michael Brown, Paul S Crozier, Pieter J in’t Veld, Axel Kohlmeyer, Stan G Moore, Trung Dac Nguyen, Ray Shan, Mark J Stevens, Julien Tranchida, Christian Trott, and Steven J Plimpton. LAMMPS – a flexible simulation tool for particle-based materials modeling at the atomic, meso, and continuum scales. Computer Physics Communications, 271:108171, 2022. [42] Philippe Tillet, H. T. Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages (MAPL), pages 10–19, 2019. [43] Oliver T Unke, Stefan Chmiela, Felix E Sauceda, Michael Gastegger, Isaac Ströhr, and Alexandre Tkatchenko. Machine learning force fields. Chemical reviews, 121(16):10142–10186, 2021. [44] Loup Verlet. Computer “experiments” on classical fluids. i. thermodynamical properties of lennard-jones molecules. Physical Review, 159(1):98–103, 1967. [45] Cheng Wan, Youjie Li, Cameron R. Wolfe, Anastasios Kyrillidis, Nam Sung Kim, and Yingyan Lin. PipeGCN: Efficient full-graph training of graph convolutional networks with pipelined feature communication. In International Conference on Learning Representations (ICLR), 2022. [46] Chu Wang, Lin Huang, Xinran Wei, Tao Qin, Arthur Jiang, Lixue Cheng, and Jia Zhang. Scalable machine learning force fields for macromolecular systems through long-range aware message passing. arXiv preprint arXiv:2601.03774, 2026. [47] Brandon M. Wood, Misko Dzamba, Xiang Fu, Meng Gao, Muhammed Shuaibi, Luis Barroso-Luque, Kareem Abdelmaqsoud, Vahe Gharakhanyan, John R. Kitchin, Daniel S. Levine, Kyle Michel, Anuroop Sriram, Taco Cohen, Abhishek Das, Ammar Rizvi, Sushree Jagriti Sahoo, Zachary Ulissi, and C. Lawrence Zitnick. UMA: A family of universal models for atoms. In Advances in Neural Information Processing Systems (NeurIPS), 2025. [48] Junfan Xia and Bin Jiang. Efficient parallelization of message passing neural network potentials for large-scale molecular dynamics. arXiv preprint arXiv:2505.06711, 2025. [49] Shiwei Yue, Maria Carolina Muniz, Marcos F Calegari Andrade, Linfeng Zhang, Roberto Car, and Athanassios Z Panagiotopoulos. When do short-range atomistic machine-learning models fall short? The Journal of Chemical Physics, 154(3):034111, 2021. [50] Linfeng Zhang, Jiequn Han, Han Wang, Roberto Car, and E Weinan. Deep potential molecular dynamics: a scalable model with the accuracy of quantum mechanics. Physical Review Letters, 120(14):143001, 2018. [51] Da Zheng, Chao Ma, Minjie Wang, Jinjing Zhou, Qidong Su, Xiang Song, Quan Gan, Zheng Zhang, and George Karypis. DistDGL: Distributed graph neural network training for billion-scale graphs. In 2020 IEEE/ACM Workshop on Irregular Applications: Architectures and Algorithms (IA3), pages 36–44. IEEE, 2020.

33

The appendices are organized one topic per letter. Appendix A records the exact model configuration behind every reported number. Appendices B–C then take the mechanisms of Section 3 one at a time—single-GPU kernels, then the implementation detail behind the partition, the neighbor search and the communication path—each supplying what the main text defers. Appendix D gives the measurement protocol, the full-inference validation, and the per-run telemetry behind Section 4, and Appendix E records the hardware and software behind every reported timing. Nothing here is needed to follow the main argument; each appendix is for a reader who wants one mechanism spelled out in full or one number reproduced. A

MODEL AND HYPERPARAMETERS

Table 9 lists the backbone hyperparameters of the E2Former V2 model [23, 13] used throughout, with the exact key/value pairs of our run configuration (keys shown are under backbone_config). The stack is a 4-layer backbone—three shortrange equivariant layers at a 5 Å cutoff followed by one long-range layer at 8 Å, for the 23 Å per-atom receptive field of Section 4. Node features carry channels up to degree 𝑙 max = 2 (irreps 256x0e+256x1e+256x2e, i.e. 𝐶 = 256 channels per degree), attention uses 128 heads with an 8-dimensional scalar head, and every run uses the Triton tensor-product path of Section 3.1, in TF32 mixed precision except where stated otherwise (the trajectory validation of Section 4.6 runs with TF32 off). The total parameter count is 23.5M. Model / Backbone setting

Symbol

Value

Backbone # layers Cutoff radius Clustering Node embedding irreps Head irreps Scalar head dim # attention heads Radial basis Neighbor cap Attention dropout Basis type Normalization Attention type Tensor product impl. Edge embedding FFN type Encoder

— 𝑅 — — 𝐶, 𝑙 max — 𝑑𝑘 𝐻 — — — — — — — — — —

e2former (E2Former V2) num_layers=4 (3 short-range + 1 long-range) short-range 5 Å (×3 layers), long-range 8 Å (×1) with_cluster=False "256x0e+256x1e+256x2e" (𝐶 = 256, 𝑙 max = 2) "2x0e+2x1e+2x2e" attn_scalar_head=8 num_attn_heads=128 number_of_basis=256 none (all in-cutoff pairs retained) alpha_drop=0 gaussiansmear rms_norm_sh so2-zero-order tp_type=QK+triton default s3 default

Table 9. E2Former V2 backbone hyperparameters used for all runs in this paper. Keys are configuration fields under backbone_config; the Symbol column gives the notation used for each in this appendix (𝐶 channels, 𝐻 attention heads, 𝑙 max maximum degree, 𝑑𝑘 scalar head dimension, 𝑅 Transformer blocks).

B

SINGLE-GPU KERNEL IMPLEMENTATION

This appendix details the kernel-level design summarized in Section 3.1. Triton Kernel Design. Two of the four axes of Section 3.1 are standard practice and detailed here rather than in the main text. The Clebsch–Gordan tensor products these kernels accelerate are the primitive introduced by Tensor Field Networks [40] and implemented generically in e3nn [10]; our kernels target the specific shapes and access patterns E2Former V2’s attention imposes on that primitive, not a general replacement for the library. Hardware-aware compute saturation. Autotuning explores the tile-size space (e.g., 𝐵𝐿𝑂𝐶𝐾_𝐻, 𝐵𝐿𝑂𝐶𝐾_𝐾) under the shared-memory constraint 𝐵𝐻 · 𝐵𝐾 + 𝐵𝐻 · 𝐵𝐷 + 𝐵𝐾 · 𝐵𝐷 ≤ 64,000; the QK kernels use 2D tiled parallelism with a 3-stage pipeline (𝑛𝑢𝑚_𝑠𝑡𝑎𝑔𝑒𝑠 = 3) to overlap FMA with irregular gathers.

34

Table 10. Exact throughput and latency of the attention kernels of Fig. 15. Each cell reports effective throughput in TFLOPS, with kernel latency in milliseconds in parentheses. The dense SDPA result is a hardware-oriented reference point rather than a deployable alternative for the sparse molecular workload. The sweep runs at 16 heads to isolate the 𝑁 and 𝐷 axes; production uses 128 (Appendix A), so these are kernel-level comparisons rather than production-shape measurements. 𝑁

head

𝐷

Sparse-Triton

Sparse-Vanilla

SDPA-Mask

SDPA-Dense

512 2048 8192 32768

16 16 16 16

8 8 8 8

0.097 (0.161) 0.192 (0.325) 0.605 (0.412) 0.692 (1.439)

0.105 (0.148) 0.141 (0.440) 0.157 (1.583) 0.163 (6.130)

0.164 (0.095) 0.088 (0.708) 0.026 (9.630) 0.007 (150.934)

2.178 (0.071) 4.028 (0.617) 5.419 (7.331) 5.783 (109.918)

512 2048 8192 32768

16 16 16 16

64 64 64 64

0.648 (0.166) 1.792 (0.240) 2.015 (0.852) 1.745 (3.936)

0.193 (0.557) 0.220 (1.955) 0.211 (8.148) 0.205 (33.472)

0.974 (0.110) 0.528 (0.813) 0.155 (11.105) 0.035 (197.232)

14.045 (0.078) 25.491 (0.687) 33.145 (8.455) 33.726 (132.952)

Deep operator fusion. Single-pass kernels fuse sparse gathers, dot products, and gated multiplications—including Í 𝑠ℎ ⊗ 𝑓 𝑒𝑎—without materializ-

the SpMV-based key-gradient backward and the outer-product reduction 𝑜𝑢𝑡 =

ing intermediates, minimizing HBM round-trips. All loads use 2D/3D tiling with stride-based indexing for memory coalescing. Allocator Policy. The recomputation interface takes a retention fraction in [0, 1]: 0 is full recomputation (classical gradient checkpointing) and 1 retains everything, so the compiler rather than the author chooses where to trade memory for recompute—the same trade selective activation recomputation [22] makes at transformer-training scale, applied here per-operator by the compiler’s own cost model rather than per-layer by the author. Alongside it we run a two-tier allocator policy. Threshold-triggered recovery: above 98% GPU occupancy, a lazy reclamation path calls gc. collect() and torch.cuda.empty_cache(), reclaiming fragmented memory before an allocation fails while keeping CPU–GPU synchronization off the common path. Deterministic deallocation: transient tensors (e.g. batched_data, step_ctx) are explicitly del’d at each loop’s end so the caching allocator can reuse their blocks immediately, which is what keeps cudaMalloc failures from appearing during continuous serving. Kernel benchmark. Fig. 15 and Table 10 give the benchmark Section 4 frames, with its exact per-configuration values. Irregular gathers and masked reductions never reach the Tensor Core path, which is why the ceiling is sparsity rather than the dense peak. What the Triton path removes is avoidable sparse-execution overhead; the intrinsic MFU penalty stays underneath it. C

IMPLEMENTATION DETAIL: PARTITION, NEIGHBOR SEARCH, AND COMMUNICATION

This appendix carries what the main text defers about three mechanisms in turn: the one caveat behind the partitioncurve comparison, the metadata descriptor the Level-1 filter publishes along with the definition of the peer and candidate-population figures, and the buffer and autograd detail behind the fused SKV exchange. The derivations those mechanisms rest on—the count-drift estimate, the three-level filter comparison, the cache exactness invariant, plan-reuse volume, sender-side projection, the overlap bound, and the local-neighbor fraction—are in Sections 3.2, 3.3 and 3.4. Morton/Hilbert Quantization Caveat. Our Morton implementation normalizes all three quantization axes by the largest extent, while Hilbert normalizes per axis. On strongly anisotropic cells the two therefore differ in effective resolution as well as in curve topology, so part of the measured Morton–Hilbert gap is attributable to quantization rather than to locality. We do not claim the two contributions are separated at every scale. The gap does persist at 64 GPUs, where the cubic benchmark cell makes the two normalizations coincide and only the topology can differ, which is the cleanest evidence that the locality effect is real.

35

(b) Throughput, 𝐷=8

(a) Latency, 𝐷=8

TFLOPS

Latency (ms)

100 10

0.1

1

0.01 512

8192

512

8192

(d) Throughput, 𝐷=64

(c) Latency, 𝐷=64

TFLOPS

Latency (ms)

100 10

0.1

1

0.01 512

8192

512

8192

𝑁

𝑁 Sparse-Triton

Sparse-Vanilla

SDPA-Mask

SDPA-Dense

Fig. 15. Single-GPU attention-kernel benchmark on H100. Latency (a,c) and effective throughput (b,d) of Sparse-Triton against a vanilla sparse baseline and two dense SDPA variants, at 𝐷=8 (a,b) and 𝐷=64 (c,d); the dashed red line marks the device’s dense FP32 peak. Dense SDPA is a hardware reference point, not a deployable alternative for this sparse workload: it computes the full 𝑁 2 attention matrix, which is why its throughput is high while its latency explodes with 𝑁 . The takeaway is the comparison within sparse implementations—Sparse-Triton dominates the vanilla baseline at practical sizes, while SDPA-Mask collapses because it still materializes dense score matrices. Exact values in Table 10.

Level-1 Periodic-Image Metadata. Each rank publishes its query AABB plus one AABB per occupied periodic image of its shard, at most 27 including the identity image, with the corresponding counts: 8 + 27 × 8 doubles, under 2 KiB per rank. Because the descriptor is this small and its size fixed, a single metadata all-gather lets every rank reconstruct the complete rank×rank×image request matrix locally, with no second exchange of request masks. For contrast, the spatial-bin variant publishes 4 + 8192 × 16 doubles (≈ 1 MiB) per rank and does need that second all-gather, which is the dominant reason it loses despite pruning the candidate side hardest. Peer Degree and Candidate Population. Fig. 11(c,d) and the peer/pool figures quoted throughout Section 4.3 are both per-step cross-rank maxima, read off the request matrix above for whichever rank is heaviest that step (not necessarily the same rank for both quantities). Peer degree is the number of source ranks it requests candidates from, self included; candidate population is the atom count it then holds after the exchange, before per-chunk filtering, periodic images included. All-gather has no sparse mask to bound either and is = 𝐾 by construction. The equal-volume comparator 36

quoted against these figures is 27, not the textbook 26: a grid cell’s 26 face-, edge-, and corner-adjacent neighbors plus the cell itself, counted the same self-inclusive way. Level-3 Auxiliary Costs. Beyond the pairwise search itself, chunking adds an 𝑂 (𝑄 loc log 𝑄 loc ) median bisection to build the chunk tree and one AABB test per chunk, linear in the gathered atoms. Both are small relative to the Í 𝑄 loc𝐶 cand − 𝑖 𝑄𝑖 𝐶𝑖 pairwise work they remove, and the per-chunk searches are issued as a single batched GPU call rather than one launch per chunk. Fused SKV Buffer Layout and Autograd. Making the SKV fusion practical requires reorganizing both data layout and operator dependencies. We redesign the communication buffer layout so subsequent attention kernels consume the received SKV data directly from the fused buffer without materializing redundant intermediate copies, and rewrite the corresponding operators’ autograd dependencies so the backward pass propagates gradients through the fused path in a well-defined manner. The result preserves end-to-end differentiability with no activation-memory footprint beyond the communication buffers the forward pass already requires. Pre-registered Buffer Pool. Send and receive buffers are carved from a pre-allocated pool whose Memory-Region registration is done in advance during plan construction, so steady-state transfers repeat neither buffer creation nor dynamic MR registration and each exchange launches with minimal bookkeeping (Fig. 7). D

MEASUREMENT PROTOCOL AND FULL-INFERENCE VALIDATION

This appendix gives the measurement protocol behind Section 4 in full. The three tables the main text uses—the neighbor-search composition, the full-inference cache validation, and the strong-scaling detail—are set in Sections 4.3 and 4.6 rather than repeated here. Single-GPU ablation timing. Each Latency entry in Table 2 is the mean over the last 50 of 100 repetitions on the fixed 10k-atom system, so the first half absorbs torch.compile tracing, autotuning, and first-touch allocation and the reported mean is the steady-state cost. Mem is the peak allocation over the same run. Neighbor-stage timing. Beyond the scope limit stated in Section 4.3, two details govern how the neighbor-pipeline numbers should be read. The sub-stage timers are each reduced across ranks with MAX independently, and the maxima of different sub-stages may originate from different ranks, so we do not decompose the pipeline additively—this is why our attribution of spatial-bin’s cost to plan overhead is described as consistent with the telemetry rather than as a measured stage decomposition. Step 0 is excluded from every series, since it carries warm-up and first-touch allocation. The 25k and 50k configurations run in neighbor-only mode; the cache result is separately confirmed in full-inference runs at 20k atoms per GPU (Table 4, Section 4.3). One official run per point, so no error bars are drawn. Neighbor-search composition and its full-inference check. Table 3 (Section 4.3) gives the exact per-mechanism numbers behind Section 4.3’s composed 15.7× reduction, including each row’s own 64-to-512-GPU growth—a scaling behavior Figure 12(b) does not show, since it plots only the 512-GPU endpoint. How the cached rows are built matters for how much weight the exact figure carries. Each is (𝑡 rebuild + 9 𝑡 hit )/10, where 𝑡 hit is the cache-hit cost measured on an interval-10 run and 𝑡 rebuild is the rebuild cost measured on a separate interval-1 run of the same configuration. The two are not interchangeable, for two reasons that both point the same way. The interval-1 arm runs with the cache off and therefore searches at each layer’s bare cutoff, while an interval-10 rebuild builds the 𝑟𝑐 +𝛿 superset the following nine steps consume—1.33× the candidate volume at the 5 Å layers and 1.20× at the 8 Å one, so 𝑡 rebuild is measured on a strictly cheaper search than the one it stands in for. Nor can we simply substitute the right quantity: the interval-10 arm’s only rebuild is step 0, which carries compile and first-touch cost and is excluded as warm-up everywhere in this study, so these runs contain no steady-state interval-10 rebuild at all. Where both quantities do exist in one place—the 20k atoms/GPU full-inference runs of Table 4 (Section 4.3), whose rebuilds fall on steps 10 and 20—a steady-state interval-10 rebuild runs 0.98–1.62× its interval-1 counterpart, the larger radius accounting for much of that. The composition therefore understates the amortized cost by up to 1.4× at those points, and the composed 15.7× should be read as an upper estimate rather than a measured speedup. Two things it 37

does not affect: the cache-hit latencies of Fig. 12(a), measured directly, and the full-step 1.87× of Section 4.5, an A/B on the complete pipeline. Table 4 likewise reports a directly timed interval-10 window; its Amortized column is the mean over that window, and the algebraic identity between that mean and (𝑡 rebuild + 9 𝑡 hit )/10 within a single 2:18 run is arithmetic, not independent confirmation that an interval-1 rebuild may stand in for an interval-10 one. Overlap A/B: what the timelines are not. Section 4.4’s overlap result is a single-flag A/B on the configuration of Section 4.5, so it inherits that protocol. Nsight timelines corroborate it qualitatively, showing the layer-wise exchange largely masked behind the local branch, but we do not count them as a second measurement: they are an aggregate, run-level view on a different denominator than the per-step timers the A/B compares. End-to-end timing. Every point in Section 4.5 is the mean over steps 300–399 of a 500-step run. The window is chosen so that it excludes the torch.compile first step and spans exactly ten rebuild periods at interval 10, which makes the mean the amortized per-step cost rather than a sample biased by where rebuilds happen to land. One official run per point, so no error bars are drawn. Trajectory Settings and Achievable Timescale. The 10.24M-atom KCl and 10.06M-atom capsid trajectories of Section 4.6 rebuild the neighbor graph every 10 steps and re-partition every 100, matching Sections 3.3 and 3.5’s defaults, with the 0.5 Å displacement skin triggering an early rebuild past half that margin (0.25 Å). At the 0.5 fs step used throughout, the capsid run is 200,000 steps, or about 12 days of wall time at its measured 5.2 s per step—≈ 8 ps/day. The 30.7Matom configuration of Section 4.5 reaches ≈ 26 ps/day at 1.684 s per step. Tens of picoseconds per day is enough for the equilibration, structural and energy-stability checks reported here, and short of the nanosecond-to-microsecond windows a binding or folding study would need; closing it is a strong-scaling problem at fixed atom count, the axis on which Fig. 14’s returns go flat past 160 GPUs at 8k atoms per GPU. MD trajectory: rebuild-reason breakdown. Section 4.6’s “cache-only rebuild” reasons condense two of the four categories the underlying per-rank-step log actually distinguishes. Over the NVT segment, scheduled_interval (the every-10-step schedule) accounts for 48.05% of rank-step records and max_displacement (a rank’s own skin crossing) for 1.77%—a 27.09:1 ratio (35.52:1 over NVE). That ratio, however, compares initiating ranks, and is not the share of rebuild work displacement is responsible for; the paragraph below resolves the two. The other two categories split along a different axis than schedule vs. displacement, not a third independent cause: scheduled_interval and dispatch (4.81%) are both unconditional—the driver’s global-sync logic invokes every rank identically on the same step, with no triggering-rank/following-rank distinction, so dispatch is every rank’s own record of a repartition event, not a subset that caused it. Only max_displacement has that triggering structure, since it is a genuinely local, per-rank measurement: whichever ranks independently cross the skin record max_displacement, and the all-reduce of Section 3.3 then forces every other rank to rebuild too, purely to keep the decision unanimous. That induced remainder is passive_follow (45.37%)—specifically max_displacement’s, not dispatch’s, since only a small minority of ranks cross the skin on any given displacement-triggered step. Attributing cost rather than blame, then, displacement is responsible for 1.77 + 45.37 = 47.14% of rank-step rebuild records against the schedule’s 48.05%: the two are near parity, which is what the main text reports and what the 90% → 79% hit-rate gap independently confirms. The 27:1 figure above is the ratio of initiating ranks only, and would badly understate displacement’s cost if read as an event frequency. The capsid trajectory of Section 4.6 splits the same four ways, and the split separates cleanly into system-independent and system-dependent causes. The two globally scheduled ones match the figures above—scheduled_interval 48.06% against 48.05%, dispatch 4.81% in both—which is what a fixed interval should do, since neither depends on what is being simulated. The one genuinely per-rank cause does not: max_displacement is 1.23% of the capsid’s rank-step records against 1.77% here, 31% lower, with the induced passive_follow share rising to match (45.89% against 45.37%). Capsid atoms are covalently constrained and tightly packed, so they displace less per step than bulk water and fewer ranks independently cross the 𝛿/2 threshold. Strong scaling: recomputation off, and the forward-pass split. Each point in Figure 14 and Table 5 (Section 4.6) is the mean step time over the last 5 logged windows of its run, the steady-state value used throughout for scaling comparisons.

38

The table extends Figure 14 two ways. Recomputation off is undefined at 32/64 GPUs (OOM, Section 4.5) but runs 5–23% faster where both do, flattening at the same 160-GPU knee. Its margin narrows with device count—22.9% at 96 GPUs, 5.0% at 256—as recomputation’s own cost follows the shrinking local batch, so on a common 96-GPU baseline it reaches only 64.0% at 256 against on’s 75.0%. The knee appears in both modes, but they do not scale identically, and the efficiency figure quoted for either is only comparable against a shared reference point. The forward-pass hit/rebuild split shows no consistent ordering (±35 ms either way), unlike the 5.6–8.4% hit-slower-than-rebuild gap Table 6 finds at the much larger per-step scale of Section 4.6. TF32: What the Precision Mode Costs. Section 3.1’s TF32 mode is characterized statically in Table 2; what it does to a trajectory we measure directly, as a single-flag A/B on 64 GPUs over the 1.28M-atom system of Fig. 14, 10 ps per arm with the drift fitted over the last 5 ps (101 samples each). Energy drift is −6.316 𝜇eV/atom/ps with TF32 on against −0.966 with it off, a factor of 6.54. Two readings say that the extra drift is systematic rather than noise, and that it does not carry the thermodynamics with it: the residual about each linear fit is the same in both arms (0.328 vs. 0.320 eV), and temperature and its rate of change are indistinguishable (201.37 ± 1.71 vs. 201.45 ± 1.72 K; +1.16 vs. +1.17 K/ps). The mode is worth 1416.9 ± 16.0 → 912.1 ± 9.4 ms per step, 1.55×. The pair runs on the same A100 cluster and with the same activation recomputation as the strong-scaling sweep, so it anchors to an already-reported point: the off arm reproduces Table 5’s 64-GPU step time to 0.9%. That is the split Section 4.1 makes: of the two observables resolved here TF32 moves only the drift, so Section 4.6, where drift is the measurement, runs without it, while the throughput and weak-scaling results, which the 6.54× does not touch, run with it. Two scope limits: this pair sits at ≈ 201 K and is still warming at 1.16 K/ps, not at Section 4.6’s equilibrated 300 K, and it is a 1.28M-atom system—so its off-arm −0.966 is not the same measurement as the 0.42 quoted there for 10.24M atoms, and the two should not be read against each other. TF32 and Structure. Drift is not the only thing a precision mode could move, so we also compare all six radial distribution functions of the same pair of runs over the same 5 ps window (Fig. 16). The solvation structure is reproduced essentially exactly: first-peak heights agree to 0.5% (K–O), 0.2% (Cl–O) and 0.0% (O–O), and, integrated to a common cutoff, both ion–water coordination numbers agree to 0.2%. The three ion–ion pairs are where the two runs visibly part—first-peak heights differ by −3.2% (K–Cl), +23.8% (K–K) and +11.6% (Cl–Cl)—and that gap is not evidence about TF32. Splitting the TF32-off trajectory alone into its two halves of 51 frames, with no precision change anywhere, moves the same three peaks by 10.3%, 14.1% and 12.5% (the gray curves in Fig. 16d–f): K–Cl and Cl–Cl sit inside that floor and K–K is 1.7× it. The reason is sample size. The system carries 1181 K+ and 1181 Cl− , so the K–K first-peak shell holds on the order of 10 pairs per frame and 103 over the window—but ion diffusion decorrelates far more slowly than 5 ps, so those 101 frames contain only 𝑂 (1–3) independent ion configurations and an effective sample of order 10–30 pairs. √ That predicts 1/ 𝑁 eff ≈ 18–32% scatter, which brackets the observed 23.8% and sits above the half-split floor—as it should, since two halves of one trajectory share an initial condition and so understate the scatter between independent runs. Resolving whether TF32 touches ion–ion structure at all needs a trajectory long enough to decorrelate the ions, not a longer averaging window over the same one; what these runs do establish is that it leaves water and the solvation shells alone. Guard and Rebuild Stability over 100 ps. Section 4.6’s capsid arm is long enough to ask whether the shape invariant survives a real trajectory rather than a steady-state window. Over its 200,000 steps the run performed 200,000 force evaluations—158,800 on a cached graph and 41,200 on a rebuild, a 79.40% hit rate—and re-cut the partition throughout on the fixed 100-step schedule. torch.compile’s guards are the mechanism that would catch a shape moving, and they held: zero graph breaks, zero fallbacks to eager, zero backend compilation failures, and a compiled-code cache that never evicted—a limit of 128 against a high-water mark of three variants. Past the 1,738 one-time compilations at step 0, the whole trajectory triggered about 41 further recompilations, all in the energy/force output head and all traced to one cause that is not shape drift: two dimensions that happened to be equal when that function was first traced (both 34,921) were specialized into an equality by duck-shaped sizing, so later steps that break the coincidence recompile. Turning duck shapes off removes them. suppress_errors was armed throughout and never fired, so the zero-fallback counts are measured rather than an artifact of silent recovery. The division of labor behind those counts is worth stating, since

39

(b) Cl–O

(a) K–O

(c) O–O

4 3

3

𝑔 (𝑟 )

4

2

0

2

2

1

1 0

0 2

4

6

8

2

4

(d) K–Cl

6

2

8

6

8

6

8

(f) Cl–Cl

(e) K–K 1.5

4

4

0.8 0.6

𝑔 (𝑟 )

1

0.4

2

0.5

0

0.2

0 2

4

6

8

0 2

Distance 𝑟 (Å)

4

6

8

2

Distance 𝑟 (Å)

TF32 off

TF32 on

4

Distance 𝑟 (Å)

halves of the off run

Fig. 16. TF32 on/off structure over the same 5 ps window. All six radial distribution functions of the 1.28M-atom KCl-in-water run of Fig. 14, TF32 off (dashed) against TF32 on (solid), 101 frames each. (a–c) Ion–water and water–water structure is indistinguishable. (d–f) The ion–ion pairs differ, but the two gray curves are the two halves of the TF32-off run by itself, 51 frames each and no precision change involved: they set a noise floor of 10.3%, 14.1% and 12.5% on these first peaks, which is the scale of the on/off gap. With 1181 ions of each sign and ion decorrelation slower than the window, the effective sample behind a K–K peak is tens of pairs, not the 103 the frame count suggests.

41,200 rebuilds might otherwise be expected to cost a recompilation apiece: the node path is count-fixed and compiled static, while the operators whose extents follow the edge set are compiled for dynamic shapes, so a rebuild moves their sizes without asking for new code. Three compiled variants and 41 attributable recompilations across 200,000 steps are the durability evidence behind Section 3.2’s claim that the equal-count cut pins node shapes for the life of a partition. Hit-Slower-Than-Rebuild: Communication Wait, Not Compute. Table 6’s cache-hit forward costs more than rebuild forward in both columns, despite skipping the search a rebuild pays for. A 256-GPU, 2000-step run with kernel-level profiling (steps 2–13, spanning 2 rebuilds and 10 hits) resolves this. The model-compute kernel is not slower on hit steps— if anything rebuild’s is 1.6% higher (1833.4 vs. 1804.3 ms, consistent with a guard check on the changed edge shape)— while the cross-rank NCCL exchange is 724 ms cheaper on rebuild (1235.4 vs. 1959.6 ms). Per-step communication time on hit steps ranges 735–2451 ms (a 3.3× spread) while compute holds within 1% (1797–1812 ms) over the same steps, so the gap is entirely in how long a rank waits at the attention All-to-All, not in what it computes. Aggregated over the full run (weighted by each window’s actual hit/rebuild count, step 0 excluded), forward-on-hit costs 3465.8 ms against forward-on-rebuild’s 3243.7 ms, a 6.85% gap consistent with the kernel-level breakdown above. We had previously ruled out cross-rank drift because Algorithm 1’s per-step collectives—the rebuild-decision all _reduce, the per-layer All-to-All, the closing scalar reduction—complete every step. That reasoning conflated completing with arriving together: a collective forces the exchange to finish before the step proceeds, not that every rank reaches it at the same wall-clock instant. A rebuild’s ghost-atom exchange is a larger, different collective that in practice re-aligns ranks; consecutive hit steps carry no equivalent, so small per-step timing differences build up until the next attention All-to-All waits on whichever rank is currently slowest. The effect is scale-dependent: across 2/4/8/32/256-GPU configurations, compute alone never favors hit over rebuild (tied at 2 and 8 GPUs, rebuild 1.3–1.6% higher at 4/32/256), but the full forward timer’s hit-slower gap needs per-step work to drift against as well as ranks: it shows at 256 GPUs here and in Table 6, and not in Table 5’s strong-scaling points, whose 256-GPU forward is a sixth as long at 5k

40

atoms/GPU and where hit is the faster path at four of the eight points. We report it as a property of the loaded regime, not a universal one. E

HARDWARE AND SOFTWARE

This appendix records the hardware and software behind every reported timing. The scaling study runs on 64 compute nodes of 8 NVIDIA H100 GPUs each (80 GB HBM3, NVLink 4.0), linked by a 400 Gbps InfiniBand NDR fabric; the largest configuration uses all 512 GPUs. The strong-scaling sweep, the TF32 A/B pair, and the two MD trajectory studies run on a separate cluster of 256 A100 GPUs (40 GB HBM2e, PCIe, no NVLink). Both stacks use CUDA 13.0, PyTorch 2.11, NCCL 2.28, Triton 3.7 and FAISS 1.14.2 with the E2Former V2 backbone of Appendix A. Every timing is the per-step value reduced across ranks with MAX, so it measures the distributed critical path rather than a rank-0 view.

41

Related documents

• Learning consistent molecular mechanics force fields from first principles#1217755
• Quantum-Accurate Vibrational and Conformational Dynamics from Small Molecules to Proteins with Machine-Learned Force Fields#1400825
• A Perspective on Explanations of Molecular Prediction Models#221898
• Multitask learning with semiempirical orbital charges enables sample-efficient MLIPs#1386288
• LightPFP: A Lightweight Route to Ab Initio Accuracy at Scale#1489313