TorchFF: A high-performance GPU-accelerated differentiable force field library.

Abstract

Molecular dynamics (MD) and the development of next-generation force fields increasingly rely on automatic differentiation for efficient simulated property prediction and parameter optimization. However, standard FFs and machine learning frameworks often suffer from significant performance bottlenecks-such as kernel launch overhead and memory bandwidth limitations-when executing the many-atom, small-kernel operations characteristic of MD simulations. Here, we present TorchFF, a high-performance, differentiable library that extends PyTorch with a suite of customized CUDA operators specifically engineered for molecular modeling. By implementing performance-critical routines-including bonded interactions, multipolar electrostatics, particle mesh Ewald, and neighbor list searches-as backend-optimized primitives, TorchFF bridges the gap between the flexible Python ecosystem and the execution speed of compiled MD engines.