Libra: Unleashing GPU Heterogeneity for High-Performance Sparse Matrix Multiplication

TL;DR

This work addresses accelerating sparse operators (SpMM and SDDMM) by exploiting GPU heterogeneity through a novel 2D-aware workload distribution. Libra analyzes the complementary strengths of Tensor Core Units and CUDA cores, enabling precise task mapping, hybrid load balancing, and occupancy-aware scheduling, with GPU-accelerated preprocessing and optimized kernels. The approach yields state-of-the-art end-to-end performance on Google-scale GNN tasks, achieving up to several-fold speedups over strong baselines while maintaining low preprocessing overhead. The results demonstrate a practical pathway to fully leverage heterogeneous GPU resources for memory-bound sparse computations in both ML and scientific computing contexts.

Abstract

Sparse matrix multiplication operators (i.e., SpMM and SDDMM) are widely used in deep learning and scientific computing. Modern accelerators are commonly equipped with Tensor Core Units (TCUs) and CUDA cores to accelerate sparse operators. The former excels at structured matrix computations, whereas the latter offers greater programming flexibility. However, how to combine these two resources to maximize sparse-operator performance remains unclear. In this work, we first identify the source of performance gains in hybrid computation and systematically analyze their complementary strengths. Motivated by this, we propose Libra, a holistic framework that efficiently leverages heterogeneous computing resources to accelerate both SpMM and SDDMM operators. Specifically, Libra introduces a 2D-aware (locality and utilization) workload distribution method to precisely identify the optimal task mapping, simultaneously leveraging the data reuse capabilities of TCUs and the flexibility of CUDA cores to minimize computational redundancy. Libra further incorporates hybrid load balancing, occupancy-aware task scheduling, and efficient kernel implementations to maximize execution efficiency. Extensive experiments on H100 and RTX 4090 GPUs demonstrate that Libra surpasses all the 12 up-to-date baselines significantly, e.g., on average 1.77x speedup over FlashSparse, 1.73x over RoDe, and 2.9x over DGL for end-to-end GNN applications. Libra opens up a new perspective for sparse operator acceleration by fully unleashing the power of heterogeneous GPU resources.

Figures (15)

• Figure 1: The ratio of NNZ-1 vectors (containing exactly one nonzero element) in all nonzero vectors of a sparse matrix (total 500 matrices from SuiteSparse). The subplot presents a case study showing how Libra’s hybrid SpMM performance varies with the TCU ratio on the matrix pkustk01.
• Figure 2: Sparse matrix partition with SGT.
• Figure 3: View of MMA operation for SpMM and SDDMM.
• Figure 4: Performance degradation of concurrent execution when increasing CUDA-core workloads. The TCU activity ratio is measured under the Hybrid-concurrent mode.
• Figure 5: Overview of Libra.