FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials

Machine Learning Interatomic Potentials (MLIPs) enable efficient molecular dynamics (MD) simulations with high accuracy. While equivariant MLIPs achieve state-of-the-art accuracy, they face significant computational bottlenecks centered around their Tensor-Product layer, which account for up to 75% of training time and cause substantial memory overhead. We present FlashTP, a highly optimized tensor-product library that addresses these inefficiencies through kernel fusion, sparse computation, and path-aggregated execution. FlashTP achieves up to 41.6× and 60.8× kernel speedups over e3nn and NVIDIA cuEquivariance, respectively. For SevenNet-l3i5, it delivers 4.2× and 3.5× speedup while reducing peak memory usage by 6.3× and 6.2× for inference and training, respectively. The code is available at https://github.com/SNU-ARC/flashTP.