FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
Authors: Seung Yul Lee
Hojoon Kim
Yutack Park
Dawoon Jeong
Seungwu Han
Yeonhong Park
Jae W. Lee
Venue: ICML 2025 spotlightposter (2025)
Keywords: Equivariant neural networks
Tensor Product
Software libraries
Efficiency
Machine-learned interatomic potential (MLIP)
Machine Learning Force Fields (MLFF)
Primary area: applications->chemistry_physics_and_earth_sciences
TL;DR: FlashTP accelerates equivariant MLIPs by optimizing Tensor-Product operations, achieving up to speedup and significantly reducing memory footprint.

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

FlashTP:
Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
Seung Yul Lee!
Hojoon Kim
Yutack Park
Dawoon Jeong'
Seungwu Han 12 Yeonhong Park
Jae W. Lee
Abstract
force errors but also real physical observables derived from MD simulations (Fu et al., 2022; Kim et al., 2023). Equivariant MLIPs typically follow the principles of graph neural networks, taking as input a graph where atoms are represented as nodes and their interactions as edges.
Machine Learning Interatomic Potentials (MLIPs) enable efficient molecular dynamics (MD) sim- ulations with high accuracy. While equivariant MLIPs achieve stte-off-theaae accuracy, they face significant computational bottlenecks centered around their Tensor-Product layer, which account for up to 75% of training time and cause substan- tial memory overhead. We present FlashTP, a highly optimized tensor-product library that ad- dresses these inefficiencies through kernel fusion, sparse computation, and path-aggregated execu- tion. FlashTP achieves up to 41.6 41.6x and 60.8x kernel speedups over e3nn and NVIDIA cuEquiv- ariance, respectively. For SevenNet-B3i5, it deliv- ers 4.2x and 3.5x speedup while reducing peak memory usage by 6.3x and 6.2x for inference and training, respectively. The code is available at https//gittub.co////UAAA//aaaaT
Equivariant MLIP models, however, suffer from slow in- ference and training speeds, as well as high memory re- quirements (Passaro & Zitnick, 2023; Luo & Krishnapriyan, 2024; Zhang et al., 2024; Xie et al., 2024). The key bottle- neck of equivariant MLIP model lies in its Tensor-Product layers, which serves as the core operation responsible for up- dating the hidden states of nodes while incorporating inter- actions with neighboring nodes. For instance, in SevenNet- 1315 (Park et al., 2024a;b), state-of-the-an of-tthe-rtt equivariant MLIP model, the Prosout layer accounts for 89% of the inference time and 75% of training time. Moreover, the Tensor-Produet layer operates on edges rather than nodes, leading to a colossal memory footprint since the number of edges far exceeds the number of nodes (e.g., 37 x more edges on average for SevenNet on MPF dataset (Chen & Ong. 2022)).
1. Introduction
The inefficiency of the Tensor-Product layer stems from three sources. First, the Tensor-Product layer consists of multiple distinct kernels, and the intermediate data ex- changed between these kernels generate substantial memory traffic. Second, the output data produced by the Tensor- Product layer significantly large, leading to memory spikes that limit the scalability of MD simulations. Lastly, large amount of ineffectual computation arises from the high sparsity of the Clebsch- Gordan (CG) coefficient ma- trix (Varshalovich et al., 1988), a constant matrix that en- sures equivariance of updated node hidden states after their tensor product with connected edges and neighbors.
Machine learning interatomic potential (MLIP) is a method to predict energy and forces given atom positions and atomic numbers with machine learning models. is gaining signif- icant attention as means to accelerate molecular dynamics (MD) simulations, which play crucial role in computa- tional materials science and chemistry. MLIP enables sim- ulations to run orders of magnitude faster than quantum mechanical methods while preserving their accuracy (Unke et al., 2021; Deringer et al., 2021; Ko & Ong, 2023).
Among the various types of MLIPs, equivariant MLIPs (Thomas et al., 2018; Batzner al., 2022: Batatia et al., 2022; Musaelian et al., 2023; Liao et al., 2024) have demonstrated superior performance across complex benchmarks that assess not only energy and
To address the aforementioned issues, we present FlashTP, specialized Tensor-Product library designed to enhance the efficiency of equivariant MLIP model by both accelerating its speed and reducing memory requirements. FlashTP fully fuses all operations within the Tensor- Product layer, signifi- cantly reducing the memory traffic caused by intermediate data. Additionally, the fully fused Tensor-Produs layer is again fused with its subsequent layer, climinating mem- ory spikes that would otherwise be caused by large output
'Seuul National University Korea Institute for Advanced Study. Correspondence to: Jac Lee <jacwlee wn..ac.ssuu............ Yeonhong Park <ilil96@snu.ac.kr an?>
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
Inference
proven to be effective in a wide range of applications, in- cluding the discovery of new materials (Hwang et al., 2023; Kruglov et al., 2023), simulations of semiconductor pro- cesses (Hong et al., 2024), and modeling of solid-state elec- trolytes ( (Wang et al., 2022; Lee et al., 2024) for safer and more energy-efficient batteries.
Backward
Energy
Energy
Forward
Force
Force
@ Backward
Bouble-Backwand
Training
Figure 1. Overview of MLIP model inference and training pipeline. Lrowe and Liamy indicates the loss of the predicted forces and energy, respectively.
MLIP Inference and Training Pipeline. Figure illus- trates the inference and training pipeline of MLIP models. Inference of MLIP models consists of two phases. The first phase involves the forward computation of the model to pre- dict the total energy which the forces are computed by taking the gradient of the energy
tensors. FlashTP also optimizes computation by skipping ineffectual operations. Lastly, as an additional optimization, FlashTP introduces path-aggregation, a technique that in- creases input reuse within the accelerating the execution speed.
The second phase follows, during
This corresponds to the backward computation of the first phase. It is important to note that this backward computation is performed with respect to the result of the first phase (energy) and not its loss, making it distinct from backpropagation used for parameter updates.
layer, further
Through these optimizations, when applied the state- of-the-art SevenNet-IBi5 model on MPF dataset, FlashTP achieves 4.2x faster inference and 3.5x faster training com- pared to e3nn (Geiger& & Smidt, 2022; Geiger et al., 2022), a widely used framework for building equivariant MLIP models, while reducing peak memory usage by 6.3x and 6.2x for inference and training, respectively.
For training, two additional phases are performed to update the model parameters. These correspond to the backward computations of the two inference phases. First, a backward- of of-backward or Nouble-backwart computation performed Fol-
with respect to the loss of the predicted forces lowing this, the backward computation is performed with While phases and 4 are computationally identical, they differ in their inputs: phase takes the energy as input, while phase takes the loss of energy as input.
Summary of Contributions.
respect to the loss of the predicted energy
We identify memory traffic, memory consumption, and ineffectual computations as the key bottlenecks in the Tensor-Product layer.
2.2. Equivariant MLIP
We develop novel optimization techniques including kernel fusion, sparse computation skipping, and path- aggregated execution to address these bottlenecks.
Equivariant MLIPs (Batzner et al., 2022: Batatia et al., 2022; Musaclian et al., 2023; Passaro & Zitnick, 2023; Liao et al., 2024; Park et al., 2024b) are class of interatomic potentials where both the predicted atomic forces and all intermediate layer outputs transform consistently with rotations of the molecular system (Thomas et al., 2018). These models have proven to be more data-efficient (Cohen & Welling, 2016; Batzner et al., 2022) and excel at predicting physical properties in MD simulations (Fu et al., 2022; Kim et al., 2023). While non-equivariant models exist (Schütt et al.,
We demonstrate that FlashTP accelerates both infer- ence (4.2x) and training (3.5x) for a state-of-the-art equivariant MLIP model.
We show that FlashTP reduces peak memory usage by 6.3x x for inference and 6.2x for training in equivariant MLIP model.
2017: Chen et al., 2019: Gasteiger et al., 2020; 2021), we focus on equivariant MLIPs due to their superior accuracy and widespread adoption in the field.
2. Background
2.1. Machine Learning Interatomic Potential (MLIP)
MLIP Basics. Energy and atomic forces estimated using interatomic potentials govern the reliability of molecular dynamics (MD) simulations, which are widely used in ma- terials science and chemistry. Recently, machine learning interatomic potentials (MLIPs) have gained significant atten- tion as promising solution for these estimations. Trained on high-accuracy data generated using quantum mechanical methods, MLIP models predict energy, from which forces can be derived by computing gradients, achieving both high accuracy and computational efficiency. MLIP models have
Model Architecture. Figure 2(a) illustrates the architec- ture of NequlP (Batzner et al., 2022), which serves as the foundation for typical equivariant MLIP models such as MACE (Batatia et al., 2022) and SevenNet (Park et al., 2024b). The model is based on the principles of graph neural networks (GNNs), taking a graph as input where nodes correspond to atoms and edges denote atomic interac- tions (Gilmer et al., 2017; Reiser et al., 2022). These models aim to generate latent representations for all nodes in the graph, from which each node's 's energy is predicted.
2

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
(a)
Energy
(b)
X X Mch. dirl
Inputs Output Feature Definitions
4-11
--
State [n_ odges, n_ ch, d_i)
MLP
Activation
a: Edge Angular Festure [n_ edges, d_a]
Xnnal
[_r: Edge Radial Feature (atter EQ-MLP)
[n edges, n ch, n pl
DouterPooduct
Interaction
Output [o [__ddge,, n_ch, d'_]
Block k 14,
Eq-Linear
Eq-Linear
iterate over edges and channels
7 Roddess Beh
1: for edge in rangein edges):
2: for channel in rangein chj:
Interaction Block i 1x,
θ
Reduce
3: p= #t Tracks valid path indices
[Redges d1
Iterate over all possible paths Lh, _e, Lout oil parh:
OP MMatrix Multiplication
Perform TensorProduct oout, max]:
valid path
Tensor-Product
test valid()
Interaction
9:
Compute (G matrix för given path
Block
[eddee Meb d,]
Medgess Rel pptthl
11:
CG matrii cG( Le, Lour]
1%
12:
Extract relevant slices h H[edge, channel, Lh]
Interaction
@
Scatter
14:
13:
e= EE aledge, I e|
r
Block
[Moodes Mety d:]
15:
r= E (edge, channel, p]
Vector Scaling
(=
x,)
16:
Embedding
Eq-Linear
2 Eq-MLP
18:
17:
# Compute output using the TP function
out TP CG matrix)
Layer
19:
Output(rdge, channel,
Laut 11
Xi [nodess Reh, d.]
Fr [Redges dr]
20:
p+=
Atomic Graph
out
Fg. [eedges do.a]
(a) Pseudocode of Tensor-Product Layer
(b) A single tensor-product
Pnodes/edgess/ch/ppatt number of modes//edges/channe///aaa
dr10.p dimension of radial/angular feature
de hidden dimension of (th layer
Figure 3. Overview of operation of Tensor-Produce layer and a single tensor- passurrrrouut..
d/: intermediate dimension of ith layer (>> d,)
sion of the hidden states is reduced in a reduce layer ( Note that after passing through Tensor-Produch layer, the hidden dimension increases from de to d,, which is an order of magnitude larger. The hidden states then pass through another equivariant linear layer with residual connection O , followed by an activation function ) producing the output of the block (or the input to the next block).
Figure Equivariant MLIP model architecture. Assumes de- grees of hidden feature to have same number of channel.
The model first passes through embedding layer, where each node and edge of the graph are given their own fea- tures. Node features (N) are trainable parameters spe- cific to each atomic number. Edge features consist of two components: one encoding distance information (E₂) and the other capturing angular relationships (Ea.s) between connected atoms. It then progresses through multiple in- teraction blocks that perform message passing, mecha- nism for updating node representations by aggregating in- formation from neighboring nodes (Gilmer et al., 2017) The first block receives the node features as hidden states, which are then iteratively updated across subsequent blocks (N X1) + X2 - ... + X₁)). Edge features are used in all blocks. The final hidden states (X (Xina)) are used to pre- dict energy, from which forces can be derived by computing its gradient.
2.3. Tensor-Product Layer in Equivariant MLIP
Operations in Tensor-Product Layer. Figure (a) presents the pseudo-code for the Tensor-Product layer. The Tensor-Product layer performs large number of tensor- products. Note that we distinguish the tensor- mansor-product (low- ercase), which is an operation within the Tensor- Product layer, from the Tensor-Product layer itself (capitalized). Specifically, the tensor-product is carried out for each edge (Redges). channel (na)) and path (7path)) resulting in a total of Redges Tich X Opaat tensor-products (Line 1-6).
A path is defined by combination of hidden, edge, and out- put degrees (4, le, l.ut). Each of these degrees can take any value from 0 to Imax- meaning that the total possible number of paths is (Amxx + 1)². 4max is a user-defined parameter that controls the trade-off between model quality and computa- tional cost. A larger (max improves accuracy (Batzner et al., 2022) but results in an exponential increase in computation due to the growing number of paths (Passaro & Zitnick, 2023; Luo & Krishnapriyan, 2024).
Figure 2(b) illustrates the operations within each interaction block. The core of the interaction block is Tensor-Prodee Product layer where the messaae--psssnng mechanism is exe- cuted. Before reaching this layer, hidden states first pass through the equivariant linear layer ()), which applies linear transformation while preserving the input-output di- mension. The hidden states then undergo a scatter layer which expands their first dimension from Rnodes to Medges Meanwhile, Eg, is directly fed into Tensor-Produst layer while E, passes through an MLP before reaching it layer, the expanded first dimen-
Meanwhile, when physical constraints are considered, some of (Amax + 1)3 path are considered to be invalid, as the tensor-product of these paths always results Only paths
After
3

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
Inference
88.9%
11.1%
Table Sparsity coefficient matrix for varying 4max
Training 75.2%
24.8%
mmx
1
2
3
4
5
0%
20%
40%
60%
80%
100%
Sparsity (%) 71
78
82
84
86
Non-TP
product, where the hidden degree dimension is 3, the edge degree dimension is 5, and the number of edges is 1K. Each phase consists of multiple separate kernels (3 for forward, for backward, and 13 for double-backward), in which intermediate data are generated. Among these, the primary bottleneck is the outer product result generated during the forward phase this outer product result in the backward phase and the
Figure 4. Portion of Pensor-Product layer in inference and training of SevenNet-I3i5.
satisfying Inequality produce non-zero output (Line 8).
hn - Lel S lout S ln 1
(1)
along with the data associated
phase (02
and
Computations in Single Path A tensor-produee for sin- gle path, defined by 3-tuple (h, le, laus) (Line 18), is illus- trated in Figure 3(b). It receives a subvector of hidden states (h) and a subvector of angular features (e) as inputs, whose sizes are (2 x ln + 1) and (2 x le + 1), respectively. It then produces output (oul) whose size is (2 x loui + 1).
These intermediate data, highlighted in red in Figure 5, scale linearly with the product of the node degree dimension and the edge degree dimension, making them significantly larger than the input degrees.
Profiling Results. Figure 6 illustrates the compute and DRAM bandwidth utilization during the forward, backward, and wouble-backward phases of a SevenNer-13i5. The backward and double-backward phases take a significantly longer time than the forward phase, and they are largely memory-bound. This emphasizes the need to reduce memory bandwidth bottlenecks by minimizing intermediate data.
The computations proceed as follows: First, outer prod- uct between h and e is performed, producing OP. This tensor is then multiplied by CGh..a.aa, a a Clebsch-Gordan (CG) coefficient matrix (Varshalovich et al., 1988) corre- sponding to this path. CG coefficient matrix encapsulates the interaction principles between node and edge tensors preserving the equivariance of the output. Finally, the result of the matrix multiplication, denoted as is scaled by r, an element of the radial edge feature (Fz) after Eq-MLP, to produce the final output, out.
layer of
3.2. Peak Memory Spikes due to Output Data
In addition to the intermediate data, the output of the Tensor- Product layer presents significant challenge, this time in terms of peak memory usage rather than memory bandwidth. This issue stems from the substantial increase in hidden di- mension after passing through the Tensor-Product layer, growing from de d,. As mentioned carlier, d, is usually order of magnitude larger than d,, making the output size Redges X Tech d; considerably large. Although the subse- quent reduce layer immediately consumes and compresses this data, the full output must still be stored in memory beforehand, leading to memory spikes. In fact, this alone can increase peak memory consumption by a factor of four, severely limiting the number of atoms that can be included in simulations and posing major scalability challenge.
Overhead of Tensor-Produe Layer. Due to its large amount of computations, which scale exponentially with mmax the Tensor-Produse layer represents the primary bottle- neck in the equivariant MLIP model architecture. Figure 4 shows the breakdown of inference and training time for SevenNet-13i5 (Park et al., 2024b) on the MPF dataset (Chen & Ong. 2022). The Tensor-Product layer accounts for ap- proximately 89% and 75% of the inference and training time, respectively.
3. Inefficiency in Tensor-Product Layer
In this section, we highlight three major sources of ineffi- ciencies in the Prosuutt layer: 1) memory traffic from intermediate data, 2) peak memory spikes due to output data and 3) sparsity in the CG coefficient matrix.
3.3. Ineffectual Computation due to Sparsity
In the forward phase of the Tensor- Product layer, the outer product of input degrees (hidden degree and edge degree) multiplied by the CG coefficient matrix 2 in Figure 5(a)). Meanwhile, this matrix is highly sparse, with sparsity rang- ing from approximately 71 % to 86%, as shown in Table 1. This implies that the matrix multiplication between the outer product result and the CG coefficient matrix involves a sig- nificant amount of ineffectual computation. These inefficien-
3.1. Memory Traffic from Intermediate Data
A substantial amount of intermediate data is generated within the Censor-Produce layer. Figure 5 illustrates how memory accesses occur during the forward, backward, and woukle-backward phases for single path of the tensor-

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
(a) Forward
(b)kacwward
X
CG
3K
3K
3K
3K
1K
3K
3K
30K
30K
8K
3K
3K
1K
DRAM
DRAM
DiDouble-Backwaae
BACKWARD
GG
BK
15K
8K
16K
30K
15K
15K
:5
3K
3K
1K
3K
3K
1K
3K
30K
8K
8K
24K
60K
8K
1K
DRAM
Figure 5. Kernels for the three phases (forward, backward, double-backward) of a tensor- pnnsor-product.. Arrows represent memory traflic, with the numbers indicating the amount of traffic (in elements), calculated for the tensor-product with IK edges. The Clebsch-Gordon (CG) coefficient matrix are reused across all edges, contributing only 45 elements to the total memory traffic. Elements placed inside the DRAM indicates the output the tensor-product.
Forward
to eliminate redundant computations, addressing third inefficiency (Section 4.3). Finally, FlashTP introduces technique reduce memory traffic associated with tensor- product input data, which emerges resolving the three aforementioned issues (Section 4.4).
Backward
Double-Backwand
Backward
100
Utilization (%6)
80
bottleneck after
60
40
4.2. Kernel Fusion
20
Kernel fusion is widely used technique that combines mul- tiple kernels into single kernel to eliminate data movement between them (Dao et al., 2022). FlashTP employs kernel fusion for the Pensor-Product layer at two levels: intra-layer kernel fusion and inter-layer kernel fusion. The former ad- dresses overhead from intermediate data, while the latter addresses overhead from output data.
0
Compute
DRAM Bandwidth
Figure 6. Compute and DRAM bandwidth utilization of Tensor- Product layer in SevenNet-13i5 on an NVIDIA A100 GPU. Time regions of forward, backward and double-backwer are labeled.
Intra-Layer Kernel Fusion. A key strategy for fusing kernels in the Tensor-Product layer-which involves execut- ing many computation kernels-is to divide them into two phases: those up to and including multiplication by the CG coefficient matrix, and those that follow.
cies propagate also the backward and double-backward phases, particularly in the backward and double-backward kernels corresponding to this matrix multiplication 4 in Figure 5(b) and in Figure 5(c)).
In the first phase, FlashTP fuses kernels by handling all operations associated with each non-zero element in the CG coefficient matrix in single pass. For instance, for the entry at index (1, 1, 0), it multiplies the A[1] by the e[1] and then by that CG coefficient. Because each pass yields only a partial sum, any operations depending on those CG coefficient matrix products must wait until the partial sums are fully accumulated. Once accumulation is complete, the second phase of fused computations can proceed. This two- phase approach eliminates redundant work and streamlines fusion in the later phase.
4. FP
4.1. Overview
Addressing the inefficiencies identified in Section 3, we introduce FlashTP, tensor-produce GPU library for equiv- ariant MLIP. To mitigate the first two inefficiencies intermediate data and output Jata-FlashTP fully fuses all kernels within the Censor-Product layer, as well as its sub- sequent layer, into a single kernel (Section 4.2). Addition- ally, it exploits the sparsity of the CG coefficient matrix
5

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
Edge [src, dat)
Tensor Product
Reduce
Edge (src, dat)
Tensor Product
Atomic Reduce
Hidden
le2
Edge
1-0/1 M
1Ng F+t I+2
Hidden
Edge
20/+3
[0,3]
Worker
0
[0,3)
Worker 0
(4
Worker
(4, 3]
Warker 1
(8) Sperate Reduce
(b) Fused Reduce
Output
Output
Figure 7. Comparison of the output stage of the Tensor- Tensor-Product layer (a) before and (h) after inter-layer kernel fusion, assuming two worker outputs share the same destination node (dst). In (a), each ouiput is stored in memory before reduction. In (b), Each output is directly accumulated using atomic reduction.
Redundant Input Load (a) Without Path-Aggregation
Maximize Reuse
(b) With Path-Aggregation
Figure 9, Visualization of the effect of path -aghegggrcg,tion.
CG
COO Format j. k val)
Unique CG Value Array
k
[0, 2, 1,
array of unique CG coefficient values (u_og) and stores only the index of each value within this array in the val field of the tuple instead of the original value. This index can be represented with fewer bits (c.g., 8 bits) rather than storing the original 32-bit value. When performing computations with this sparse-formatted CG coefficient matrix, FlashTP iteratively processes each tuple one by one, illustrated in Figure 8(b).
vol
2,
[0, 00,4,000 0
V4_g
1.3.1.3
12,0,011
(a) Sparse-formatted CG coefficient matrix
J. k, vol]
Execution
20,
10, 1,
h
z
u_cg
2,
00.4.0.0
4.4. Path Aagh--ggregttion
1,,
z += h[2]
e[0]
[1cg[1]
02,0,
FlashTP introduces path-aggregation, technique designed to mitigate the new bottleneck load traffic in tensor products-thal arises after the integration of kernel fusion and sparse computation. By grouping tensor-product paths that share the same input and executing them within single kernel, this method reduces memory traffic.
(b) Computation with sparse-formatted CG coefficient matrix
Figure 8. Implementation of sparse tensor-product in FlashTP.
Figure 9 illustrates the path-aggregation by visualizing con- nection between inputs and outputs involved in each path. Without path-aggregation, multiple paths may redundantly read the same input data (Figure 9a)). For example, the degree-1 subvector of hidden states is read five times, and the degree-1 subvector edge features read three times. On the other hand, with aath-aggregation, hidden states and edge features are read from memory only once, reducing memory traffic by factors of five and three, respectively, in this example (Figure 9b)). In FlashTP, paths are grouped by the hidden subvector they use, as shown in Figure 9(b).
Inter-Layer Kernel Fusion. Figure 7 illustrates how the Pensor-Product layer is fused with subsequent layer, the re- tensor-product
duce layer. The reduce layer accumulates tensor product re- sults for edges that share the same destination nodes. Since rensor-products for different edges are likely to be executed by different workers on GPUs (e.g., thread blocks in CUDA), synchronization issues may arise during fusion when mul- tiple edges map to the same destination node. To address this, Fl employs atomic add operations such cases to ensure correct reduction.
4.3. Applying Sparsity in Tensor-Product
5. Implementation
In order to exploit the sparsity of CG coefficient matrix, FlashTP stores CG coefficient matrix in sparse data format. Figure 8(a) shows sparse data format used in FlashTP. The coordinate format (COO) is employed, where each nonzero value is represented as a 4-tuple three indices correspond- ing to each dimension (i, k) and the index of non-zero value (val). To further enhance storage efficiency, FP exploits the fact that the CG coefficient matrix contains many duplicated values. Specifically, FlashTP creates an
Programming Interface. FlashTP provides programming interface that is easily compatible with e3nn (Geiger Smidt, 2022; Geiger al., 2022), MyTorch-based frame- work for equivariant operations which is widely used for building equivariant MLIP models. Specifically, the inter- face of FlashTP well aligns with the implementation of Ten- sorProduct class in e3nn as can be seen in Figure 10. This allows programmers to integrate FlashTP with a minimal
6

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
Initialize
Initialize
Imax =
2
Imax
tp c3nn.oa. TensarProduct(
flashtp flashTP cannn.uv TP( i_int, i_in2, s_out, inst_ tuple)
i ini, i_in2, i_out, inst_tupley
20x
2.6x
25x
Speedup(vs e3nn)
16x
Forward input[edge_snC) y= tp(x, M) 3: output ryduce(y, edge_dst
Forward
20x
1: output (input e, edge_src, edge_ edge_dst)
12x
2.3x
15x
2.9x
FlashTP
8x 2.0x
10x
2.5x
Figure 10. Code change necessary to integrate FlashTP into the existing implementation based on the e3nn framework.
4x
5x-
1x
6x
Forward
Backward
Double Backward
Forward
Backward
Double
Backward
Tused+sparsity
sparsity path-aggregation
code change. In addition, FlashTP preserves the same ma- trix storage ordering and model structure as e3nn, ensuring consistency and ease of adoption.
Figure II. Ablation study result of three optimization techniques introduced by FlashTP. The speedup over eănn is reported.
Preprocessing. Before performing inference or training, during model initialization, FlashTP requires preprocessing for two purposes. First, identifies paths to aggregate by examining model configurations. Second, it constructs the CG coefficient matrix in sparse format.
6.1. Kernel Microbenchmark
Table 2 presents latency for three phase (forward, back- ward, and double-backward) of tensor-product across vary- ing mmx configuration. For both single and double preci- sion, FlashTP significantly outperforms e3nn for all cases. On average, for single precision, FlashTP achieves 6.75x, 14.37x, and 26.54x speedups for the forward, backward, and double-backward phase, respectively. Similarly, for double precision (fp64). it achieves 6.38x, 16.26x, and 25.98x, respectively. The speedup is particularly notable in the backward and wouble-backward phases, as they benefit the most from reduced memory traffic due kernel fusion and path aggregation, given their highly memory-bound nature.
GPU Parallelization Scheme. The tensor-product opera- tion FlashTP is parallelized across three dimensions: the number of edges, the number of aggregated paths, and the channel dimension associated with the node degree for each path. For a single warp, FlashTP assigns tensor-product computations with the same path but of different channel and edge index.
This ensures that all threads access the same CG coeffi- cient matrix metadata simultaneously allowing FlashTP to fully leverage constant memory for storing the metadata for sparse CG coefficient matrix. However, if multiple paths execute on single compute unit using different CG co- efficient matrices, cache evictions occur in constant memory, leading to long cache miss stalls. FlashTP pre- vents this by assigns single path to each thread block and maximizes its size. FlashTP automatically selects the block size by maximizing it while ensuring efficient shared memory utilization, thereby improving GPU occupancy.
Compared to cuEq, FlashTP is also consistently faster in all cases. While cuÊq demonstrates impressive results (e.g., cases and 2) and occasionally comes close to FlashTP in performance, its efficiency diminishes lmax increases, limiting its applicability to high-precision MD simulations. This is because cuEq's efficiency is primarily attributed to its effective use of shared memory, which becomes infeasible for large max due to its substantial memory footprint.
6. Evaluation
We evaluate FlashTP against two comparison base- e3m (Geiger et al.. 2022) and cuEquivariance (cuEq) (Geiger et al., 2024). e3nn is a widely used frame- work for building equivariant MLIP models, while cuEq is a recently released CUDA library by NVIDIA that accelerates cquivariant operations, such as tensor products. We first evaluate how much FlashTP accelerates each tensor consor-product operations (Section 6. and present how it translates into end-to-end speedup (Section 6.2). All evaluations, except for multi-GPU training, are conducted on a single NVIDIA A100 80GB GPU. Further details on the evaluation setup are provided in Appendix A.
Ablation Study. Figure 11 presents how kernel perfor- mance improves with the introduction of each of the three optimization techniques of FlashTP: 1) kernel fusion, 2) sparse computation and 3) aath-aggreggtion. For this evalu- ation, we use the case of laax = and 4max = 3. Results for other Imax values are provided in Appendix A.3. All three optimization techniques significantly enhance kernel performance.
lines:
Numerical Stability. Since kernel fusion can alter the order of floating-poin operations, it may produce slightly different results compared to unfused kernels, especially
7

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
Table 2. Speedup comparison of different tensor-product configurations relative eánn haseline for both single (fp32) and double precision (1p64), measured across varying (max. Entries marked with indicate speedup relative to cuEq due to e3nn encountering memory (OOM) errors.
Forward
Backward
Latency (ms) (Speedup over c3nn)
Latency (ms) (Speedup over ednmn)
Latency (ms) (Speedup over e3nn)
max
e33n
cuFq
FlashTP
e:3mm
cuRq
FlashTP
e:3mn
cuFq
FlashTP
2,05
0.61 (3.4x)
0.22 (9.4x)
6.18
(10.3x)
0.54 (11.4x)
9.13
0.92 (9.9%)
0.83 (11.0x)
fp32
6.20
1.24 (5.0x)
0.81 (7.7x)
32.26
2.28 (14.2x)
1.86 (17.3x)
63.86
4.28 (14.9x)
2.50 (25.6x)
3
18.40
16,06 (1.1x)
2.54 (7.3x)
85,00
91.68 (0.9x)
6,15 (13,8x)
254.13
204.07 (1.2x)
8.73 (29.1x)
4
45.58
122.20 (0.4x)
9.29 (4.9x)
268.48
835,63 (0.3x)
17.89 (15.0%)
1196.89
1822.67 (0.7x)
29.99 (39.9x)
5
99.01
121.02 (0.8%)
18.22 (5.4x)
779.14
858,95 (0.9x)
51.96 (15.0%)
3752.63
1835.29 (2.0x)
93.11 (40.3x)
Forward
Backward
Double-Bacte Backward
Latency (ms) (Speedup over ednn)
Latency (ms) (Speedup over e3nn)
Latency (ms) (Speedup over ešnn)
max
e3nn
cuEq
FlashTP
63nm
euRq
FlashTP
63nn
cuFq
FlashTP
2.39
0.76 (3.1x)
0.44 (5.4x)
12.21
0,94
0.74 (16.4x)
17.96
1.55 (11.6x)
1.19 (15.1x)
fp64
2
9.50
2.13 (4.5%)
1.44 (6.6%)
43.63
3.18
(13.7k)
2.86 (15.3%)
95.40
6.66 (14.3x)
4.36 (21.9x)
3
30.10
41.38 (0.7x)
4,00 (7.5x)
178.75
243.59 (0.7x)
10.81 (16.5x)
594.36
532.85 (1.1x)
17.97 (33.1x)
4
74.33
56.33 (1.3x)
11.93 (6,2x)
577.02
(1.3x)
34.04 (16.9x)
2566. 10
T158 80 (2.2x)
61.72 (41.6x)
5
NIA
134.47
30.18 (4.5x)*
N/A
1103.31
114.85 (9.6x)*
N/A
2861.31
199.18 (14,4x)*
MD Step time (sec)
3
e3nn
Table 3. Per-epoch training time of SevenNet models on a GPU
cuEq
NevenNet-
Seven Net-I3i5
SevenNet-A4iS
2
FlashTP
Time (min)
Peak Mem.
Time
Peak Men.
Time
Peak Mem.
(GB)
(min)
(GB)
(mine
(GB)
1-
e3nn
31
3.78
76
8.52
213
17.43
cuFq
23
2.56
97
5.23
358
9.97
FlashTP
20
0.89
22
1.37
32
2.07
0K
5K
10K
15K
20K
25K
Number of Simulated Atoms
of the Tensor- Prosor-Product layer. In terms of simulation time, FlashTP also outperforms the comparison baselines. For a system with 4K atoms, FlashTP achieves 4.2x speedup over e3nn and 6.2x speedup over cuEq, while using 6.3x and 4.3x less memory compared to e3nn and cuEq, respec- tively.
Figure 12. Comparison of average MD simulation step time across varying numbers of copper atoms, scaling up until simulations encounter an Out-of-Memory (OOM) error.
case. To quantify this potential error, we assume the fp64 results from e3nn as the golden refer- ence and compare the deviations of e3nn fp32 results and Fl fp32 results from it. Across the forward, back- ward, and double-backward phases, the average root mean squared error for e3nn fp32 and FlashTP fp32 are compa- rable 25.64110-6 vs. 4.60x10 5), indicating that kernel fusion does not appear to have an adverse impact on numer- ical stability. Detailed results are in Appendix B.
for
Single GPU Training. Table 3 presents the training time per-epoch and peak memory usage for three variants of SevenNet: 12i5, 13i5, and 14i5, each corresponding to max values of 2. 3, and 4, respectively. FlashTP achieves speedups of 1.6x, 3,5x, and 6.7x over e3nn for the 12i5, 13i5, and 14i5 configurations, respectively. Notably, as the proportion of computation spent on the Tensor--roduuct layer increases with (max, the benefits of FlashTP become more pronounced. Compared to cuEq, FlashTP is 1.15x, 4.41 x, and 1 11.19x faster. Consistent with the observations in Sec- tion 1, cuEq demonstrates impressive performance for Lmax = 2 but falls significantly behind for larger (max val- ues. In terms of peak memory usage, FlashTP requires significantly less memory than both comparison baselines, further emphasizing its efficiency in handling large systems.
6.2. End-to-end Speedup
Inference. Figure 12 presents the average time required for a single step of an MD simulation, which consists a single model inference and an atomic position update, using pre-trained SevenNet-13i5 model. We evaluate FlashTP alongside comparison baselines, increasing the number of atoms until an out-of-memory error occurs. Notably, FlashTP supports up to 28K atoms, 4.7x and 3.1x more than the limitations of e3mn and cuEq, respectively. This im- provement is primarily attributed to FlashTP's kernel fusion, which eliminates memory spikes caused by the output data
Multi- Training. Table 4 shows the per-epoch train- ing time of the SevenNet-13i5 model on a multi-GPU sys- tem. Each GPU node consists of eight NVIDIA 1 80GB GPUs, interconnected via NVLink and NVSwitch, with inter-node communication handled over a 100GB/s network.
8

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
2021) and Matbench (Riebesell al., 2024) continue rely on CGTP, underscoring the critical need for further acceleration of CGTP
Table Per-epoch training time of SevenNet-I3is on multi-GPUs # of Nodes (f of GPUs)
1
(8)
(16)
(32)
8 (64)
e3nn (sec)
735
399
218
123
FlashTP (sec)
164
30
Speedup
4.5x
4.6x
4.4x
4.1x
8. Conclusion
In this paper, we presented FlashTP, solution to two key in- efficiencies in the Tensor-Product layer of equivariant MLIP model models: the memory footprint from intermediate data and the sparsity of the CG coefficient matrix. Through kernel fusion, sparsity utilization, and path-aggregated ex- ecution, FlashTP achieved significant speedups in both in- ference and training of SevenNet-13i5, a state-of- state-of-the-arr the--artt equivariant MLIP model.
FlashTP demonstrates consistent speedup with an increas- ing number of GPUs. The slight reduction in speedup at higher GPU counts is likely due to inter-node communica- tion overhead, which, in turn, reduces the effective region of interest for FlashTP -ie., the portion of the workload that benefits from its acceleration.
7. Related Work
A Acknowledgments
Acceleration of Clebsch-Gordan Tensor Products. Sev- eral approaches have been proposed to accelerate the Clebsch-Gordan Tensor Product (CGTP). One notable ap- proach is the SO(2) Tensor-Produce (Passaro & Zitnick, 2023), which reduces the number of nero elements in the Clebsch-Gordan coefficient matrix by aligning the axis via rotation, thereby lowering the computational complexity of CGTP The SO(2) method is orthogonal to FlashTP and can be complementary, FIashTP also benefits from duced number of non-zero elements in the Elebsch-Gordan coefficient matrix. Another approach, Fused Tensor Product (FTP) or FusedTensor (Unke& & Maennel, 2024), acceler- ates CGTP by fusing all irreducible representations (irreps) into single tensor that can be processed using standard matrix multiplication. The output is then decomposed to recover the original irreps. While this method reduces com- putational cost compared to traditional CGTPs, it comes at the expense of reduced expressivity (Xie et al., 2024).
This work was supported by research grant from Sam- sung Advanced Institute of Technology (SAIT) and the Na- tional Research Foundation of Korea (NRF) grant funded by the Korean Government (MSIT) (RS 202400020003)000). We thank Saerom Choi and Yongdeok Kim from the Materials AI Lab at Samsung AI Center for providing GPU resources and supporting the multi-GPU evaluation of FlashTP on their server infrastructure.
Impact Statement
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
References
Batatia, I., Kovacs, D. P., Simm, G., Ortner, C., and Csányi, G. MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields. Ad- vances in Neural Information Processing Systems, 35: 11423-11436, 2022.
Alternatives to Clebsch-Gordant Tensor Products. Al- ternative approaches have explored using different sets of bases beyond spherical harmonics to accelerate tensor- The Gaunt Tensor Product (GTP) (Luo& & Krishnapriyan, 2024) changes the basis the frequency domain and uses Fast Fourier Transforms to speed up tensor-product computations. While computationally ef- ficient, this approach trades off some expressivity compared to CGTP. Additionally, due to the symmetric nature of its operations, GTP is unable to capture chiral features in 3D structures (Xie et al., 2024). Another notable direction in- volves Cartesian-basee models, which use Irreducible Carte- sian Tensors (ICT) instead of spherical harmonics. Recent studies have demonstrated promising performance in both accuracy and latency. particularly at lower ranks, when ap- plied to MLIP tasks (Simeon & De Fabritiis, 2023; Zaverkin et al., 2024). Nonetheless, CaTe-based models remain the dominant paradigm for lop-performing entries, as widely used MLIP benchmarks such as OC20 (Chanussot et al.,
product computations.
Batzner, S., Musaelian, A., Sun, L.,, Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E., and Kozinsky, B. Equivaranr Graph Neural Networks for Data- Efficient and Accurate Interatomic Potentials. Nature communications, 13(1):2453, 2022.
Chanussot, L.,, Das, A., Goyal, S., Lavril, T., Shuaibi, M., Riviere, M., Tran, K., Heras- soras-Dominggo, J., Ho, C., Hu, W., al. Open Catalyst 2020 (OC20) Dataset and Com- munity Challenges. Acs Catalysis, 11(10)605596072 2021.
Chen, C. and Ong, S. P. A universal graph deep learning

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
interatomic potential for the periodic table. Nature Com- putational Science, 2(11::1187728 2022.
Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and Dahl, G. E. Neural Message Passing for Quantum Chem- istry. In Proceedings of the 34th International Conference on Machine Learning, volume 70, pp. 1263-1272. PMLR, 2017.
Chen, C., Ye, W., Zuo, Y., Zheng, C., and Ong, S. P. Graph Networks as Universal Machine Learning Framework for Molecules and Crystals. Chemistry of Materials, 31 (9):3544-3572, 2019,
Hong. C., Oh, S.,, An, H., Kim, P.-h., Kim, Y., Ko. J.-H., Sue, J.. Oh, D., Park, S., and Han, S. Atomistic Simulation of HF Etching Process of Amorphous SizN, Using Machine Learning Potential. ACS Applied Materials & Interfaces. 16(36):4457-4333 48469, 2024.
Cohen, T. and Welling, M. Group Equivariant Convolutional Networks. In Proceedings of the 33rd International Con- ference on Machine Learning, volume 48, pp. 2990-2999. PMLR. 2016.
Hwang, S., Jung, J., Hong, C., Jeong, W., Kang, S., and Han, S. Stability and Equilibrium Structures of Unknown Ternary MetalOxides Explored by Machine-Learned Po- tentials. Journal of the American Chemical Society, 145 (35):13378- 193866 2023.
Dao, T., Fu, D. Y., Ermon, S., Rudra, A., and Ré, C. FlashAt- tention: Fast and memory-efficieen exact attention with 10-awareness. In Advances in Neural Information Pro- cessing Systems (NeurIPS), 2022.
Kim, G., Na, B., Kim, G., Cho, H., Kang, S., Lee, H. S., Choi, S., Kim, H., Lee, S., and Kim, Y. Benchmark of Machine Learning Force Fields for Semiconductor Sim- ulations: Datasets, Metries, and Comparative Analysis. Advances in Neural Information Processing Systems, 36, 2023.
Deringer, V. L.. Bartók, A. P., Bernstein, N., Wilkins, D. M., Ceriotti, M., and Csányi, G. Gaussian Process Regression for Materials and Molecules. Chemical Reviews, 121(16): 10073-10141, 2021.
Fu, X., Wu, Z., Wang, W., Xie, T., Keten, S., Gómez- Bombarelli, R., and Jaakkola, T. Forces are not Enough: Benchmark and Critical Evaluation for Machine Learning Force Fields with Molecular Simulations. Transactions on Machine Learning Research, 2022.
Ko, T. W. and Ong, S. P. Recent advances and outstanding challenges for machine learning interatomic potentials. Nature Computational Science, 3(12):998-1000 2023.
Kruglov, I. A.,, Yanilkin, A. V., Propad, Y., Mazitov, B., Rachitskii, P., and Oganov, A. R. Crystal structure predic- tion at finite temperatures. npj Computational Materials, 9(1): 2023.
Gasteiger, J., GroB. J., and Günnemann. S. Directional Message Passing for Molecular Graphs. In The Eighth International Conference on Learning Representations, 2020.
Larsen, A. H., Mortensen, J. J., Blomqvist, J., Castelli, 1. E., Christensen, R., Dutak, M., Friis, J., Groves, M. N., Hammer, B.,, Hargus, C., Hermes, E. D., Jennings, P. C., Jensen, P. B., Kermode, J., Kitchin, J. R., Kolsbjerg, E. L., Kubal, J., Kaasbjerg, K., Lysgaard, S., Maronsson, J. B., Maxson, T., Olsen, T., Pastewka, L.., Peterson, A., Rost- gaard, C., Schiøtz, J., Schütt, O., Strange, M., Thygesen, K. S., Vegge, T., Vilhelmsen, L., Walter, M., Zeng, Z., and Jacobsen, K. W. The atomic simulation environ- ment- Python library for working with atoms. Journal of Physics: Condensed Matter, (972):277302, 2017.
Gasteiger, J., Becker, F., and Günnemann, S. GemNet: Uni- versal Directional Graph Neural Networks for Molecules. Advances in Neural Information Processing Systems, 34: 6790-6802, 2021.
Geiger, M. and Smidt, T. e3nn: Euclidean Neural Networks. arXiv preprint arXiv:2207.09455. 2022.
Geiger, M.. Smidt, T., M., A., Miller, B. K., Boomsma, W., Dice, B., Lapchevskyi, K., Weiler, M., Tyszkiewicz, M.. Batzner, S., Madisetti, D., Uhrin, M. Frellsen, J.. Jung, N., Sanborn, S., Wen, M., Rackers, J., Rød, M., and Bailey, M. Euclidean neural networks: e3nn, April 2022. URL https::/ org/10. 521/1/1..111///// zenodo. 6459381.
Lee, J., Ju, S., Hwang, S., You, J., Jung, J., Kang, Y., and Han, S. Risorder-Depender Li Diffusion in LigPS, Investigated by Machine- Potential. ACS Applied Materials & Interfaces, 16135):46442-44555 2024.
Geiger, M., Kucukbenli, E., Zandstein, B., and Tretina, K. Accelcrate Drug and Material Discovery with New Math Library NVIDIA cuEquivariance, Novem- ber 2024. URL https: : /tt://gttttt github...t... com/ /VIIII// cuEquivarian
Liao, Y-I.., Wood, B. M., Das, A., and Smidt, T. EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degre Representations. In The Twelfth International Conference on Learning Representations, 2024.
10

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
Luo, S., C. T. and Krishnapriyan, A. S. Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products. In The Twelfth International Conference on Learning Representations, 2024.
Unke, O. T., Chmiela, S., Sauceda, H. E., Gastegger, M., Poltavsky, I., Schutt, K. T., Tkatchenko, A., and Muller, K.-R. Machine Learning Force Fields. Chemical Reviews, (1(16)):0142-10188 2021.
Musaelian, A., Batzner, S., Johansson, A., Sun, L.,, Owen, J., Kornbluth, M., and Kozinsky, B. Learning Local Equivariant Representations for Large-Scale Atomistic Dynamics, Nature Communications, 14(1):579, 2023.
Varshalovich, D. A., Moskalev, A. N., and Khersonskii, V. K. Quantum Theory of Angular Momentum. World Scientific, 1988.
Wang. J., Panchal, A. A., Gautam, S.S and Canepa, The resistive nature of decomposing interfaces of solid elec- trolytes with alkali metal electrodes. Journal of Materials 2022.
Park, Y., Kim, J., Hwang. S., and Han, S. SevenNet graph neural network interatomic potential package supporting efficient multi-GPU parallel molecular dy- namics simulations, December 2024a. URL https: //github. cOm/MDIL SIi---NNUSSeeeeeee
Chemistry A, 19(37)
Xie, Y., Daigavane, A., Kotak, M., and Smidt, The Price of Freedom: Exploring Tradeoff's between Expressivity and Computational Efficiency Equivariant Tensor Prod- ucts. In ICML 2024 Workshop
Park, Y., Kim, J., Hwang, S., and Han, S. Scalable Parallel Algorithm for Graph Neural Network Interatomic Poten- tials in Molecular Dynamics Simulations. Journal of Chemical Theory and Computation, 20(1 (881))888744888 2024b. doi: 10.
grounded Representation Learning and Generative Modeling, 2024.
Zaverkin, V., Alesiani, F., Maruyama, T., Errica, F., Chris- tiansen, H., Takamoto, M., Weber, N., and Niepert, M. Higher-Ram Irreducible Cartesian Tensors for Equivari- ant Message Passing. Advances in Neural Information Processing Systems, 38, 2024.
Passaro, S. and Zitnick, C. Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs. In Proceed- ings of the 40th International Conference on Machine Learning, volume 202, pp. 17420-274333 PMLR, 2023.
Zhang, Y.-W., Sorkin, V., Aitken, H., Politano, A., Behler, J., Thompson, A., Ko, T. W., Ong, S. P., Chalykh, O., Korogod, D., et al. Roadmap for the development of machine learning- based interatomic potentials. Modelling and Simulation in Materials Science and Engineering, 2024.
Reiser, P. Neubert, M., Eberhard, A., Torresi, L., Zhou, C., Shao, C., Metni, H., van Hoesel, C., Schopmans, H., Sommer, T., and Friederich, P. Graph neural networks for materials science and chemistry. Communications Materials, 3(1)993, 2022.
Ricbesell, J., Goodall, R. A., Benner, P., Chiang, Y., Deng, B., Ceder, G., Asta, M., Lee, A. A., Jain, A., and Persson, K. A. Matbench Discovery A framework to evaluate machine learning crystal stability predictions, 2024. URL https:/ /ttp://aaavv../.ii org/ wttggaa//////oo..ee/ggo///eeaeb/ /2008 14920.
Schüitt, K., Kindermans, P.-J., Sauceda Felix, E., Chmicla, S., Tkatchenko, A., and Müller, K.-R. SchNet: A continuous- filter convolutional neural network for model- ing quantum interactions. Advances in Neural Informa- tion Processing Systems, 30, 2017.
Simeon, G. and De Fabritiss, G. TensorNet: Cartesian Ten- sor Representations for Efficient Learning of Molecular Potentials. Advances in Neural Information Processing Systems, 36:37334 334-3334-37333, 2023.
Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation- and translation-equivarian neural networks for 3D point clouds. arXiv preprint arXiv:1802.08221. 2018.
Unke, O. T. and Maennel, H. E3x: Eq3i-aauu Deep Learning Made Easy. arXiv preprint 5rX1v:2401.07595, 2024.

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
A. Evaluation Details
We used pytorch 2.5. +cu12.4, e3nn 0.5.4 (Geiger et al., 2022), cuEquivariance 0.2.0 (Geiger et al., 2024), Atomic Simulation Environment (ASE) 3.24.0 (Larsen al., 2017) for evaluation.
GPU kernel runtime and memory usage measurements conducted using the PyToorc profiler, leveraging its capability to track forward and backward connections for precise measurement of backward and double-backward operations. All evaluations, except for multi-GPU training, are performed on an NVIDIA A100 GPU with 80GB of memory. For multi-GPU training, GPU nodes that consist of cight NVIDIA A 100 GPUs (80GB), interconnected via NVLink and NVSwitch, with inter-node communication handled over a 100GB/s network was used.
A.1. Kernel Microbenchmark
We benchmark the Tensor-Product layer 4 from NequlP (Batzner et al., 2022) across different values of maxx keeping the channel of hidden feature fixed to 32 channels.
Our synthetic input graph has 512 nodes, each with outgoing edges (32K edges total). To generate connectivity, we assign each node 64 distinct random destinations drawn uniformly from the other 511 nodes. This exact graph topology reused for all kernel implementations (e3nn, cuEq and FlashTP) ensure fair comparison. All hidden, edge and weight values are randomly generated.
Table Tensor-Product configurations for kernel microbenchmark
Imar
Hidden
Edge
Output
32x0e+32200o333xxeee333A
Ix0e+Ixlo
64x00=:64x0e-99x1oo9001111
32x
96x0e+96x0e=192xto
1x0e+Ix1o+Ix2e
+32xxce+32x1o+32222
+192xlc=19222+~192222>
32x0e+:32x0o++22xcee33222
125x06-112x00-22881102222211
1xDe=1xloo4x2e<1xx3o
A.2. Channel Scaling
We repeated the kernel microbenchmark for mmx 3 using hidden-channel dimensions of 64 and 128, reporting kernel runtimes in milliseconds. Across all channel sizes, FlashTP shows consistent speedup.
Table Kerel microbenchmark results for Imax 3 across different channel sizes
Forward
Backward
Youble-Backwee
Channel Size
32
64
128
32
64
128
32
64
128
e3nn
18.4
34.25
66.8
85
163.96
324.73
254.13
503.28
1003.95
FlashTP
2.54
5.02
9.98
6.15
12.14
24.19
8.73
17.41
34.81
Speedup
7.2x
6.8x
6.7x
13.8x
13.5x
13.4x
29.1x
28.9x
28.8x
A.3. Ablation Study
Table 7 reports the ablation results on our kernel microbenchmark, showing the speedup over e3nn for each of the three optimization strategies kernel fusion (Fused), sparsity optimization (Sparse), and path aggregation (Path)-as as well their combined effect (All).
Because the Crehsch-Gordan (CG) coefficient matrix for the non-sparse variants exceeds the capacity of constant memory, we disabled constant-memory storage for all ablation configurations, including the Sparse variant, to ensure a fair comparison. Consequently, the cumulative "All" configuration in the ablation study achieves slightly lower speedup than the full Pl implementation.
12

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
Table 7. Ablation results the kernel microhenchmark
Forward
Backward
Fmax
Fused
Path
Sparse
All
eônn
Fused
Path
Spurse
All
eônn
Fused
Puth
Sparse
All
1
3.46x
4.37x
4.46x
5.25x
1x
x0,95
15.85%
12.16x
18.41x
1x
9.70x
12.80x
10.87x
14.77x
2
3.30x
4.16x
4,68x
6.60x
1x
5.03x
7.26x
6.63x
11.35x
1x
8.45x
11.85x
11.66x
19.13x
x
x
4.79x
6.80x
1x
4.48x
6.32x
7.24x
13,01x
1x
9.79
12.45x
15,46x
25.82x
2.32x
x
4.51x
6.02x
1x
3,97x
4.99x
7
13.16x
11.66x
14.58x
19.89x
30.73x
A.4. Roofline Analysis
We estimate the theoretical performance limits of FlashTP using roofline analysis, based on the peak capabilities of the A100 GPU (19.5 TFLOPS and 1.9 TB/s memory bandwidth). According this analysis, the estimated latencies for the fp32 kernel microbenchmark with fmax 3 are 0.92 ms for the forward pass, 1.78 ms for the backward pass, and 2.70 ms for the wouble-backward pass. The gap between these theoretical bounds and our measured performance in Table 2 (Section 6.1) indicates that PlashTP has room for further optimization.
A.5. End-to-end Inference
We evaluate end-to-end inference performance by integrating pre-trained SevenNet-13i5 model (Park et al., 2024a) into the ASE framework to run molecular-dynamic (MD) simulations of copper atoms. To gauge memory demands, we also determine the maximum number of atoms supportable, which correlates directly with peak memory usage. Each simulation spans 500 timesteps, and our reported runtimes are averaged over the final 100 timesteps for stability.
In experiments, the out-of-memory (OOM) failures observed when using FlashTP originate not from the Tensor- Product layer but from the default e3nn implementation of the Linear layer. This suggests that by optimizing the Linear layer implementation, FlashTP could scale to even larger atomic systems, highlighting its robustness and scalability for large-scale MD simulations.
A.6. End-to-end Training
The models SevenNet-I2i5 (which is SevenNet-O) (Park et al., 2024b), SevenNet- 13i5, and SevenNet- 14i5 are trained using the same training configuration. The training was conducted using the MPF dataset (Chen & Ong, 2022), which comprises 1 168,921 samples, with each sample containing an average of 29 atoms. Training time is reported as the duration a single epoch, and peak-memory reduction ratios are computed per batch and then averaged across multiple batches. For multi-GPU experiments, only SevenNet-13i5 was used, maintaining a per-GPU batch size of to match the single-GPu setup.
Table 8. Training configuration
Training Parameters
Model Parameters
Dataset
MPF
Precision
FP32
Batch Size
16
Cutoff
6
Optimizer
Adam
Cutoff Function
XPLOR (5.5)
Learning Rate
0.001
Radial Basis
Bessel (8 basis)
Loss Function
Huber Loss
Activation
SiLU (e), tanh (o)
Scheduler
Linear
Table 9. SevenNet-12i5 Tensor-Produs layer configuration
Seven Net-Iiet
Hidden
Edge
Output
Layer 1
128x0e
1x0e+1x1e++ 1x2e
128x0e++28x1e+1228x2
Layer 2
128x0c+64xlc++222ee
1x0e+Ixlc+1x2 1x2c
224x0e++384xlc4352222
Layer 3
128x0e+6641e+42222e
1x0e+1xle+1x3)
224x0e+3884xle+355222
Layer 4
128x0c+64x1e+32222e
1x0e+1xle+1x2e
224x0e+384xle4332222
Layer 5
128x0c+64xlc++2222e
1x0e+Ixlc+1x2e
224x0c
13

---

FlashTP: Fused, Sparsity-Aware Tensor Product for Machine Learning Interatomic Potentials
Table 10. SevenNet-13i5 Tensor-Produce layer configuration
SevenNet-13i5
Hidden
Edge
Output
Layer
128x0c
Ix0e+Ixl++Ix2++Il 1x3c
Layer 2
12800=-64xJec455X24455XXX3
Lx0++Ixlc+1M2++1k33
266x0=+4B0Xlc<554422244400
Layer 3
1BXX0=466x1e433x22443XX
Ix0e+1xle+1x2++133
Layer 4
128X0c+6kxlc+42222+333ee
Ix0c+1xlc+1x2c+413ce
25xx0++-80010:334222+44000
Layer 5
128x0e+66x1e452222eH77
1x0e41xle+1x2++1333
256x0e
Table SevenNet-14i5 Tensor-Product layer configuration
SevenNet-14i5
Hidden
Edge
Output
Layer 1
128x0e
128x0c+128x
Layer 2
Layer 3
X0e+Ixxle+122++1133441444
Layer 4
Layer
288x0e
B. Numerical Stability
We assessed the numerical stability FlashTP against the default e3nn implementation as follows. First, we sampled FP32 tensors for the hidden features (h), edge features (e), and weights (r) from uniform distribution over the interval (0,1). We then computed the "reference" outputs by rumning e3nn in fp64 precision. Next, re-computed the layer outputs using both FlashTP and e3nn in fp32 precision. Finally, calculated the root-mean-square error (RMSE) between each fp32 output (and its corresponding gradients) and the fp64 reference to quantify numerical stability.
Overall, results shown in Table 12 indicate that FlashTP offers improved or comparable accuracy in most cases, making it a reliable alternative.
Table 12, RSME of e3nn, and FlashTP
Stage
Tensor
e3nn
FlashTP
Forward
9.95651-07
9.9171E-07
Backward
h grad
1.4393E-05
1.4393E-05
Backward
grad
7.2380E-06
5,4978E-06
Backward
grad
6.3440E-08
6.5808E-08
muuble-Bakkward
grad
2.5270E-07
3.5300E-07
nouble-Backwaar
grad
1.6240B-05
9.8454E-06
grad
1.4850E-07
1.3050E-07 07
Double- Barkwdd
grad
5.7579E-06
5.5367B-06
14


Reviews:
[{"review_id": "5eGNzcy33I", "reviewer": "Reviewer_hLSR", "summary": "In this paper, the authors develop FlashTP, an optimized tensor-product library that uses kernel fusion, sparse computation, and path-aggregated execution. FlashTP achieves significant performance improvement in terms of increasing throughput and decreasing memory usage compared to common libraries, e3nn and cuEquivariance. They authors validation the speed improvements using the SevenNet-l3i5 model.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback.\n# R1. Evaluation on multi-GPU training [W1, Q1]\n- Multi-GPU training with FlashTP can be performed using PyTorch's Distributed Data Parallel (DDP). The table below shows the one-epoch training time (in seconds) for SevenNet-l3i5 on the MPF dataset, using varying numbers of GPU nodes. Each node consists of eight NVIDIA A100 GPUs (80 GB), interconnected with NVLink and NVSwitch, while inter-node communication is handled via a 100 GB/s network. A per-GPU batch size of 16 was used for all runs. FlashTP exhibits strong scalability across multiple GPUs. The modest reduction in speedup at higher GPU counts is likely due to increased inter-node communication overhead.\n\n| # of Nodes (# of GPUs) |1 (8)|2 (16)|4 (32)|8 (64)|\n|-|-|-|-|-|\n|e3nn|735|399|218|123|\n|FlashTP|164|87|49|30|\n|Speedup|4.5|4.6|4.4|4.1|\n\n# R2. Essential references not discussed\n- FlashAttention is a well-known work that optimizes self-attention via kernel fusion, and it served as an inspiration for the naming of FlashTP. We will include a reference to FlashAttention in Section 4.2 of the final draft.\n\n# R3. Acceleration of Clebsch-Gordan Tensor Product (CGTP)\nWe appreciate the reviewers' insightful comments regarding alternative acceleration methods for CGTP. Below, we summarize several relevant approaches. We will incorporate this discussion into the final draft.\n- **SO(2) Tensor Product [1]**: This technique reduces the number of non-zero elements in the Clebsch–Gordan coefficient matrix by aligning the axis via rotation, thereby lowering computational complexity. This method is orthogonal to FlashTP and can be synergistic—FlashTP can also benefit from a reduced number of non-zero elements in the input.\n- **Gaunt Tensor Product (GTP) [2]**: GTP leverages Fast Fourier Transforms to accelerate CGTP computations by operating in the frequency domain. While computationally efficient, it trades off some expressivity compared to CGTP. Due to the symmetric nature of its operations, GTP is unable to capture chiral features in 3D structures [3]. \n- **Fused Tensor Product (FTP) [4]**: FTP aggregates all irreducible representations (irreps) into a single tensor, applies native matrix multiplication, and subsequently decomposes the result to recover the original irreps. This reduces computational cost compared to CGTPs, but similar to GTP, it exhibits reduced expressivity [3]. While it has slightly higher computational complexity than GTP, it is asymmetric operations, making it more suitable for capturing parity-sensitive interactions.\n- **Cartesian-based models**: These models employ Irreducible Cartesian Tensors (ICT) instead of spherical harmonics as the basis. Recent studies have demonstrated promising performance in both accuracy and latency, particularly at lower ranks, when applied to MLIP tasks [5,6]. Nonetheless, CGTP-based models remain the dominant paradigm—top-performing entries in widely used MLIP benchmarks such as OC20 and Matbench continue to rely on CGTP, underscoring the critical need for further acceleration of CGTP [7,8].\n\n# R4. Future plans for FlashTP [W2, Q2]\nFlashTP is already deployed in production at a major tech company to deliver SOTA performance for their simulation workloads. We are also in active discussions with another major tech company to integrate FlashTP into their widely used open-source library for equivariant neural networks.\n- We plan to release FlashTP as a standalone Python package and integrate it into the e3nn codebase by implementing a variant of the `TensorProduct` class that utilizes FlashTP under the hood.\n- To facilitate reproducibility and adoption by the broader community, we will include a benchmarking script and a comprehensive integration example in the repository.\n\n# R5. Roofline model analysis [W3]\n- Please refer to our response R2 to Reviewer GZbV for the theoretical performance of FlashTP estimated from roofline analysis.\n- The roofline analysis from Nsight Compute shows that FlashTP kernels are compute-bound, whereas kernels of cuEq are L1 memory-bound. For more analysis on the differences between FlashTP and cuEq, please see our response R3 to Reviewer LtnP.\n\n# References\n[1] Passaro et al., Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs, ICML 2023.\\\n[2] Luo et al., Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products, ICLR 2024.\\\n[3] Xie et al., The price of freedom: Exploring tradeoffs between expressivity and computational efficiency in equivariant tensor products, ICML 2024 Workshop GRaM.\\\n[4] Unke et al., E3x: E(3)-equivariant deep learning made easy, CoRR 2024.\\\n[5] Simeon et al., TensorNet: Cartesian tensor representations for efficient learning of molecular potentials, NeurIPS 2023.\\\n[6] Zaverkin et al., Higher-Rank Irreducible Cartesian Tensors for Equivariant Message Passing, NeurIPS 2024.\\\n[7] https://opencatalystproject.org/leaderboard.html \\\n[8] https://matbench-discovery.materialsproject.org", "claims_and_evidence": "All of the claims made in the paper are directly supported by evidence. The authors make claims about the inefficiencies of current TP operations, and support these claims in Figures 4, 6 and in Table 1. All of the claims about speed-ups relative to e3nn and cuEquivariance are well supported by extensive experimental results. Additionally, the authors provide a good ablation study to demonstrate the contributions of each part of FlashTP.", "methods_and_evaluation_criteria": "The two main efficiency metrics of interest to downstream researchers are throughput/latency and GPU memory usage, which are the primary metrics reported in the paper. There exist more sophisticated ways to quantify performance such as the Roofline model, and the evaluation/discussion might be better if such metrics were used, but the reality is that MLFF researchers only care about how fast the model runs and how much memory they will need to run it. \n\nFor downstream model evaluations, the authors use SevenNet-l3i5, which is a fair representative model for equivariant MLFFs as it shares the same backbone (NequIP) as many other models (MACE, GNoME, etc.)\n\nThe authors use a reasonable dataset to evaluate on (MPF). The choice of the dataset doesn't really matter and the authors include a study on scaling the number of atoms in the system,  thereby demonstrating applicability to nearly all other MLFF datasets.", "theoretical_claims": "N/A -- no theoretical claims are made.", "experimental_designs_or_analyses": "As mentioned in the Methods And Evaluation Criteria section, all the experiments/analysis are well designed are the models/datasets used are good representative samples of the MLFF research community as a whole. I would not expect to see noticeably different results on different models/datasets, at least in terms of the relative performance of FlashTP to e3NN/cuEquivariance.", "relation_to_broader_scientific_literature": "MLFF foundation models have recently become incredibly powerful and useful for a variety of downstream tasks. Most of the top performing foundation models heavily rely on TP operations, and as such are incredibly slow. Additionally, the TP operations are known to be bounded by memory latency and often have poor GPU utilization. This work takes a massive step in this direction towards speeding up these foundation models.", "essential_references_not_discussed": "On the MLFF side, one potential reference that is missing is [1], which is an orthogonal approach used to improve the speed of TP operations. I am not very familiar with low-level ML systems research and likely would not be aware of any missing references. However, one paper that comes to mind that maybe should have been included is FlashAttention [1], which famously used kernel fusion to accelerate LLMs.\n\n[1] Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs, Passaro et al, https://arxiv.org/abs/2302.03655\n[2] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, Dao et al, https://arxiv.org/abs/2205.14135", "other_strengths_and_weaknesses": "Strengths:\n- S1: The paper is well written and provides good explanations of why existing TP libraries have huge performance bottlenecks. Figure 6 is a fantastic illustration of this. \n- S2: FlashTP has very strong performance improvements over existing libraries\n- S3: The experiments in the paper are well designed and representative of MLFF research/use cases.\n- S4: The authors provide an initial (and minimal) codebase for FlashTP. \n- S5: The paper includes a good ablation study to further analyze the creation of their fused kernels. \n\nWeaknesses:\n- W1: It is impractical to train top MLFF foundation models on a single GPU. Does FlashTP currently work on multi-GPU and can the authors provide such experiments?\n- W2: This paper has the potential to be massively impactful, but only if the authors continue to work to build out their FlashTP library including writing good tutorials and adding more of the existing functionality in e3nn.\n- W3: The authors evaluation of GPU performance evaluation could be more nuanced using something like the Roofline model", "other_comments_or_suggestions": "### Update After Rebuttal\n\nThe authors include some additional experiments to show that the proposed method works in a more realistic multi-GPU training setup and improved the performance analysis using the Roofline model. This is a strong work the potential to have a large impact on the MLFF community.", "questions_for_authors": "- Q1: The main thing preventing me from giving this work a 5 is that they do not evaluate on multi-GPU training. Can the authors provide such an evaluation and demonstrate performance scaling on as many GPUs as possible?\n- Q2: What is the authors plan to deploy the proposed code? There are many different operations/utilities in e3nn, do the authors plan to release a more developed software library to replace e3nn? For MLFF researchers, the ideal case would be for FlashTP kernels to be integrated into the existing e3nn codebase to maximize code reusability. In my view, this work is only a strong contribution to the community if the authors can continue to develop the software for people to use.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oPqXyQVfz9", "reviewer": "Reviewer_LtnP", "summary": "This paper presents FlashTP, an optimized tensor-product library designed to improve the computational efficiency of equivariant machine-learning models that employ spherical tensors. The authors identify three key inefficiencies in existing tensor-product layers: excessive memory traffic from intermediate data, memory spikes from large output tensors, and lack of sparsity exploitation in the Clebsch-Gordan coefficients. FlashTP addresses these challenges through kernel fusion, sparse tensor-product computation, and path-aggregated execution. Experimental results show significant speedups over e3nn and NVIDIA cuEquivariance.\n\n## update after rebuttal\n\nThe authors addressed all my questions. I increased my score to 4 and recommend this work for publication.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback.\n# R1. Discussion regarding Cartesian-based models\n- Please refer to the last paragraph of our response R3 to Reviewer hLSR.\n\n# R2. Choice of SevenNet over MACE for end-to-end evaluation\n- We chose SevenNet over MACE for end-to-end evaluation because 1) we believe SevenNet more accurately reflects the current SOTA, as demonstrated by its strong performance on a widely used benchmark [1]; and 2) its architecture is more broadly representative than that of MACE. SevenNet is primarily composed of CGTP-based interactions—a design principle shared by many other leading models, such as GNoME and Allegro. In contrast, MACE introduces a unique operation called _symmetric contraction_ on top of CGTP, which is specific only to MACE and limits its generality as a benchmark.\n- Meanwhile, FlashTP can still be applied to accelerate MACE through a hybrid approach with cuEq. Specifically, the CGTP components within MACE can benefit from FlashTP, while the symmetric contraction operation—which falls outside the scope of FlashTP—can continue to be handled by cuEq. In our evaluation, this complementary use of FlashTP and cuEq reduces the per-epoch training time of the MACE-MP Large model (fp32) by 1.2× compared to using cuEq alone.\n\n# R3. Speculation on the performance differences between FlashTP and cuEq  \nSince cuEq’s CUDA implementation is not open-source, identifying the exact performance bottlenecks of cuEq is infeasible. As a result, we refrained from including potentially premature speculation in the manuscript. Nonetheless, we have conducted an in-depth analysis of cuEq to the extent possible. Below, we share our informed speculations regarding the observed performance differences between FlashTP and cuEq.\n\n## Performance degradation of cuEq for higher values of $l_{max}$.\n- cuEq launches different CUDA kernels depending on the CGTP configuration. In particular, the kernel used for CGTPs with $l_{max} \\leq 2$ (cuEq-fast) differs from the one used for $l_{max} > 2$ (cuEq-slow). \n\n- Nsight Compute profiling indicates that while both cuEq-fast and cuEq-slow are L1-bandwidth-bound, their L1 utilization differs markedly: cuEq-fast achieves high L1 utilization, whereas cuEq-slow exhibits significantly lower L1 utilization. We believe this disparity is a primary contributor to the degraded performance of cuEq at higher $l_{\\text{max}}$ values. Note that our references to L1 bandwidth/utilization refer to the combined bandwidth/usage of the L1 cache and shared memory, which share the same underlying hardware.\n- We speculate that the difference in L1 utilization arises from how operands are managed. Specifically, cuEq-slow shows substantially higher shared memory usage compared to cuEq-fast. Our hypothesis is that cuEq-fast handles operands primarily using registers, while cuEq-slow relies more on shared memory due to operand sizes exceeding register capacity. This increased reliance on shared memory may reduce kernel occupancy, which in turn lowers L1 utilization and ultimately impacts performance.\n## Source of speedup for FlashTP \n- FlashTP employs a shared memory–based design regardless of the $l_{max​}$ value, yet it does not experience performance degradation due to low occupancy and consistently outperforms cuEq. This is because FlashTP substantially reduces overall memory traffic—including L1 traffic—such that low occupancy and the resulting lower L1 utilization do not become performance bottlenecks. A suite of memory traffic optimization techniques, including effective kernel fusion and path aggregation, are the key differentiator that sets FlashTP apart from cuEq in both efficiency and scalability.\n\n\n# R4. On the practical benefits of higher $l_{max}$\n- While it is true that increasing $l_{max}$​ can lead to diminishing returns in accuracy relative to the added computational cost, numerous empirical studies show that modest increases can yield meaningful improvements. For example, models such as MACE (using $l_{max}$​ from 0 to 2), SevenNet (2–3), NequIP (0–3), eSCN (6), and EquiformerV2 (6) demonstrate improved accuracy when employing higher $l_{max}$​ settings. The range of $l_{max}$​ used in our main evaluation (up to 5) is therefore well-aligned with standard practice in MLIP research and reflects a practically useful regime.\n\n# R5. Architectural details of SevenNet variants and NequIP usage\n- The tensor-product configurations used for the kernel microbenchmarks are detailed in Table 4 of Appendix A.1.\n- **NequIP** was _not_ used in the evaluation.\n- We apologize for the omission of configuration details for the three SevenNet variants. This information will be added to Appendix A.4 in the final draft.\n\n# References\n[1] Dunn et al., Benchmarking Materials Property Prediction Methods: The Matbench Test Set and Automatminer Reference Algorithm, npj Computational Materials 2020.", "claims_and_evidence": "The paper claims that FlashTP significantly accelerates tensor-product operations while reducing memory usage. The experimental results support this claim, showing notable performance improvements across different benchmarks. However, the comparison with cuEquivariance is primarily numerical, lacking a discussion of key differences. It remains unclear whether FlashTP's speedup stems from algorithmic advancements or benefits from being more constrained to architectures similar to NequIP and SevenNet.", "methods_and_evaluation_criteria": "The proposed library is well-aligned with the problem, and the evaluation consists of runtime and memory benchmarks on an NVIDIA A100 GPU. While these benchmarks provide valuable insight, the paper could benefit from additional evaluations on different MLIP architectures, particularly models incorporating symmetric contraction (e.g., MACE).", "theoretical_claims": "The paper suggests that higher-rank tensors (larger $l_\\mathrm{max}$ values) are important to consider when improving the numerical efficiency of tensor-product layers in models designed for atomistic simulations, using this argument to support the advantages over NVIDIA cuEquivariance. However, this claim is not entirely accurate—higher-rank tensors are generally only necessary in environments with higher local symmetry, typically relaxed in atomistic simulations. A more general argument beyond atomistic simulations might provide a stronger justification.", "experimental_designs_or_analyses": "The experiments convincingly show that FlashTP improves computational efficiency but are focused solely on models using spherical tensors. A broader comparison would strengthen the evaluation by mentioning Cartesian-based models and architectures and incorporating symmetric contraction for spherical models.", "relation_to_broader_scientific_literature": "The paper addresses an important issue in the context of equivariant machine learning models. However, it focuses exclusively on models in the spherical basis and does not mention other bases, such as the Cartesian one. Including a discussion of these alternatives would offer a more comprehensive understanding of how FlashTP fits within the broader landscape of MLIP development. Some references for Cartesian models include: https://arxiv.org/abs/2306.06482, https://arxiv.org/abs/2405.14253, and https://arxiv.org/abs/2412.18263 (and references therein).", "essential_references_not_discussed": "The paper does not mention works that explore machine-learned force fields in bases other than spherical. Given the growing body of research on Cartesian representations (see, e.g., above references), mentioning these studies could offer valuable context for the paper’s claims.", "other_strengths_and_weaknesses": "Please refer to the comments in previous sections for all strengths and weaknesses of the presented work.", "other_comments_or_suggestions": "Please refer to the comments in previous sections for all comments and suggestions.", "questions_for_authors": "My questions would relate to the issues or comments raised above, so addressing them would suffice to change my evaluation of the paper.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PU2WDs2Hrd", "reviewer": "Reviewer_GZbV", "summary": "The paper presents FlashTP, a highly optimized tensor-product library designed to address computational inefficiencies through kernel fusion, sparse computation, and path-aggregated execution. The proposed approach significantly accelerates tensor-product operations in equivariant neural networks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable feedback.\n\n# R1. Why Diffdock/MACE was not selected for end-to-end evaluation [Q1]\n- DiffDock is a diffusion-based molecular docking model, not a MLIP model. Its CGTP configuration is different from the configurations used in MLIPs, so it falls outside the scope of FlashTP. \n- Please refer our response R2 to Reviewer LtnP for why SevenNet was chosen over MACE.\n\n# R2. Roofline analysis of FlashTP [Q2]\n- We estimate the theoretical performance limits of FlashTP using a roofline analysis based on the peak capabilities of the A100 GPU (19.5 TFLOPS and 1.9 TB/s memory bandwidth). Under this model, the estimated latency for the kernel microbenchmark with $l_{\\text{max}} = 3$ is 0.46 ms for the forward pass, 0.89 ms for the backward pass, and 1.35 ms for the double-backward pass.\n- While the exact derivation of the \"Speed of Light\" (SoL) latencies reported in the cuEq blog [1] is not publicly documented, FlashTP’s exploitation of sparsity in the Clebsch–Gordan coefficient matrix reduces both theoretical memory traffic and floating-point operations. Consequently, the SoL latency for FlashTP is expected to be lower than that of cuEq.\n\n# R3. Channel scaling in FlashTP [Q3, Q4]\n- FlashTP supports various channel sizes without requiring any modifications, and the channel size is **not** limited by CUDA kernel register storage. The SevenNet variants used in the end-to-end evaluation employ multiple channel sizes (32, 64, and 128) for their features.\n- Below, we report the kernel microbenchmark result (in milliseconds) for $l_{\\text{max}} = 3$ with channel sizes of 32, 64, and 128. These results demonstrate that the performance improvement of FlashTP remains consistent across different channel sizes. \n\n|Phase|Channel|e3nn|FlashTP (Speedup)|\n|-|-|-|-|\n|Forward|32|18.40|2.54 (7.2x)|\n||64|34.25|5.02 (6.8x)|\n||128|66.80|9.98 (6.7x)|\n|Backward|32|85.00|6.15 (13.8x)|\n||64|163.96|12.14 (13.5x)|\n||128|324.73|24.19 (13.4x)|\n|Double Backward|32|254.13|8.73 (29.1x)|\n||64|503.28|17.41 (28.9x)|\n||128|1003.95 |34.81 (28.8x)|\n\n\n# R4. Discussion on cuEq performance characteristics [Q5, Q6]\n- Please view our response R3 to Reviewer LtnP.\n\n# R5. Contributions regarding sparsity\n- First, it is important to clarify that the sparsity discussed in [2] refers to pruning-based sparsity in CG tensor-product (CGTP) paths, not the sparsity in the Clebsch–Gordan (CG) coefficient matrix exploited by FlashTP. These are orthogonal techniques.\n- While the idea of leveraging sparsity in CG matrices has been explored (e.g., [3]), prior attempts have not yielded practical speedups. In fact, the sparse CGTP implementation in [3] was slower than its dense counterpart. This is often due to the overhead from metadata and irregular control flow present in sparse GPU computations, which can outweigh computational gains—a well-known issue in deep learning [4].\n- In contrast, we believe FlashTP makes a significant contribution by successfully leveraging sparsity to obtain measurable performance improvements as demonstrated in our ablation study. This is enabled by introduction of efficient sparse data structure and careful kernel design that leverages constant memory.\n\n# R6. Discrepancy between Figure 4 and reported inference speedup [Q7]\nWe apologize for the confusion and would like to clarify the source of the discrepancy:\n- The difference stems from the number of atoms (or more precisely, edges) used in the two settings. Inference breakdown in Figure 4 is obtained from the forward pass during training (~500 atoms), whereas the reported inference speedup is based on a MD simulation on a larger system with 4,096 atoms. In the 4K-atom setting, the tensor product accounts for approximately 88% of the total inference time, suggesting a maximum achievable speedup of around 8×. We will update Figure 4 in our final draft.\n\n# R7. Discussion about other acceleration methods\n- Please view our response R3 to Reviewer hLSR.\n\n# R8. Will FlashTP be open-sourced [Q8]\n- Yes, the code will be open-sourced. For more details on our future plans, please refer to our response R4 to Reviewer hLSR.\n\n# R9. Removing the inter-layer fusion [Q9]\n- Without the use of message-passing, inter-layer fusion becomes unnecessary. Removing this fusion requires only a minimal code change, which we have implemented for debugging purposes.\n\n# References\n[1] https://developer.nvidia.com/blog/accelerate-drug-and-material-discovery-with-new-math-library-nvidia-cuequivariance \\\n[2] Wang et al., Towards Flexible, Efficient, and Effective Tensor Product Networks, NeurIPS 2023 GLFrontiers Workshop.\\\n[3] Xie et al., The price of freedom: Exploring tradeoffs between expressivity and computational efficiency in equivariant tensor products, ICML 2024 Workshop GRaM.\\\n[4] Wen, et al., Learning structured sparsity in deep neural networks, NeurIPS 2016.", "claims_and_evidence": "Overall, this paper provides clear evidence to support its claims. \n\nIn Section 3, the authors conduct a detailed empirical analysis of the computational bottlenecks in tensor-product operations. In Section 4, they explicitly explain how their proposed methods address these issues one by one. This one-to-one correspondence in problem identification and solution design, to some extent, substantiates the potential effectiveness of their approach. Furthermore, in the evaluation section, the paper presents comprehensive ablation studies, including comparisons with baseline methods in forward, backward, and double-backward propagation scenarios, as well as end-to-end experiments within neural networks. \n\nThese results collectively demonstrate that the proposed approach successfully accelerates tensor-product computations, thereby validating the claim that it enhances the efficiency of equivariant neural networks.", "methods_and_evaluation_criteria": "Overall, the experimental design in this paper is suitable for addressing the proposed problems, though further optimization or clarification is needed in the benchmarking aspect.\n\nThe benchmarks chosen for the single tensor-product (tp) tests are e3nn and cuEquivariance, while for the end-to-end network tests, SevenNet-l3i5 is used. To my knowledge, cuEquivariance uses DiffDock and MACE for its end-to-end experiments, also demonstrating the difference between its implementation and theoretical limits. Based on this, \n\n1. why were DiffDock and MACE not chosen or included as additional end-to-end benchmarks for a more direct comparison of performance between the two methods? \n\n2. Does your method have any theoretical limit differences when compared to cuEquivariance?\n\nIn the experiments, the chosen node feature dimension is 32, which is common, but it shows a noticeable gap compared to the typical dimensions (128 or 256) used in other popular equivariant networks, such as the Equiformer series. \n\n3. Does your method scale well to tensor features with a larger channel size? If so, is this expansion a simple extension, or does it require additional specialized design or implementation? \n\n4. Is there a channel size limit based on the CUDA kernel register storage limits?\n\n[1]Liao Y L, Smidt T. Equiformer: Equivariant graph attention transformer for 3d atomistic graphs[J]. arXiv preprint arXiv:2206.11990, 2022.\n\n[2]Liao Y L, Wood B, Das A, et al. Equiformerv2: Improved equivariant transformer for scaling to higher-degree representations[J]. arXiv preprint arXiv:2306.12059, 2023.\n\n[3]https://developer.nvidia.com/blog/accelerate-drug-and-material-discovery-with-new-math-library-nvidia-cuequivariance/", "theoretical_claims": "The paper does not contain much content regarding theoretical claims. It would be beneficial to add a discussion about the theoretical maximum speed of the proposed method, as well as a scaling analysis of the tensor product with respect to the tensor order in comparison to the baseline methods.", "experimental_designs_or_analyses": "5. In Tables 2 and 3, CuEq shows significant performance degradation when l_max > 2, even falling below e3nn. This suggests that CuEq may not have implemented adaptive modifications or optimizations for these cases. Have you checked the correctness of CuEq's execution in these cases or analyzed potential reasons for this behavior? In this context, does the comparison with CuEq still hold significant value?\n\n6. If we focus only on the cases where l_max ≤ 2, the acceleration of FlashTP relative to CuEq seems limited, with only the Double-Backward acceleration being significantly improved. Is this behavior related to the specific design of your method?\n\n7. According to Figure 4, assuming that tensor-product operations account for 60.3% and 75.2% of the inference and training time in SevenNet-l3i5, does this imply an upper limit on the speedup (i.e., when the tensor-product time is optimized to 0) of approximately 2.5x for inference and 4x for training? How does this conclusion reconcile with the reported speedups of \"4.2× and 3.5×\" for inference and training, respectively?", "relation_to_broader_scientific_literature": "I believe the paper has sufficiently discussed the connection between tensor product computations and various prediction tasks in scientific fields.", "essential_references_not_discussed": "The authors should consider including a discussion on acceleration methods beyond cuEquivariance, particularly in the context of other approaches that improve efficiency.\n\nFor instance:  \n- Acceleration based on sparsity:  \n[1] Wang N, Lin C, Bronstein M, et al. \"Towards Flexible, Efficient, and Effective Tensor Product Networks\" (NeurIPS 2023 Workshop: New Frontiers in Graph Learning).  \n\n- Acceleration based on frequency domain:  \n[2] Luo S, Chen T, Krishnapriyan A S. \"Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products\" (arXiv preprint arXiv:2401.10216, 2024).  \n[3] Xie Y Q, Daigavane A, Kotak M, et al. \"The price of freedom: Exploring tradeoffs between expressivity and computational efficiency in equivariant tensor products\" (ICML 2024 Workshop on Geometry-grounded Representation Learning and Generative Modeling, 2024).", "other_strengths_and_weaknesses": "Strengths:\n1. The paper clearly describes the background and necessity of tensor-product acceleration, providing a detailed analysis of the speed bottlenecks in current tensor-product computations. This leads naturally to the proposed solution and evaluations, presenting a coherent story with a very smooth and readable writing style.\n2. The method section offers clear motivation and solutions, and the source code is provided in the appendix, which enhances the credibility and reproducibility of the work.\n3. The experimental design in the method section is comprehensive, particularly with the inclusion of the Double-Backward time tests, ensuring its applicability to MD tasks. The results also demonstrate significant speedups, confirming the potential of the method in real-world applications.\n\n Weaknesses:\n1. The paper lacks a discussion and validation on whether the proposed method can be applied to other general equivariant models, which would significantly impact the overall quality of the paper and its real influence in the community.\n2. The experimental section still contains aspects that need further clarification and optimization to strengthen the credibility of the results and to establish the speed advantages over existing methods. Further detailed discussions in other sections can serve as a reference.\n3. While the method presented in the paper is effective, there seems to be limited claim about its originality and innovation. Additionally, since the paper does not have a section discussing related works, the innovations in FlashTP, as described in Chapter 4, are not directly compared to existing methods. For example, Chapter 4.3 discusses \"Applying Sparsity in Tensor-Product,\" but the application of CG sparsity has already been discussed and applied earlier, as seen in the referenced papers mentioned above.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Most of the questions have already been raised in the sections above. Additionally:\n\n8. Do you have plans to open-source the work presented in this paper?\n\n9. Regarding Inter-Layer Kernel Fusion, is the fusion heavily tied to the message-passing function, making its extension to other networks a non-trivial task?\n\nOverall, I am inclined to believe that there is significant engineering work behind this paper, but I still have concerns about the reliability of the experimental results, the innovation and generalization of the method. If the authors are unable to clarify or address these issues, I may reconsider my evaluation.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: All reviewers agree that this is an important engineering contribution in the domain of machine learning for molecular applications.  The analogy to flash attention for transformers appears fair, and this work has the potential to inspire follow-up work or re-use in adjacent domains.
Decision: Accept (spotlight poster)