Track papers across labs, topics, and trends.
Search papers, labs, and topics across Lattice.
This paper addresses the computational bottleneck of Clebsch--Gordan tensor product (CGTP) convolutions in equivariant machine learning interatomic potentials (MLIPs) by introducing LunaTP, a novel framework for parallel execution. LunaTP leverages node-centric execution, path-group parallelism, and multi-level reuse to significantly accelerate CGTP convolutions and higher-order derivative computations. The result is substantial kernel and end-to-end training speedups (up to 250.0x and 28.7x respectively) over existing libraries, enabling more scalable and efficient high-accuracy equivariant MLIPs.
Equivariant MLIPs just got a 250x speedup: LunaTP unlocks efficient tensor products by rethinking parallel execution.
Equivariant machine learning interatomic potentials (MLIPs) promise atomistic simulations with accuracy approaching that of quantum-mechanical calculations. However, the high computational cost of their core Clebsch--Gordan tensor product (CGTP) convolutions severely limits model scalability and efficiency. In this paper, we rethink the parallel execution of CGTP convolutions and introduce LunaTP. Built on node-centric execution, LunaTP exploits path-group parallelism and multi-level reuse to accelerate CGTP convolutions and higher-order derivative computation while preserving the underlying mathematical formulation. On an NVIDIA H100 GPU, LunaTP achieves kernel speedups of up to 250.0$\times$ over e3nn and 17.1$\times$ over OpenEquivariance. It also delivers end-to-end training throughput gains of up to 28.7$\times$ over e3nn and 3.11$\times$ over OpenEquivariance. Finally, LunaTP demonstrates scalability and robust performance across models, workload scales, and GPUs, delivering substantial acceleration for training and atomistic simulation with high-accuracy equivariant MLIPs.