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.


翻译:稀疏图神经网络聚合运算(CSR格式的SpMM/SDDMM)的性能受节点度分布偏斜、特征维度宽度及GPU微架构差异影响显著。本文提出AutoSAGE,一种基于输入感知的CUDA调度器,通过轻量级预估模型结合设备端微探针优化,为每个输入动态选择分块策略与内存映射方案。该系统配备安全防护机制:可自动回退至厂商优化内核,并采用持久化缓存确保确定性计算重现。AutoSAGE覆盖SpMM与SDDMM两类核心算子,并可组合构建CSR注意力计算流水线(SDDMM → 行级Softmax → SpMM)。在Reddit和OGBN-Products数据集测试中,该系统在带宽受限的特征宽度场景下与厂商基线性能持平,而在小特征宽度条件下实现性能提升;在合成稀疏性与偏斜分布压力测试中,其内核级加速比最高达4.7倍。我们同步开源了CUDA源码、Python接口绑定、可复现测试框架及可重放的缓存日志。

0
下载
关闭预览

相关内容

[IEEE TPAMI 2024] 将CLIP模型转换为端到端文本识别器
专知会员服务
22+阅读 · 2024年4月12日
【CVPR2022】EDTER:基于Transformer的边缘检测(CVPR2022)
专知会员服务
33+阅读 · 2022年3月18日
在TensorFlow中对比两大生成模型:VAE与GAN
机器之心
12+阅读 · 2017年10月23日
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
46+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
VIP会员
相关基金
国家自然科学基金
4+阅读 · 2015年12月31日
国家自然科学基金
46+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
Top
微信扫码咨询专知VIP会员