From 13 Hours to 4.6: Profiling-Driven Kernel Fusion for TensorNet in Molecular Simulation

Wall-clock time to simulate 1 ns of molecular dynamics 10⁶ steps at 1 fs · one NVIDIA A100

2.82×end-to-end speedup

PyTorcheager

13.0 h

Fused Tritonthis work

4.6 h

Figure 1. Fusing TensorNet's memory-bound operations into single-launch Triton kernels cuts a 1 ns trajectory from 13 h to 4.6 h on one A100, about 1.8 → 5.2 ns/day. Same weights, same equivariance structure; energies agree with the PyTorch reference to <10⁻⁶ kcal/mol.

2.82×end-to-end inference speedup on MD17 / MD22

3.14×geometric-mean speedup on fused ops, 1K to 64K atoms

5.2 ns/dayon a single A100, up from 1.8

75–88%fewer kernel launches per fused block

Abstract

Machine learning force fields such as TensorNet accelerate molecular dynamics relative to ab initio methods, but their PyTorch inference is itself bottlenecked by memory-bound, kernel-launch-heavy operations. Profiling TensorNet on an NVIDIA A100 shows that element-wise tensor operations and index-based message passing together account for 60.8% of inference time. Fusing these into single-launch Triton kernels gives a 3.14× geometric-mean speedup on isolated operations (1K to 64K atoms) and a 2.82× speedup on end-to-end TensorNet inference on MD17/MD22, reducing a 1M-step, 1 ns trajectory from 13 hours to 4.6 hours with numerically verified physical accuracy. The scope is explicit and bounded: end-to-end measurements cover small organic molecules (21 to 42 atoms), operation-level benchmarks establish scaling to 64K atoms, and regimes where PyTorch eager remains competitive are reported.

1Where the time goes

TensorNet represents atomic environments with equivariant Cartesian tensors. Its reference implementation in TorchMD-NET runs each small operation as a separate CUDA kernel. Because neighbour lists change every step, gather and scatter access memory in an irregular order, so these kernels are bound by memory traffic, not arithmetic.

36.0%

24.8%

12.4%

26.8%

60.8% memory-bound, fused in this work

Index-based message aggregation · gather / scatter36.0%
Element-wise tensor algebra24.8%
GEMM · already cuBLAS, left untouched12.4%
Everything else26.8%

Figure 2. Share of GPU time in TensorNet inference, PyTorch profiler on an A100 SXM4 80GB. 4,096-atom system, 5.0 Å cutoff, ~32 neighbours per atom, 131,072 edges.

torch.compile gave negligible improvement on these graph-structured index operations, so every number on this page is measured against PyTorch 2.7.1 eager mode on the same GPU. That is the fair baseline for this workload, and it is stated up front because it changes how the speedups should be read.

2Fusing the memory-bound path

Each eager-mode kernel reads its inputs from global memory and writes its result back. Fusing a chain of them into one Triton kernel keeps the intermediate tile in on-chip SRAM, so the chain pays one round trip instead of eight. Nothing about the model changes: no weights, no equivariance structure, no numerics beyond floating-point-safe reordering.

PyTorch eagerradial cutoff + message passing · 8 launches

op 1

op 2

op 3

op 4

op 5

op 6

op 7

op 8

global memory (HBM)8 round trips

Fused Triton kernel1 launch

cutoff · gather · multiply · scatter

tile held in on-chip SRAM

global memory (HBM)1 round trip

Figure 3. The largest fusion: a smooth radial cutoff folded into message passing, 8 launches → 1. It is worth up to 4.89× on its own, and it matters because cutoffs are evaluated billions of times over a trajectory.

Five kernel families cover the two non-GEMM bottlenecks. Each is tuned by grid search over block size (128 to 256 in 1D, 16 to 32 in 2D) and uses atomic writes only where scatter conflicts are unavoidable. The kernels and benchmark scripts are on GitHub.

Fused kernel family	Launches
Vector → symmetric tensor	5 → 1
Vector → skew tensor	5 → 1
Tensor decomposition	6 → 1
Gather · multiply · scatter message passing	4 → 1
Radial cutoff + message passing	8 → 1

3Results

Every fused operation beats PyTorch eager at every system size from 1K to 64K atoms. Speedups peak around 4K atoms and taper at 64K, but none falls below 2×.

Operation / atoms	1K	4K	16K	64K
Vector → symmetric tensor	3.77×	3.82×	3.88×	3.04×
Vector → skew tensor	3.08×	3.06×	3.09×	2.33×
Tensor decomposition	2.20×	2.21×	2.09×	2.18×
Message passing	3.45×	3.51×	3.10×	2.61×
Fused cutoff + message passing	4.02×	4.89×	3.45×	2.95×
Geometric mean	3.21×	3.38×	3.05×	2.60×

Table 1. Speedup over PyTorch eager for isolated operations on an A100, median of 5 runs × 100 iterations after 20 warm-up. Shading scales with speedup. Geometric mean across all operations and sizes: 3.14×.

End-to-end inference	Atoms	Batch 1	Batch 32
MD17 aspirin	21	2.54×	2.85×
MD22 Ac-Ala3-NHMe	42	2.96×	2.85×

Table 2. Full TensorNet forward pass on the two datasets standard in the TorchMD-NET literature, 2.82× on average. The gap to the 3.14× micro-benchmark mean is the GEMM share (22.1% of time), which already runs on cuBLAS, plus residual Python dispatch.

At 4.6 hours per nanosecond, one A100 produces about 5.2 ns of trajectory per day. A practical target for accurate ML force fields in drug-discovery screening is roughly 1 to 10 ns/day, which puts a single GPU inside that range.

4Same physics

Fusion reorganizes the computation; it does not approximate it. Every fused kernel was validated against the PyTorch reference with torch.allclose (rtol = atol = 10⁻⁴) over 1,000 inference steps on MD17 / MD22.

Energies< 10⁻⁶ kcal/mol

Forces< 10⁻⁵ kcal/mol/Å

Tensor symmetry and tracelessness< 10⁻⁷

5Scope, and where PyTorch still wins

On an individual high-contention scatter (4D index-add, [E,C,3,3] → [N,C,3,3]), atomic contention makes Triton up to 2.3× slower than PyTorch's hand-tuned atomics.
Isolated operations below ~512 atoms are dominated by launch overhead. End-to-end, fusion still wins at 21 to 42 atoms because many small operations are fused together.
The recommended deployment is therefore a hybrid dispatch: route fusion-favourable operations to Triton, and leave high-contention scatters and simple reductions to PyTorch.
End-to-end validation covers 21 to 42 atom molecules. Protein-scale evidence is operation-level only, and there is no throughput comparison to native MD engines (ACEMD, GROMACS, OpenMM kernels).

BibTeX

@inproceedings{singh2026accelerating,
  title     = {Accelerating Molecular Simulations with Open{AI} Triton:
               Fused {GPU} Kernels for TensorNet Neural Potentials},
  author    = {Manpreet Singh},
  booktitle = {EurIPS 2025 Workshop on SIMBIOCHEM},
  year      = {2026},
  url       = {https://openreview.net/forum?id=2UUI4uPNjM}
}

More research

This is one piece of a broader line of work at Embedded LLM on portable, fused GPU kernels for biological and clinical foundation models, measured on both NVIDIA and AMD.

All research →

Let's stay in touch

I'm an undergraduate working on GPU kernels for scientific ML. Happy to talk about this work; questions and feedback are always welcome.

Email LinkedIn About me