# On Efficient Scaling of GNNs via IO-Aware Layers Implementations

arXiv: https://arxiv.org/abs/2605.31500 · HTML: https://mlsanity.com/paper/2605.31500/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 29 of its 42 references that are in the map.
- First posted on arXiv 29 May 2026.

## Builds on (most cited first, 29 of 29)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Neural Message Passing for Quantum Chemistry](https://mlsanity.com/paper/1704.01212/) (2017)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Revisiting Semi-Supervised Learning with Graph Embeddings](https://mlsanity.com/paper/1603.08861/) (2016)
- [Learning to Simulate Complex Physics with Graph Networks](https://mlsanity.com/paper/2002.09405/) (2020)
- [Learning Mesh-Based Simulation with Graph Networks](https://mlsanity.com/paper/2010.03409/) (2020)
- [How Attentive are Graph Attention Networks?](https://mlsanity.com/paper/2105.14491/) (2021)
- [FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision](https://mlsanity.com/paper/2407.08608/) (2024)
- [Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification](https://mlsanity.com/paper/2009.03509/) (2020)
- [Data Movement Is All You Need: A Case Study on Optimizing Transformers](https://mlsanity.com/paper/2007.00072/) (2020)
- [Transformer for Graphs: An Overview from Architecture Perspective](https://mlsanity.com/paper/2202.08455/) (2022)
- [DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusion](https://mlsanity.com/paper/2108.13342/) (2021)
- [PyG 2.0: Scalable Learning on Real World Graphs](https://mlsanity.com/paper/2507.16991/) (2025)
- [GNNAdvisor: An Adaptive and Efficient Runtime System for GNN Acceleration on GPUs](https://mlsanity.com/paper/2006.06608/) (2020)
- [The Deep Learning Compiler: A Comprehensive Survey](https://mlsanity.com/paper/2002.03794/) (2020)
- [GE-SpMM: General-purpose Sparse Matrix-Matrix Multiplication on GPUs for Graph Neural Networks](https://mlsanity.com/paper/2007.03179/) (2020)
- [Understanding GNN Computational Graph: A Coordinated Computation, IO, and Memory Perspective](https://mlsanity.com/paper/2110.09524/) (2021)
- [MaxK-GNN: Extremely Fast GPU Kernel Design for Accelerating Graph Neural Networks Training](https://mlsanity.com/paper/2312.08656/) (2023)
- [GraphLand: Evaluating Graph Machine Learning Models on Diverse Industrial Data](https://mlsanity.com/paper/2409.14500/) (2024)
- [TC-GNN: Bridging Sparse GNN Computation and Dense Tensor Cores on GPUs](https://mlsanity.com/paper/2112.02052/) (2021)
- [Cached Operator Reordering: A Unified View for Fast GNN Training](https://mlsanity.com/paper/2308.12093/) (2023)
- [Can Graph Reordering Speed Up Graph Neural Network Training? An Experimental Study](https://mlsanity.com/paper/2409.11129/) (2024)
- [Can Tensor Cores Benefit Memory-Bound Kernels? (No!)](https://mlsanity.com/paper/2502.16851/) (2025)
- [Fused3S: Fast Sparse Attention on Tensor Cores](https://mlsanity.com/paper/2505.08098/) (2025)

## Abstract (arXiv)

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