基于张量核心稀疏内核加速深度神经网络

《IEEE Transactions on Parallel and Distributed Systems》：Based on Tensor Core Sparse Kernels Accelerating Deep Neural Networks

针对大型语言模型参数多导致的存储和计算挑战，提出基于Tensor Core的TC-CSR稀疏矩阵格式，设计块稀疏矩阵乘法（SpMM）和采样密集矩阵乘法（SDDMM）内核，在中等稀疏场景实现141%加速，大规模高稀疏矩阵达138%加速，并在稀疏Transformer推理中实现175%加速并节省内存。

　　

摘要：

深度学习中的大型语言模型拥有大量参数，需要大量的存储空间和计算资源。压缩技术在应对这些挑战方面非常有效。随着图形处理单元（GPU）等硬件的发展，Tensor Core能够加速低精度矩阵乘法运算，但加速稀疏矩阵运算仍然具有挑战性。由于稀疏性，Tensor Core的利用率相对较低。为了解决这个问题，我们提出了一种基于Tensor Core的压缩稀疏行格式（TC-CSR），该格式有助于数据在GPU上的加载以及Tensor Core上的矩阵运算。基于这种格式，我们设计了块状稀疏矩阵乘法（SpMM）和采样密集矩阵乘法（SDDMM）内核，这两种运算在深度学习中非常常见。利用这些设计，我们在中等稀疏度场景下实现了1.41×的加速效果；而在大规模高稀疏矩阵场景下，实现了1.38×的加速效果。得益于我们的设计，在使用稀疏Transformer进行端到端推理时，我们实现了1.75×的加速，并节省了内存。