Fused3S融合三步稀疏计算,让显卡张量核跑得更快更省力。
Fused3S: Fast Sparse Attention on Tensor Cores
- 把三步稀疏运算合为一炉,直接利用显卡张量核加速
- 在H100和A30上比现有方法快1.6到16.3倍
- 适合做图神经网络或稀疏序列建模的开发者用
稀疏注意力是众多先进神经网络模型的核心组件,可分解为三个稀疏矩阵操作(3S):采样稠密-稠密矩阵乘法(SDDMM)、Softmax归一化和稀疏矩阵乘法(SpMM)。由于非结构化稀疏性与专为稠密运算优化的张量核心不匹配,且数据移动开销高,高效执行3S模式在现代GPU上仍具挑战。此前工作仅分别优化各操作或解决单一问题。本文提出Fused3S,首个联合优化张量核心利用率与数据移动的3S融合算法。在真实图数据集上,Fused3S在H100和A30 GPU上分别实现1.6–16.3倍和1.5–14倍加速。将其集成至图变压器推理后,端到端性能提升1.05–5.36倍,跨不同数据集与GPU架构均优于所有3S基线。
原文摘要 · Abstract (English)
Sparse attention is a core building block in many leading neural network models, from graph-structured learning to sparse sequence modeling. It can be decomposed into a sequence of three sparse matrix operations (3S): sampled dense-dense matrix multiplication (SDDMM), softmax normalization, and sparse matrix multiplication (SpMM). Efficiently executing the 3S computational pattern on modern GPUs remains challenging due to (a) the mismatch between unstructured sparsity and tensor cores optimized for dense operations, and (b) the high cost of data movement. Previous works have optimized these sparse operations individually or addressed one of these challenges. This paper introduces Fused3S, the first fused 3S algorithm that jointly maximizes tensor core utilization and minimizes data movement. Across real-world graph datasets, Fused3S achieves $1.6- 16.3\times$ and $1.5-14\times$ speedup over state-of-the-art on H100 and A30 GPUs. Furthermore, integrating Fused3S into Graph Transformer inference accelerates end-to-end performance by $1.05-5.36\times$, consistently outperforming all 3S baselines across diverse datasets (single and batched graphs) and GPU architectures.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。