Batmobile: 10-20x Faster CUDA Kernels for Equivariant Graph Neural Networks
Summary
Batmobile presents hand-tuned CUDA kernels to accelerate equivariant GNN primitives (spherical harmonics and Clebsch-Gordan tensor products) for L_max=3, achieving 10-20x speedups over generic implementations. It explains bottlenecks in e3nn, describes compile-time constants, register-based intermediates, and fused SH+TP kernels, with benchmark results on RTX 3090 and a Python API example.