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.

翻译：摘要：在深度神经网络训练的加速过程中，GPU已成为主流平台。然而，GPU在图神经网络（GNN）中面临工作负载不平衡和内存访问不规则等重大挑战，导致硬件利用率不足。现有解决方案（如PyG、结合cuSPARSE的DGL以及GNNAdvisor框架）部分解决了这些问题，但内存流量仍然显著。我们认为，只有通过算法与系统创新的垂直优化才能实现显著的性能提升，而非将加速优化视为“事后考虑”（即：(i) 给定GNN算法后设计加速器，或(ii) 给定硬件后主要优化GNN算法）。本文提出了MaxK-GNN，一个融合算法与系统创新的先进高性能GPU训练系统。(i) 我们引入了MaxK非线性激活函数，并从理论上证明了MaxK非线性作为通用逼近器的特性，同时设计了压缩平衡稀疏行（CBSR）格式，用于存储非线性激活后特征矩阵的数据和索引；(ii) 我们设计了一种基于CBSR的逐行乘积SpGEMM内核的合并增强前向计算，用于输入特征矩阵的获取，并在共享内存中策略性地放置稀疏输出累积缓冲区；(iii) 我们开发了基于外积和SSpMM内核的优化反向计算。我们对MaxK-GNN进行了广泛评估，并报告了端到端系统运行时间。实验表明，根据Amdahl定律，MaxK-GNN系统可接近理论加速极限。我们实现了与最先进GNN相当的精度，但速度显著提升：在Reddit数据集上，与DGL和GNNAdvisor实现相比，加速比达到3.22/4.24倍（理论极限为5.52/7.27倍）。