Sparse Matrix Dense Matrix Multiplication plays a vital role in modern AI and high-performance computing. It is key to scientific computing, graph neural networks, and high-dimensional data analysis. However, as GPU compute power grows, cross-node network bandwidth has fallen behind. This makes communication overhead the main bottleneck for multi-GPU operator scalability. Current mainstream algorithms fail to account for hierarchical topologies, where intra-node bandwidth is high but inter-node bandwidth is low. They also suffer from heavy graph partitioning overhead and rigid asynchronous communication pipelines. Furthermore, they struggle to balance coarse-grained block transfers with fine-grained sparsity-aware transfers. These bottlenecks prevent multi-GPU clusters from reaching their full hardware potential in distributed unstructured sparse workloads.

To address these challenges, researchers from the AI Departments of CNIC developed DistSpMM. This efficient distributed sparse computing framework targets multi-GPU platforms. It mitigates communication bottlenecks through deep co-design across data layout, computing pipelines, and communication strategies. First, DistSpMM introduces a lightweight, hierarchy-aware dense matrix allocation algorithm. By optimizing the physical distribution of dense matrix sub-blocks, it eliminates heavy graph partitioning and cuts preprocessing time to sub-milliseconds. Second, DistSpMM co-designs an "aggregate-exchange" topology-aware pipeline. Using multiple CUDA streams and hierarchical double buffering, it achieves deep three-way overlap among cross-node transfers, intra-node transfers, and local GPU compute. Finally, the framework includes an adaptive communication selector based on the LogGP model. It uses a non-zero column density threshold function to match the optimal transfer granularity. Experiments across 48 real-world ultra-large datasets and synthetic graphs demonstrate clear performance gains. Compared to the benchmark algorithm SUMMA, DistSpMM achieves average speedups of 1.6×–2.6× on single nodes and 4.0×–5.1× across multiple nodes. It also substantially outperforms existing state-of-the-art frameworks like CoLa and TileSpMM.

Overall Architecture of DistSpMM

This research has been accepted by ACM Transactions on Architecture and Code Optimization (CCF A). The work was supported by the National Key R&D Program of China (2023YFB3001900) and the Strategic Priority Research Program of the Chinese Academy of Sciences (XDB0500103). The first author is Junyu Gu, a Ph.D. student at CNIC. Senior Engineer Jue Wang is the corresponding author.

Related publication

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.