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.