广义邻域注意力：多维稀疏注意力以光速般的速度
计算机视觉与模式识别
2025-04-24 v1 人工智能
机器学习
摘要
许多稀疏注意力机制，如邻域注意力，通常难以持续地实现对自注意力基准的加速。这在很大程度上是由于注意力基础设施的复杂度以及 AI 硬件架构的快速演化所致。与此同时，许多领先的基础模型，特别是在计算机视觉中，都严重受限于注意力机制，需要可靠的稀疏性才能摆脱 O(n^2) 的复杂度。本文研究了一类具有前景的稀疏注意力机制，这些机制聚焦于局部性，旨在开发更好的分析模型来衡量其性能提升。我们首先引入广义邻域注意力 (Generalized Neighborhood Attention, GNA)，其可以描述滑动窗口、抽屉滑动窗口和阻塞注意力。随后，我们考虑在实现这些方法时可能的设计选择，并创建了一个仿真器，可为给定设置提供 much more 真实的加速上限。最后，我们在为 NVIDIA Blackwell 架构设计的领先的融合多头注意力 (FMHA) 内核上实现 GNA。我们的实现在许多理想的完全阻塞稀疏情况下能够完全实现理论上可能的最大加速速度，在 FP16 下实现了 1.3 PetaFLOPs/秒 的有效利用率。此外，我们将 various GNA 配置插入面向生成模型中，如 Cosmos-7B、HunyuanVideo 和 FLUX，展示在不进行微调的情况下可在 B200 上实现 28%至46%的端到端加速。我们将通过 NATTEN 项目直接开源我们的仿真器和 Blackwell 内核。
引用
@article{arxiv.2504.16922,
title = {Generalized Neighborhood Attention: Multi-dimensional Sparse Attention at the Speed of Light},
author = {Ali Hassani and Fengzhe Zhou and Aditya Kane and Jiannan Huang and Chieh-Yun Chen and Min Shi and Steven Walton and Markus Hoehnerbach and Vijay Thakkar and Michael Isaev and Qinsheng Zhang and Bing Xu and Haicheng Wu and Wen-mei Hwu and Ming-Yu Liu and Humphrey Shi},
journal= {arXiv preprint arXiv:2504.16922},
year = {2025}
}
备注
https://github.com/SHI-Labs/NATTEN/