On Efficient Scaling of GNNs via IO-Aware Layers Implementations

arXiv:2605.31500v1 Announce Type: cross Abstract: 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 extbf{3.9}imes speedup for Graph Transformer (median extbf{1.6}imes), with Tensor Core (block-sparse) variants up to extbf{7.3}imes on locally dense graphs; for GATv2 we reach up to extbf{8.5}imes speedup (median extbf{2.0}imes) while reducing peak memory by up to extbf{76}imes (median extbf{6}imes). Our degree-aware reduction kernels achieve up to extbf{10}imes speedup (median extbf{2.6}imes). For SpMM-based layers, properly cached cuSPARSE achieves up to extbf{8}imes 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.

Source: arXiv cs.AI | 2026-06-01