Thank you for your valuable feedback.

R1. Why Diffdock/MACE was not selected for end-to-end evaluation [Q1]

• 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.
• Please refer our response R2 to Reviewer LtnP for why SevenNet was chosen over MACE.

R2. Roofline analysis of FlashTP [Q2]

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

R3. Channel scaling in FlashTP [Q3, Q4]

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

R4. Discussion on cuEq performance characteristics [Q5, Q6]

• Please view our response R3 to Reviewer LtnP.

R5. Contributions regarding sparsity

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

R6. Discrepancy between Figure 4 and reported inference speedup [Q7]

We apologize for the confusion and would like to clarify the source of the discrepancy:

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

R7. Discussion about other acceleration methods

• Please view our response R3 to Reviewer hLSR.

R8. Will FlashTP be open-sourced [Q8]

• Yes, the code will be open-sourced. For more details on our future plans, please refer to our response R4 to Reviewer hLSR.

R9. Removing the inter-layer fusion [Q9]

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

References

[1] https://developer.nvidia.com/blog/accelerate-drug-and-material-discovery-with-new-math-library-nvidia-cuequivariance
[2] Wang et al., Towards Flexible, Efficient, and Effective Tensor Product Networks, NeurIPS 2023 GLFrontiers Workshop.
[3] Xie et al., The price of freedom: Exploring tradeoffs between expressivity and computational efficiency in equivariant tensor products, ICML 2024 Workshop GRaM.
[4] Wen, et al., Learning structured sparsity in deep neural networks, NeurIPS 2016.

Thank you for your valuable feedback.

R1. Discussion regarding Cartesian-based models

• Please refer to the last paragraph of our response R3 to Reviewer hLSR.

R2. Choice of SevenNet over MACE for end-to-end evaluation

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

R3. Speculation on the performance differences between FlashTP and cuEq

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

Performance degradation of cuEq for higher values of $l_{max}$.

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

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

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

Source of speedup for FlashTP

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

R4. On the practical benefits of higher $l_{max}$

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

R5. Architectural details of SevenNet variants and NequIP usage

• The tensor-product configurations used for the kernel microbenchmarks are detailed in Table 4 of Appendix A.1.
• NequIP was not used in the evaluation.
• 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.

References

[1] Dunn et al., Benchmarking Materials Property Prediction Methods: The Matbench Test Set and Automatminer Reference Algorithm, npj Computational Materials 2020.

Thank you for your valuable feedback.

R1. Evaluation on multi-GPU training [W1, Q1]

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

R2. Essential references not discussed

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

R3. Acceleration of Clebsch-Gordan Tensor Product (CGTP)

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

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

R4. Future plans for FlashTP [W2, Q2]

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

• 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.
• To facilitate reproducibility and adoption by the broader community, we will include a benchmarking script and a comprehensive integration example in the repository.

R5. Roofline model analysis [W3]

• Please refer to our response R2 to Reviewer GZbV for the theoretical performance of FlashTP estimated from roofline analysis.
• 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.

References

[1] Passaro et al., Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs, ICML 2023.
[2] Luo et al., Enabling Efficient Equivariant Operations in the Fourier Basis via Gaunt Tensor Products, ICLR 2024.
[3] Xie et al., The price of freedom: Exploring tradeoffs between expressivity and computational efficiency in equivariant tensor products, ICML 2024 Workshop GRaM.
[4] Unke et al., E3x: E(3)-equivariant deep learning made easy, CoRR 2024.
[5] Simeon et al., TensorNet: Cartesian tensor representations for efficient learning of molecular potentials, NeurIPS 2023.
[6] Zaverkin et al., Higher-Rank Irreducible Cartesian Tensors for Equivariant Message Passing, NeurIPS 2024.
[7] https://opencatalystproject.org/leaderboard.html
[8] https://matbench-discovery.materialsproject.org