cuSOAP: a GPU-accelerated Generator of Smooth Overlap of Atomic Positions Descriptor

RESUMO

The Smooth Overlap of Atomic Positions (SOAP) descriptor is one of the most widely adopted representations of atomic environments in molecular machine learning, but the cost of evaluating it and its derivatives remains a principal bottleneck of SOAP-based interatomic potentials, particularly at the large radial and angular basis sizes demanded by complex, multi-species condensed-phase environments. We present cuSOAP, a GPU-accelerated generator of atom-wise SOAP vectors and their analytic derivatives. Built on PyTorch, cuSOAP evaluates closed-form expressions or quadratures for the projection coefficients and their Cartesian gradients for Gaussian-type-orbital and polynomial radial bases, through fused CUDA and Triton kernels that eliminate the multi-gigabyte intermediates of a naive tensor formulation. The package is a drop-in replacement for the CPU-based reference DScribe, reproducing its constructor signature, feature ordering, and output to within ${\sim}10^{-6}$, and accepts structures directly as Atomistic Simulation Environment (ASE) Atoms objects. On an NVIDIA Grace--Blackwell (GB200) node, a single Blackwell GPU generates the full descriptor-plus-Jacobian workload for a 1000-molecule water cluster up to two orders of magnitude faster than DScribe, i.e., from $13\times$ to $227\times$ across the entire $(n_{\max}, l_{\max})$ hyperparameter map, with the speedup growing with the angular band limit $l_{\max}$. Descriptor-only generation scales as $t \propto n^{1.34}$, close to linear, up to a million-atom water cluster, which is processed in 19.5s on one GPU and, through a shared-memory multiprocessing driver, in 5.7s on the four GPUs of the node, demonstrating $\sim$90% parallel efficiency with no sign of saturation as devices are added. These results bring on-the-fly SOAP evaluation for large-scale condensed-phase simulation within reach.