Results for Graph Transformer speedups on selected graphs (heads H = 4, head dim D = 128). N/A means DF-GNN encountered illegal memory access on this setup and failed to launch; higher is better.
On efficient scaling of GNNs via IO-aware layer implementations
We introduce a set of IO-Aware kernels for Graph Neural Networks, bridging the gap between large-scale Graph Machine Learning and hardware-efficient deep learning. The paper was accepted at ICML 2026 as a Spotlight Paper.
Graph Neural Networks (GNNs) power a wide range of applications, from recommendation systems and social network analysis to molecular modeling and fraud detection. Yet despite their success, GNNs remain notoriously difficult to run efficiently on GPUs. Moreover, one of the most popular libraries for Graph-ML — Deep Graph Library (DGL) — is becoming unmaintained as it hasn’t received any meaningful updates over the past year. This further highlights the need for efficient, easy-to-maintain implementations of core graph operations.
At the heart of this challenge is the way graph computations interact with modern hardware. The core operation behind most graph models is aggregating information from neighboring nodes. Unlike dense tensors used in transformers, graph neighborhoods are highly irregular. One node may have three neighbors, another three thousand, and their features can be scattered throughout the entire tensor in GPU memory. As a result, the bottleneck is often not computation itself, but moving data between high-bandwidth memory (HBM) and fast on-chip GPU memory.
As part of a research project course, we set out to address this challenge. The project brought together Yandex researchers and engineers with students from the Yandex School of Data Analysis (YSDA) and HSE University. The team included Fedor Velikonivtsev as project supervisor, Daniil Krasilnikov, Alexey Boykov, and Andrey Dolgovyazov from YSDA, Vyacheslav Zhdanovskiy from Yandex’s LLM Inference Platform team, and Daria Fomina from Yandex’s ML Infrastructure team. Together, we developed a collection of specialized GPU kernels for the most widely used families of graph neural networks.
In our Spotlight paper accepted at ICML 2026, we revisit the implementation of popular GNN layers from an IO-awareness perspective. In many ways, our goal is similar to what FlashAttention achieved for transformers: reducing memory movement rather than focusing solely on arithmetic efficiency.
Latency speedup (tDGL/t) and Peak-memory ratio (memDGL/mem) for GATv2 on selected graphs (heads H = 2, head dim D = 64); higher is better.
We identify three major families of graph layers, each suffering from memory bottlenecks in different ways:
- Convolutional layers which can be represented as Sparse-Dense Matrix Multiplication (SpMM), like Graph Convolutional Networks or Mean-Aggregation;
- Aggregation layers (such as min/max pooling);
- Attention-based layers, including GATv2 and Graph Transformers.
For attention-based architectures, we introduce IO-aware kernels inspired by FlashAttention. Intermediate edge-level attention values are never materialized in slow global memory. Instead, they are computed and consumed on the fly using an online softmax algorithm directly in SRAM. During the backward pass, attention weights are recomputed from compact statistics rather than stored explicitly, following ideas introduced in FlashAttention-2. We also propose a block-sparse graph representation that enables the use of tensor cores, providing the largest gains on locally dense graphs.
For aggregation layers, the key challenge comes from the heavy-tailed degree distribution commonly observed in real-world graphs. A small number of high-degree nodes can dominate execution time and limit parallelism. To address this, we introduce degree-aware tiling and additional parallelization strategies specifically designed for such nodes.
Interestingly, we find that for graph convolution layers, modern NVIDIA implementations combined with effective caching strategies already outperform many specialized alternatives. This suggests that custom kernels are not always necessary — a point that has often been overlooked in previous work.
Across a wide range of datasets and architectures, our implementations deliver up to 8.5× speedups while reducing peak memory consumption by more than an order of magnitude in some configurations. Importantly, we release all kernels as drop-in replacements for popular GNN frameworks, allowing practitioners to benefit from these improvements without rewriting existing models.
Results for reduction-based convolutions (with min-aggregation as the operator) on selected graphs (hidden dim D = 512).
We hope this work helps make large-scale graph learning more practical and efficient, while highlighting the importance of IO-aware design for the next generation of graph neural network systems.