AFH-SpMM: Auto-Fit Heterogeneous Block Sparse-Dense Matrix Multiplication on Tensor Core GPUs

Abstract

Sparse-dense matrix multiplication (SpMM), a fundamental computational kernel in graph analytics and scientific computing, can be substantially accelerated on modern GPUs by leveraging both dense and sparse Tensor Core units, thereby enabling high-throughput computation. However, existing methods often fail to exploit these hardware units effectively when applied to real-world sparse matrices that exhibit strong local structural heterogeneity. In particular, methods that either (i) reorganize all regions into dense-like tiles or (ii) aggressively convert them into strict 2:4 structured sparsity blocks typically incur low effective block density, substantial padding overhead, and nontrivial preprocessing costs. To address these challenges, we propose AFH-SpMM, a novel Auto-Fit Heterogeneous SpMM framework designed for adaptively parallelizing sparse-dense matrix multiplication on Tensor Core-equipped GPUs. Using a 16-row window as the basic processing granularity, AFH-SpMM adaptively maps local regions to two hardware-efficient computation paths: 16 × 16 dense tiles targeted to Dense Tensor Cores and 16 × 8 row-wise 2:4 structured-sparse tiles targeted to Sparse Tensor Cores. For the sparse computation path, AFH-SpMM further exploits local column proximity to mitigate subsequent memory-access and address-generation overheads. At runtime, the two block types are executed within a single fused kernel, while preserving distinct operand layouts and specialized MMA pipelines for each path. Experiments on 600 SuiteSparse matrices across NVIDIA RTX PRO 6000, H100, and A800 show that AFH-SpMM achieves average speedups of 1.33 × (up to 5.72 ×), and often leads cuSPARSE, ASpT, Sputnik, RoDe, Acc-SpMM, and MP-SpMM, with especially strong gains on medium and large matrices.