提出GNA稀疏注意力机制,实现闪电级加速且无需微调。
Generalized Neighborhood Attention: Multi-dimensional Sparse Attention at the Speed of Light
- 设计通用邻域注意力(GNA),统一多种局部稀疏模式。
- 在NVIDIA Blackwell上实现1.3 petaFLOPs/秒的峰值算力利用率。
- 在多个生成模型中实现28%-46%端到端加速,适合视觉大模型部署。
许多稀疏注意力机制如邻域注意力通常未能持续超越自注意力基线,主要源于注意力基础设施复杂度及AI硬件架构快速演进。当前主流基础模型,尤其在计算机视觉领域,严重依赖注意力机制,亟需可靠稀疏性以突破O(n²)复杂度瓶颈。本文研究一类聚焦局部性的有前景稀疏注意力机制,旨在建立更优性能分析模型。首先提出广义邻域注意力(GNA),可描述滑动窗口、步进滑动窗口与分块注意力。接着分析实现中的设计选择,构建模拟器以提供更真实的速度提升上限。最后,在NVIDIA Blackwell架构的CUTLASS优化融合多头注意力(FMHA)内核上实现GNA。该实现可在诸多理想分块稀疏场景下完全达到理论最大加速比,FP16下有效算力达1.3 petaFLOPs/秒。将多种GNA配置嵌入Cosmos-7B、HunyuanVideo和FLUX等现成生成模型,无需微调即可在B200上实现28%至46%的端到端加速。代码与内核将通过NATEN项目开源。
原文摘要 · Abstract (English)
Many sparse attention mechanisms such as Neighborhood Attention have typically failed to consistently deliver speedup over the self attention baseline. This is largely due to the level of complexity in attention infrastructure, and the rapid evolution of AI hardware architecture. At the same time, many state-of-the-art foundational models, particularly in computer vision, are heavily bound by attention, and need reliable sparsity to escape the O(n^2) complexity. In this paper, we study a class of promising sparse attention mechanisms that focus on locality, and aim to develop a better analytical model of their performance improvements. We first introduce Generalized Neighborhood Attention (GNA), which can describe sliding window, strided sliding window, and blocked attention. We then consider possible design choices in implementing these approaches, and create a simulator that can provide much more realistic speedup upper bounds for any given setting. Finally, we implement GNA on top of a state-of-the-art fused multi-headed attention (FMHA) kernel designed for the NVIDIA Blackwell architecture in CUTLASS. Our implementation can fully realize the maximum speedup theoretically possible in many perfectly block-sparse cases, and achieves an effective utilization of 1.3 petaFLOPs/second in FP16. In addition, we plug various GNA configurations into off-the-shelf generative models, such as Cosmos-7B, HunyuanVideo, and FLUX, and show that it can deliver 28% to 46% end-to-end speedup on B200 without any fine-tuning. We will open source our simulator and Blackwell kernels directly through the NATTEN project.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。