LunaTP: Enabling Efficient and Scalable Tensor Product for Equivariant MLIPs

arXiv:2610.18393v1 Announce Type: cross Abstract: 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.0imes over e3nn and 17.1imes over OpenEquivariance. It also delivers end-to-end training throughput gains of up to 28.7imes over e3nn and 3.11imes 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.

Source: arXiv cs.LG | 2026-10-14