DistSpMM: Accelerating Sparse Matrix Dense Matrix Multiplication on GPUs
TL;DR
The proposed DistSpMM proposes DistSpMM, a co-design framework integrating data layout, pipelining, and communication strategies, which introduces HSDMA, a lightweight algorithm that reduces communication by optimizing the dense matrix allocation.
Abstract
Sparse matrix-dense matrix multiplication (SpMM) is a core operation in scientific computing and deep learning. On multi-GPU platforms, its scalability is limited by communication bottlenecks. To address this, we propose DistSpMM, a co-design framework integrating data layout, pipelining, and communication strategies. DistSpMM introduces HSDMA, a lightweight algorithm that reduces communication by optimizing the dense matrix allocation. DistSpMM features a topology-aware two-stage pipeline that manages the IB/NVLink bandwidth disparity to maximize the overlap of computation and communication. Finally, DistSpMM employs an adaptive selector that uses a performance model to dynamically choose the optimal communication granularity (coarse vs. fine-grained) based on data sparsity and network tier. Experiments on diverse real-world datasets demonstrate the superior performance of our method. It achieves average speedups of 1.6 × to 2.6 × in single-node multi-GPU environments and 4.0 × to 5.1 × in multi-node multi-GPU environments over the baseline. Compared to the best state-of-the-art implementations, our method delivers up to 2.0 × speedup.