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非线性函数，通过理论分析证明其具备通用逼近能力，并提出压缩平衡稀疏行（CBSR）格式，用于存储非线性变换后特征矩阵的数据与索引；（ii）设计基于行向量乘积SpGEMM内核的融合增强前向计算，借助CBSR实现输入特征矩阵获取，并在共享内存中策略性部署稀疏输出累加缓冲区；（iii）开发基于外积与SSpMM内核的优化反向计算。我们通过广泛评估MaxK-GNN系统端到端运行时间，实验表明该系统可根据Amdahl定律逼近理论加速极限。与SOTA GNN相比，MaxK-GNN在保持相当精度的同时实现了显著加速：在Reddit数据集上，相比DGL与GNNAdvisor实现分别获得3.22/4.24倍加速比（理论极限为5.52/7.27倍）。