Based on Tensor Core Sparse Kernels Accelerating Deep Neural Networks

Abstract

Large language models in deep learning have numerous parameters, requiring significant storage space and computational resources. Compression techniques are highly effective in addressing these challenges. With the development of hardware like Graphics Processing Unit (GPU), Tensor Core can accelerate low-precision matrix multiplication but achieve acceleration for sparse matrices is challenging. Due to its sparsity, the utilization of Tensor Cores is relatively low. To address this, we propose the based on Tensor Core Compressed Sparse Row format (TC-CSR), which facilitates data loading on GPUs and matrix operations on Tensor Cores. Based on this format, we designed block Sparse Matrix-Matrix Multiplication (SpMM) and Sampled Dense-Dense Matrix Multiplication (SDDMM) kernels, which are common operations in deep learning. Utilizing these designs, we achieved a $\mathbf {1.41\times }$ speedup on Sputnik in scenarios of moderate sparsity and a $\mathbf {1.38\times }$ speedup with large-scale highly sparse matrices. Benefit from our design, we achieved a $\mathbf {1.75\times }$ speedup in end-to-end inference with sparse Transformers and save memory.