TRANSMISSION // 2611.06122v1
Tessera scales molecular AI to 30.7 million atoms across 512 GPUs
Source published
An IQuest Research team reports a distributed runtime for equivariant molecular Transformers that preserves reusable execution shapes while atoms move. The preprint demonstrates large-scale force evaluation and separate live trajectories, while identifying precision, load-balance and simulation-timescale limits.

For molecular AI, distributing a model means distributing an evolving physical system: neighboring atoms exchange information, move apart and acquire new neighbors. Tessera tackles that changing workload by combining geometry-aware partitioning with reusable computation graphs, reaching 30.7 million atoms in its largest force-evaluation benchmark.
Predicting molecular motion with an AI potential requires more than a fast neural network. When one molecular system spans hundreds of GPUs, information must cross device boundaries repeatedly—and the atoms defining those boundaries keep moving. In a preprint submitted to arXiv on November 4, 2026, Allex Kan and colleagues at IQuest Research describe Tessera, a runtime designed around that problem. The supplied source reports no accepted peer-reviewed publication status.
The team reports force evaluation for 30.7 million atoms on 512 NVIDIA H100 GPUs, alongside separate moving-atom simulations on 256 A100 GPUs. The contribution is distributed execution, not a newly established chemistry model: the authors use a variant of UBio-MolFM built on E2Former V2 and refer to separate work for its underlying molecular accuracy.
Why molecular Transformers are difficult to distribute
The deployed model represents atoms using features that transform predictably under rotations. This equivariance allows directional quantities to follow the molecular geometry. Attention passes information between neighboring atoms through four layers: three with five-ångström cutoffs and one with an eight-ångström cutoff. Together, they provide a nominal per-atom receptive-field bound of 23 ångströms.
That extends information propagation beyond one layer’s cutoff; it does not establish unlimited-range physical interactions. Each layer needs features from neighboring devices. Forces are calculated as the negative gradient of the predicted energy, so reverse-mode differentiation also requires cross-device communication. This backward pass computes forces with respect to coordinates, rather than training the model’s parameters.
Meanwhile, motion changes the neighbor graph. Its changing edge counts interfere with CUDA Graphs, which accelerate execution by replaying recorded GPU operations with fixed buffer addresses and launch geometry.
Keeping shapes stable without freezing the atoms
Tessera sorts atoms along a Hilbert space-filling curve, converting three-dimensional positions into a locality-preserving order, then divides that order into equal-count shards. Each GPU retains a fixed atom count while nearby atoms tend to remain grouped. This stabilizes node-tensor shapes without abandoning spatial locality.
The neighbor cache supplies the second half of the design. At rebuild time, Tessera searches beyond each physical cutoff using a 0.5-ångström protective skin. Between rebuilds, it updates distances and masks inactive pairs without shrinking the cached edge tensors. A displacement exceeding 0.25 ångström triggers a coordinated rebuild; scheduled rebuilds occur every ten steps, and ownership is repartitioned every hundred.
Three-level neighbor filtering, custom Triton kernels and fused feature transfers reduce work further. Communication overlaps local computation in both the forward and backward passes. Each GPU also integrates its own atoms, avoiding centralized position scattering and force collection every step.
What the benchmarks establish
In weak scaling, the problem grows with GPU count while atoms per GPU remain fixed. From 32 to 512 H100 GPUs, step time increased by factors of 1.29, 1.18 and 1.11 at three increasing loads. Corresponding efficiencies were 77.3%, 84.7% and 89.7%.
The largest configuration reached 30.7 million atoms at 1.684 seconds per force-evaluation step. Its 89.7% efficiency needs qualification: this load alone used activation recomputation, trading extra computation for lower memory consumption. These scaling measurements used repeated independent evaluations, not a live trajectory at the largest size.
A controlled full-pipeline comparison found that rebuilding every step cost 1.87 times as much as interval-ten reuse at the largest load. The paper’s larger, neighbor-only 15.7-fold reduction is a composed upper estimate, not a directly measured end-to-end speedup.
For live validation, a 10.24-million-atom potassium chloride solution showed approximately 0.42 microelectronvolts per atom per picosecond of energy drift over ten picoseconds without a thermostat. Selected structural distributions closely matched a smaller serial reference. Separately, a solvated 10.06-million-atom norovirus capsid completed 100 picoseconds at roughly eight picoseconds per day, with no reported broken bonds, anomalously close pairs, graph breaks or eager-execution fallbacks.
Where the evidence stops
The throughput benchmarks enabled TF32 arithmetic; the trajectory validation disabled it. A separate precision comparison found about 6.5 times greater energy-drift magnitude with TF32, while its effect on sparsely sampled ion–ion structure remained unresolved.
Equal atom counts also do not guarantee equal work. The heterogeneous capsid cost 13–33% more per step than the near-uniform solution across sampled windows, consistent with differing edges per atom. Neighbor-search completeness is conditional too: FAISS’s 2,048-entry selection limit can constrain results when saturated.
Every scaling point came from one official run, and comparisons were internal ablations rather than external production-system benchmarks. Billion-atom execution remains extrapolation, while eight picoseconds per day remains short of typical binding or folding timescales. For AI researchers, Tessera’s demonstrated value is narrower but substantial: coordinating placement, graph reuse and differentiation can make a moving scientific graph tractable across hundreds of GPUs.
- Tessera
- molecular dynamics
- equivariant Transformers
- machine-learning force fields
- distributed GPU computing
- CUDA Graphs
- Hilbert space-filling curves
- AI for science
Related reading
Automatically matched by shared topic keywords. These links do not imply that the sources agree or establish a consensus.