Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MaxK-GNN: Extremely Fast GPU Kernel Design for Accelerating Graph Neural Networks Training

Hongwu Peng, Xi Xie, Kaustubh Shivdikar, MD Amit Hasan, Jiahui Zhao, Shaoyi Huang, Omer Khan, David Kaeli, Caiwen Ding

TL;DR

MaxK-GNN targets the core bottleneck in GNN training on GPUs: irregular memory access and memory traffic during SpMM-based graph aggregation. By introducing the MaxK nonlinearity (top-$k$ feature selection per node) and the Compressed Balanced Sparse Row (CBSR) format, the authors create regular sparsity patterns that enable a row-wise SpGEMM forward path and an outer-product-based SSpMM backward path, both optimized for on-chip buffering and coalesced memory accesses. The CUDA-based kernels and memory-system design achieve substantial memory-traffic reductions (around 90% in key paths) and significant speedups over baselines like cuSPARSE SpMM and GNNAdvisor, while preserving competitive accuracy across GraphSAGE, GCN, and GIN on multiple datasets. The results indicate the practical potential of algorithm-system co-design for scalable GNN training on GPUs, with room to extend the MaxK paradigm to other DNN architectures and hardware platforms.

Abstract

In the acceleration of deep neural network training, the GPU has become the mainstream platform. GPUs face substantial challenges on GNNs, such as workload imbalance and memory access irregularities, leading to underutilized hardware. Existing solutions such as PyG, DGL with cuSPARSE, and GNNAdvisor frameworks partially address these challenges but memory traffic is still significant. We argue that drastic performance improvements can only be achieved by the vertical optimization of algorithm and system innovations, rather than treating the speedup optimization as an "after-thought" (i.e., (i) given a GNN algorithm, designing an accelerator, or (ii) given hardware, mainly optimizing the GNN algorithm). In this paper, we present MaxK-GNN, an advanced high-performance GPU training system integrating algorithm and system innovation. (i) We introduce the MaxK nonlinearity and provide a theoretical analysis of MaxK nonlinearity as a universal approximator, and present the Compressed Balanced Sparse Row (CBSR) format, designed to store the data and index of the feature matrix after nonlinearity; (ii) We design a coalescing enhanced forward computation with row-wise product-based SpGEMM Kernel using CBSR for input feature matrix fetching and strategic placement of a sparse output accumulation buffer in shared memory; (iii) We develop an optimized backward computation with outer product-based and SSpMM Kernel. We conduct extensive evaluations of MaxK-GNN and report the end-to-end system run-time. Experiments show that MaxK-GNN system could approach the theoretical speedup limit according to Amdahl's law. We achieve comparable accuracy to SOTA GNNs, but at a significantly increased speed: 3.22/4.24 times speedup (vs. theoretical limits, 5.52/7.27 times) on Reddit compared to DGL and GNNAdvisor implementations.

MaxK-GNN: Extremely Fast GPU Kernel Design for Accelerating Graph Neural Networks Training

TL;DR

MaxK-GNN targets the core bottleneck in GNN training on GPUs: irregular memory access and memory traffic during SpMM-based graph aggregation. By introducing the MaxK nonlinearity (top- feature selection per node) and the Compressed Balanced Sparse Row (CBSR) format, the authors create regular sparsity patterns that enable a row-wise SpGEMM forward path and an outer-product-based SSpMM backward path, both optimized for on-chip buffering and coalesced memory accesses. The CUDA-based kernels and memory-system design achieve substantial memory-traffic reductions (around 90% in key paths) and significant speedups over baselines like cuSPARSE SpMM and GNNAdvisor, while preserving competitive accuracy across GraphSAGE, GCN, and GIN on multiple datasets. The results indicate the practical potential of algorithm-system co-design for scalable GNN training on GPUs, with room to extend the MaxK paradigm to other DNN architectures and hardware platforms.

Abstract

In the acceleration of deep neural network training, the GPU has become the mainstream platform. GPUs face substantial challenges on GNNs, such as workload imbalance and memory access irregularities, leading to underutilized hardware. Existing solutions such as PyG, DGL with cuSPARSE, and GNNAdvisor frameworks partially address these challenges but memory traffic is still significant. We argue that drastic performance improvements can only be achieved by the vertical optimization of algorithm and system innovations, rather than treating the speedup optimization as an "after-thought" (i.e., (i) given a GNN algorithm, designing an accelerator, or (ii) given hardware, mainly optimizing the GNN algorithm). In this paper, we present MaxK-GNN, an advanced high-performance GPU training system integrating algorithm and system innovation. (i) We introduce the MaxK nonlinearity and provide a theoretical analysis of MaxK nonlinearity as a universal approximator, and present the Compressed Balanced Sparse Row (CBSR) format, designed to store the data and index of the feature matrix after nonlinearity; (ii) We design a coalescing enhanced forward computation with row-wise product-based SpGEMM Kernel using CBSR for input feature matrix fetching and strategic placement of a sparse output accumulation buffer in shared memory; (iii) We develop an optimized backward computation with outer product-based and SSpMM Kernel. We conduct extensive evaluations of MaxK-GNN and report the end-to-end system run-time. Experiments show that MaxK-GNN system could approach the theoretical speedup limit according to Amdahl's law. We achieve comparable accuracy to SOTA GNNs, but at a significantly increased speed: 3.22/4.24 times speedup (vs. theoretical limits, 5.52/7.27 times) on Reddit compared to DGL and GNNAdvisor implementations.
Paper Structure (29 sections, 2 theorems, 4 equations, 10 figures, 5 tables, 2 algorithms)

This paper contains 29 sections, 2 theorems, 4 equations, 10 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background and Related work
  3. Graph Convolution Network
  4. GNN Acceleration
  5. Introducing Sparsity in GNN Training
  6. MaxK-GNN Dataflow
  7. MaxK Nonlinearity as a Universal Approximator
  8. MaxK-GNN Training Dataflow
  9. GPU System Support for MaxK-GNN
  10. Forward SpGEMM GPU Kernel
  11. Backward SSpMM GPU Kernel
  12. Memory System
  13. Kernel Profiling
  14. Evaluation
  15. Experimental Setup
  16. ...and 14 more sections

Key Result

proposition 1

Given any positive integers $s$ (input dimension) and $t$ (output dimension), two parameter groups $W$ and $W'$ are determined such that $g(X)$ is expressed as a linear combination of $r$ convex PWL functions: $g(X)$ operates as the neural network approximator and denotes the continuous PWL function with $r$ locally affine regions on $\mathbb{R}^s$. Thus, $g(X)$ (could be any continuous PWL funct

Figures (10)

  • Figure 1: GraphSAGE structure analysis: latency breakdown of full-batch GraphSAGE training on the ogbn-proteins dataset over 30 epochs, with 256 hidden dimensions. GPU platform: Nvidia A100.
  • Figure 2: GraphSAGE layer example with (a) ReLU (b) MaxK nonlinearity. SSpMM: sampled sparse matrix dense matrix multiplication.
  • Figure 3: Computational workflow of a GCNConv layer.
  • Figure 4: MLP with MaxK and ReLU non-linearity. $y = x^2$ function approximation example with different number of hidden units.
  • Figure 5: Training dataflow of single MaxK based GNN layer. In the backward computation, the transposed CSC format is equal to original CSR format.
  • ...and 5 more figures

Theorems & Definitions (2)

  • proposition 1
  • theorem 1