Research map / Architectures and efficiency / RNNs and normalization
FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Presents FlashTP, a fused, sparsity-aware GPU tensor-product library for equivariant machine learning interatomic potentials that speeds up SevenNet-l3i5 inference and training and cuts peak memory.
Problem
The Tensor-Product layer in equivariant MLIPs dominates inference and training time and causes large memory traffic and peak memory usage because of intermediate data, large outputs, and ineffectual computation from sparse Clebsch-Gordan coefficients.
Approach
Fuse all kernels of the tensor-product layer and its following reduce layer into one GPU kernel, skip zero entries of the CG coefficient matrix via a sparse COO format with deduplicated values, and aggregate paths sharing the same hidden input to reduce input loads.
Key results
- FlashTP achieves up to 41.6x (fp64, lmax=4 double-backward) kernel speedup over e3nn and up to 60.8x over cuEquivariance per the abstract.
- For SevenNet-l3i5 on a 4K-atom system, FlashTP gives a 4.2x inference speedup over e3nn with 6.3x less peak memory.
- Single-GPU SevenNet-l3i5 training is 3.5x faster than e3nn with 6.2x lower peak memory (0.89 vs 3.78 GB for l2i5 etc.; l3i5 1.37 vs 8.52 GB).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use FlashTP instead of e3nn for the tensor-product implementation | SevenNet-l3i5 · MPF inference | speedup: 4.2× | Helped |
| Use FlashTP instead of e3nn for the tensor-product implementation | SevenNet-l3i5 · MPF training | speedup: 3.5× | Helped |
| Fuse tensor-product output with the following reduce layer | FlashTP · MD simulation, SevenNet-l3i5 | maximum supported atoms: 28K atoms; 4.7× and 3.1× more than the limitations of e3nn and cuEq | Helped |
| Apply FlashTP instead of e3nn | SevenNet-l3i5 · 4K-atom MD simulation | simulation speed: 4.2× speedup | Helped |
| Apply FlashTP instead of cuEq | SevenNet-l3i5 · 4K-atom MD simulation | simulation speed: 6.2× speedup | Helped |
| Use FlashTP instead of e3nn | SevenNet-l3i5 · single-GPU training | time per epoch: 76 -> 22 min | Helped |
| Use FlashTP instead of e3nn | SevenNet-l3i5 · single-GPU training | peak memory: 8.52 -> 1.37 GB | Helped |
| Add kernel fusion (ablation) | FlashTP tensor-product kernel · Tensor-product microbenchmark | kernel performance | Helped |
| Add sparse computation using the CG coefficient matrix (ablation) | FlashTP tensor-product kernel · Tensor-product microbenchmark | kernel performance | Helped |
| Add path-aggregation (ablation) | FlashTP tensor-product kernel · Tensor-product microbenchmark | kernel performance | Helped |
| Replace separate tensor-product and reduce kernels with fused atomic reduction (ablation) | FlashTP · Tensor-product output stage | peak memory usage | Helped |
| Group paths sharing the same input and execute them in one kernel (ablation) | FlashTP · Tensor-product layer | memory traffic: factors of five and three, respectively, in this example | Helped |
Surprising findings
- cuEquivariance is slower than e3nn for some high-lmax forward cases (e.g., 0.3x-0.9x at lmax 4-5 in forward/backward), showing its efficiency degrades as lmax grows.
Limitations
- Evaluation is focused on SevenNet-l3i5 and a single dataset (MPF), with limited generality across other equivariant MLIPs.
- Multi-GPU speedup slightly decreases at higher GPU counts due to inter-node communication overhead.
- Fused kernels reorder floating-point operations, causing small numerical differences in single precision.
Experimental setup
Datasets: MPF dataset
Benchmarks: tensor-product kernel microbenchmark (lmax 1-5) SevenNet-l3i5 MD step time (copper) SevenNet per-epoch training time and peak memory
Metrics: latency (ms) speedup over e3nn peak memory (GB) MD step time (s) root mean squared error
Baselines: e3nn cuEquivariance FusedTensor SO(2) Tensor-Product
Builds on: e3nn NequIP SevenNet
Closest papers by meaning
- The Importance of Being Scalable: Improving the Speed and Accuracy of Neural Network Interatomic Potentials Across Chemical Domains (NeurIPS 2024 · similarity 0.89)
- MLIP Arena: Advancing Fairness and Transparency in Machine Learning Interatomic Potentials via an Open, Accessible Benchmark Platform (NeurIPS 2025 · similarity 0.88)
- Stein Kernelized Molecular Dynamics for Active Learning of Interatomic Potentials (NeurIPS 2026 · similarity 0.88)
- Higher-Rank Irreducible Cartesian Tensors for Equivariant Message Passing (NeurIPS 2024 · similarity 0.88)
- Physics-Informed Weakly Supervised Learning For Interatomic Potentials (ICML 2025 · similarity 0.88)
- Towards Fast, Specialized Machine Learning Force Fields: Distilling Foundation Models via Energy Hessians (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.