Batmobile: 10-20x Faster CUDA Kernels for Equivariant Graph Neural Networks
4 months ago
- #CUDA
- #GNN
- #Machine Learning
- Batmobile通过定制CUDA内核加速等变图神经网络（如MACE、NequIP和Allegro）中的球谐函数和张量积运算
- 等变图神经网络遵循物理对称性（旋转/平移/反射）但计算成本高昂，导致实际应用受限
- 球谐函数编码3D方向信息，张量积则在保持等变性的前提下融合特征
- 标准库e3nn因Python/PyTorch开销、内存带宽浪费、缺乏算子融合和动态形状导致性能低下
- Batmobile采用编译时常量、寄存器中间变量和融合运算实现性能优化
- 基准测试显示Batmobile在球谐函数和张量积运算上比e3nn快10-20倍
- 该工具专为L_max=3场景设计，34条Clebsch-Gordan路径全部展开且系数作为编译时常量
- 项目命名Batmobile（蝙蝠车）寓意其专为分子模拟设计的高性能特性
- 已在GitHub开源并提供基准测试和示例，支持快速集成