# AutoSAGE: Input-Aware CUDA Scheduling for Sparse GNN Aggregation (SpMM/SDDMM) and CSR Attention

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 12 of its 24 references that are in the map.
- First posted on arXiv 17 Nov 2025.

## Builds on (most cited first, 12 of 12)

- [Fast Graph Representation Learning with PyTorch Geometric](https://mlsanity.com/paper/1903.02428/) (2019)
- [Open Graph Benchmark: Datasets for Machine Learning on Graphs](https://mlsanity.com/paper/2005.00687/) (2020)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [Design Space for Graph Neural Networks](https://mlsanity.com/paper/2011.08843/) (2020)
- [Sparse GPU Kernels for Deep Learning](https://mlsanity.com/paper/2006.10901/) (2020)
- [GNNAdvisor: An Adaptive and Efficient Runtime System for GNN Acceleration on GPUs](https://mlsanity.com/paper/2006.06608/) (2020)
- [BGL: GPU-Efficient GNN Training by Optimizing Graph Data I/O and Preprocessing](https://mlsanity.com/paper/2112.08541/) (2021)
- [GE-SpMM: General-purpose Sparse Matrix-Matrix Multiplication on GPUs for Graph Neural Networks](https://mlsanity.com/paper/2007.03179/) (2020)
- [FusedMM: A Unified SDDMM-SpMM Kernel for Graph Embedding and Graph Neural Networks](https://mlsanity.com/paper/2011.06391/) (2020)
- [TC-GNN: Bridging Sparse GNN Computation and Dense Tensor Cores on GPUs](https://mlsanity.com/paper/2112.02052/) (2021)
- [GraphBLAST: A High-Performance Linear Algebra-based Graph Framework on the GPU](https://mlsanity.com/paper/1908.01407/) (2019)
- [AdaptGear: Accelerating GNN Training via Adaptive Subgraph-Level Kernels on GPUs](https://mlsanity.com/paper/2305.17408/) (2023)

## Abstract (arXiv)

Sparse GNN aggregations (CSR SpMM/SDDMM) vary widely in performance with degree skew, feature width, and GPU micro-architecture. We present AutoSAGE, an input-aware CUDA scheduler that chooses tiling and mapping per input using a lightweight estimate refined by on-device micro-probes, with a guardrail that safely falls back to vendor kernels and a persistent cache for deterministic replay. AutoSAGE covers SpMM and SDDMM and composes into a CSR attention pipeline (SDDMM -> row-softmax -> SpMM). On Reddit and OGBN-Products, it matches vendor baselines at bandwidth-bound feature widths and finds gains at small widths; on synthetic sparsity and skew stress tests it achieves up to 4.7x kernel-level speedups. We release CUDA sources, Python bindings, a reproducible harness, and replayable cache logs.