RA-SpMM：面向GPU上GNN工作负载的机制感知稀疏矩阵乘法

《Future Generation Computer Systems》：RA-SpMM: Regime-aware sparse matrix multiplication for GNN workloads on GPUs

　　

## 摘要

稀疏矩阵-稠密矩阵乘法（SpMM）是GPU上图神经网络（GNN）训练和推理的主要瓶颈，然而没有任何单一SpMM内核能在所有图结构上都取得最优性能。RA-SpMM以O(M)复杂度计算三个CSR衍生描述符：矩阵规模M、平均行度d?和度变异系数CVd。随后将其与特征宽度N结合，通过七条可解释的规则和一个cuSPARSE兜底机制，在五个专门的内核家族中进行选择。在51个图产生的192种可行的（图，N）配置上，路由器实现了相对于FP32 cuSPARSE的1.702倍几何平均加速，达到所测预言值的99.76%；每个配置上路由器与预言值比率的最小值为0.868倍。对于首次调用，原生CSR模式避免了格式构建开销，在分量级计量下，在N=128时相对于冷启动cuSPARSE达到2.269倍加速，平均描述符计算与路由开销仅为0.17毫秒。在50个被两种评估的DTC-SpMM模式都正确完成的配置上，RA-SpMM分别快1.598倍和1.773倍，同时避免了DTC平均38.5秒的变体扫描成本。微架构分析表明，在这类内存受限的工作负载上，减少内存访问流量而非峰值Tensor Core吞吐量是加速的主要来源。在不重新调优的情况下，同一路由器在RTX 4090上保留了Ada专属预言值的0.942倍性能。实现代码和测量CSV数据发布在 https://github.com/tariqaf/RA-SpMM ，处理后的51图测试套件归档在Zenodo上。

## 引言

稀疏矩阵-稠密矩阵乘法（SpMM，C=A·B，其中A为稀疏矩阵）是GNN推理和训练中的重要计算瓶颈[1],[2]。NVIDIA的cuSPARSE提供了强大的通用基线，但固定的厂商路径无法针对每种图结构进行特化。DTC-SpMM[3]和RoDe[4]等高吞吐方法通过专门格式、重排序或变体选择来利用特定结构。因此它们的收益随矩阵而异，其图特定的设置通常需要跨多次调用进行摊销。

这种差异促使我们进行 regime-aware（区域感知）的选择：使用廉价的CSR描述符对输入进行表征，然后将其分发到针对该结构区域专门化的内核。

一个可解释的路由器能否在异构GNN图上有效地在专门的SpMM内核之间进行选择，而无需计时候选内核或执行逐矩阵自动调优？

RA-SpMM通过以下四项贡献回答了这个问题：

1. **结构区域。** 我们使用矩阵规模M、平均行度d?和度变异系数CVd将图分为六个类别：稀疏均匀、稀疏倾斜、稠密小、稠密大规模、社区/局部性和混合/不规则。

2. **五内核组合。** RA-SpMM将两个原生CSR家族（CSR_DIRECT和ZERO_OVERHEAD_CSR）与三个分块Tensor Core家族（TC_DIRECT、COMMUNITY_TC和SEGMENT_HYBRID）相结合。分块家族共享一个带路由器选择的FP16和TF32路径的分块MMA引擎。cuSPARSE兜底机制处理没有任何专家内核获益的输入。没有任何家族执行逐矩阵变体选择或自动调优。

3. **可解释路由。** 一个七规则决策树将三个CSR衍生描述符与调用时的特征宽度N结合，无需计时候选内核即可选择路径。在校准套件上，它达到所测预言值几何均值的99.76%。我们在留出协议下单独评估学习型替代方案。

4. **复用感知执行。** 原生CSR模式为首次调用提供无需格式构建的服务。对于重复调用，稳态路由器可以在一次性构建成本摊销后选择更高吞吐的分块路径。我们同时测量了构建成本和复用盈亏平衡点。

我们评估26个真实图，特征宽度最多四个N∈{64,128,256,512}；受内存限制，留下92个可行的真实图配置。25个合成图额外增加100个压力测试配置，总计192个。

图1展示了端到端架构。在此套件上，路由器实现相对于FP32 cuSPARSE的1.702倍预热几何平均加速，并且在每个配置上都保持在所测预言值的0.868倍以上。评估还比较了RA-SpMM与DTC-SpMM和其他近期系统的性能，测量了集成GNN训练步骤，并在不重新调优的情况下在RTX 4090上应用固定路由器（第7.7节）。

## SpMM公式化

