arXiv:2605.31500cs.LGcs.AI2026-05

优化图神经网络的内存访问,提升大规模图计算效率

On Efficient Scaling of GNNs via IO-Aware Layers Implementations

  • 按计算与访存特性分类核心层,设计减少数据移动的专用GPU内核
  • 图注意力层加速达8.5倍,内存峰值降低76倍,局部稠密图上提速7.3倍
  • 代码可直接替换现有框架,适合追求高效图模型部署的研究者

图神经网络受稀疏不规则内存访问限制。主流框架如DGL和PyTorch Geometric虽支持通用消息传递,但复杂层常生成边级中间结果,增加内存流量并限制大规模图扩展性。本文从访存与算术强度角度出发,发现常用层可分为三类:基于SpMM的卷积、基于规约的聚合、注意力型层(GATv2/Graph Transformer)。针对每类,开发了减少数据移动、提升局部性的GPU内核,并研究图重排影响:其对邻居并行(聚集主导)内核更有效。实验表明,融合注意力内核在Graph Transformer上实现最高3.9倍加速(中位1.6倍),使用张量核心(块稀疏)版本在局部稠密图上达7.3倍;GATv2最高提速8.5倍(中位2.0倍),峰值内存降低76倍(中位6倍)。度感知规约内核最高提速10倍(中位2.6倍)。SpMM类层通过缓存cuSPARSE实现最高8倍加速,优于多数自研基线。代码已开源,可作为即插即用组件支持可复现的硬件感知图神经网络加速。

原文摘要 · Abstract (English)

Graph 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 intermediates, increasing memory traffic and limiting scalability on large graphs. We take an I/O- and arithmetic-intensity--centric view and show that widely used layers fall into three kernel families: SpMM-based convolutions, reduction-based aggregations, and attention-based layers (GATv2/Graph Transformer). For each family, we develop GPU kernels that reduce data movement, improve locality, and remain robust across realistic graphs. We also study graph reordering and find that its impact depends on the kernel mapping: it benefits neighbor-parallel (gather-dominated) kernels more consistently than feature-parallel designs. Empirically, our fused attention kernels reach up to $\textbf{3.9}\times$ speedup for Graph Transformer (median $\textbf{1.6}\times$), with Tensor Core (block-sparse) variants up to $\textbf{7.3}\times$ on locally dense graphs; for GATv2 we reach up to $\textbf{8.5}\times$ speedup (median $\textbf{2.0}\times$) while reducing peak memory by up to $\textbf{76}\times$ (median $\textbf{6}\times$). Our degree-aware reduction kernels achieve up to $\textbf{10}\times$ speedup (median $\textbf{2.6}\times$). For SpMM-based layers, properly cached cuSPARSE achieves up to $\textbf{8}\times$ speedup over DGL and outperforms evaluated custom baselines in the majority of evaluations. We release our implementations as drop-in replacements to support reproducible, hardware-aware GNN acceleration.

图神经网络GPU优化内存效率深度学习

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。