arXiv CS open access

LunaTP: Enabling Efficient and Scalable Tensor Product for Equivariant MLIPs

Yuanchang Zhou et al. · arxiv_cs

arXiv CS · Papers

cloud distributed-computing parallel-computing

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

LunaTP: Enabling Efficient and Scalable Tensor Product for Equivariant MLIPs Yuanchang Zhou1,2 , Mingzhen Li1,2,∗ , Chen Wang1,2 , Xiangyu Zhang1,2 , Guangming Tan1,2 , Weile Jia1,2,∗ 1 2

Corresponding author

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× over e3nn and 17.1× over OpenEquivariance. It also delivers end-to-end training throughput gains of up to 28.7× over e3nn and 3.11× 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. Email: zhouyuanchang23s@ict.ac.cn, limingzhen@ict.ac.cn, jiaweile@ict.ac.cn Code: https://github.com/HPC-AI-Team/LunaEquivariance

Introduction

Scaling equivariant machine learning interatomic potentials (MLIPs) is challenging, since the Clebsch-Gordan tensor products (CGTP) at their heart have computational cost that scales as O(L6max ) in the maximum degree Lmax (Passaro and Zitnick, 2023; Luo et al., 2024). Consequently, modern equivariant MLIPs are often restricted to Lmax to small values (typically no greater than 3), limiting the use of high-degree equivariant features (Batatia et al., 2022; Liao and Smidt, 2023; Park et al., 2024). Ideally, we would like to use higher-degree features to better capture angular resolution and directional information, which are crucial for accurate potential energy surface (PES) modeling (Batzner et al., 2022; Liao et al., 2024; Duval et al., 2024; Xie et al., 2025). In modern equivariant MLIPs (Batzner et al., 2022; Kim et al., 2026; Xu et al., 2026; Kavanagh et al., 2026), CGTP convolutions can account for over 70% of the training time, making them a computational bottleneck for model scaling. A series of pioneering works have explored ways to optimize CGTP. Their approaches fall into two categories (see Figure 1). Algorithm-level methods reformulate the tensor product or reduce its degrees of freedom, resulting in a trade-off between expressivity and efficiency (Unke and Maennel, 2024; Luo et al., 2024; Passaro and Zitnick, 2023; Xie et al., 2025). By contrast, operator-level methods improve the efficiency of CGTP on GPUs without changing its mathematical formulation or sacrificing expressivity through kernel optimizations (NVIDIA, 2024;

Operator-level optimization

(L 3)

Expressivity

arXiv:2610.18393v1 [cs.DC] 13 Oct 2026

∗

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

tio

Ours M=4

(L 2)

M=2

iza

lev

thm

m pti

M=1 e3nn eSCN GTP(grid) MTP

(L) 0

FlashTP OpenEq cuEq LunaTP

Latency (ms) Figure 1 Expressivity–efficiency landscape of tensor product

convolutions. Forward latency is measured on an NVIDIA H100 with 768 atoms, Lmax = 6, and 64 channels per degree.

Lee et al., 2025; Bharadwaj et al., 2025). Our work follows this direction and introduces LunaTP to further improve the efficiency of CGTP convolutions. However, further accelerating CGTP convolutions while preserving the full CGTP formulation remains challenging. Although kernel fusion and sparse CG computation reduce computational cost and memory traffic (NVIDIA, 2024; Lee et al., 2025; Bharadwaj et al., 2025), edge-parallel execution still suffers from atomic write contention when multiple edges accumulate messages at the same destination node. Node-centric reduction can eliminate this contention, but does not necessarily improve overall performance. Variations in node degree and per-path computational cost complicate load balancing, while increasing data reuse across paths raises on-chip resource requirements and can limit parallelism. Further improving CGTP convolution performance therefore requires jointly optimizing reduction, data reuse, and parallel execution. To bridge this performance gap, we introduce LunaTP, which improves the execution of CGTP convolutions without sacrificing expressivity. Our main contributions are: • We rethink the parallel execution of CGTP convolutions and introduce a node-centric execution scheme that accumulates edge messages on chip, reducing global memory updates. • Building on this execution scheme, we develop path-group parallelism to balance input reuse and register usage, and reduce redundant computation through sparse expansion and multi-level reuse. • We extend this execution scheme to higher-order backward passes, enabling efficient force- and Hessiansupervised training. • We evaluate LunaTP through kernel benchmarks and end-to-end experiments on representative MLIP workloads. On an NVIDIA H100, LunaTP achieves up to a 250.0× kernel speedup and 28.7× higher training throughput than e3nn.

Preliminaries

2.1

Clebsch–Gordan Tensor Product

The Clebsch–Gordan tensor product (CGTP) computes each degree-l component in the direct-sum decomposition. Given two irreducible representations (irreps) x(l1 ) ∈ R2l1 +1 and y(l2 ) ∈ R2l2 +1 , the m-th component of the degree-l output is given by: h i(l) x(l1 ) ⊗ y(l2 ) = m

l1 X

l2 X

(l,m)

(l1 ) (l2 ) C(l1 ,m1 )(l2 ,m2 ) xm ym2 , 1

(1)

m1 =−l1 m2 =−l2

(l,m)

where C(l1 ,m1 )(l2 ,m2 ) are the CG coefficients and are nonzero only when |l1 − l2 | ≤ l ≤ |l1 + l2 |, making the coefficient matrices more than 70% sparse (Bharadwaj et al., 2025; Lee et al., 2025). We refer to each coupling from input degrees l1 and l2 to an output degree l as a path, denoted by (l1 , l2 ) → l. Each path can be assigned a separate learnable weight while preserving equivariance. For irreps up to degree L, the full tensor product contains O(L3 ) paths. In a naive implementation, each path requires O(L3 ) multiply–add operations, yielding O(L6 ) total complexity. To understand how per-path CGTPs are combined in equivariant MLIPs, we next describe the tensor product convolution layer.

2.2

Tensor Product Convolution

A tensor product convolution computes one CGTP for every combination of edge, channel, and path (Geiger and Smidt, 2022; Batzner et al., 2022; Batatia et al., 2022; Park et al., 2024). With E edges, C channels, and P paths, the layer performs a total of E × C × P such computations. The resulting edge messages are then aggregated at their destination nodes. Algorithm 1 presents the pseudocode for the tensor product convolution. For each edge ji, the CGTP (l ,c) (l ) couples the source-node irrep xj 1 with the angular edge irrep yji2 . The result is a full degree-l output, 2

Algorithm 1 Standard Clebsch–Gordan Tensor Product Convolution (l ,c)

Input: xj 1

: node irreps

(l )

Input: yji2 : angular edge irreps (l ,l ,l,c)

Input: wji1 2 : path weights generated by the radial MLP Input: edges ji, from source node j to destination node i (l,c)

Output: x′ i

: updated node irreps

1: Initialize x′ i ← 0 for all nodes i 2: for ji in the E edges do e ji ← 0 3: Initialize m 4: for c in the C channels do 5: for (l1 , l2 ) → l in the P paths do (l,c)

e ji m

end for end for

7: 8:

(l,c)

e ji ←m

(l ,l ,l,c)

+ wji1 2

▷ Compute edge messages ▷ Process each feature channel ▷ Compute a CGTP for each path (l ,c)

xj 1

(l )

⊗ yji2

i(l)

▷ Aggregate at destination node i

e ji 9: x′ i ← x′ i + m 10: end for

(l ,l ,l,c)

whose m-th component is given by Eq. 1. This output is scaled by the corresponding path weight wji1 2 and accumulated into the edge message (Line 6). The same operation is performed for every channel–path e ji (Lines 2–8). The message is then aggregated at destination combination to construct the edge message m node i (Line 9). In particular, training conservative MLIPs typically requires higher-order derivatives, and we provide the corresponding pseudocode in Appendix F. While the standard implementation is mathematically straightforward, its performance on modern GPUs depends on how the computation is organized, as we discuss next.

Rethinking Parallelism for CGTP Convolution (a) Edge-parallel execution worker 1

xj1

𝑗" 𝑗!

𝑖

worker 2

xj2

CGTP

worker 3

xj3

𝑗#

CGTP

𝑗$

