AutoSAGE：面向稀疏 GNN 汇聚 (SpMM/SDDMM) 与 CSR 注意力的输入感知 CUDA 调度器
机器学习
2025-11-25 v1 性能
摘要
稀疏 GNN 汇聚 (CSR SpMM/SDDMM) 在度不均、特征宽度以及 GPU 微体系结构不同时性能差异显著。我们提出 AutoSAGE，一种输入感知 CUDA 调度器，通过轻量估计器配合设备内微探针进行细化，设定安全回退机制以适配厂商内核，并配备持久化缓存以实现确定性重放。AutoSAGE 涵盖 SpMM 和 SDDMM，并可组合成 CSR 注意力流水线 (SDDMM -> 行软最大 -> SpMM)。在 Reddit 和 OGBN-Products 基准上，AutoSAGE 在带宽受限特征宽度下与厂商基准持平，并在小宽度场景实现收益；在合成稀疏性和偏斜压力测试中实现最高 4.7 倍内核级加速。我们发布 CUDA 源码、Python 绑定、可复现的测试基准以及可回放的缓存日志。
引用
@article{arxiv.2511.17594,
title = {AutoSAGE: Input-Aware CUDA Scheduling for Sparse GNN Aggregation (SpMM/SDDMM) and CSR Attention},
author = {Aleksandar Stankovic},
journal= {arXiv preprint arXiv:2511.17594},
year = {2025}
}
备注
10 pages, several figures. Code and artifacts: https://github.com/SV25-22/AutoSAGE