通过优化稀疏矩阵乘法减少计算冗余，显著提升显卡性能。
FlashSparse: Minimizing Computation Redundancy for Fast Sparse Matrix Multiplications on Tensor Cores
- 采用交换-转置策略最小化稀疏粒度，降低计算冗余。
- 在H100和RTX 4090上实现平均5.5倍加速，优于现有方法。
- 适合需要高效稀疏计算的科学计算与深度学习场景。
稀疏矩阵-矩阵乘法（SpMM）和采样稠密-稠密矩阵乘法（SDDMM）在科学计算与深度学习中至关重要。张量核心单元（TCUs）虽具备强大算力，但因非结构化稀疏数据的不规则性，难以在TCUs上获得实际加速。为此，我们提出FlashSparse，一种新型方法，通过创新的交换-转置矩阵乘法策略，最小化SpMM与SDDMM在TCUs上的稀疏粒度。该策略显著减少计算冗余，同时充分释放TCU算力。此外，FlashSparse配备内存高效的线程映射策略以实现合并访问，并采用节省内存的稀疏矩阵存储格式。在H100与RTX 4090 GPU上的大量实验表明，FlashSparse在稀疏矩阵乘法上达到新基准，几何平均速度比DTC-SpMM快5.5倍，比RoDe快3.22倍。
原文摘要 · Abstract (English)
Sparse Matrix-matrix Multiplication (SpMM) and Sampled Dense-dense Matrix Multiplication (SDDMM) are important sparse operators in scientific computing and deep learning. Tensor Core Units (TCUs) enhance modern accelerators with superior computing power, which is promising to boost the performance of matrix operators to a higher level. However, due to the irregularity of unstructured sparse data, it is difficult to deliver practical speedups on TCUs. To this end, we propose FlashSparse, a novel approach to bridge the gap between sparse workloads and the TCU architecture. Specifically, FlashSparse minimizes the sparse granularity for SpMM and SDDMM on TCUs through a novel swap-and-transpose matrix multiplication strategy. Benefiting from the minimum sparse granularity, the computation redundancy is remarkably reduced while the computing power of TCUs is fully utilized. Besides, FlashSparse is equipped with a memory-efficient thread mapping strategy for coalesced data access and a sparse matrix storage format to save memory footprint. Extensive experimental results on H100 and RTX 4090 GPUs show that FlashSparse sets a new state-of-the-art for sparse matrix multiplications (geometric mean 5.5x speedup over DTC-SpMM and 3.22x speedup over RoDe).
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。