给定以压缩稀疏行（CSR）格式存储的稀疏矩阵A∈?^M×K和稠密矩阵B∈?^K×N，SpMM计算C=A·B∈?^M×N。在GNN工作负载中，A是邻接矩阵或消息传递矩阵（每个图固定），B包含节点特征（每层变化），N是特征维度（通常为64–512）。

## GPU内存层次与Tensor Core

相关的GPU存储层次包括全局DRAM（RTX 3090上约900 GB/s）、L2缓存（约2 TB/s有效带宽）以及片上共享内存和寄存器（约19 TB/s）。在RTX 3090上，这些层次

## 负载均衡分解

RoDe[4]将每行分离为稠密块前缀和稀疏残差，然后通过块拆分和子块流水线平衡工作。Sputnik[5]通过1D分块、子warp执行、反向偏移对齐和行洗牌在更细粒度上解决负载不均衡和内存吞吐问题。基于合并的稀疏计算[6]则沿着CSR流上的虚拟合并路径划分工作。

## 带逐矩阵预处理的Tensor Core SpMM

Tensor Core SpMM方法通常通过构建

## 结构分类体系

我们使用三个CSR衍生描述符组织测试套件：矩阵规模M、平均行度d?=nnz/M，以及度变异系数CVd（行度的变异系数）。表1列出了26个真实图。每个图最多评估四个特征宽度，省略内存不可行的情况；真实和合成套件分别贡献92和100个配置。附录S1列出了被省略的图-宽度对。

1. **稀疏均匀**（d?<20，CVd<1.0，M>100K）：道路网络、Web图、

## 内核组合

RA-SpMM包含五个覆盖第4节结构区域的内核家族。六个类别组织评估但与这些家族不是一一对应。两个家族直接在CSR上执行；三个使用带FP16和TF32路径的共享分块Tensor Core引擎。路由器还可以退回到cuSPARSE。本节描述每个家族的实现；第7节报告性能。

## 路由规则

路由器使用M、平均行度d?、度变异系数CVd和调用时宽度N。前三个值通过对CSR行指针数组的一次O(M)遍历计算。七条规则从上到下评估，第一条匹配的规则确定路径；未匹配的输入使用TC_DIRECT。阈值拟合于完整的192配置校准套件。第7.6节在单独的留出协议下报告学习型替代方案。规则如下：

- **极端倾斜稀疏尾部**

## 硬件与套件

主要平台是NVIDIA GeForce RTX 3090（Ampere，SM 86）配CUDA 11.8；第7.7节在RTX 4090（Ada，SM 89）上验证。套件由表1的26个真实图和25个合成图组成：92个真实加100个合成可行（图，N）配置，共192个。

## 匹配生命周期计时

我们报告预热和冷启动生命周期区域。对于预热测量，所有可复用计划、描述符、缓冲区和暂存区在计时前构建；然后运行50次预热和200次计时的CUDA事件迭代。

## 何时RA-SpMM有帮助？

RA-SpMM暴露两个针对不同复用模式的运行点。原生CSR首次调用模式针对推理、动态图、流式图等低复用场景，在这些场景中格式构建无法摊销。预热稳态模式针对固定图上的重复训练或推理，可以在构建成本回收后使用分块家族。

盈亏平衡点是配置相关的。在169个分块路径是预热预言值的情况中，中位

## 结论

RA-SpMM将两个原生CSR家族和三个分块Tensor Core家族置于一个可解释的七规则路由器之下。在192个配置上，它实现相对于FP32 cuSPARSE的1.702倍预热几何平均加速，达到所测预言值的99.76%，而原生CSR首次调用模式在N=128时实现2.269倍加速且无需分块构建。集成GNN研究和外部比较表明该组合在孤立内核基准之外同样有用，性能剖析将分块增益归因于

## CRediT作者贡献声明

Tariq Habib Afridi：写作—原始草稿、可视化、验证、软件、方法论、调查、数据整理、概念化。Young-Koo Lee：写作—审阅与编辑、监督、资金获取、概念化。

## 生成式AI和AI辅助技术在稿件准备过程中的声明

在本工作准备过程中，作者使用了Claude（Anthropic）和ChatGPT（OpenAI）来辅助稿件组织并提高写作的清晰度和一致性。使用这些工具后，作者根据需要审阅和编辑了内容，并对文章的内容承担全部责任。

## 资助

本工作由韩国国家研究基金会（NRF）资助，由韩国科学技术信息通信部（MSIT）提供资金，资助编号RS-2026-25489654。

## 利益冲突声明

作者声明他们没有已知的可能影响本文报告工作的竞争财务利益或个人关系。

---

Tariq Habib Afridi | Young-Koo Lee

庆熙大学计算机科学与工程系，全球校区，永仁市，大韩民国