An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks

TL;DR

This work tackles the computational bottleneck of the Clebsch-Gordon tensor product in $O(3)$-equivariant deep networks by presenting an open-source GPU sparse kernel generator. Through JIT compilation, static analysis, and warp-focused scheduling, it efficiently handles the highly structured yet block-sparse CG contraction, and it fuses this primitive with graph convolutions to cut memory traffic. The authors demonstrate substantial speedups over both open-source and closed-source baselines across forward, backward, and higher-derivative computations, including notable gains for Nequip and MACE in FP32 and FP64. Overall, the approach delivers practical, scalable acceleration for large-scale equivariant chemistry models, enabling more efficient end-to-end training and inference. Potential future work includes extending to lower precisions and further integration with new equivariant architectures.

Abstract

Rotation equivariant graph neural networks, i.e. networks designed to guarantee certain geometric relations between their inputs and outputs, yield state of the art performance on spatial deep learning tasks. They exhibit high data efficiency during training and significantly reduced inference time for interatomic potential calculations compared to classical approaches. Key to these models is the Clebsch-Gordon (CG) tensor product, a kernel that contracts two dense feature vectors with a highly-structured sparse tensor to produce a dense output vector. The operation, which may be repeated millions of times for typical equivariant models, is a costly and inefficient bottleneck. We introduce a GPU sparse kernel generator for the CG tensor product that provides significant speedups over the best existing open and closed-source implementations. Our implementation achieves high performance by carefully managing the limited GPU shared memory through static analysis at model compile-time, minimizing reads and writes to global memory. We break the tensor product into a series of smaller kernels with operands that fit entirely into registers, enabling us to emit long arithmetic instruction streams that maximize instruction-level parallelism. By fusing the CG tensor product with a subsequent graph convolution, we reduce both intermediate storage and global memory traffic over naive approaches that duplicate input data. We also provide optimized kernels for the gradient of the CG tensor product and a novel identity for the higher partial derivatives required to predict interatomic forces. Our kernels offer up to 1.3x speedup over NVIDIA's closed-source cuEquivariance package, as well as 10x speedup over the widely-used e3nn package. In FP64 precision, we offer up to 6.2x inference-time speedup for the MACE chemistry foundation model over the original unoptimized version.

Figures (10)

• Figure 1: The CG tensor product, which contracts a block-sparse tensor $\mathcal{P}$ with two dense vectors to output a new vector. It is usually followed by multiplication with a structured weight matrix $\bm W$, and by convention we use "CG tensor product" to refer to both operations in sequence. Each blue block is itself sparse (see Figure \ref{['fig:cg_tensor_structure']}), and several blocks may share identical structure.
• Figure 2: Three examples of coefficient tensors depicted by blue cubes in Figure \ref{['fig:cg_tensor_product']}. The fraction of zero entries increases with the tensor order, and the full CG tensor contains several copies of the blocks pictured here.
• Figure 3: Fundamental subkernels that compose to implement the CG tensor product in Figure \ref{['fig:cg_tensor_product']}. In (A), $\bm x$, $\bm y$, and $\bm z$ refer to segments of the longer vectors in Figure \ref{['fig:cg_tensor_product']}, and $\bm W$ contains entries from the larger weight matrix rearranged appropriately. In (B) and (C), $\bm X$ and $\bm Z$ refer to segments from $\bm x$ and $\bm z$ that have been reshaped into matrices to exploit repeating sparse tensor structure. $b$ and $b'$ are multiples of 32 in many models.
• Figure 4: Applications of the CG tensor product. The simplest and most general use case (A) calls the kernel repeatedly with distinct $\bm x$, $\bm y$, and $\bm W$ inputs. Interatomic potential models embed the operation in a graph convolution (B), where the tensor product combines node features with edge features.
• Figure 5: Throughput of CG tensor products (batch size 50K), kernel B configurations without SpMM kernel fusion. On difficult configurations like Nequip-benzene with massive output vector lengths, we exhibit more than 10x improvement over e3nn.