一句话：提出 MaxK-GNN 系统，通过算法与系统协同设计大幅加速 GPU 上的 GNN 训练。
GPU 是深度神经网络训练的主流平台，但 GNN 训练面临负载不均衡与不规则内存访问，SpMM 内存流量尤其显著，PyG、DGL/cuSPARSE、GNNAdvisor 等仅部分解决。本文提出 MaxK-GNN，强调算法与系统垂直协同优化：引入 MaxK 非线性并理论分析其通用逼近能力，设计 CBSR 格式存储非线性后的特征矩阵；前向使用基于 CBSR 的合并增强 SpGEMM 核，并在共享内存中放置稀疏输出累加缓冲；反向使用外积式 SSpMM 核。评估显示其在 Reddit 上相较 DGL 与 GNNAdvisor 达 3.22×/4.24× 加速，精度与 SOTA GNN 相当，接近 Amdahl 定律加速上限。
In the acceleration of deep neural network training, the graphics processing unit (GPU) has become the mainstream platform. GPUs face substantial challenges on Graph Neural Networks (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. However, the memory traffic involved with Sparse-Dense Matrix Matrix Multiplication (SpMM) 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 Sparse Matrix-Matrix Multiplication (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 Sampled Sparse Matrix Dense Matrix Multiplication (SSpMM) Kernel. We conduct extensive evaluations of MaxK-GNN and report the system training time. Experiments show that MaxK-GNN system could approach the speedup limit according to Amdahl's law. We achieve comparable accuracy to SOTA GNNs, but at a significantly increased speed: 3.22×/4.24× speedup (vs. 5.52×/7.27×) on Reddit compared to DGL and GNNAdvisor implementations. Our implementation can be found on GitHub1.
DOI 原文 · arXiv · PDF（开放获取） · DBLP
@inproceedings{PengXSHZHKKD24,
title = {MaxK-GNN: Extremely Fast GPU Kernel Design for Accelerating Graph Neural Networks Training},
author = {Hongwu Peng and Xi Xie and Kaustubh Shivdikar and Md Amit Hasan and Jiahui Zhao and Shaoyi Huang and Omer Khan and David Kaeli and Caiwen Ding},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3620665.3640426}
}