LunaTP: Enabling Efficient and Scalable Tensor Product for Equivariant MLIPs

Yuanchang Zhou; Mingzhen Li; Chen Wang; Xiangyu Zhang; Guangming Tan; Weile Jia

cs.DCOct 13, 2026

LunaTP: Enabling Efficient and Scalable Tensor Product for Equivariant MLIPs

Authors: Yuanchang Zhou, Mingzhen Li, Chen Wang, Xiangyu Zhang, Guangming Tan, Weile Jia

Organizations: State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences

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.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.

Figures & tables

Figure 1: Expressivity–efficiency landscape of tensor product convolutions. Forward latency is measured on an NVIDIA H100 with 768 atoms, $L_{\max}=6$ , and 64 channels per degree.

Figure 2: CGTP convolution parallelism and accumulation latencies across the memory hierarchy.

Figure 3: Overview of LunaTP. (a) Node-centric execution accumulates edge messages on chip before storing results. (b) Paths are grouped under an output component budget and mapped to warps by channel block. (c) Products are reused across outputs, and matching coefficients are factored.

Figure 4: LunaTP performance in kernel microbenchmarks. Solid lines show runtime speedup over e3nn, while hollow markers show the reduction in peak memory relative to e3nn. Evaluations are performed in FP32 on NVIDIA H100 80GB GPUs. The raw data are provided in Appendix B.1 .

Figure 5: End-to-end training performance. We benchmark EFS (energy, forces, and stress) and EFSH (EFS plus a Hessian objective). Matched uses equal batch sizes across backends, and MaxBS selects each backend’s maximum batch size. Bars show throughput speedup over e3nn, while bar color indicates the CGTP convolution time fraction. The raw data are provided in Appendix B.2 .

Figure 6: Inference scaling. MD speed on Si with a lattice of 5.43 Å and a 1-fs time step.

Figure 7: Forward scalability of CGTP convolution.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Case ID	$L_{\max}$	Hidden irreps	Edge irreps	Output irreps	Representative MLIP
1	2	128x0e+64x1e+32x2e	1x0e+1x1e +1x2e	224x0e+384x1e +352x2e	SevenNet-0 , GNoME
2	2	128x0e+128x1o+128x2e	1x0e+1x1o +1x2e+1x3o	384x0e+640x1o +640x2e+512x3o	MACE-MP-L
3	3	128x0e+64x1o+32x2e +32x3o	1x0e+1x1o +1x2e+1x3o	256x0e+352x1o +384x2e+320x3o	Nequix-MP/OAM , SevenNet-MF/Omni
4	3	256x0e+256x1o+256x2e +256x3o	1x0e+1x1o +1x2e+1x3o	1024x0e+1536x1o +1792x2e+1536x3o	EquFlash-OMat/OAM
5	4	320x0e+96x1o+64x2e +32x3o+32x4e	1x0e+1x1o+1x2e +1x3o+1x4e	544x0e+736x1o+832x2e +800x3o+704x4e	NequIP-XL series
6	4	128x0e+128x1o+128x2e +128x3o+128x4e	1x0e+1x1o+1x2e +1x3o+1x4e	640x0e+1024x1o+1280x2e +1280x3o+1152x4e	EquFlashV2-OMat/OAM

Appendix

Table 1: Irrep configurations used in the kernel microbenchmarks.

Atoms	SevenNet-Omni	MACE-L	NequIP-XL	EquFlashV2	TACE-L	NequIX	Prophet	BAM
Matched
EFS	2340	2450	1630	290	1030	8900	430	1280
EFSH	870	800	600	90	300	3200	120	580
MaxBS: EFS
e3nn	2340	2450	1630	290	1030	8900	430	1280
FlashTP	14600	3850	9550	2550	3940	30400	2220	3440

Appendix

Table 2: Training batch sizes (atoms per GPU). Matched uses the same batch size across all backends, while MaxBS uses the maximum batch size supported by each backend.

