# MaxK-GNN: Extremely Fast GPU Kernel Design for Accelerating Graph Neural Networks Training

arXiv: https://arxiv.org/abs/2312.08656 · HTML: https://mlsanity.com/paper/2312.08656/

## Key facts

- Cited by 25 papers in the map (data through 14 Oct 2026).
- Builds on 46 of its 76 references that are in the map.
- First posted on arXiv 14 Dec 2023.

## Builds on (most cited first, 30 of 46)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Graph Attention Networks](https://mlsanity.com/paper/1710.10903/) (2017)
- [Inductive Representation Learning on Large Graphs](https://mlsanity.com/paper/1706.02216/) (2017)
- [How Powerful are Graph Neural Networks?](https://mlsanity.com/paper/1810.00826/) (2018)
- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://mlsanity.com/paper/1803.03635/) (2018)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Graph Convolutional Neural Networks for Web-Scale Recommender Systems](https://mlsanity.com/paper/1806.01973/) (2018)
- [GraphSAINT: Graph Sampling Based Inductive Learning Method](https://mlsanity.com/paper/1907.04931/) (2019)
- [Cluster-GCN: An Efficient Algorithm for Training Deep and Large Graph Convolutional Networks](https://mlsanity.com/paper/1905.07953/) (2019)
- [Rigging the Lottery: Making All Tickets Winners](https://mlsanity.com/paper/1911.11134/) (2019)
- [Graph Neural Networks in Recommender Systems: A Survey](https://mlsanity.com/paper/2011.02260/) (2020)
- [Scalable Training of Artificial Neural Networks with Adaptive Sparse Connectivity inspired by Network Science](https://mlsanity.com/paper/1707.04780/) (2017)
- [Graph Neural Network for Traffic Forecasting: A Survey](https://mlsanity.com/paper/2101.11174/) (2021)
- [Sparse Networks from Scratch: Faster Training without Losing Performance](https://mlsanity.com/paper/1907.04840/) (2019)
- [Maxout Networks](https://mlsanity.com/paper/1302.4389/) (2013)
- [A Systematic DNN Weight Pruning Framework using Alternating Direction Method of Multipliers](https://mlsanity.com/paper/1804.03294/) (2018)
- [A Unified Lottery Ticket Hypothesis for Graph Neural Networks](https://mlsanity.com/paper/2102.06790/) (2021)
- [Computing Graph Neural Networks: A Survey from Algorithms to Accelerators](https://mlsanity.com/paper/2010.00130/) (2020)
- [Molecular graph generation with Graph Neural Networks](https://mlsanity.com/paper/2012.07397/) (2020)
- [HyGCN: A GCN Accelerator with Hybrid Architecture](https://mlsanity.com/paper/2001.02514/) (2020)
- [EnGN: A High-Throughput and Energy-Efficient Accelerator for Large Graph Neural Networks](https://mlsanity.com/paper/1909.00155/) (2019)
- [BNS-GCN: Efficient Full-Graph Training of Graph Convolutional Networks with Partition-Parallelism and Random Boundary Node Sampling](https://mlsanity.com/paper/2203.10983/) (2022)
- [Dissecting the NVIDIA Volta GPU Architecture via Microbenchmarking](https://mlsanity.com/paper/1804.06826/) (2018)
- [AutoReP: Automatic ReLU Replacement for Fast Private Network Inference](https://mlsanity.com/paper/2308.10134/) (2023)
- [LinGCN: Structural Linearized Graph Convolutional Network for Homomorphically Encrypted Inference](https://mlsanity.com/paper/2309.14331/) (2023)
- [ReConTab: Regularized Contrastive Representation Learning for Tabular Data](https://mlsanity.com/paper/2310.18541/) (2023)
- [A Length Adaptive Algorithm-Hardware Co-design of Transformer on FPGA Through Sparse Attention and Dynamic Pipelining](https://mlsanity.com/paper/2208.03646/) (2022)
- [Accel-GCN: High-Performance GPU Accelerator Design for Graph Convolution Networks](https://mlsanity.com/paper/2308.11825/) (2023)

## Built on by (most cited first, 25 of 25)

- [SwitchTab: Switched Autoencoders Are Effective Tabular Learners](https://mlsanity.com/paper/2401.02013/) (2024)
- [Time Series Modeling for Heart Rate Prediction: From ARIMA to Transformers](https://mlsanity.com/paper/2406.12199/) (2024)
- [Regional Style and Color Transfer](https://mlsanity.com/paper/2404.13880/) (2024)
- [Enhance Image-to-Image Generation with LLaVA-generated Prompts](https://mlsanity.com/paper/2406.01956/) (2024)
- [Optimizing Automated Picking Systems in Warehouse Robots Using Machine Learning](https://mlsanity.com/paper/2408.16633/) (2024)
- [Identification of Prognostic Biomarkers for Stage III Non-Small Cell Lung Carcinoma in Female Nonsmokers Using Machine Learning](https://mlsanity.com/paper/2408.16068/) (2024)
- [LoR-VP: Low-Rank Visual Prompting for Efficient Vision Model Adaptation](https://mlsanity.com/paper/2502.00896/) (2025)
- [Enhancing Visual Question Answering through Ranking-Based Hybrid Training and Multimodal Fusion](https://mlsanity.com/paper/2408.07303/) (2024)
- [Research on reinforcement learning based warehouse robot navigation algorithm in complex warehouse layout](https://mlsanity.com/paper/2411.06128/) (2024)
- [Autonomous Navigation of Unmanned Vehicle Through Deep Reinforcement Learning](https://mlsanity.com/paper/2407.18962/) (2024)
- [Cross-border Commodity Pricing Strategy Optimization via Mixed Neural Network for Time Series Analysis](https://mlsanity.com/paper/2408.12115/) (2024)
- [Deep Analysis of Time Series Data for Smart Grid Startup Strategies: A Transformer-LSTM-PSO Model Approach](https://mlsanity.com/paper/2408.12129/) (2024)
- [Advanced AI Framework for Enhanced Detection and Assessment of Abdominal Trauma: Integrating 3D Segmentation with 2D CNN and RNN Models](https://mlsanity.com/paper/2407.16165/) (2024)
- [Advanced Payment Security System:XGBoost, LightGBM and SMOTE Integrated](https://mlsanity.com/paper/2406.04658/) (2024)
- [Recording Brain Activity While Listening to Music Using Wearable EEG Devices Combined with Bidirectional Long Short-Term Memory Networks](https://mlsanity.com/paper/2408.12124/) (2024)
- [GTA-Net: An IoT-Integrated 3D Human Pose Estimation System for Real-Time Adolescent Sports Posture Correction](https://mlsanity.com/paper/2411.06725/) (2024)
- [Intelligent logistics management robot path planning algorithm integrating transformer and GCN network](https://mlsanity.com/paper/2501.02749/) (2025)
- [FlashSparse: Minimizing Computation Redundancy for Fast Sparse Matrix Multiplications on Tensor Cores](https://mlsanity.com/paper/2412.11007/) (2024)
- [Data-Algorithm-Architecture Co-Optimization for Fair Neural Networks on Skin Lesion Dataset](https://mlsanity.com/paper/2407.13896/) (2024)
- [Key Information Retrieval to Classify the Unstructured Data Content of Preferential Trade Agreements](https://mlsanity.com/paper/2401.12520/) (2024)
- [A Self-guided Multimodal Approach to Enhancing Graph Representation Learning for Alzheimer's Diseases](https://mlsanity.com/paper/2412.06212/) (2024)
- [International Trade Flow Prediction with Bilateral Trade Provisions](https://mlsanity.com/paper/2407.13698/) (2024)
- [Cardiovascular Disease Detection By Leveraging Semi-Supervised Learning](https://mlsanity.com/paper/2412.10567/) (2024)
- [On Efficient Scaling of GNNs via IO-Aware Layers Implementations](https://mlsanity.com/paper/2605.31500/) (2026)
- [SparseDitto: An Agentic Sparse Compilation Framework through Architecture-Aware Synthesis on GPUs](https://mlsanity.com/paper/2608.05033/) (2026)

## Abstract (arXiv)

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.