摘要
arXiv:2504.16922v1 Announce Type: cross
摘要:许多稀疏注意机制,如邻域注意,通常未能在自我注意基线之上提供稳定的速度提升。这主要是由于注意基础设施的复杂程度以及AI硬件架构的快速演变。与此同时,许多最先进的基础模型,尤其是计算机视觉领域中的模型,高度依赖注意力机制,并需要可靠的稀疏性来摆脱O(n^2)的复杂性。在本文中,我们研究了一类专注于局部性的有希望的稀疏注意机制,并旨在开发出对他们性能改进的更好分析模型。我们首先介绍了通用邻域注意(GNA),它可以描述滑动窗口、跳跃滑动窗口以及块状注意力。然后,我们考虑了实施这些方法时可能的设计选择,并创建了一个模拟器,可以为任何给定的设置提供更现实的速度上限。最后,我们在NVIDIA Blackwell架构中为CUTLASS设计的最先进的融合多头注意力(FMHA)内核上实现了GNA。我们的实现可以在许多完美块稀疏的情况下充分利用理论上的最大速度提升,并在FP16中实现了每秒1.3拍浮点运算。此外,我们将在现成的生成模型,如Cosmos-7B、HunyuanVideo和FLUX中插入各种GNA配置,并表明它可以在不进行微调的情况下,在B200上实现28%到46%的端到端速度提升。我们将直接通过NATTEN项目开源我们的模拟器和Blackwell内核。