Table 3: MLIP configurations for training.

Setting	Model	Throughput (CGTP time fraction)
Setting	Model	e3nn	FlashTP	OpenEq	cuEq	LunaTP (ours)
Matched	SevenNet-Omni	936.6 (95.9%)	8470.3 (51.2%)	9160.7 (47.5%)	358.7 (76.5%)	18024.6 (21.3%)
	MACE-L	4023.5 (53.5%)	7084.6 (10.7%)	7435.6 (7.9%)	774.6 (59.7%)	8455.9 (1.4%)
	NequIP-XL	625.8 (93.8%)	4696.5 (57.5%)	5225.1 (48.9%)	370.7 (70.3%)	7703.4 (20.2%)
	NequIX	4929.0 (92.8%)	26052.0 (46.1%)	40017.6 (41.0%)	4989.0 (88.9%)	52397.8 (20.9%)
	EquFlashV2	144.0 (78.4%)	547.1 (26.1%)	538.0 (19.8%)	504.6 (9.5%)	659.7 (6.0%)

Appendix

Table 5: EFS training performance. Each entry reports throughput in “atoms/s”, with the CGTP convolution time fraction in parentheses.

Setting	Model	Throughput (CGTP time fraction)
Setting	Model	e3nn	FlashTP	OpenEq	cuEq	LunaTP (ours)
Matched	SevenNet-Omni	277.5 (93.4%)	–	2036.4 (52.4%)	156.1 (71.7%)	3232.3 (18.0%)
	MACE-L	1204.6 (56.1%)	–	2371.2 (13.5%)	486.3 (48.7%)	2775.5 (2.3%)
	NequIP-XL	182.5 (92.0%)	–	1100.4 (52.9%)	–	1480.7 (22.6%)
	NequIX	1465.7 (89.9%)	–	8352.2 (56.3%)	1506.9 (91.8%)	14369.5 (25.0%)
	EquFlashV2	47.0 (55.4%)	–	85.6 (11.2%)	72.5 (11.6%)	91.4 (9.2%)

Appendix

Table 6: EFSH training performance. Each entry reports throughput in “atoms/s”, with the tensor product convolution time fraction in parentheses.

Configuration	Path groups	Execution units	Registers	Spill stores (B)	Latency (ms)
All compatible paths	5	17	64	2976	0.5393
Default ( $\beta=16$ )	17	45	64	0	0.3449
$\beta=8$	29	73	56	0	0.3851
$\beta=4$	37	96	56	0	0.4167

Appendix

Table 7: Output component budget with automatic warp selection (C5, FP32, forward). Latency is the median of two independent measurements. Spill stores are kernel-reported statistics, not total runtime traffic.

Configuration	Path groups	Execution units	Latency (ms)
Baseline (grouped by edge irrep)	25	85	0.3688
$+$ Cross-path input reuse	17	45	0.3446
$+$ Work partitioning (full)	17	45	0.3447

Appendix

Table 8: Cross-path input reuse and work partitioning (C5, FP32, forward). All configurations use node-centric execution.

CG-term count order	Node degree order	Latency (ms)	vs. both descending
Descending	Descending	0.3446	—
Descending	Ascending	0.3895	$+13.0\%$
Ascending	Descending	0.3645	$+5.8\%$
Ascending	Ascending	0.4036	$+17.1\%$

Appendix

Table 9: Task ordering with a fixed fusion scope and warp count (C5, FP32, forward). Each ordering axis is reversed independently; the task set, instruction set, and output are unchanged.

Product reuse	Coefficient reuse	Latency (ms)	vs. both enabled
Enabled	Enabled	0.3449	—
Enabled	Disabled	0.3578	$+3.7\%$
Disabled	Enabled	0.3526	$+2.2\%$
Disabled	Disabled	0.3565	$+3.4\%$

Appendix

