AutoSAGE动态调度稀疏GNN计算，自动适配不同输入性能更优。
AutoSAGE: Input-Aware CUDA Scheduling for Sparse GNN Aggregation (SpMM/SDDMM) and CSR Attention
- 根据输入特征动态选择最优分块策略，结合设备微探测优化
- 在小特征宽度下最高提速4.7倍，大宽度下媲美厂商内核
- 支持稀疏GNN聚合与注意力计算，适合需要高效推理的场景
稀疏图神经网络聚合（CSR SpMM/SDDMM）的性能受度偏斜、特征宽度和GPU微架构影响显著。本文提出AutoSAGE，一种输入感知的CUDA调度器，通过轻量级估计结合设备端微探测精细调整分块与映射策略，并设有安全回退至厂商内核的保障机制及持久化缓存实现确定性重放。该方案覆盖SpMM与SDDMM，并可组成CSR注意力流水线（SDDMM → row-softmax → SpMM）。在Reddit与OGBN-Products数据集上，于带宽受限的特征宽度下表现匹配厂商基线，在小特征宽度下实现提升；在合成稀疏性与偏斜压力测试中，最高获得4.7倍内核级加速。我们开源了CUDA源码、Python绑定、可复现测试框架及可重放缓存日志。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。