GNN 加速新方法:IO 感知的 GPU 内核实现,最高 8.5 倍提速

On Efficient Scaling of GNNs via IO-Aware Layers Implementations

精选理由

做图神经网络训练或推理的开发者,可以直接用这些内核替换现有层,无需改模型结构就能获得数倍加速和大幅内存节省,值得一试。

AI 摘要

图神经网络(GNN)因稀疏、不规则的内存访问而性能受限。该研究从 I/O 和计算强度角度出发,将常用 GNN 层分为三类(SpMM 卷积、归约聚合、注意力层),并为每类开发了减少数据移动、提升局部性的 GPU 内核。实验显示,融合注意力内核在 Graph Transformer 上最高提速 3.9 倍(中位数 1.6 倍),GATv2 最高提速 8.5 倍(中位数 2.0 倍),峰值内存降低最多 76 倍。研究还发现图重排序对邻居并行内核更有效。所有实现作为即插即用替代方案开源,便于复现。

AI 翻译 · 中文

图神经网络(GNN)因稀疏、不规则的内存访问而性能受限。该研究从 I/O 和计算强度角度出发,将常用 GNN 层分为三类(SpMM 卷积、归约聚合、注意力层),并为每类开发了减少数据移动、提升局部性的 GPU 内核。实验显示,融合注意力内核在 Graph Transformer 上最高提速 3.9 倍(中位数 1.6 倍),GATv2 最高提速 8.5 倍(中位数 2.0 倍),峰值内存降低最多 76 倍。研究还发现图重排序对邻居并行内核更有效。所有实现作为即插即用替代方案开源,便于复现。

arXiv cs.LGGraph Neural Networks (GNNs) are bottlenecked by sparse, irregular memory access. Popular frameworks such as DGL and PyTorch Geometric support general message passing, but complex layers often materialize edge-wise inter