UltraGNN: A Sparse-Operator-Aware Framework for Accelerating Graph Neural Networks on Tensor Cores

Abstract

Graph Neural Networks (GNNs) have achieved widespread success from social networks to AI-for-Science. Most existing GNN frameworks adopt scatter-first (edge-centric) or gather-first (vertex-centric) scheduling paradigms for message passing. However, these paradigms are closely tied to traditional CUDA-core execution model in terms of data access patterns, scheduling granularity, and compute mapping, limiting their ability to fully leverage the superior matrix computation capabilities of Tensor Core Units (TCUs) on modern GPUs. To this end, we propose UltraGNN, the first sparse-operator-aware framework designed to efficiently accelerate both topology-driven and dynamic-weight GNN models on TCUs. UltraGNN introduces a novel subgraph-centric scheduling paradigm, which reformulates message generation and aggregation as TCU-friendly sparse matrix multiplications. To enable efficient support for dynamic-weight GNN architectures on TCUs, UltraGNN incorporates a transpose-friendly sparse storage format explicitly designed to facilitate fast gradient computation and memory-efficient training and inference. UltraGNN adopts PyTorch as its frontend, delivering an out-of-the-box user experience. Extensive experimental results on the H100 and RTX 4090 GPUs demonstrate that UltraGNN sets a new state-of-the-art in GNN acceleration, achieving geometric mean speedups of 2.47× and 5.46× over DGL on GAT and GATv2 models, respectively.