TileSpMM: A Variable-Size Tiled Algorithm for Sparse Matrix-Matrix Multiplication on Tensor Cores
Sparse Matrix-Matrix Multiplication (SpMM) is a fundamental primitive in scientific computing and artificial intelligence applications. Modern hardware, notably Tensor Core Units (TCUs), offers immense computational power, creating promising opportunities for SpMM acceleration. However, it is difficult to map sparse matrices with irregular structures onto TCUs due to their requirement for regular operands. Existing fixed-granularity tiling methods frequently face a trade-off between massive zero-padding in sparse regions and poor spatial data locality in dense regions. To bridge this gap, we propose TileSpMM, which breaks the static-granularity bottleneck through a variable-size tiling algorithm that dynamically adapts to local sparsity patterns. Furthermore, TileSpMM is equipped with an adaptive load-balancing strategy and customized granularity-specific kernels to improve hardware utilization and mitigate computation redundancy. Experiments on NVIDIA H100 and RTX 5090 GPUs with a diverse range of benchmark matrices show that TileSpMM delivers overall better performance than existing SpMM methods across the evaluated platforms and datasets. Compared with cuSPARSE, SSpMM, Acc-SpMM and FlashSparse, TileSpMM achieves geometric mean speedups of 4.76 × , 2.74 × , 2.38 × and 1.58 × , respectively.