An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks
作者: Vivek Bharadwaj, Austin Glover, Aydin Buluc, James Demmel
分类: cs.LG, cs.AI
发布日期: 2025-01-23 (更新: 2025-05-08)
备注: To appear in the Proceedings of the 2025 SIAM Conference on Applied and Computational Discrete Algorithms. 15 pages, 10 figures, 4 tables
💡 一句话要点
提出高效稀疏核生成器，加速O(3)-等变深度网络的Clebsch-Gordan张量积运算。
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 等变神经网络 Clebsch-Gordan张量积 GPU优化 稀疏核生成 分子动力学 原子间势 高性能计算
📋 核心要点
- 旋转等变图神经网络在空间深度学习中表现出色，但Clebsch-Gordan张量积运算是其性能瓶颈。
- 论文提出一种GPU稀疏核生成器，通过静态分析管理GPU共享内存，优化指令流，加速CG张量积运算。
- 实验结果表明，该方法在推理速度上优于现有方案，最高可达10倍加速，尤其在FP64精度下提升显著。
📝 摘要（中文）
本文提出了一种用于Clebsch-Gordan (CG) 张量积的GPU稀疏核生成器，旨在加速旋转等变图神经网络。该网络在空间深度学习任务上表现出色，与传统方法相比，在训练期间具有更高的数据效率，并显著减少了原子间势计算的推理时间。CG张量积是该模型的关键，它使用高度结构化的稀疏张量来收缩两个密集特征向量，以产生密集输出向量。该操作在典型的等变模型中可能会重复数百万次，是一个代价高昂且效率低下的瓶颈。通过在模型编译时进行静态分析，仔细管理有限的GPU共享内存，从而最大限度地减少对全局内存的读写，实现高性能。将张量积分解为一系列更小的内核，其操作数完全适合寄存器，从而能够发出最大化指令级并行性的长算术指令流。通过将CG张量积与后续的图卷积融合，减少了中间存储和全局内存流量。此外，还为CG张量积的梯度以及预测原子间力所需的高阶偏导数提供优化的内核。与NVIDIA的cuEquivariance软件包相比，内核提供了高达1.3倍的加速，与广泛使用的e3nn软件包相比，提供了10倍的加速。在FP64精度下，对于MACE化学基础模型，与原始未优化的版本相比，提供了高达6.2倍的推理时间加速。
🔬 方法详解
问题定义：旋转等变图神经网络依赖于Clebsch-Gordan (CG) 张量积，但其计算复杂度高，成为性能瓶颈。现有实现方案效率低下，限制了模型在实际应用中的速度。
核心思路：通过生成优化的GPU稀疏核来加速CG张量积运算。核心在于充分利用GPU的共享内存和寄存器，减少全局内存的访问，并最大化指令级并行性。
技术框架：该方法包含以下几个关键阶段：1) 模型编译时进行静态分析，确定最佳的内存管理策略；2) 将CG张量积分解为一系列更小的内核，每个内核的操作数都能放入寄存器；3) 生成优化的GPU内核代码，利用长算术指令流最大化指令级并行性；4) 将CG张量积与后续的图卷积融合，减少中间存储和全局内存流量。
关键创新：主要创新在于针对CG张量积的特性，设计了高效的GPU稀疏核生成器。该生成器能够根据模型结构，自动生成优化的内核代码，从而避免了手动优化带来的困难和局限性。此外，还针对CG张量积的梯度和高阶偏导数进行了优化，使其能够应用于更广泛的场景。
关键设计：关键设计包括：1) 静态分析模块，用于确定最佳的共享内存使用策略；2) 内核分解模块，用于将CG张量积分解为更小的、适合寄存器操作的内核；3) 代码生成模块，用于生成优化的GPU内核代码，包括指令调度和内存访问优化；4) 融合模块，用于将CG张量积与后续的图卷积融合，减少内存访问。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在CG张量积运算速度上显著优于现有方案。与NVIDIA的cuEquivariance软件包相比，实现了高达1.3倍的加速；与广泛使用的e3nn软件包相比，实现了10倍的加速。在FP64精度下，对于MACE化学基础模型，与原始未优化的版本相比，推理时间加速高达6.2倍。
🎯 应用场景
该研究成果可广泛应用于分子动力学模拟、材料科学、药物发现等领域。通过加速原子间势的计算，可以更高效地进行分子模拟，从而加速新材料的开发和药物的筛选。此外，该方法还可以应用于其他需要旋转等变性的深度学习任务，例如三维物体识别和场景理解。
📄 摘要（原文）
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.