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，一个高性能的 GPU 训练系统，通过集成算法和系统创新来实现。在 MaxK-GNN 系统的广泛评估中，实验结果显示其接近 Amdahl 法则的理论极限，相对于 DGL 和 GNNAdvisor 的实现，在 Reddit 上实现了 3.22/4.24 倍的速度提升（相对于理论极限的 5.52/7.27 倍），并且达到了与 SOTA GNN 相当的准确性。