arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2505.08098cs.DCcs.LG
收藏

Fused3S:Tensor Core 上的快速稀疏注意力

Fused3S: Fast Sparse Attention on Tensor Cores

University of California, Irvine(加州大学 Irvine 分校)

Zitong Li, Aparna Chandramowlishwaran

AI总结:

本文提出融合式 3S 算法 Fused3S,通过协同提升 tensor core 利用率并减少数据移动,在多种 GPU 和图数据集上显著加速稀疏注意力及 Graph Transformer 推理。

摘要:

稀疏注意力是从图结构学习到稀疏序列建模等许多领先神经网络模型中的核心构建模块。它可分解为三个稀疏矩阵操作(3S)的序列:采样稠密-稠密矩阵乘法(SDDMM)、softmax 归一化和稀疏矩阵乘法(SpMM)。在现代 GPU 上高效执行 3S 计算模式仍具挑战,原因在于:(a)非结构化稀疏性与为稠密操作优化的 tensor core 之间不匹配;(b)高昂的数据移动成本。以往工作分别优化这些稀疏操作,或只解决其中一项挑战。本文提出 Fused3S,这是首个融合式 3S 算法,可同时最大化 tensor core 利用率并最小化数据移动。在真实图数据集上,Fused3S 在 H100 和 A30 GPU 上相比最先进方法实现 1.6–16.3 倍和 1.5–14 倍加速。此外,将 Fused3S 集成到 Graph Transformer 推理中,可使端到端性能提升 1.05–5.36 倍,并在不同数据集(单图和批处理图)及 GPU 架构上持续优于所有 3S 基线。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