稀疏矩阵-稠密矩阵乘法在科学计算、图神经网络以及高维数据分析等现代人工智能与高性能计算领域扮演着至关重要的角色。但通信开销成为制约多GPU分布式算子可扩展性的核心瓶颈。

    中心人工智能技术与应用发展部门研究团队提出了面向多GPU平台的高效分布式稀疏计算框架DistSpMM。该框架通过“数据布局-计算流水线-通信策略”的深层协同设计大幅缓解通信瓶颈。该研究提出了轻量级的层次化稀疏感知稠密矩阵分配算法，通过优化稠密矩阵子块的物理分布避开了繁重的图划分预处理，将预处理开销降至亚毫秒级；协同设计了“聚合-交换”拓扑感知两阶段流水线，利用多计算流与层次化双缓冲区实现了跨节点通信、节点内高带宽通信交换与本地GPU计算的三重深度重叠；同时基于LogGP通信模型构建了自适应通信策略选择器，通过非零列密度临界值函数动态匹配最优通信粒度。实验表明，DistSpMM在48个真实超大规模数据集与合成图谱测试中展现出显著优势，相比作为基准的经典分布式算法SUMMA在单节点与跨节点场景下分别取得了 1.6×~2.6× 和 4.0×~5.1× 的显著平均加速效果，并大幅超越CoLa、TileSpMM等现有先进的分布式稀疏计算框架。

DistSpMM整体架构

该成果已被ACM Transactions on Architecture and Code Optimization (CCF A)录用。该研究得到国家重点研发计划和中国科学院先导专项的支持。论文第一作者为中心博士研究生顾峻瑜，中心正高级工程师王珏为通信作者。

相关成果：Junyu Gu, Jue Wang, Zhikuang Xin, Chunbao Zhou, Zhiqiang Liang, Yucheng Pang, Rongqiang Cao, Zongguo Wang, Fang Liu, Jing Wang, Yangang Wang. DistSpMM: Accelerating Sparse Matrix Dense Matrix Multiplication on GPUs. ACM Trans. Archit. Code Optim.

责任编辑：郎杨琴