Table 10: Compile-time reuse switches with fixed fusion scope and warp count (C5, FP32, forward). Coefficient reuse retains shared products even when the product-reuse switch is disabled.

Figure 8: LunaTP performance in kernel microbenchmarks (FP64, NVIDIA H100-80GB).

Figure 9: LunaTP performance in kernel microbenchmarks (NVIDIA A100-80GB PCIe).

Figure 10: LunaTP performance in kernel microbenchmarks (RTX 3090-24GB).

Stage	Output	e3nn (FP32)	LunaTP (FP32)	LunaTP (FP64)
Forward	$\mathbf{x}^{\prime}$	9.910429e-07	9.978108e-07	9.517594e-16
Backward	$\bar{\mathbf{x}}$	3.582043e-06	6.545904e-06	6.437515e-15
	$\bar{\mathbf{y}}$	3.516555e-06	2.934879e-06	8.966707e-15
	$\bar{\mathbf{w}}$	1.696597e-07	1.712254e-07	4.281635e-16
Double bwd	$\mathbf{h}_{x}$	5.083432e-06	7.955895e-06	1.274438e-14
	$\mathbf{h}_{y}$	5.002968e-06	4.300956e-06	5.290353e-15

Appendix

Table 11: Numerical error comparison between e3nn and LunaTP.

Figure 11: Loss convergence on MPTrj-PFT dataset.

Explore similar work

BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials
Oct 1, 2026Laura Zichi, Gil Harari, Chuin Wei Tan +6Machine Learning Interatomic Potentials Equivariant Neural Networks
Harvard University · Robert Bosch LLC Research and Technology Center
MatRIS: Toward Reliable and Efficient Pretrained Machine Learning Interatomic Potentials
Mar 2, 2026Yuanchang Zhou, Siyu Hu, Xiangyu Zhang +3Machine Learning Interatomic Potentials Computational Materials Science
State Key Lab of Processors, Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences · School of Advanced Interdisciplinary Sciences, UCAS
FlashCart: Fast Cartesian Tensor Products for Equivariant Interatomic Potentials
Oct 5, 2026Viktor Zaverkin, Payman Goodarzi, Sergey V. Sukhomlinov +4Machine Learning Interatomic Potentials GPU Kernel Optimization
INM – Leibniz Institute for New Materials · Saarland University · German Research Center for Artificial Intelligence (DFKI) +2

Oct 1, 2026J/K move · Enter open · S save

BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials

Laura Zichi, Gil Harari, Chuin Wei Tan, Marc L. Descoteaux+5

Harvard University · Robert Bosch LLC Research and Technology Center

Machine Learning Interatomic Potentials Equivariant Neural Networks Adaptive Computation Computational Materials Science

Open paper AskPDF arXiv

Equivariant machine learning interatomic potentials (MLIPs) have revolutionized atomistic modeling, but accurate treatment of complex materials and molecular systems demands expensive models. This limits simulation length- and time-scales, with tensor products a key computational bottleneck. The recent emergence of foundation-scale MLIPs further exacerbates this challenge. We present Branch Interatomic Potential (BranchIP), a single-model framework for learned adaptive tensor product computation, trained with a novel distillation loss. In our experiments on two systems of physical interest, a heterogeneous catalysis system and a proton-conducting solid acid electrolyte, BranchIP accelerates MLIPs across model sizes by up to

2.4\times

while reducing memory usage by up to

2.6\times

. This is achieved while maintaining physical fidelity. Furthermore, the learned adaptive computation provides model interpretability by revealing which interactions demand deeper computation and showing how computational depth relates to chemical complexity and dynamics.

Abstract

Explore similar work

BranchIP: Learning Adaptive Equivariant Computation for Interatomic Potentials

MatRIS: Toward Reliable and Efficient Pretrained Machine Learning Interatomic Potentials

FlashCart: Fast Cartesian Tensor Products for Equivariant Interatomic Potentials