𝑁(𝑖) = {𝑗! , 𝑗" , 𝑗# , 𝑗$ }

CGTP

worker 4

xj4

CGTP

(b) Node-centric execu6on Thread block ( node i )

mj1i mj2i

for 𝑗 ∈ 𝑁(𝑖)

atomicAdd

𝑥!𝑖

CGTP Register accumulation

mj3i

Block reducIon (if needed)

mj4i 𝑥!𝑖

global store

SRAM On-chip

Shared Memory ~31 cycles (228 KiB)

Global Memory ~280 cycles (80 GiB)

HBM Oﬀ-chip

Memory Hierarchy with Accumulation Latency & Memory Size

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

Parallelism across edges is a common execution scheme in CGTP convolutions (Figure 2(a)) (Lee et al., 2025; Bharadwaj et al., 2025). Since CGTP on each edge are independent, this scheme exploits data parallelism by assigning per-edge computations to different workers. However, during message aggregation (Algorithm 1, line 9), many workers attempt to update the same node concurrently, causing atomic contention and hardware-level serialization. Our profiling on an NVIDIA H100 GPU shows that aggregation can account for more than 30% of the total runtime. Thus, although kernel fusion and sparse CG evaluation reduce intermediate memory accesses and tensor product computation, aggregation still incurs substantial overhead. On modern GPUs, the latency of accumulation depends on where the data resides. As Figure 2(c) shows, accumulation in registers and shared memory typically has lower latency than global atomic updates. This raises a natural question: can we use these low-latency on-chip memories to perform neighbor aggregation? 3

(b) Path grouping and mapping

(a) Node-centric execution 𝑙=2

𝒀: 𝑬×𝑳 𝑙=0

𝑙=1

𝑙 ,𝑐

𝑙

𝑥𝑗#

𝑦𝑗𝑖"

On-chip

𝑙

𝑝 ,𝑐

𝑤𝑗𝑖!

𝑮 = {𝒑𝟐 , … , 𝒑𝟏 }

𝑝! : (1,1) → 0

𝑦𝑗𝑖"

𝑾: 𝑬×𝑷×𝑪

𝑙=2

…

𝑝" : (1,1) → 2

𝑝 ,𝑐

𝑙

𝑦𝑗𝑖"

𝑤𝑗𝑖!

𝑝 ,𝑐

𝑤𝑗𝑖"

𝑙 ,𝑐

…

CGTP

Output size 2𝑙𝑝 +1

𝑝4 𝑝3 𝑝2 𝑝1 𝑝0

𝑞𝑚₁𝑚₂

3 1

First-Fit Decreasing

𝐺1

9 3

warp0

warp1

𝐺1

0–31

32–63

𝐺2

0–31

⋯

32–63 thread block

𝐴𝑚

𝐴 𝑚′ Coefficient reuse

⋯

𝑞𝑚 1 𝑚 2

𝑞𝑚 1′ 𝑚 2′

×𝐂 𝑝𝑚 1 𝑚 2

×𝐂 𝑝𝑚 ′1 𝑚 ′2 +

One task → one warp ⋯

×𝐶𝑚₁𝑚₂𝑚ʹ

Mapping Strategy ⋯

Multiply once 𝑝

𝑝

×𝐶𝑚₁𝑚₂𝑚

𝑝,𝑐

~ 𝑦 𝑗𝑖,𝑚₂

×𝐂 𝑝𝑚 1 𝑚 2

One channel → one thread

𝑞𝑚 1 𝑚 2 + 𝑞𝑚 ′1 𝑚 ′2

Outer Loop

𝑿′: 𝑵×𝑳×𝑪

Product reuse (𝑙 ,𝑐)

1 𝑥 𝑗,𝑚₁

Budget size: 𝛽

𝑝 ,𝑐

CGTP

𝐴𝑖 "

𝑝 ,𝑐 𝐴𝑖 !

𝑝 ,𝑐

𝑤𝑗𝑖"

Grouping Strategy

𝑝𝑙 : (2,2) → 𝑙

𝐺2

𝑥𝑗#

𝑝 ,𝑐

𝑤𝑗𝑖#

Inner Loop

𝑿: 𝑵×𝑳×𝑪 𝑙=0 𝑙=1

𝐀 𝑝,𝑐 𝑖

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.

Node-centric execution provides a natural way to do so (Figure 2(b)). This scheme aggregates messages for each node on chip and writes the aggregated result to global memory. A straightforward way to implement this is to assign one thread block to each node. However, our experiments show that this direct mapping eliminates per-edge atomic updates but does not improve runtime, as it does not fully exploit intra-node parallelism. Thus, efficient node-centric execution still faces the following primary challenges: (1) Work partitioning limited by on-chip reduction. Node-centric execution requires each thread block to traverse

all edges of a node and aggregate messages on chip. This makes work partitioning difficult, since the workload depends on both node degree and path computation. On the one hand, node degree determines how many edges a block traverses (e.g., degree distributions are often skewed in atomistic datasets (Barros-Luque et al., 2026; Levine et al., 2026; Gharakhanyan et al., 2026)). On the other hand, the amount of CG computation also varies considerably across paths. Thus, blocks computing the same number of output features may still have different workloads.

(2) Fusion scope limited by register capacity. In node-centric execution, we aim to exploit shared inputs across

paths to reduce redundant memory accesses and computation. However, each fused path needs to keep its output accumulators live throughout neighbor reduction, so a larger fusion scope increases register demand and may lower occupancy or cause spilling. Under identical thread configuration and register constraints, fusing all compatible paths runs 35.4% slower than limiting the fusion scope and incurs register spilling.

LunaTP

In this section, we introduce LunaTP’s node-centric execution (Section 4.1). Building on this execution scheme, we describe path-group parallelism (Section 4.2), followed by sparse expansion and multilevel reuse (Section 4.3). Finally, we present higher-order differentiation and GPU implementation details (Section 4.4).

4.1

Node-Centric Execution

LunaTP organizes CGTP computation around destination nodes, accumulates edge messages on chip, and stores the results in global memory after the inner loop. Figure 3(a) shows this execution scheme. The outer loop runs in parallel across destination nodes, paths, and channels, while the inner loop traverses the edges of each destination node. To support this traversal, LunaTP stores the adjacency list in CSR format by destination node, with the edges of node i in the range [rowptr[i], rowptr[i + 1]). (l ,c)

In each inner-loop iteration, a thread loads the source-node features xj 1 4

(l )

, edge features yji2 , and path

p,c weight wji for edge ji. It computes CGTP between the source-node features and the weighted edge (l )

p,c p,c eji features y = wji yji2 . The thread computes only component products corresponding to nonzero CG terms, multiplies them by their CG coefficients, and accumulates the results directly in registers. For path p = (l1 , l2 ) → l, the accumulator Api,c is defined as

Ap,c = i

i(l) X h (l ,c) p,c eji xj 1 ⊗ y ,

(2)

j∈N (i) p,c where N (i) denotes the neighbors of node i, wji is the path weight generated by the radial MLP, and the CG tensor product is defined in Eq. 1.

After traversing all edges of the destination node, each thread stores its accumulated result once at the corresponding output location. Computation across channels and paths remains parallel, while global updates decrease from O(E) to O(N ) for each fixed path and channel. However, the accumulators Ap,c need to remain i in registers throughout the inner loop. Computing more paths that share inputs can improve reuse but also increases register demand. Work partitioning therefore needs to account for both input reuse and accumulator storage. The next section describes how path-group parallelism controls group size and organizes parallel computation.

4.2

Path-Group Parallelism

Multiple paths for the same destination node can share source features. Prior work aggregates these paths to reduce repeated loads (Lee et al., 2025). However, the aggregation scope affects both reuse and register usage. Aggregating all paths that share inputs keeps their output accumulators live until the inner loop ends, which may cause register spilling and limit parallelism. Processing each path separately reduces register usage but repeatedly loads the same features, limiting input reuse. LunaTP therefore uses a grouping strategy to balance input reuse and register usage. For a path group G, let lp denote the output degree of path p. The number of output components retained per channel is X RG = (2lp + 1), (3) p∈G

LunaTP uses RG ≤ β as the grouping constraint, where β is the per-channel output component budget (β = 16 in our experiments). We use First-Fit Decreasing (FFD) to group paths in descending order of output size, as shown in Figure 3(b). Since each path group can be computed independently, we use path groups as units of parallel work, which we call path-group parallelism. Specifically, LunaTP partitions work across destination nodes, path groups, and channels. Each task is identified by (i, G, B), where i is the destination node, G is the path group, and B is a contiguous channel block. In the mapping strategy shown in Figure 3(b), one warp processes one task covering 32 consecutive channels, with each thread handling one channel and computing all paths in the group to reuse the loaded source features. Each task traverses all edges of its destination node and completes its accumulation on chip, avoiding an additional reduction to combine partial results from splitting the edge loop. This partitioning provides more parallel tasks, but differences in node degree and per-path cost can still cause workload imbalance within a thread block. When some warps finish early, the block must wait for the remaining warps before releasing all its resources. To reduce this tail effect, LunaTP sorts path groups by their total number of nonzero CG terms in descending order, then sorts nodes within each group by degree in descending order, placing heavier tasks first and helping assign tasks with similar workloads to the same thread block. Task ordering.

4.3

Sparse Expansion and Multi-Level Reuse

LunaTP exploits sparsity and reuse in the CG structure to reduce memory accesses and computation. Path grouping enables input reuse across paths, while product and coefficient reuse further reduce computation within each path. 5

CG coefficients and their component indices are determined by the path and do not depend on the graph or input features. Like OpenEq (Bharadwaj et al., 2025), LunaTP embeds normalized nonzero coefficients and their indices directly in the generated code, avoiding runtime table lookups and checks for zero terms. Sparse expansion.

Beyond CG sparsity, LunaTP uses path grouping to reuse loaded inputs across paths. In each inner-loop iteration, a thread loads each source-feature and edge-feature component needed by the group once and shares it across the paths that use it. (l ,c)

p,c 1 The same input product qm1 m2 = xj,m yeji,m may be used by multiple output components. 1 2 LunaTP therefore organizes CG terms by input component pair, computes each product once, and reuses it p for all terms with Cm ̸= 0, as shown in the upper part of Figure 3(c). Here, m1 and m2 follow the input 1 m2 m component indices in Eq. 1.

Product reuse.

Different input component pairs may have coefficient vectors that are identical up to an overall sign. Let Cpm1 m2 ∈ R2l+1 collect the CG coefficients over all output components. When Cpm′ m′ = 1 2 ±Cpm1 m2 , we combine the products before applying the shared coefficients, Coefficient reuse.

qm1 m2 Cpm1 m2 + qm′1 m′2 Cpm′ m′ = (qm1 m2 ± qm′1 m′2 )Cpm1 m2 . 1

(4)

This reduces coefficient multiplications, as shown in the lower part of Figure 3(c). In particular, we identify these matches only within each path at operator construction time.

4.4

Implementation Details

LunaTP extends the node-centric execution scheme and reuse methods above to backward, double-backward, and triple-backward passes by adapting the reduction direction to the differentiation target. Computing the gradient of node features xj requires accumulating results over all edges with source node j, so reduction shifts from destination nodes to source nodes. To support this traversal, LunaTP organizes a CSR by source node while retaining the partitioning across nodes, path groups, and channels and accumulating gradients on chip. Derivatives with respect to the incoming gradient instead reduce at destination nodes, as in the forward pass. Edge-feature and path-weight derivatives are computed per edge, with edge-feature derivatives also requiring a reduction across channels. Double-backward and triple-backward use the same approach, organizing computation around the reduction required by each output. The full pseudocode is provided in Appendix F. Extension to higher-order derivatives.

Higher-order backward passes also allow reuse across derivative branches. In double-backward, for example, computing node-feature and edge-feature derivatives requires the same incoming gradient for a given edge. LunaTP combines branches that share inputs to avoid repeated loads and applies sparse expansion, product reuse, and coefficient reuse within each branch. LunaTP uses __ldg to load read-only features, weights, and adjacency indices, and assigns consecutive channels to threads within each warp to coalesce path-weight loads. We express accumulation as multiply-add operations so the compiler can generate fused multiply-add instructions. In higher-order backward passes, we assign independent derivative branches to different CUDA streams to overlap their execution. Meanwhile, we provide a programming interface for LunaTP to improve ease of use (see Appendix G). Low-level optimization and interface.

Experiments

We evaluate LunaTP across four dimensions: computational efficiency, memory efficiency, scalability, and accuracy. Our baselines include e3nn (Geiger and Smidt, 2022), FlashTP (Lee et al., 2025), OpenEquivariance (OpenEq) (Bharadwaj et al., 2025), and cuEquivariance (cuEq) (NVIDIA, 2024). We report kernel microbenchmark results (Section 5.1), end-to-end training and inference performance (Section 5.2), and scalability analysis (Section 5.3). The experimental setup is detailed in Appendix A and ablation studies are described in Appendix C. 6

5.1

Kernel Microbenchmark

The microbenchmarks cover ten irrep configurations with Lmax ranging from 2 to 6. Most are based on configurations used in existing equivariant MLIPs (Batzner et al., 2022; Batatia et al., 2022; Park et al., 2024; Kim et al., 2026; Kavanagh et al., 2026). The hidden, edge, and output irreps and their representative MLIPs are summarized in Appendix Table 1.

Gain over e3nn (×)

FlashTP 256 Runtime 128 64 32 16 8 4 2 1

OpenEq

(a) Forward

cuEq (uniform)

cuEq (naive)

(b) Backward

58.9×

70.7×

15.0×

51.0× 14.8×

10.2×

LunaTP (ours)

Avg. speedup (d) Triple-bwd

151.8×

103.6× 14.0×

35.3×

9.1×

30.8×

6.0×

7.8×

21.3×

1.2×

1.0×

1.1×

1 2 3 4 5 6 7 8 9 10

30 Memory 20 10 1

Case ID

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 4 summarizes the FP32 microbenchmark results on H100 (FP64 results are provided in Appendix D.1). LunaTP achieves average speedups of 58.9×, 70.7×, 151.8×, and 103.6× over e3nn for forward, backward, double-backward, and triple-backward, respectively. Its corresponding speedups over OpenEq are 3.93×, 5.93×, 4.78×, and 13.31×. Although FlashTP and OpenEq already use kernel fusion to reduce memory traffic, LunaTP further improves performance through node-centric execution, path-group parallelism, and multilevel reuse. The gains are most pronounced for double- and triple-backward, which offer more opportunities to share data and computation across derivative branches. cuEq also performs well with uniform channel counts across degrees, but uses a naive implementation with limited speedups otherwise. Runtime and peak memory.

In addition to runtime improvements, LunaTP uses only 4.0%–16.2% of e3nn’s peak memory. Its memory footprint is comparable to that of FlashTP and OpenEq from forward through double-backward, since all three avoid materializing per-edge message tensors. For triple-backward, LunaTP uses approximately 48% of OpenEq’s peak memory while also achieving lower runtime. Since kernel fusion and parallel reductions can alter floating-point operation order, LunaTP and e3nn may exhibit small numerical discrepancies. Using e3nn’s FP64 results as the reference, we measure the RMSEs of their FP32 outputs and corresponding derivatives. From the forward pass through third-order derivatives, both implementations maintain FP32 RMSEs on the order of 10−7 –10−6 , while LunaTP’s FP64 RMSEs remain between 10−16 and 10−14 (see Appendix E.1). Numerical stability and training convergence.

To assess training stability, we train NequIP-XL (Batzner et al., 2022; Kavanagh et al., 2026) on MPTrj-PFT (Deng et al., 2023; Koker et al., 2026) with the same settings. LunaTP’s validation MAEs for E, F, S, and H differ from those of e3nn by less than 0.3% across all targets, while showing similar convergence trends. (see Appendix E.2). We evaluate LunaTP on NVIDIA A100 and RTX 3090 GPUs to assess its portability across data-center and consumer hardware. These GPUs differ in compute performance and memory bandwidth. LunaTP also has a performance advantage over the baselines on both GPUs. Full results are reported in Appendix D.2. Hardware portability.

(b) Matched EFSH

(d) MaxBS EFSH

Speedup over e3nn (×)

19.2 11.6

8.3

9.8

5.3

15 4.4

4.9

3.7

7.5 5.3

5.1

4.6

4.8

6.2

1.0 0.6

0.4

6.0

2.2

10.6 9.9

25 20

2.4 1.9

16.5

13.3

5.7

5.2

5.0

4.5

7.3

8.1 9.0

8.1

25.4

23.0 20.5

9.8 12.3 10.6

100

28.7

25.6

1.0

0.6

9.6

5.6

7.3 6.0

8.7

11.9

11.1 11.0 9.7

9.0 14.9

6.7 5.0

0.6

0.4

9.2

11.2 8.2 10.4 7.2 10.1 6.2 6.5 7.6 5.2

40 5.8

0.6

1.8 10 10 5 2.3 1.8 5.6 5.5 S S S S 5.0 1.0 3.5 1.7 O ev O ev O ev O ev 1.0 1.0 2.3 5.0 Ne mn enN Ne mn enN 5 Ne mn enN 5 Ne mn enN 2.0 0.9 1.1 qu i etqu i etqu i etqu i et0.9 2.5 1.5 Eq Eq 1.9 Ne IX IX I IP P 1.2 uF uF -X -X qu qu 0.5 1.8 BA BA Ne Ne las las L L 0 0 0 0 1.9 IP IP 0.9 1.2 -X -X qu qu 1.0 hV hV M M 0.5 Pr Pr Pr Pr L L IX IX P 2 2 op op op op TP ) TP ) TP ) 0.4 0.6 aT s) q a a a 0.2 0.3 E E T T h h h h s s s n n n n q q q A A q q et et et et u ur nE u ur nE u ur nE uF uF CE CE Lu (our enE L L L T T B B o o o la la ( pe ( pe ( pe -L -L A A A A sh sh TP TP TP TP V2 V2 Op O O O MA CE-L MA CE-L MA M MA M sh sh sh sh CE CE CE CE Eq Eq Eq Eq Fla Fla Fla Fla -L -L -L -L cu cu cu cu

3.7

2.1

1.8

3.4

3.8

3.5 0.9

1.1

Percentage (%)

(a) Matched EFS

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.

5.2

End-to-end Performance

We evaluate LunaTP’s end-to-end performance in training and inference. For training, we measure throughput on eight state-of-the-art (SOTA) equivariant MLIPs with EFS and EFSH objectives. For inference, we measure MD speed and the maximum system size each backend can run. Experimental details are in Appendix A.3. Training. As shown in Figure 5(a) and (b), LunaTP achieves speedups of 2.10–19.24× over e3nn for EFS and 1.81–11.6× for EFSH at matched batch sizes (Matched). The speedup depends on the fraction of training time spent in CGTP convolutions. For example, MACE-L achieves a 2.10× speedup under Matched EFS. With e3nn, CGTP convolutions account for 53.5% of runtime, limiting the overall speedup achievable by optimizing this operator. In contrast, SevenNet-Omni spends a larger fraction of runtime in CGTP convolutions and achieves a 19.2× speedup. Under Matched EFSH, LunaTP achieves speedups of 1.06–1.72× over OpenEq, with further gains also limited by work outside CGTP convolutions.

To evaluate larger workloads, we increase the batch size to the maximum each backend can run on a single GPU (MaxBS). As shown in Figure 5(c) and (d), LunaTP achieves speedups of 2.30–25.6× over e3nn for EFS and 3.50–28.7× for EFSH. The corresponding speedups over OpenEq are 1.14–2.10× and 1.34–3.11×. These gains over OpenEq generally exceed those under Matched, showing that LunaTP provides greater end-to-end benefits on larger workloads. We evaluate inference scalability on TACE-L and NequIP-XL, as shown in Figure 6. LunaTP achieves an 8.70× speedup on TACE-L and a 15.27× speedup on NequIP-XL at e3nn’s maximum system size. Since kernel fusion avoids storing per-edge message tensors, LunaTP increases the maximum atom count by 4.63× and 5.04× for TACE-L and NequIP-XL, respectively. OpenEq and FlashTP also use kernel fusion, so all three backends reach the same maximum system size. At these workloads, LunaTP remains 1.8× faster on TACE-L and 2.1× faster on NequIP-XL than OpenEq. This shows that LunaTP’s optimizations further improve inference performance even with kernel fusion already in place. Inference.

e3nn

cuEq

LunaTP (ours) 1.8×

1.8×

MD speed (ns/day)

0.1 0.01

(a) TACE-L

2.1×

1 0.1

(b) NequIP-XL 64

Scalability Analysis

OpenEq

0.01

5.3

FlashTP

128 256 512 1K 2K 4K Number of atoms

16K

We evaluate LunaTP’s forward scalability across four Figure 6 Inference scaling. MD speed on Si with a lattice of 5.43 Å and a 1-fs time step. dimensions: atom count, average in-degree, maximum degree, and channel size. As Figure 7 shows, latency generally increases with workload size, but LunaTP 8

e3nn

FlashTP

Forward latency (ms)

(a) System scale

OpenEq

cuEq (uniform)

LunaTP (ours)

(b) Graph density

(d) Feature width

100 10 1

2.6×

2.4×

10.3×

5.8×

0.1 64

256

16k

128k 8

Atoms

128

512

Lmax

Avg in-degree

128

512

Channel dim

Figure 7 Forward scalability of CGTP convolution.

achieves the lowest latency at nearly all measured points. At the largest comparable point in each sweep, it is 2.4–10.3× faster than the fastest available baseline, with larger advantages in the maximum-degree and channel-size sweeps. As Lmax increases, the number of coupling paths and CG terms grows. LunaTP controls accumulator storage through path grouping and reduces repeated computation through multilevel reuse. As channels become wider, LunaTP partitions the additional computation into more independent channel-block tasks that run across warps without increasing the per-thread accumulator count under the basic mapping. These results demonstrate LunaTP’s scalability to large graphs, higher-degree features, and wider channels. The experimental setup is provided in Appendix A.2.

Related Work

SO(2)-based methods align a local coordinate system with each edge direction to reduce SO(3) convolutions to simpler SO(2) operations (Passaro and Zitnick, 2023). Other approaches reformulate feature representations and coupling operations to exploit efficient computational primitives. Fused Tensor Product organizes computations across irreps into matrix multiplications (Unke and Maennel, 2024), while Gaunt Tensor Product uses a Fourier basis and FFTs to accelerate products of spherical functions (Luo et al., 2024). These approaches offer efficient alternatives for high-degree equivariant models, although their parameterizations can involve trade-offs between expressivity and efficiency (Xie et al., 2025). LunaTP instead improves execution efficiency while preserving the original CGTP formulation and expressivity. Algorithm-level optimization.

A complementary direction improves tensor product implementations on GPUs. e3nn provides general-purpose equivariant operators (Geiger and Smidt, 2022), while cuEquivariance offers specialized GPU kernels (NVIDIA, 2024). FlashTP exploits CG sparsity, aggregates paths sharing source features, and fuses convolution operations to reduce repeated loads and intermediate storage (Lee et al., 2025). OpenEq generates sparse kernels just in time for CG contractions and supports fused graph convolutions and higher-order derivatives. Its deterministic implementation accumulates over sorted edges to reduce global writes, then merges node partial sums across warps with a fixup kernel (Bharadwaj et al., 2025). Building on these advances, LunaTP focuses on parallel execution and reuse within node-centric computation. It groups paths under an accumulator budget to expose intra-node parallelism and reuses common products and coefficients to reduce computation. Our experiments demonstrate performance gains over these baselines. Operator-level optimization.

Conclusion

In this work, we rethink the parallelism of CGTP convolutions and present LunaTP, a GPU kernel for CGTP convolutions in equivariant MLIPs. By organizing node reduction, data reuse, and work partitioning, LunaTP improves computational efficiency without sacrificing expressivity, and extends to third-order derivatives, making efficient Hessian-based training practical. Experiments show that LunaTP accelerates kernels by up to 250.0× over e3nn and 17.1× over OpenEq, while delivering up to 28.7× the end-to-end training throughput of e3nn. Future work will extend these optimizations to other operators in equivariant MLIPs.

AI use statement In this work, we did not use generative AI tools for any tasks requiring disclosure. We used generative AI tools only for language editing to improve the clarity and readability of the manuscript and for literature search, which are tasks with recommended disclosure. We reviewed all AI-assisted edits and independently verified the identified literature and references. We take responsibility for the final content of this work, including all AI-assisted language edits and literature-related content.

Reproducibility statement We describe our methods in Section 4. Appendix A provides the experimental setup, including the hardware and software environments, benchmarking protocols, and model configurations. The raw benchmark data underlying the results reported in the main text are provided in Appendix B. We will make the complete source code publicly available. Our code is available at https://github.com/HPC-AI-Team/LunaEquivariance

Acknowledgments This work is supported by the following funding: National Science Foundation of China (92270206, 62372435, T2125013, 62502501), Beijing Natural Science Foundation (4254087), National Key Research and Development Program of China (2025YFB4507000), and the Innovation Funding of ICT, CAS. Part of the training is performed on the robotic AI-Scientist platform of Chinese Academy of Sciences.

References Luis Barros-Luque, Muhammed Shuaibi, Xiang Fu, Brandon M. Wood, Misko Dzamba, Meng Gao, Ammar Rizvi, Matt Uyttendaele, C. Lawrence Zitnick, and Zachary W. Ulissi. The Open Materials 2024 (OMat24) inorganic materials dataset and models. Nature Computational Science, 6(6):642–652, June 2026. ISSN 2662-8457. doi: 10.1038/s43588-026-00996-w. https://doi.org/10.1038/s43588-026-00996-w. Ilyes Batatia, David Peter Kovacs, Gregor N. C. Simm, Christoph Ortner, and Gabor Csanyi. MACE: Higher order equivariant message passing neural networks for fast and accurate force fields. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. https://openreview.net/forum?id=YPpSngE-ZU. Ilyes Batatia, Philipp Benner, Yuan Chiang, Alin M. Elena, Dávid P. Kovács, Janosh Riebesell, Xavier R. Advincula, Mark Asta, Matthew Avaylon, William J. Baldwin, Fabian Berger, Noam Bernstein, Arghya Bhowmik, Filippo Bigi, Samuel M. Blau, Vlad Cărare, Michele Ceriotti, Sanggyu Chong, James P. Darby, Sandip De, Flaviano Della Pia, Volker L. Deringer, Rokas Elijošius, Zakariya El-Machachi, Edvin Fako, Fabio Falcioni, Andrea C. Ferrari, John L. A. Gardner, Mikołaj J. Gawkowski, Annalena Genreith-Schriever, Janine George, Rhys E. A. Goodall, Jonas Grandel, Clare P. Grey, Petr Grigorev, Shuang Han, Will Handley, Hendrik H. Heenen, Kersti Hermansson, Cheuk Hin Ho, Stephan Hofmann, Christian Holm, Jad Jaafar, Konstantin S. Jakob, Hyunwook Jung, Venkat Kapil, Aaron D. Kaplan, Nima Karimitari, James R. Kermode, Panagiotis Kourtis, Namu Kroupa, Jolla Kullgren, Matthew C. Kuner, Domantas Kuryla, Guoda Liepuoniute, Chen Lin, Johannes T. Margraf, Ioan-Bogdan Magdău, Angelos Michaelides, J. Harry Moore, Aakash A. Naik, Samuel P. Niblett, Sam Walton Norwood, Niamh O’Neill, Christoph Ortner, Kristin A. Persson, Karsten Reuter, Andrew S. Rosen, Louise A. M. Rosset, Lars L. Schaaf, Christoph Schran, Benjamin X. Shi, Eric Sivonxay, Tamás K. Stenczel, Christopher Sutton, Viktor Svahn, Thomas D. Swinburne, Jules Tilly, Cas van der Oord, Santiago Vargas, Eszter Varga-Umbrich, Tejs Vegge, Martin Vondrák, Yangshuai Wang, William C. Witt, Thomas Wolf, Fabian Zills, and Gábor Csányi. A foundation model for atomistic materials chemistry. The Journal of Chemical Physics, 163(18):184110, 11 2025. ISSN 0021-9606. doi: 10.1063/5.0297006. https://doi.org/10.1063/5.0297006. Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Kornbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13(1), May 2022. ISSN 2041-1723. doi: 10.1038/s41467-022-29939-5. http://dx.doi.org/10.1038/s41467-022-29939-5. Vivek Bharadwaj, Austin Glover, Aydin Buluc, and James Demmel. An efficient sparse kernel generator for o(3)equivariant deep networks, 2025. https://arxiv.org/abs/2501.13986.

Bowen Deng, Peichen Zhong, KyuJung Jun, Janosh Riebesell, Kevin Han, Christopher J. Bartel, and Gerbrand Ceder. CHGNet as a pretrained universal neural network potential for charge-informed atomistic modelling. Nature Machine Intelligence, 5(9):1031–1041, September 2023. ISSN 2522-5839. doi: 10.1038/s42256-023-00716-3. https://doi.org/10.1038/s42256-023-00716-3. Alexandre Duval, Simon V. Mathis, Chaitanya K. Joshi, Victor Schmidt, Santiago Miret, Fragkiskos D. Malliaros, Taco Cohen, Pietro Liò, Yoshua Bengio, and Michael Bronstein. A hitchhiker’s guide to geometric gnns for 3d atomic systems, 2024. https://arxiv.org/abs/2312.07511. Xiang Fu, Brandon M Wood, Luis Barroso-Luque, Daniel S. Levine, Meng Gao, Misko Dzamba, and C. Lawrence Zitnick. Learning smooth and expressive interatomic potentials for physical property prediction. In Forty-second International Conference on Machine Learning, 2025. https://openreview.net/forum?id=R0PBjxIbgm. Johannes Gasteiger, Janek Groß, and Stephan Günnemann. Directional message passing for molecular graphs, 2022. https://arxiv.org/abs/2003.03123. Mario Geiger and Tess Smidt. e3nn: Euclidean neural networks, 2022. https://arxiv.org/abs/2207.09453. Vahe Gharakhanyan, Luis Barroso-Luque, Yi Yang, Muhammed Shuaibi, Kyle Michel, Daniel S. Levine, Misko Dzamba, Xiang Fu, Meng Gao, Xingyu Liu, Haoran Ni, Keian Noori, Brandon M. Wood, Matt Uyttendaele, Arman Boromand, C. Lawrence Zitnick, Noa Marom, Zachary W. Ulissi, and Anuroop Sriram. Open Molecular Crystals 2025 (OMC25) dataset and models. Scientific Data, 13(1):354, February 2026. ISSN 2052-4463. doi: 10.1038/s41597-026-06628-2. https://doi.org/10.1038/s41597-026-06628-2. Seán R. Kavanagh, Chuin Wei Tan, Menghang Wang, Marc L. Descoteaux, Gabriel de Miranda Nascimento, Ulrik Unneberg, Laura Zichi, Francesco Libbi, Norma Rivano, Austin Glover, Vivek Bharadwaj, Anders Johansson, William C. Witt, Albert Musaelian, and Boris Kozinsky. Fast and accurate equivariant foundation models for atomistic simulation, 2026. https://arxiv.org/abs/2607.28461. Jaesun Kim, Jinmu You, Yutack Park, Yunsung Lim, Yujin Kang, Jisu Kim, Haekwan Jeon, Suyeon Ju, Deokgi Hong, Seung Yul Lee, Saerom Choi, Yongdeok Kim, Jae W. Lee, and Seungwu Han. Optimizing cross-domain transfer for universal machine learning interatomic potentials. Nature Communications, 17(1):3432, March 2026. ISSN 2041-1723. doi: 10.1038/s41467-026-70195-8. https://doi.org/10.1038/s41467-026-70195-8. Teddy Koker, Abhijeet Sadashiv Gangan, Mit Kotak, Jaime Marian, and Tess Smidt. PFT: Phonon fine-tuning for machine learned interatomic potentials. In Forty-third International Conference on Machine Learning, 2026. https://openreview.net/forum?id=QC9S8gUOYc. Seung Yul Lee, Hojoon Kim, Yutack Park, Dawoon Jeong, Seungwu Han, Yeonhong Park, and Jae W. Lee. FlashTP: Fused, sparsity-aware tensor product for machine learning interatomic potentials. In Forty-second International Conference on Machine Learning, 2025. https://openreview.net/forum?id=wiQe95BPaB. Daniel S. Levine, Muhammed Shuaibi, Evan Walter Clark Spotte-Smith, Michael G. Taylor, Muhammad R. Hasyim, Kyle Michel, Ilyes Batatia, Gábor Csányi, Misko Dzamba, Peter Eastman, Nathan C. Frey, Xiang Fu, Vahe Gharakhanyan, Aditi S. Krishnapriyan, Joshua A. Rackers, Sanjeev Raja, Ammar Rizvi, Andrew S. Rosen, Zachary Ulissi, Santiago Vargas, C. Lawrence Zitnick, Samuel M. Blau, and Brandon M. Wood. The open molecules 2025 (omol25) dataset, evaluations, and models, 2026. https://arxiv.org/abs/2505.08762. Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs. In The Eleventh International Conference on Learning Representations, 2023. https://openreview.net/forum?id= KwmPfARgOTD. Yi-Lun Liao, Brandon M Wood, Abhishek Das, and Tess Smidt. Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations. In The Twelfth International Conference on Learning Representations, 2024. https://openreview.net/forum?id=mCOBKZmrzD. Yi-Lun Liao, Alexander J. Hoffman, Sabrina C. Shen, Alexandre Duval, Sam Walton Norwood, and Tess Smidt. Equiformerv3: Scaling efficient, expressive, and general se(3)-equivariant graph attention transformers, 2026. https://arxiv.org/abs/2604.09130. Shengjie Luo, Tianlang Chen, and Aditi S. Krishnapriyan. Enabling efficient equivariant operations in the fourier basis via gaunt tensor products. In The Twelfth International Conference on Learning Representations, 2024. https://openreview.net/forum?id=mhyQXJ6JsK. NVIDIA. cuEquivariance. https://github.com/NVIDIA/cuEquivariance, 2024.

Yutack Park, Jaesun Kim, Seungwoo Hwang, and Seungwu Han. Scalable parallel algorithm for graph neural network interatomic potentials in molecular dynamics simulations. Journal of Chemical Theory and Computation, 20(11): 4857–4868, May 2024. ISSN 1549-9626. doi: 10.1021/acs.jctc.4c00190. http://dx.doi.org/10.1021/acs.jctc.4c00190. Saro Passaro and C. Lawrence Zitnick. Reducing so(3) convolutions to so(2) for efficient equivariant gnns. ICML’23. JMLR.org, 2023. Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, and Tri Dao. Flashattention-3: Fast and accurate attention with asynchrony and low-precision. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. https://openreview.net/forum?id=tVConYid20. Oliver T. Unke and Hartmut Maennel. E3x: E(3)-equivariant deep learning made easy, 2024. https://arxiv.org/abs/ 2401.07595. Soohaeng Yoo Willow, Tae Hyeon Park, Gi Beom Sim, Sung Wook Moon, Seung Kyu Min, D. ChangMo Yang, Hyun Woo Kim, Juho Lee, and Chang Woo Myung. Bayesian e(3)-equivariant interatomic potential with iterative restratification of many-body message passing, 2026. https://arxiv.org/abs/2510.03046. YuQing Xie, Ameya Daigavane, Mit Kotak, and Tess Smidt. The price of freedom: Exploring expressivity and runtime tradeoffs in equivariant tensor products. In Forty-second International Conference on Machine Learning, 2025. https://openreview.net/forum?id=EvIwwGYTLc. Zemin Xu, Wenbo Xie, and P. Hu. Spectral/spatial tensor atomic cluster expansion with universal embeddings in cartesian space, 2026. https://arxiv.org/abs/2509.14961. Yuanchang Zhou, Siyu Hu, Xiangyu Zhang, Hongyu Wang, Guangming Tan, and Weile Jia. MatRIS: Toward reliable and efficient pretrained machine learning interatomic potentials. In The Fourteenth International Conference on Learning Representations, 2026. https://openreview.net/forum?id=5xBT5Ziute.

Appendix A Evaluation setup A.1 Kernel microbenchmark settings A.2 Scalability analysis settings A.3 Training setup details B Detailed Benchmark Data B.1 Raw data for kernel microbenchmarks B.2 Raw data for training benchmarks C Ablation Studies C.1 Path grouping C.2 Task ordering C.3 Compile-time reuse D Additional kernel microbenchmark results D.1 Kernel microbenchmark (FP64) D.2 Kernel microbenchmark (A100 & RTX 3090) E Details on numerical stability and training convergence E.1 Operator-level numerical stability E.2 End-to-end training convergence F Differentiation of CGTP convolution G Programming interface

Evaluation setup

All evaluations use FP32 with PyTorch 2.11.0+cu128, e3nn v0.6.0 (Geiger and Smidt, 2022), FlashTP v0.1.1 (Lee et al., 2025), OpenEquivariance v0.7.0 (Bharadwaj et al., 2025), and cuEquivariance v0.11.1 (NVIDIA, 2024). Notably, cuEquivariance uses uniform_1d when all segments have the same shape and naive otherwise, which we denote as cuEq(uniform) and cuEq(naive), respectively. Software and hardware.

Unless otherwise specified, all experiments are conducted on NVIDIA H100-80GB SXM GPUs with Intel Xeon Platinum 8468 CPUs. NVLink provides up to 900 GB/s of bidirectional bandwidth per GPU. To construct workloads that reflect realistic atomic data distributions, we randomly shuffle the OMat24 dataset (Barros-Luque et al., 2026) and sample structures to form batches with specified atom counts. For each structure, we construct a neighbor graph under periodic boundary conditions with rcut = 6.0 Å, thereby preserving the neighbor distribution of the original dataset. Workload construction protocol.

For all microbenchmarks, we warm up each configuration for 10 iterations and then measure GPU kernel execution time over the next 20 iterations using the PyTorch profiler. We report the mean runtime, excluding initial JIT compilation, graph construction, and host-to-device data transfers. Measurement protocol.

A.1

Kernel Microbenchmark Settings

All kernel microbenchmarks in Section 5.1 use the same workload unless otherwise specified. The workload contains 768 nodes and 37304 edges, with an average of 48.6 neighbors per atom. To cover different angularmomentum orders and channel sizes, we use irrep configurations from several SOTA equivariant MLIPs, as summarized in Table 1. Table 1 Irrep configurations used in the kernel microbenchmarks. Case ID Lmax Hidden irreps

Edge irreps

Output irreps

Representative MLIP

224x0e+384x1e +352x2e

SevenNet-0 , GNoME

128x0e+64x1e+32x2e

1x0e+1x1e +1x2e

128x0e+128x1o+128x2e

1x0e+1x1o +1x2e+1x3o

384x0e+640x1o +640x2e+512x3o

MACE-MP-L

128x0e+64x1o+32x2e +32x3o

1x0e+1x1o +1x2e+1x3o

256x0e+352x1o +384x2e+320x3o

Nequix-MP/OAM , SevenNet-MF/Omni

256x0e+256x1o+256x2e +256x3o

1x0e+1x1o +1x2e+1x3o

1024x0e+1536x1o +1792x2e+1536x3o

EquFlash-OMat/OAM

320x0e+96x1o+64x2e +32x3o+32x4e

1x0e+1x1o+1x2e +1x3o+1x4e

544x0e+736x1o+832x2e +800x3o+704x4e

NequIP-XL series

128x0e+128x1o+128x2e +128x3o+128x4e

1x0e+1x1o+1x2e +1x3o+1x4e

640x0e+1024x1o+1280x2e +1280x3o+1152x4e

EquFlashV2-OMat/OAM

128x0e+64x1o+32x2e +32x3o+32x4e+32x5o

1x0e+1x1o+1x2e +1x3o+1x4e+1x5o

320x0e+480x1o+576x2e +608x3o+608x4e+512x5o

–

8∗

96x0e+96x1o+96x2e +96x3o+96x4e+96x5o

1x0e+1x1o+1x2e +1x3o+1x4e+1x5o

576x0e+960x1o+1248x2e +1344x3o+1344x4e+1152x5o

–

128x0e+64x1o+32x2e+32x3o 1x0e+1x1o+1x2e+1x3o 352x0e+544x1o+672x2e+736x3o +32x4e+32x5o+32x6e +1x4e+1x5o+1x6e +768x4e+736x5o+640x6e

–

10∗

64x0e+64x1o+64x2e+64x3o +64x4e+64x5o+64x6e

–

1x0e+1x1o+1x2e+1x3o 448x0e+768x1o+1024x2e+1152x3o +1x4e+1x5o+1x6e +1216x4e+1152x5o+1024x6e

“*” e3nn encounters OOM under the default workload; these cases therefore use 448 atoms and 20386 edges. “–” Since no existing MLIP uses the corresponding settings, we constructed these configurations for evaluation.

A.2

Scalability Analysis Settings

We evaluate LunaTP’s FP32 forward scalability along four dimensions: atom count, graph density, maximum degree, and channel size. In each test, we vary one dimension while holding the others fixed, and continue 14

until OOM or compilation timeout. Scaling with atom count. We combine distinct OMat structures and double the atom count starting from 16. We use a 6.0 Å neighbor cutoff, Lmax = 4, and 64 channels per degree. Scaling with graph density. Using a fixed 768-atom batch, we increase the neighbor cutoff from 3.5 Å to 30.5 Å while keeping Lmax = 4 and 64 channels per degree.

For this test and the channel-size test, we use the same fixed 768-atom graph. We hold the channel count at 64 per degree and increase Lmax from 0 in steps of one. Scaling with degree.

We fix the 768-atom graph and Lmax = 4, and increase the number of channels per degree from 32 by successive doubling. Scaling with channel size.

A.3

Training Setup Details

We benchmark training in FP32 on NVIDIA H100 GPUs, using OMat-1M (Barros-Luque et al., 2026) for EFS and MPTrj-PFT (Deng et al., 2023; Koker et al., 2026) for EFSH. All models are trained using their official PyTorch implementations (Batatia et al., 2025; Kim et al., 2026; Koker et al., 2026; Xu et al., 2026; Kavanagh et al., 2026; Willow et al., 2026). For EFSH, we follow the training protocol of Koker et al. (2026) and sample Hessian columns during training. Tables 2 and 3 summarize per-GPU batch sizes and the model configurations, respectively. 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. Atoms

SevenNet-Omni

MACE-L

NequIP-XL

EFS EFSH

2340 870

2450 800

1630 600

e3nn FlashTP OpenEq cuEq LunaTP

2340 14600 14600 14600 14600

2450 3850 3850 3850 3850

e3nn FlashTP OpenEq cuEq LunaTP

870 – 3830 3830 5130

800 – 2600 3140 6200

EquFlashV2

TACE-L

NequIX

Prophet

BAM

290 90

1030 300

8900 3200

430 120

1280 580

1630 9550 9550 3300 9550

MaxBS: EFS 290 2550 2550 2550 2550

1030 3940 3940 2340 3940

8900 30400 37600 14500 37600

430 2220 2220 510 2220

1280 3440 3440 1600 3440

600 – 2750 – 3730

MaxBS: EFSH 90 – 660 660 860

300 – 1380 1340 1650

3200 – 9870 3060 13200

120 – 600 120 680

580 – 1700 580 1700

Matched

Table 3 MLIP configurations for training. (a) MACE-L (10.10M) Hyper-parameter

Value

Layers Cutoff radius

2 5.0 Å

Stage

(b) NequIP-XL (32.14M) Hyper-parameter

Irreps

Edge Input / Output(L2)

128x0e / 128x0e

Hidden (L1)

128x0e +128x1o +128x2e

Hyper-parameter

Value

5 6.0 Å

Stage

Irreps

Edge Input / Output(L6)

Input / Output (L5)

Hidden (L1–4)

Irreps

Hidden (L1–9)

128x0e / 128x0e 128x0e+128x1o +128x2e+128x3o +128x4e

(f) Prophet (62.30M)

Value

Hyper-parameter

4 6.0 Å

Layers Cutoff radii Irreps

Value

10 3.0, 7.0 Å

Stage

Irreps

Edge

1x0e+1x1o +1x2e+1x3o

Edge

64x0e / 64x0e

Input / Output (L4)

89x0e / 128x0e

Input / Output(L10) 256x0e / 256x0e

64x0e+64x1o +64x2e+64x3o +64x4e+64x5o

Hidden (L1–3)

128x0e+64x1o +32x2e +32x3o

Hidden (L1–9)

(g) SevenNet-Omni (25.76M) Hyper-parameter

Layers Cutoff radius Stage

Hyper-parameter

5 6.0 Å

Layers Cutoff radius Irreps

1x0e+1x1e+1x2e+1x3e

Input / Output (L5)

128x0e / 128x0e

Hidden (L1, L4)

128x0e+64x1o+32x2e+32x3o

Hidden (L2)

128x0e+64x1o+64x1e +32x2+32x2e+32x3o+32x3e

1x0e+1x1o+1x2e +1x3o+1x4e

256x0e+96x1o +64x2e+32x3o +32x4e

(h) BAM (40.15M)

Value

Edge

Hidden (L3)

Input / Output(L10)

(e) NequIX (0.71M)

Stage

10 6.0 Å

1x0e+1x1o+1x2e +1x3o+1x4e

Edge

320x0e+96x1o +64x2e+32x3o +32x4e

Hyper-parameter

Value

Stage

32x0e / 320x0e

Hidden (L1–5)

Irreps

Edge

Layers Cutoff radius

1x0e+1x1o+1x2e +1x3o+1x4e

Layers Cutoff radius

1x0e+1x1o+1x2e +1x3o+1x4e+1x5o

Hyper-parameter

6 6.0 Å

Stage

(d) TACE-L (82.90M)

Layers Cutoff radius

Value

Layers Cutoff radius

1x0e+1x1o +1x2e+1x3o

Value

5 5.0 Å

Stage

128x0e+128x0e+64x1o+64x1e +32x2o+32x2e+32x3o+32x3e

Irreps

Edge

1x0e+1x1o +1x2e+1x3o

Input

128x0e

Hidden (L1–5)

64x0o+64x0e +64x1o+64x1e +64x2o+64x2e +64x3o+64x3e

Detailed Benchmark Data

B.1

Raw Data for Kernel Microbenchmarks

We evaluate LunaTP using the experimental setup described in Section A. Table 4 reports the raw results of our kernel microbenchmarks. Table 4 Kernel measurements from forward through third-order differentiation. Parentheses report runtime speedup

and peak-memory reduction relative to e3nn. Forward

Case ID e3nn

FlashTP

Backward

OpenEq

cuEq

LunaTP (ours)

e3nn

FlashTP

OpenEq

cuEq

LunaTP (ours)

Runtime (ms) C1 C2 C3 C4 C5 C6 C7 C8 C9 C10

3.56 (1.00x) 0.563 (6.32x) 0.248 (14.4x) 4.40 (0.809x) 0.0763 (46.7x) 9.85 (1.00x) 1.59 (6.19x) 0.681 (14.5x) 1.12 (8.81x) 0.163 (60.6x) 5.77 (1.00x) 0.947 (6.09x) 0.439 (13.1x) 6.80 (0.849x) 0.127 (45.4x) 28.3 (1.00x) 4.70 (6.03x) 1.99 (14.3x) 2.16 (13.1x) 0.449 (63.2x) 19.8 (1.00x) 3.10 (6.39x) 1.36 (14.6x) 14.7 (1.35x) 0.344 (57.6x) 39.5 (1.00x) 6.13 (6.45x) 2.46 (16.1x) 2.93 (13.5x) 0.586 (67.4x) 26.0 (1.00x) 4.33 (6.00x) 1.54 (16.9x) 17.0 (1.53x) 0.435 (59.7x) 34.7 (1.00x) 5.90 (5.88x) 2.23 (15.5x) 3.84 (9.05x) 0.544 (63.8x) 46.8 (1.00x) 8.67 (5.40x) 2.85 (16.4x) 27.3 (1.71x) 0.755 (62.0x) 44.5 (1.00x) 8.42 (5.29x) 3.05 (14.6x) 6.97 (6.39x) 0.713 (62.4x)

C1 C2 C3 C4 C5 C6 C7 C8 C9 C10

1.50 (1.00x) 0.157 (9.60x) 0.156 (9.62x) 1.51 (0.997x) 4.19 (1.00x) 0.359 (11.7x) 0.359 (11.7x) 0.391 (10.7x) 2.50 (1.00x) 0.218 (11.5x) 0.218 (11.5x) 2.47 (1.01x) 12.0 (1.00x) 0.972 (12.3x) 0.972 (12.3x) 1.06 (11.3x) 8.15 (1.00x) 0.607 (13.4x) 0.606 (13.5x) 8.00 (1.02x) 14.3 (1.00x) 0.907 (15.8x) 0.907 (15.8x) 1.01 (14.2x) 9.09 (1.00x) 0.535 (17.0x) 0.535 (17.0x) 8.86 (1.03x) 12.8 (1.00x) 0.628 (20.4x) 0.628 (20.5x) 0.713 (18.0x) 16.7 (1.00x) 0.781 (21.4x) 0.781 (21.4x) 15.1 (1.11x) 16.4 (1.00x) 0.658 (25.0x) 0.658 (25.0x) 0.758 (21.7x)

9.21 (1.00x) 1.29 (7.16x) 0.331 (27.8x) 9.59 (0.960x) 0.214 (43.1x) 24.0 (1.00x) 3.28 (7.32x) 1.21 (19.8x) 0.514 (46.7x) 0.366 (65.6x) 16.3 (1.00x) 1.95 (8.35x) 0.610 (26.8x) 17.3 (0.944x) 0.320 (50.9x) 76.9 (1.00x) 9.59 (8.03x) 9.40 (8.18x) 1.56 (49.3x) 1.02 (75.8x) 67.3 (1.00x) 7.08 (9.51x) 6.40 (10.5x) 64.6 (1.04x) 1.05 (63.8x) 130 (1.00x) 13.4 (9.74x) 13.0 (10.0x) 2.02 (64.5x) 1.53 (85.0x) 101 (1.00x) 9.76 (10.4x) 9.53 (10.6x) 92.2 (1.10x) 1.38 (73.3x) 133 (1.00x) 13.1 (10.2x) 11.8 (11.3x) 2.37 (56.1x) 1.54 (86.2x) 212 (1.00x) 21.3 (9.99x) 18.6 (11.4x) 181 (1.17x) 2.79 (76.3x) 199 (1.00x) 20.2 (9.87x) 17.0 (11.7x) 5.14 (38.8x) 2.28 (87.4x)

Peak memory (GB) 0.157 (9.59x) 0.360 (11.7x) 0.218 (11.5x) 0.972 (12.3x) 0.611 (13.4x) 0.907 (15.8x) 0.539 (16.9x) 0.629 (20.4x) 0.787 (21.3x) 0.659 (25.0x)

2.41 (1.00x) 0.331 (7.27x) 0.312 (7.72x) 1.52 (1.59x) 0.313 (7.70x) 7.38 (1.00x) 0.775 (9.53x) 0.718 (10.3x) 0.749 (9.85x) 0.718 (10.3x) 3.76 (1.00x) 0.471 (7.99x) 0.435 (8.65x) 2.31 (1.63x) 0.436 (8.64x) 20.2 (1.00x) 2.10 (9.63x) 1.94 (10.4x) 2.03 (9.93x) 1.94 (10.4x) 13.2 (1.00x) 1.33 (9.87x) 1.21 (10.9x) 7.00 (1.88x) 1.22 (10.8x) 28.0 (1.00x) 1.98 (14.2x) 1.81 (15.5x) 1.91 (14.6x) 1.81 (15.4x) 16.4 (1.00x) 1.18 (13.8x) 1.07 (15.3x) 7.18 (2.28x) 1.07 (15.3x) 22.5 (1.00x) 1.39 (16.2x) 1.26 (17.9x) 1.34 (16.8x) 1.26 (17.9x) 28.5 (1.00x) 1.64 (17.4x) 1.56 (18.3x) 11.9 (2.40x) 1.57 (18.2x) 27.4 (1.00x) 1.37 (20.0x) 1.32 (20.8x) 1.42 (19.4x) 1.32 (20.8x)

Double-bwd

Case ID e3nn

FlashTP

OpenEq

Triple-bwd

cuEq

LunaTP (ours)

e3nn

FlashTP

OpenEq

cuEq

LunaTP (ours)

47.2 (1.00x) 132 (1.00x) 89.5 (1.00x) 429 (1.00x) 401 (1.00x) 758 (1.00x) 629 (1.00x) 813 (1.00x) 1630 (1.00x) 1420 (1.00x)

– – – – – – – – – –

6.38 (7.41x) 19.9 (6.64x) 12.0 (7.48x) 82.1 (5.22x) 53.5 (7.49x) 109 (6.93x) 73.8 (8.53x) 104 (7.85x) 148 (11.0x) 145 (9.80x)

45.5 (1.04x) 4.77 (27.7x) 86.7 (1.03x) 23.7 (18.1x) 362 (1.11x) 112 (6.77x) 555 (1.13x) 243 (3.35x) 1230 (1.33x) OOM

0.683 (69.2x) 1.74 (75.9x) 1.17 (76.2x) 5.14 (83.3x) 3.98 (101x) 6.39 (119x) 5.13 (123x) 7.10 (114x) 11.4 (143x) 10.7 (132x)

– – – – – – – – – –

1.98 (3.20x) 4.52 (4.20x) 2.74 (3.63x) 12.2 (4.48x) 7.58 (4.60x) 11.3 (6.68x) 6.64 (6.87x) 7.77 (8.23x) 9.39 (8.50x) 7.97 (9.93x)

5.27 (1.20x) 2.16 (8.79x) 8.24 (1.21x) 5.87 (9.34x) 25.0 (1.39x) 5.47 (13.8x) 29.0 (1.57x) 3.79 (16.9x) 48.3 (1.65x) OOM

0.947 (6.69x) 2.18 (8.71x) 1.33 (7.49x) 5.93 (9.24x) 3.72 (9.38x) 5.56 (13.6x) 3.28 (13.9x) 3.84 (16.7x) 4.67 (17.1x) 3.97 (20.0x)

Runtime (ms) C1 C2 C3 C4 C5 C6 C7 C8 C9 C10

25.7 (1.00x) 1.95 (13.2x) 74.4 (1.00x) 4.71 (15.8x) 52.8 (1.00x) 2.78 (19.0x) 251 (1.00x) 14.1 (17.8x) 262 (1.00x) 10.9 (24.1x) 467 (1.00x) 21.0 (22.2x) 427 (1.00x) 15.9 (26.8x) 537 (1.00x) 21.5 (25.0x) 1010 (1.00x) 39.8 (25.4x) 895 (1.00x) 38.4 (23.3x)

1.12 (23.0x) 3.61 (20.6x) 2.23 (23.7x) 12.1 (20.8x) 7.77 (33.7x) 16.5 (28.4x) 10.6 (40.3x) 16.3 (33.0x) 22.2 (45.5x) 22.6 (39.5x)

25.3 (1.02x) 1.43 (52.1x) 52.0 (1.01x) 5.11 (49.2x) 245 (1.07x) 12.9 (36.3x) 392 (1.09x) 22.9 (23.5x) 926 (1.09x) 57.5 (15.6x)

0.403 (63.9x) 0.802 (92.8x) 0.651 (81.2x) 2.36 (106x) 1.96 (134x) 2.83 (165x) 2.20 (194x) 2.84 (189x) 4.04 (250x) 3.70 (242x)

C1 C2 C3 C4 C5 C6 C7 C8 C9 C10

3.80 (1.00x) 0.614 (6.19x) 0.604 (6.29x) 3.92 (0.969x) 0.614 (6.18x) 10.2 (1.00x) 1.41 (7.25x) 1.38 (7.41x) 1.44 (7.10x) 1.39 (7.35x) 5.71 (1.00x) 0.855 (6.68x) 0.836 (6.84x) 6.35 (0.900x) 0.855 (6.68x) 30.7 (1.00x) 3.81 (8.07x) 3.73 (8.24x) 3.90 (7.88x) 3.75 (8.20x) 20.7 (1.00x) 2.37 (8.75x) 2.31 (8.99x) 20.2 (1.03x) 2.38 (8.72x) 47.1 (1.00x) 3.52 (13.4x) 3.44 (13.7x) 3.64 (13.0x) 3.52 (13.4x) 28.9 (1.00x) 2.07 (14.0x) 2.01 (14.3x) 22.8 (1.27x) 2.08 (13.9x) 41.1 (1.00x) 2.41 (17.1x) 2.35 (17.5x) 2.52 (16.3x) 2.42 (17.0x) 52.5 (1.00x) 2.99 (17.5x) 2.91 (18.0x) 38.9 (1.35x) 3.02 (17.4x) 52.3 (1.00x) 2.50 (20.9x) 2.44 (21.4x) 2.64 (19.8x) 2.51 (20.8x)

Peak memory (GB)

“–” denotes an unsupported setting.

6.33 (1.00x) 19.0 (1.00x) 9.94 (1.00x) 54.8 (1.00x) 34.9 (1.00x) 75.6 (1.00x) 45.6 (1.00x) 63.9 (1.00x) 79.8 (1.00x) 79.2 (1.00x)

B.2

Raw Data for Training Benchmarks

Tables 5 and 6 report throughput and the CGTP convolution time fraction for all eight models under EFS and EFSH training. Table 5 EFS training performance. Each entry reports throughput in “atoms/s”, with the CGTP convolution time

fraction in parentheses. Setting

Throughput (CGTP time fraction)

Model e3nn

FlashTP

OpenEq

cuEq

LunaTP (ours)

Matched

SevenNet-Omni MACE-L NequIP-XL NequIX EquFlashV2 TACE-L Prophet BAM

936.6 (95.9%) 4023.5 (53.5%) 625.8 (93.8%) 4929.0 (92.8%) 144.0 (78.4%) 477.8 (79.0%) 160.2 (84.3%) 531.4 (82.3%)

8470.3 (51.2%) 7084.6 (10.7%) 4696.5 (57.5%) 26052.0 (46.1%) 547.1 (26.1%) 1625.5 (28.6%) 724.5 (31.5%) 2654.6 (16.3%)

9160.7 (47.5%) 7435.6 (7.9%) 5225.1 (48.9%) 40017.6 (41.0%) 538.0 (19.8%) 1748.3 (23.3%) 782.7 (24.8%) 2538.7 (15.7%)

358.7 (76.5%) 774.6 (59.7%) 370.7 (70.3%) 4989.0 (88.9%) 504.6 (9.5%) 245.5 (69.5%) 149.2 (50.3%) 601.7 (80.0%)

18024.6 (21.3%) 8455.9 (1.4%) 7703.4 (20.2%) 52397.8 (20.9%) 659.7 (6.0%) 2099.9 (4.6%) 824.7 (11.0%) 2824.5 (3.1%)

MaxBS

SevenNet-Omni MACE-L NequIP-XL NequIX EquFlashV2 TACE-L Prophet BAM

936.6 (95.1%) 4023.5 (53.5%) 625.8 (93.8%) 4929.0 (92.8%) 144.0 (78.4%) 477.8 (79.0%) 160.2 (84.3%) 531.4 (82.3%)

11136.3 (71.6%) 7521.5 (11.3%) 5648.9 (69.1%) 24715.1 (43.4%) 1390.4 (65.6%) 2614.6 (46.6%) 1076.0 (43.9%) 2659.0 (15.7%)

12469.2 (61.6%) 7759.8 (8.6%) 6961.3 (64.0%) 42825.2 (44.1%) 1577.2 (60.5%) 2858.3 (40.5%) 1174.9 (38.6%) 2640.0 (15.9%)

361.1 (82.8%) 1019.9 (78.0%) 368.8 (71.1%) 5046.9 (66.8%) 2147.8 (38.1%) 242.6 (69.5%) 169.0 (53.0%) 618.5 (80.0%)

23933.1 (26.1%) 9289.7 (1.5%) 12844.3 (26.2%) 52424.9 (20.9%) 3311.9 (19.5%) 4563.7 (9.1%) 1583.3 (14.2%) 2998.9 (3.2%)

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

time fraction in parentheses. Setting

Throughput (CGTP time fraction)

Model e3nn

FlashTP

OpenEq

cuEq

LunaTP (ours)

Matched

SevenNet-Omni MACE-L NequIP-XL NequIX EquFlashV2 TACE-L Prophet BAM

277.5 (93.4%) 1204.6 (56.1%) 182.5 (92.0%) 1465.7 (89.9%) 47.0 (55.4%) 122.6 (47.0%) 34.0 (53.9%) 182.9 (83.7%)

– – – – – – – –

2036.4 (52.4%) 2371.2 (13.5%) 1100.4 (52.9%) 8352.2 (56.3%) 85.6 (11.2%) 202.5 (12.5%) 76.0 (6.3%) 949.9 (26.7%)

156.1 (71.7%) 486.3 (48.7%) – 1506.9 (91.8%) 72.5 (11.6%) 116.3 (52.2%) 29.8 (64.7%) 179.7 (86.1%)

3232.3 (18.0%) 2775.5 (2.3%) 1480.7 (22.6%) 14369.5 (25.0%) 91.4 (9.2%) 221.5 (3.5%) 80.7 (4.2%) 1138.2 (6.2%)

MaxBS

SevenNet-Omni MACE-L NequIP-XL NequIX EquFlashV2 TACE-L Prophet BAM

277.5 (93.4%) 1204.6 (56.1%) 182.5 (92.0%) 1465.7 (89.9%) 42.7 (55.4%) 122.6 (46.8%) 39.1 (89.3%) 182.9 (83.7%)

– – – – – – – –

2562.0 (72.2%) 2979.7 (9.5%) 1490.5 (78.0%) 9082.6 (60.0%) 307.1 (64.3%) 635.7 (50.6%) 256.0 (43.3%) 1029.6 (30.1%)

166.9 (77.8%) 755.92 (87.7%) – 1401.0 (92.9%) 247.5 (64.1%) 150.3 (69.3%) 36.0 (64.4%) 179.7 (86.1%)

7956.0 (30.2%) 4248.1 (2.7%) 4627.4 (33.4%) 16450.9 (27.0%) 704.6 (17.6%) 1233.1 (9.8%) 408.5 (15.8%) 1386.7 (5.4%)

“–” denotes an unsupported setting.

Ablation Studies

We evaluate the effects of path fusion scope, cross-path input reuse, task ordering, and compile-time reuse on LunaTP’s forward performance. All experiments use the C5 configuration in Table 1 with FP32 precision. This configuration includes channel counts of 320, 96, 64, and 32.

C.1

Path Grouping

Path grouping needs to balance input reuse and accumulator storage. We first vary the per-channel output component budget β to examine the effect of fusion scope (Table 7). Among the tested settings, the default budget β = 16 gives the lowest latency. Fusing all compatible paths increases latency by 56.3%, while reducing the budget to 8 and 4 increases latency by 11.6% and 20.8%, respectively. Smaller groups require more repeated edge traversals, while full fusion keeps more accumulators live at once and is the only tested setting with reported register spilling. Since the default policy automatically adjusts the number of warps as the grouping changes, we also compare the fusion scopes with fixed warp counts. With 8 and 16 warps, full fusion is still 44.3% (0.5001 vs. 0.3467 ms) and 53.0% (0.5423 vs. 0.3545 ms) slower, respectively. In both cases, the reported spill stores are 2976 B for full fusion and zero for the default budget. 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 All compatible paths Default (β = 16) β=8 β=4

Path groups

Execution units

Registers

Spill stores (B)

Latency (ms)

5 17 29 37

17 45 73 96

64 64 56 56

2976 0 0 0

0.5393 0.3449 0.3851 0.4167

Performance also depends on which paths can share inputs within a group. Compared with the baseline that groups paths by edge irrep, broader cross-path input reuse reduces the number of path groups from 25 to 17 and lowers latency by about 6.6% (Table 8). This change both expands input reuse and reduces the number of tasks, so the measured gain reflects the overall effect of the grouping change. With this grouping, enabling the work partitioning option shown in the table leaves latency nearly unchanged, with no additional benefit observed on C5. Table 8 Cross-path input reuse and work partitioning (C5, FP32, forward). All configurations use node-centric

execution.

Configuration

Path groups

Execution units

Latency (ms)

25 17 17

85 45 45

0.3688 0.3446 0.3447

Baseline (grouped by edge irrep) + Cross-path input reuse + Work partitioning (full)

C.2

Task Ordering

We next vary the task order while keeping the fusion scope and number of warps fixed. We independently reverse the ordering by nonzero CG-term count and by node degree, while keeping the task set unchanged (Table 9). Compared with descending order on both axes, changing only the node order to ascending increases latency by 13.0%. Changing only the CG-term count order to ascending increases latency by 5.8%, while changing both increases it by 17.1%. Sorting both in descending order, as used in the main text, gives the lowest latency among these settings.

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.

C.3

CG-term count order

Node degree order

Latency (ms)

vs. both descending

Descending Descending Ascending Ascending

Descending Ascending Descending Ascending

0.3446 0.3895 0.3645 0.4036

— +13.0% +5.8% +17.1%

Compile-Time Reuse

Finally, we compare all four combinations of product reuse and coefficient reuse with a fixed fusion scope and number of warps (Table 10). Enabling both gives the lowest latency of 0.3449 ms. Disabling coefficient reuse increases latency by 3.7%, while disabling both increases it by 3.4%. These two switches are not fully independent. The code generation path for coefficient reuse retains shared products, so disabling the product reuse switch does not necessarily remove product sharing from these computations. When coefficient reuse is disabled, disabling product reuse also does not increase latency. 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.

Product reuse

Coefficient reuse

Enabled Enabled Disabled Disabled

Enabled Disabled Enabled Disabled

Latency (ms)

vs. both enabled

0.3449 0.3578 0.3526 0.3565

— +3.7% +2.2% +3.4%

Additional Kernel Microbenchmark Results

D.1

Kernel microbenchmark (FP64)

We evaluate FP64 performance using the same irrep configurations as in the FP32 benchmarks. Since FP64 requires more computation and memory, we reduce the atom count by half to 384 for most cases and further reduce it to 192 for C8 and C10. All runtime speedups and peak-memory reduction factors are relative to FP64 e3nn on the same hardware. Figure 8 summarizes the results. LunaTP achieves arithmetic-mean FlashTP

OpenEq

Gain over e3nn (×)

(a) Forward 128 Runtime 64 32 16 8 4 2 1

cuEq (uniform)

cuEq (naive)

(b) Backward

LunaTP (ours)

Avg. speedup (d) Triple-bwd

133.7×

81.9×

9.1×

8.5×

36.5× 19.9× 19.5×

1.0×

1 2 3 4 5 6 7 8 9 10

30 Memory 20 10 1

45.6×

58.4× 34.2×

10.4× 6.9× 6.1×

15.7×

8.6×

Case ID

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

speedups of 45.6×, 58.4×, 133.7×, and 81.9× over FP64 e3nn for forward, backward, double backward, and triple backward, respectively. The corresponding speedups over OpenEq are 4.42×, 4.40×, 3.56×, and 9.61×.

D.2

Kernel microbenchmark (A100 & RTX 3090)

We evaluate LunaTP on NVIDIA A100-80GB and RTX 3090-24GB GPUs using FP32 microbenchmarks. On A100, we use the same workloads as in the main evaluation. To fit within the RTX 3090’s memory capacity, we use 192 atoms by default and reduce C8 and C10 to 122 and 97 atoms, respectively. For each case, we keep the workload fixed across all backends and all four stages. On A100, LunaTP achieves average speedups of 49.7×, 62.4×, 122.7×, and 90.5× over e3nn for forward, backward, double backward, and triple backward, respectively. On RTX 3090, the corresponding speedups are 75.3×, 83.5×, 203.5×, and 151.2×. Figures 9 and 10 show the full results.

Gain over e3nn (×)

FlashTP 256 Runtime 128 64 32 16 8 4 2 1

OpenEq

(a) Forward

cuEq (uniform)

cuEq (naive)

(b) Backward

LunaTP (ours)

Avg. speedup (d) Triple-bwd 90.5×

49.7×

62.4×

17.0×

44.7× 15.2×

122.7×

6.6×

1 2 3 4 5 6 7 8 9 10

8.3×

8.0×

19.1×

1.0×

1.1×

1 2 3 4 5 6 7 8 9 10

1.2×

30 Memory 20 10 1

11.9×

35.3× 32.2×

13.7×

Case ID

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

FlashTP 512

OpenEq

(a) Forward

cuEq (uniform)

cuEq (naive)

(b) Backward

Avg. speedup (d) Triple-bwd

Runtime

203.5×

128

Gain over e3nn (×)

LunaTP (ours)

151.2×

75.3×

23.3×

24.8×

83.7× 83.5× 29.0×

7.9× 6.8×

59.8× 59.8× 22.5×

9.9×

14.1×

0.6×

0.9×

1 2 3 4 5 6 7 8 9 10

0.5 30 Memory 20 10 1

Case ID

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

LunaTP delivers performance gains on both data-center and consumer GPUs, although the speedups vary across hardware. Across all four stages, its average speedups over e3nn are highest on RTX 3090, followed by H100 and A100. This ordering matches that of the GPUs’ peak FP32 compute-to-memory-bandwidth ratios. Since e3nn’s tensor-product computations are primarily limited by memory bandwidth (Lee et al., 2025), LunaTP’s reductions in intermediate memory traffic through kernel fusion and data reuse yield larger gains on GPUs with greater compute capacity relative to memory bandwidth.

Details on Numerical Stability and Training Convergence

E.1

Operator-Level Numerical Stability

Following the numerical validation protocol of Shah et al. (2024), we use the e3nn FP64 as the reference. We run e3nn in FP32 and LunaTP in both FP32 and FP64 on the same inputs, and report RMSE for the forward and each derivative stage. The inputs are constructed as described in Appendix A. Table 11 reports the full results, averaged across all cases in Table 1. Table 11 Numerical error comparison between e3nn and LunaTP.

E.2

Stage

Output

e3nn (FP32)

LunaTP (FP32)

LunaTP (FP64)

Forward

x′

9.910429e-07

9.978108e-07

9.517594e-16

Backward

x̄ ȳ w̄

3.582043e-06 3.516555e-06 1.696597e-07

6.545904e-06 2.934879e-06 1.712254e-07

6.437515e-15 8.966707e-15 4.281635e-16

Double bwd

hx hy hw hg

5.083432e-06 5.002968e-06 2.546925e-07 1.769849e-06

7.955895e-06 4.300956e-06 2.582805e-07 2.792170e-06

1.274438e-14 5.290353e-15 6.404447e-16 7.214677e-15

Triple bwd

tx ty tw tg tvx tvy tvw

5.818900e-06 6.808703e-06 3.152390e-07 2.530354e-06 5.368344e-06 5.896732e-06 2.677879e-07

6.511971e-06 6.200321e-06 2.198522e-07 3.285050e-06 5.646177e-06 5.317266e-06 2.514537e-07

9.811956e-15 9.310777e-15 6.591994e-16 5.920629e-15 8.875174e-15 9.918370e-15 6.436796e-16

End-to-End Training Convergence

We evaluate end-to-end training convergence by training NequIP-XL on MPTrj-PFT with either e3nn or LunaTP as the CGTP backend. Both runs use FP32 precision with the same model initialization, data, and training settings. The loss weights for energy, forces, stress, and Hessian are set to 20, 20, 100, and 100, respectively. We train for 50 epochs with a cosine learning-rate schedule, starting at 5 × 10−4 and ending at 1 × 10−6 . Figure 11 shows the training and validation MAE curves for both backends. e3nn train

e3nn validation

LunaTP train

(b) Force (meV/Å)

(a) Energy (meV/atom)

LunaTP validation 3

(d) Hessian (eV/Å2)

Validation MAE

3.067 × 10

Validation MAE

103

MAE

Validation MAE

2.644 × 10−3

7.042 × 101

2.651 × 10−3

7.030 × 10

3.069 × 101

Validation MAE

2.297 × 10−2 2.302 × 10−2

10−1

10−2 102

102

101

10−3 0K

20 K

40 K

60 K

80 K

20 K

40 K

60 K

80 K

20 K

40 K

60 K

Step

Figure 11 Loss convergence on MPTrj-PFT dataset.

80 K

20 K

40 K

60 K

80 K

Differentiation of CGTP convolution

In this section, we derive the first-, second-, and third-order derivatives of tensor product convolution. The resulting formulas evaluate the gradients needed at each order without forming full Jacobian, Hessian, or thirdorder derivative tensors, and provide the basis for our derivative kernels. Throughout, we present the derivatives (l ,c) (l ) componentwise, following Eq. 1. For each channel–path combination, xj 1 ∈ R2l1 +1 , yji2 ∈ R2l2 +1 , and (l ,l ,l,c)

wji1 2

∈ R. (l,c)

(l,c)

For the convolution output x′ i , let gi = ∂L/∂x′ i denote its gradient. Applying the chain rule to the weighted CGTP gives the first-order derivatives in Algorithm 2. The first-order derivative.

Algorithm 2 First-Order Derivative of Clebsch–Gordan Tensor Product Convolution (l ,c)

Input: xj 1

: node irreps

(l ) Input: yji2 : angular edge irreps (l ,l ,l,c) Input: wji1 2 : path weights (l,c) Input: gi : gradients with respect to the convolution output

Input: edges ji, from source node j to destination node i Output: x̄, ȳ, and w̄

1: Initialize all three outputs to zero 2: for ji in the E edges do 3: for c in the C channels do 4: for (l1 , l2 ) → l in the P paths do 5:

(l1 ,c) (l ,l ,l,c) x̄j,m += wji1 2 1

(l )

(l ,l ,l,c)

2 += wji1 2 ȳji,m 2

l X

l2 X

m=−l m2 =−l2 l1 l X X

(l,m)

(l )

(l,m)

(l ,c) (l,c)

1 gi,m C(l1 ,m1 )(l2 ,m2 ) xj,m 1

m=−l m1 =−l1

(l ,l ,l,c)

w̄ji1 2

8: end for 9: end for 10: end for

l X m=−l

(l,c)

2 C(l1 ,m1 )(l2 ,m2 ) yji,m g 2 i,m

i(l) h (l,c) (l ,c) (l ) gi,m xj 1 ⊗ yji2 m

The second-order derivative. Let vx , vy , and vw denote the gradients with respect to x̄, ȳ, and w̄, respectively. Applying the product rule to the first-order formulas gives the second-order derivatives in Algorithm 3.

Algorithm 3 Second-Order Derivative of Clebsch–Gordan Tensor Product Convolution (l ,c)

Input: xj 1

: node irreps

(l ) Input: yji2 : angular edge irreps (l ,l ,l,c) Input: wji1 2 : path weights (l,c) Input: gi : gradients with respect to the convolution output (l1 ,c) (l2 ) (l1 ,l2 ,l,c) Input: vx,j , vy,ji , and vw,ji : gradients with respect to the first-order outputs

Input: edges ji, from source node j to destination node i Output: hx , hy , hw , and hg

1: Initialize all four outputs to zero 2: for ji in the E edges do 3: for c in the C channels do 4: for (l1 , l2 ) → l in the P paths do 5:

l2 X

(l ,c)

l X

(l )

m=−l m2 =−l2 l1 l X X

1 += hx,j,m 1

(l,m) (l,c) (l2 ) (l1 ,l2 ,l,c) (l2 ) (l1 ,l2 ,l,c) C(l1 ,m1 )(l2 ,m2 ) gi,m vy,ji,m w v + y ji ji,m2 w,ji 2 (l ,l ,l,c) (l1 ,c) (l1 ,l2 ,l,c) (l,m) (l,c) (l1 ,c) vw,ji wji1 2 + xj,m C(l1 ,m1 )(l2 ,m2 ) gi,m vx,j,m 1 1

2 += hy,ji,m 2

m=−l m1 =−l1 ! l h i(l) h i(l) X (l1 ,l2 ,l,c) (l,c) (l1 ,c) (l2 ) (l1 ,c) (l2 ) hw,ji += gi,m vx,j ⊗ yji + xj ⊗ vy,ji m m m=−l ! (l,c)

(l ,l ,l,c)

hg,i += wji1 2

(l ,c)

(l )

vx,j1

⊗ yji2

i(l)

h i(l) (l ,c) (l2 ) + xj 1 ⊗ vy,ji

8: (l ,l ,l,c)

1 2 + vw,ji

(l ,c)

xj 1

(l )

i(l)

9: end for 10: end for 11: end for

Let qx , qy , qw , and qg denote the gradients with respect to hx , hy , hw , and hg , respectively. Differentiating the second-order formulas gives the third-order derivatives in Algorithm 4. The third-order derivative.

Algorithm 4 Third-Order Derivative of Clebsch–Gordan Tensor Product Convolution (l ,c)

Input: xj 1

: node irreps

(l )

Input: yji2 : angular edge irreps (l ,l ,l,c)

Input: wji1 2

: path weights

(l,c) Input: gi : gradients with respect to the convolution output (l1 ,c) (l2 ) (l1 ,l2 ,l,c) Input: vx,j , vy,ji , and vw,ji : gradients with respect to the first-order outputs (l1 ,c) (l2 ) (l1 ,l2 ,l,c) (l,c) Input: qx,j , qy,ji , qw,ji , and qg,i : gradients with respect to the second-order outputs

Input: edges ji, from source node j to destination node i Output: tx , ty , tw , tg , tvx , tvy , and tvw

1: Initialize all seven outputs to zero 2: for ji in the E edges do 3: for c in the C channels do 4: for (l1 , l2 ) → l in the P paths do l2 X

l X

(l ,c)

1 tx,j,m += 1

h (l,m) (l,c) (l2 ) (l1 ,l2 ,l,c) (l2 ) (l1 ,l2 ,l,c) C(l1 ,m1 )(l2 ,m2 ) gi,m qy,ji,m v + v q y,ji,m2 w,ji 2 w,ji

m=−l m2 =−l2

i (l,c) (l2 ) (l ,l ,l,c) (l2 ) (l ,l ,l,c) + qg,i,m vy,ji,m wji1 2 + yji,m v 1 2 2 2 w,ji l1 l h X X (l1 ,c) (l1 ,l2 ,l,c) (l,m) (l,c) (l1 ,c) (l1 ,l2 ,l,c) (l2 ) q v + v C g q += ty,ji,m x,j,m1 w,ji x,j,m1 w,ji (l1 ,m1 )(l2 ,m2 ) i,m 2 6:

m=−l m1 =−l1

i (l,c) (l1 ,c) (l1 ,l2 ,l,c) (l1 ,c) (l1 ,l2 ,l,c) + qg,i,m vx,j,m w + x v ji j,m1 w,ji 1 ! l h i(l) h i(l) X (l1 ,l2 ,l,c) (l,c) (l1 ,c) (l2 ) (l1 ,c) (l2 ) tw,ji += gi,m qx,j ⊗ vy,ji + vx,j ⊗ qy,ji m

m=−l

h i(l) (l ,c) (l2 ) + xj 1 ⊗ vy,ji m m ! h i(l) h i(l) (l ,c) (l2 ) (l ,c) (l2 ) qx,j1 ⊗ vy,ji + vx,j1 ⊗ qy,ji

(l,c) + qg,i,m (l,c)

(l ,l ,l,c)

tg,i += wji1 2 8:

(l ,c) (l ) vx,j1 ⊗ yji2

i(l)

h i(l) h i(l) (l ,c) (l ) (l ,c) (l2 ) qx,j1 ⊗ yji2 + xj 1 ⊗ qy,ji

(l1 ,l2 ,l,c) + vw,ji

i(l) i(l) h h (l ,c) (l2 ) (l ,c) (l ) + xj 1 ⊗ vy,ji vx,j1 ⊗ yji2

(l1 ,l2 ,l,c) + qw,ji

(l ,c)

tvx1,j,m1 += 9:

l2 X

l X

h (l,m) (l,c) (l ,l ,l,c) (l2 ) (l2 ) (l1 ,l2 ,l,c) C(l1 ,m1 )(l2 ,m2 ) gi,m wji1 2 qy,ji,m2 + yji,m q 2 w,ji

m=−l m2 =−l2 (l,c)

(l ,l ,l,c) (l2 ) yji,m2

+ qg,i,m wji1 2 (l )

tvy2,ji,m2 += 10:

l1 X

l X

h (l ,l ,l,c) (l1 ,c) (l1 ,c) (l1 ,l2 ,l,c) (l,m) (l,c) C(l1 ,m1 )(l2 ,m2 ) gi,m wji1 2 qx,j,m1 + xj,m q 1 w,ji

m=−l m1 =−l1

i (l,c) (l ,l ,l,c) (l1 ,c) + qg,i,m wji1 2 xj,m1 ! l i(l) h i(l) h X (l1 ,l2 ,l,c) (l,c) (l1 ,c) (l2 ) (l1 ,c) (l2 ) qx,j ⊗ yji + xj ⊗ qy,ji tvw ,ji += gi,m 11:

m=−l

h i(l) (l,c) (l ,c) (l ) + qg,i,m xj 1 ⊗ yji2 m

end for end for 14: end for 12: 13:

Programming Interface

LunaTP replaces the gather–TP–reduce sequence with a single convolution call, using the same e3nn irreps, instructions, and per-edge weights. It also supports precision selection and optional envelope smoothing of the messages (Gasteiger et al., 2022; Fu et al., 2025; Zhou et al., 2026; Liao et al., 2026). The code below illustrates the change in usage. e3nn

from e3nn import o3 # 1.Initialize e3nn_tp = o3.TensorProduct( irreps_in, irreps_edge, irreps_out, instructions, internal_weights=False, shared_weights=False, ) # 2.Forward x_edge = node.index_select(0, src) # Gather messages = e3nn_tp(x_edge, edge_attr, weight) # Tensor product if use_envelope: messages = messages * envelope # Apply envelope out = node.new_zeros((node.shape[0], e3nn_tp.irreps_out.dim)) out.index_add_(0, dst, messages) # Reduce

LunaTP

import LunaTP # 1.Initialize from an e3nn TP or custom instructions. lunatp_conv = LunaTP.initialize(e3nn_tp) # 2.Forward out = lunatp_conv( node, edge_attr, weight, env=envelope, # Default: ’none’ (no envelope) dtype=’fp32’ )

Beyond the convolution formulation considered in this work, LunaTP supports a broader range of tensor products and their convolutional forms, including those using the uuu, uvu and uvw connection modes in e3nn.

LunaTP: Enabling Efficient and Scalable Tensor Product for Equivariant MLIPs

Analyze & explore related research in Conceptio

Related documents