Scaling Computational Performance of Spherical Harmonics Kernels with Triton

Kin Long Kelvin Lee, Mikhail Galkin, Santiago Miret

Submission Track: Findings

Submission Category: AI-Guided Design

Keywords: spherical harmonics, e3nn, computational bottlenecks, equivariant neural networks

TL;DR: We implement Triton-based e3nn kernels and show 5x improved compute performance.

Abstract:

Spherical harmonics are the key ingredient in equivariant neural networks often used in Geometric Deep Learning applications to molecules, proteins, and materials where it is crucial to imbue the model with rotational and translational symmetries.
However, computing spherical harmonics at higher orders and on larger inputs often constitutes a major performance bottleneck in equivariant models.
In this work, we propose a set of efficient forward and backward kernels implemented in the Triton language for computing spherical harmonics on systems of up to 100 million atoms.
Experimentally, the Triton implementation brings improvements of up to $5\times$ in time and up to $3\times$ in memory compared to the popular e3nn library while being portable to any GPU accelerator with the Triton backend.

Submission Number: 11