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