点击下方卡片，关注「集智书童」公众号
本文主要解决了什么问题
Sparse注意力机制的性能瓶颈 ：许多现有的Sparse注意力方法未能在自注意力Baseline之上提供显著的速度提升，主要是由于基础设施复杂性和硬件架构快速演进导致的实现开销。
局部性注意力的统一框架需求 ：当前的局部性Sparse注意力方法（如滑动窗口、分块注意力）缺乏一个统一的定义和分析模型，难以比较不同方法的性能潜力。
多维Token布局的效率挑战 ：针对视觉任务的多维Token布局（如图像和视频），现有方法在实现高效Sparse注意力时面临额外的计算和内存开销。
本文的核心创新是什么
广义邻域注意力（GNA）框架 ：提出了一种扩展的邻域注意力模式家族，通过引入“步长”参数，将滑动窗口注意力、分块注意力等多种方法统一到一个框架下，为灵活性和效率提供了新的权衡空间。
NATTENSim分析工具 ：开发了一个模拟工具，能够计算不同设计选择下的加速上限，帮助研究者优化GNA配置并找到最有效的参数组合。
Blackwell架构上的高效实现 ：基于CUTLASS FMHA核实现了GNA，能够在FP16和FP8精度下分别达到1.2 petaFLOPs/s和1.7 petaFLOPs/s的有效性能，并在完全块Sparse配置下实现理论上的加速比。
结果相较于以前的方法有哪些提升
速度提升 ：
- • 在Cosmos-7B上，通过引入56%的Sparse度，实现了1.26倍的整体加速比，接近理论FLOP加速比1.28倍。
- • 在HunyuanVideo上，通过引入91%的Sparse度，完全实现了1.63倍的理论FLOP加速比。
- • 在FLUX-1.dev上，通过引入90%的Sparse度，实现了1.45倍的整体加速比，接近理论FLOP加速比1.46倍。
质量保持 ：在所有实验中，使用GNA生成的输出质量与自注意力Baseline相当，无需额外训练或微调。
适用性扩展 ：GNA框架不仅涵盖了现有的滑动窗口和分块注意力方法，还支持更灵活的配置（如步长调整），适用于更多应用场景。
局限性总结
部分场景下的性能差距 ：在非完美块Sparse配置下，细粒度Mask的成本可能导致实际加速比低于理论值。
Token置换的效率问题 ：当前实现中的Token置换操作仅利用了约1/8的内存带宽，未来可以通过专门的拷贝Kernel进一步优化。
依赖特定硬件特性 ：虽然GNA框架本身具有通用性，但其高效实现依赖于特定硬件（如Blackwell架构）的支持，可能不适用于所有平台。
复杂性增加 ：GNA引入了更多的参数（如步长），虽然增加了灵活性，但也使得参数选择变得更加复杂，需要借助NATTENSim等工具进行优化。
导读
许多Sparse注意力机制，如邻域注意力机制，通常未能始终在自注意力 Baseline 之上提供速度提升。这主要归因于注意力基础设施的复杂性以及人工智能硬件架构的快速演进。与此同时，许多最先进的基座模型，特别是在计算机视觉领域，严重依赖注意力机制，并需要可靠的Sparse性以摆脱
的复杂度。在本文中，作者研究了一类注重局部性的有前景的Sparse注意力机制，旨在为其性能提升开发更好的分析模型。作者首先介绍了广义邻域注意力机制（GNA），该机制能够描述滑动窗口、步进滑动窗口和分块注意力。然后，作者考虑了实现这些方法时可能的设
- 引言
快速Sparse注意力机制长期以来备受关注[5, 9, 33, 35, 37, 43, 53]，但很少能避免复杂问题。基础设施仍然是注意力机制普遍面临的挑战，因为注意力机制的实施很少能充分利用现代GPU的计算能力，至少与稠密线性代数原语（如广义矩阵-矩阵乘法（GEMM））相比，后者通常能利用约80%的峰值FLOPs/秒。迄今为止最成功的例子是Flash Attention 3[39]，在NVIDIA Hopper架构中以半精度运行时，能利用高达75%的峰值FLOPs/秒。然而，更低精度的实现仍落后于预期，随着每一代新架构和编程模型的改变，基础设施方面的新挑战不断涌现。因此，需要修改核心注意力实现的Sparse方法进展缓慢。一个例子是针对注意力权重Sparse性的特定方法，如滑动窗口[5,21, 33, 37]和分块注意力[31, 43]，这些方法只能通过利用分块Sparse性来加速计算。分块Sparse性主要包含跳过在预定义注意力 Mask 中完全被屏蔽的计算块。实现分块Sparse性可能并不简单，有时会带来显著的开销，如果基础实现已经是当前最佳水平，这些开销可能会抵消性能提升。结果，这些方法通常未能充分发挥性能潜力，其利用率低于稠密注意力。
尽管存在这些挑战，滑动窗口和分块注意力机制已成功集成到大语言模型中[3, 10, 25]，其中块Sparse性和细粒度 Mask 处理的是单维度的token布局。这使得它们的实现相比图像和视频等进入二维和三维token布局的应用要简单得多。直接结果是，专注于多维token布局的实现通常落后得更远。一个这样的例子是融合邻域注意力（FNA）[19]，它实现了多维滑动窗口注意力，但由于随着维度升高而增加的开销（即软件预测和细粒度 Mask ）导致显著的性能差距。同时，FNA针对的是NVIDIA Ampere架构，这使得其与Hopper等最新架构相比存在更大的差距。虽然最近出现了一些成功的多维局部注意力实现[54]，但在模式和参数方面它们不够灵活，仅提供全块Sparse Mask ，并且局限于非常特定的用例。
目前存在多种提供线性代数和张量布局原语（primitives）的框架，例如CuTe和CULTASS [41]、Triton [42]，以及专门为块Sparse注意力设计的编译器，如Flex Attention [14]。尽管这些框架对研究进展产生了巨大影响，但不同方法之间差异巨大的实现方式使得分析这些多维Sparse注意力方法的确切性能差异，并找出其根本原因变得越来越困难。
为此，作者提出了广义邻域注意力（GNA），这是对邻域注意力（NA）[21]的扩展，而NA本身是一种滑动窗口方法，旨在将许多现有方法统一到一个定义下。GNA能够在效率、质量和归纳偏置（即平移等变性）之间提供权衡空间。使用GNA实现的方法可分为滑动窗口[21, 33, 37]、步长滑动窗口[43, 54]和分块[31, 43]注意力。作者重新审视了HaloNet[43]中讨论的滑动窗口与步长滑动窗口及分块注意力的方法，但重点在于加速而非内存操作。不同GNA模式的示意图展示在图2中。
作者进一步创建了一个仿真工具，该工具能够在不同的使用案例和实现设计选择下，计算GNA定义下任何方法可达到的上限加速比。通过抽象实现细节，并以tile/block粒度定义加速比，而不是以FLOP/FMA粒度定义，作者最终能够对这些不同方法进行公平的分析比较。
最后，作者在NVIDIA Blackwell架构上基于CUTLASS的注意力 Kernel 实现了GNA，该架构是目前性能最优的选择之一，并表明在许多GNA配置中，作者能够成功实现高达100%的FLOP效率。
尽管注意力机制本身在光速（SoL）性能方面通常无法接近标准GEMM，但鉴于作者实现的性能开销可以忽略不计，作者相信NATTEN可以为Sparse和局部注意力方法提供一个追赶标准FMHA实现性能的方案。
作者识别出能够实现完美加速的GNA变体，并将它们与标准邻域注意力一同引入生成模型和基础模型，例如Cosmos[1]、HunyuanVideo[27]和Flux[28]。作者展示，通过直接集成到NATT EN项目中、作者开发的Blackwell核，可以在不进行微调的情况下，将这些模型加速高达63%。除了改进和加速GNA，并支持更多Sparse注意力模式之外，作者希望提供一套有用的工具，例如作者的模拟器，帮助研究行人找到模型的最有效配置，并轻松地对一组有用参数进行消融实验，而不是迷失在GNA所提供的庞大参数空间中。作者所有的工作都将开源。
- 相关工作
注意力机制因其与token数量相关的二次复杂度而被广泛认为是一项计算密集型操作。这通常是研究Sparse注意力形式的主要动机之一。点积注意力机制，即最常用的注意力形式，主要包含两次矩阵乘法，这意味着它可以利用矩阵乘法中存在的多种Sparse性形式，例如结构化Sparse性[8]和低秩近似[7, 46]。通过选择更粗粒度的Sparse目标，也可以将Sparse性引入注意力机制，这有时具有应用特定性。例如，某些方法专门设计用于将Sparse性引入大语言模型推理工作负载，其中KV缓存是必要的，因此需要某种形式的Sparse性或压缩[4, 30, 40]。Token合并[6]等方法尝试直接减少 Query 和上下文token的数量，并在某些视觉任务中显示出有效性。
Sparse注意力方法也可以分为静态方法和动态方法。动态方法[6, 8, 38]无需进一步微调即可更有效，而静态方法[5, 7, 9, 33, 37]更有可能实现更好的加速效果[10, 25, 31]。某些方法可以归类为混合方法[6, 48, 49]，其中运行时收集的统计数据指导计算重新布局，然后使用静态方法（如块Sparse性[35]）。有些方法甚至使用完全密集的计算[6, 31]。例如Native Sparse Attention[52]，结合了多种Sparse方法，包括静态和动态方法。
在本文中，作者特别关注静态方法，其中Sparse化的目标是注意力权重矩阵本身，因此 Token 坐标空间决定了某个区域/权重是否被 Mask 。局部注意力是最典型的例子，但通过非连续上下文窗口引入全局上下文的局部注意力变体[5, 20]也属于此类。局部方法[11, 16, 23, 26, 29, 31]的成功可归因于它们对Transformer[44]架构所需改动极少，且空间局部性是自然界中存在的固有倾向。过去的[25]和最近的[3, 10]大语言模型除了标准自注意力外，也将其架构中采用了局部方法。
2.1. 局部Sparse注意力
多年来，许多研究者提出了将局部性引入注意力机制，有些是为了利用归纳偏差，有些是为了实现亚二次复杂度，还有些两者兼顾。在本节中，作者将这些方法总结为三类，具体描述如下。
滑动窗口注意力机制。最早提出这一机制的论文包括Image Transformer [33]和StandAlone Self-Attention (SASA) [37]，其中它们特别设计了一种注意力机制，该机制包含卷积中存在的一些归纳偏差。滑动窗口模式本身与卷积非常相似，有时也受其启发。然而，由于缺乏基础设施，它们最初并未被广泛采用，尤其是在视觉领域。Ramachandran等人 [37]指出，尽管他们的模型在精度上优于同等规模的CNN，但训练该模型需要更多的资源和时间。后来，在单维因果域中使用了相同的概念，例如Longformer [5]，并将其扩展为一种“扩张”形式，这种形式可以在与标准局部注意力相同的复杂度下引入全局上下文。
多年后，邻域注意力（NA）[21]重新审视了这一概念，用于层次化视觉Transformer。SASA和NA之间的关键区别在于处理坐标空间中的边界情况。SASA采用零填充特征图，类似于卷积中的填充。NA将窗口“反射”回去，因此每个 Query 都保证关注固定数量的上下文 Token ，无论其在坐标空间中的位置和窗口大小如何。这种行为还保证了当窗口大小等于输入大小时，NA在数值上与自注意力机制相匹配。类似于Longformer，但在视觉领域，后来引入了一种扩张形式的NA [20]。在层次化视觉Transformer中结合标准和扩张NA，在各种任务上的精度都超越了原始模型 [20, 45]，且没有任何额外成本。
步长滑动窗口注意力。也称为分块局部注意力[43]，这种方法有效地将延迟步骤引入滑动窗口注意力中。这最初是由滑动窗口注意力的典型实现[33, 37]需要显式内存操作来从上下文中提取滑动窗口这一事实所驱动的，这些操作除了会增加内存占用外，还会迅速抵消计算上的任何节省。这种形式的滑动窗口注意力有一个极端情况，称为分块注意力，其中步长与窗口大小相同，导致局部窗口不重叠。最近，这种方法在滑动Tile注意力（STA）[54]中得到了重新审视，其目标是实现完美的分块Sparse Mask ，以此作为最小化滑动窗口方法中 Mask （因此也是浪费的FLOPs）的一种手段[21, 33, 37]。作者将在第2.3节中进一步阐述这种方法。
阻塞注意力。这种方法有效地将 Query 集和上下文集进行分区，并在每个分区上独立（且并行）地执行自注意力操作，以处理其余部分。除了HaloNet [43]之外，Swin Transformer [31] 中的窗口自注意力（WSA）是阻塞注意力的另一个实例。阻塞注意力易于实现，并且在块 Level 上具有明显的并行性，但这以牺牲块间交互为代价。这可以通过引入全局注意力、跨层变化窗口（块）大小（Swin的移位wSA）或引入卷积来弥补。这种方法已被许多视觉模型采用，例如Segment Anything [26] 和 ViTDet [29]。
2.2. （Sparse）注意力基础设施
滑动窗口注意力机制，特别是在视觉领域，通常被认为效率低下[31, 43]，但这基于一个假设，即对上下文 Token （ Key和Value ）的任何形式变换都必须在全局内存中进行显式操作。邻域注意力Transformer[21]的一个关键贡献是一组简单的CUDA Kernel ，它们直接计算向量矩阵乘法问题，而无需在全局内存中进行这种显式复制。这些 Kernel 被封装为一个PyTorch扩展，称为NATTEN。然而，在初始发布的同时，注意力机制的一般实现即将迎来重大变革。
截至2022年[13, 36]，大多数点积注意力机制的实现，尤其是那些未针对特定推理使用场景绑定的实现，均采用两次矩阵乘法，并在其中插入softmax算子。在大多数深度学习框架中，前者通常针对强大的密集线性代数库中的通用矩阵-矩阵乘法（GEMM）例程，即NVIDIA GPU的cuBLAS和CUTLASS[41]，这些例程通常提供优异的性能，通常可达峰值FLOPs/秒利用率的80%。然而，这种方法的缺点在于中间矩阵——注意力权重矩阵的尺寸呈平方级增长，导致内存占用呈平方级增加，进而使得GEMM中的内存操作次数也呈平方级增长。
因此，这种实现方式严重受限于内存带宽，而现代GPU的内存带宽通常远低于计算能力（FLOPs/秒）。Flash Attention[13]表明，通过将两次GEMM和softmax融合到同一 Kernel 中，并利用在线softmax[24, 32]，作者可以在SRAM中持续累积上下文，而无需完全生成
的权重矩阵。这种方法也被称为融合多头注意力（FMHA）。Flash Attention 2[12]随后对原始版本进行了改进，并成为NVIDIA Ampere架构的当前最优方案。Flash Attention 3[39]通过利用CuTe和CUTLASS[41]中的新编程模型和硬件特性，将该方法扩展到NVIDIA Hopper架构。
由于自注意力机制是所有Sparse注意力机制的基础，并且该基础通过FMHA得到了显著改进，因此Sparse注意力方法不得不随之跟进。Mistral-7B [25]是首批采用一维滑动窗口注意力的模型之一，它通过直接将其实现为块Sparse Mask ，在Flash Attention 2 [12]等最先进的实现中直接实现了这一机制。最近，像Command A [10]这样的语言模型也采用了这种方法，并结合全局自注意力机制使用滑动窗口注意力。
然而，像
这样的实现面临着额外的挑战，因为处理token的多维布局增加了额外的负担，这使得它们的高效块Sparse性变得不平凡。
2.3. 多维诅咒
涉及注意力权重Sparse性的研究在LLM（大语言模型）应用和视觉应用之间存在一个关键且往往未被讨论的差异，即视觉数据的维度布局（图像、视频）。作者认为这是一个额外的负担，因为用于高效实现的方案选择有限，且通常伴随着明显的开销。为说明这一问题，作者建议读者参考图3。例如，当考虑 Token 的二维布局（即图像模型中的特征图）时，如果FMHA（特征映射高斯混合注意力）实现采用单维度分块来覆盖这些 Token 在内存中的典型布局，那么“浪费的计算”的可能性会增加，并且与分块大小相关。
融合邻域注意力（FNA）[19]提出通过采用多维Tile技术来解决这个问题，该技术将FMHA中的GEMM转换为张量收缩的一个实例。这种实现方式显著提升了
中的原始核，但同时也可能因多维Tile导致的软件预测开销而受到很大限制。虽然硬件预测，例如通过NVIDIA的张量内存加速器（TMA）2等组件，可能有助于最小化这种开销，但它们也可能影响某些对实现最佳加速效果至关重要的设计选择，如动态KVTile。
多维铺砖的一种更易于实现的替代方案是通过 Token 重布局来模拟行为。最早演示这一方法的是Flex Attention中FNA的实现[14]3。Flex Attention是一个PyTorch API，用于评估用户定义的注意力 Mask ，并将其编译成块SparseTriton Kernel 。这种方法增加了额外的内存操作开销，这些操作对于重布局来说是不可避免的。它仍然需要对原始注意力 Kernel 进行一些修改，但与融合多维铺砖相比，修改次数要少得多。在Flex Attention的情况下，这可以通过在用户定义的 Mask 中直接将一维坐标转换为 Token 的新布局坐标来处理。
滑动铺砖注意力（STA）[54]试图进一步最小化由于细粒度 Mask 造成的FLOPs浪费（见图3），并提出在铺砖/块上定义邻域/滑动窗口注意力，而不是在单个 Token 上进行，铺砖/块的大小与底层FMHA Kernel 相匹配。这与HaloNet[43]等步进滑动窗口方法非常相似，但动机是计算密度和因此的速度提升，而不是内存操作的成本，对于块SparseFMHA Kernel 来说，内存操作成本是不存在的。STA的实现采用了Flex Attention的 Token 重布局，而不是实时多维铺砖，并且在相同的FLOPSparse度下成功超越了FNA。值得注意的是，该实现针对Hopper架构，而FNA针对Ampere张量核心。除此之外，STA的主要限制是，与Blocked Attention类似，它假设 Query 和上下文张量总是根据相同的铺砖大小进行铺砖。这种假设在实践中并不总是成立，即使某些实现支持它，也没有保证这种配置在特定用例中达到最佳性能。
鉴于存在多种局部Sparse注意力方法，且每种方法在实现中都有不同的设计选择，加之不同硬件架构具有各自的特性和性能水平，作者旨在阐明它们之间的差异与相似之处。为此，作者首先在新的框架下将它们统一，该框架称为“广义邻域注意力”（GNA）。在该框架下，作者为滑动窗口注意力[33, 37] / 邻域注意力[21]、步长滑动窗口注意力[43]、分块注意力[31, 43]以及介于其中的方法（包括STA[54]等特殊情况）创建了统一的实现。这包括扩展NATTEN以支持所有这些方法。
作者进一步构建了一个分析工具，名为NATTENSim，该工具在不同实现设计选择下计算这些方法更细粒度的上界加速比。这对于在现有实现下优化速度或精度，以及调查新的设计选择/实现是否基于其对端到端加速比的影响而合理，都很有帮助。
最后，作者受原始FNA[19]的启发，基于CUTLASS的Blackwell FMHA Kernel [41]，为Blackwell架构创建了一个新的FNA Kernel ，并表明在许多情况下，作者可以在FLOPs和NATTENSim方面完全匹配分析加速比。
- 方法论
在本节中，作者定义了广义邻域注意力机制，该机制为静态注意力模式家族引入了一个新的维度，其唯一目标是在准确性和平移等变性以及效率之间提供一个权衡空间。随后，作者描述了此类方法实现中的设计选择，这是作者创建分析工具NATT ENSim的动机之一。作者阐述了模拟器的工作原理，以及它如何根据用例信息、设计选择甚至特定硬件的详细信息来指导选择GNA参数的过程。最后，作者转向针对Blackwell架构的实现，该实现部分受到作者分析工具研究的启发。此实现基于Blackwell架构中性能最佳的FMHA Kernel 之一，并在强大的CuTe和CUTLASS [41]框架中完成。
3.1. 广义邻域注意力
邻域注意力（NA）[21]以及更一般地，滑动窗口注意力模式是细粒度的 Mask ，允许每个 Query token拥有一个固定大小的独特上下文窗口，该窗口大小基于其坐标确定。这种效果类似于卷积，其中固定大小的滤波器应用于滑动窗口区域，并在输出中收缩为单个点。这两种方法之间的关键区别在于卷积中的滤波器是静态的，通常通过梯度下降学习，而在滑动窗口注意力中，“滤波器”是动态的，基于上下文。先前的工作将NA扩展以支持扩张窗口[20]，模拟卷积中的扩张，以及因果 Mask [19]，这在视频应用中可能很有用，其中可能需要原子时空注意力 Mask 。
在这项工作中，作者从放宽标准NA的定义开始，允许窗口大小为偶数。NA最初是在窗口大小为奇数的情况下定义的，因为目标是使 Query 在其邻域中完全居中。作者通过将窗口大小分为两部分来实现这一点：窗口大小左和窗口大小右。在标准NA的定义中，这两部分都等于窗口大小除以二取整。对于偶数大小的窗口，作者可以选择其中一边比另一边大1。默认情况下，作者选择左侧为较大的一侧，这意味着对于窗口大小为8的情况，非角落的token会关注其左侧的4个token、自身以及右侧的3个token。
作者进一步为NA添加了一个第四个参数，称为"步长"，因为它可以让人联想到卷积中的步长。与卷积类似，步长为NA中的滑动窗口添加了一个延迟步骤，这意味着步长为1相当于标准的NA，即每个 Query Token 都移动滑动窗口，而步长为2则每两个 Query Token 移动一次滑动窗口。另一种理解方式是，步长将 Query 分组，并迫使它们共享同一个邻域（上下文窗口）。每个组中的 Query 静态地选择一个领导者 Query ，并关注其邻域，该邻域由标准的NA定义。领导者的选择，类似于窗口左和窗口右，可以是用户定义的，但在GNA中的默认选择是组中最中心的 Query ，以与原始NA工作保持一致。如果组的大小均匀，则会有偏袒右侧的趋势。步长可以取任何小于或等于窗口大小的正整数值。大于窗口大小的步长会在注意力中引入"空洞"，其中一些上下文 Token 永远不会被关注，因此是不允许的。作者在图2中展示了固定窗口大小下不同步长的可视化。步长NA表现出与HaloNet [43]中的阻塞局部注意力相似的特性，其"块大小参数"为滑动窗口注意力引入了与步长在邻域注意力 [21]中相同的延迟步骤效果。这两种情况中还有一个有趣的特性是，当步长等于窗口大小时，模式将等同于（完全）阻塞注意力，也称为窗口自注意力（WSA）[31]。在某些情况下，步长NA也可以实现STA [54]，但作者指出，这只有在对底层实现做出某些假设时才得到保证。作者将在第3.2节和第3.3节中进一步阐明这一点。
步长是Sparse注意力模式NA家族中的一个关键参数，其动机主要源于滑动窗口注意力无法实现速度提升这一事实。HunyuanVideo中91%的Sparse度与其Sparse率完全成正比。这是因为滑动窗口注意力本质上是一个向量矩阵乘法问题，通常通过矩阵矩阵乘法（GEMM）[19]和点积上的 Mask 操作来实现。 Mask 操作确保了正确性，但同时也浪费了已经完成的工作（FLOPs）。
另一方面，步长迫使同一组的 Query 共享其整个邻域，从而将它们的窗口边界对齐。这意味着步长可以弥合滑动窗口方法[21, 33, 37]与分块方法[31, 43, 54]之间的差距。分块方法的优势在于，在满足某些假设的情况下，它们可以主要执行密集计算，无需对注意力权重进行细粒度 Mask ，进而不会浪费FLOPs。更具体地说，任何被定义为完全分块Sparse[35]的方法，本质上只涉及密集计算，并可能导致与FLOPs节省成正比的速度提升。因此，作者的目标是回答诸如：具体浪费了多少工作，以及更大的步长是否保证更少的浪费FLOPs等问题。为此，作者创建了一个分析工具来模拟此类方法常见实现的行为，并使用该工具计算每种组合节省了多少工作块。
3.2. GNA分析工具
步幅扩展了已相当庞大的参数空间。这，连同实现中的不同设计选择和配置，特别是FMHA Kernel 的配置，使得为最终用户提供具有最佳权衡的有用参数变得困难。因此，作者创建了一个名为NATT ENSim的分析工具，它能够针对这一点提供洞察。NATT ENSim计算实现访问的上下文（KV）Tile数量，考虑了动态与静态KVTile、一维与多维Tile等不同的设计选择，当然还有 Query （Q）和上下文（KV）张量的Tile形状。
设计选择：Tile尺寸。对于任何GEMM或基于GEMM的 Kernel 而言，Tile尺寸是最重要的配置之一，它们直接决定了工作负载如何被分配和调度到各个工作单元。在大多数融合两个GEMM的FMHA Kernel 中，它们被迫共享大部分Tile尺寸，但存在排列。在第一个GEMM中， Query 张量（Q）沿 Query 序列模式以某种Tile尺寸
进行Tile化，键张量（K）沿上下文序列模式以某种Tile尺寸
进行Tile化。共享头维度模式（即收缩模式）以某种Tile尺寸
进行Tile化。在第二个GEMM中，左侧操作数是一个形状为
的注意力权重TileA，这意味着右侧操作数，即值张量的Tile，必须在收缩维度（现在是上下文序列模式）上匹配Tile尺寸。
因此，
和V采用相同的Tile尺寸，作者将其今后称为
。
的头维度模式通常也以
进行Tile化。Tile尺寸
和
因此直接影响被掩盖的FLOPs数量（浪费的）。许多因素可以决定这些Tile尺寸的有效选择，例如可用的SRAM量、GEMM流水线中的阶段数量、操作数所需的布局变换，当然还有底层矩阵乘累加（MMAs）的形状，或在NVIDIA GPU的情况下Tensor Core指令形状。在现代架构如Hopper和Blackwell中，对于GEMM而言，有许多不同的Tile尺寸选择，它们在不同的问题形状上表现出不同的性能水平。像STA [54]这样的工作中的关键限制在于，其方法假设
始终等于
，从而限制了选择数量。而高效的FMHA Kernel 在挑选Tile尺寸方面并不总是能像标准密集GEMM那样提供相同的灵活性。
设计选择：单维度与多维度平铺。大多数FMHA Kernel 假设token采用单维度布局，考虑到注意力操作的置换等变性，这是合理的。然而，这为token采用多维度布局的情况带来了挑战，例如在视觉模型中，token代表图像和视频中的像素块。这种情况会导致更多不必要的FLOPs浪费，如图3所示。一个自然的解决方案是在多个维度上平铺，并尊重原始的多维度布局，这本质上将GEMM问题转换为张量收缩（GETT^6）。这种解决方案被FNA[19]采用，其中基础FMHA Kernel 中的
和
被重新解释为多维度平铺（即
）。这种方法的缺点是可能因额外的软件预测逻辑引入显著的开销，并且由于该 Kernel 是为Ampere架构设计的，硬件预测并非可选方案。一个实用的解决方案是将多维度平铺从 Kernel 中移除，而是将其实现为重新布局操作。这种解决方案被Flex Attention[14]在其多维度NA Mask 实现（称为"Tiled NATTEN"）中采用。这种方法唯一的潜在缺点是内存操作不可避免的总固定成本，该成本与Sparse程度无关，仅取决于Q、K、V和输出张量的大小。作者将这种方法称为token置换，因为它主要由token空间的重新布局组成，并且与批次、注意力头数和头维度无关。
设计选择：KV分块。分块通常是静态的，在许多块SparseFMHA核中，静态KV分块根据 Mask 定义要么被访问，要么被跳过。然而，某些实现，如FNA [19]，首先在KV Token 空间中切出需要访问的区域，然后动态地对该区域进行分块。这可以节省一些额外的计算，并实现可能的最小FLOPs浪费。另一方面，这种方法并不总是可以实现，特别是在依赖硬件预测的设计中。例如，Hopper张量内存加速器（TMA）需要在核启动之前确定分块副本（张量映射）的参数。虽然可以在运行时修改/替换张量映射，但这并非没有开销。
用例。问题形状（token布局）、窗口大小、扩张、步长和因果 Mask 都是用户指定的，并在确定以Tile为单位的计算节省方面发挥作用。
Tile模拟。NATTENSim的主要目标是模拟根据实现中的设计选择如何对给定的用例进行Tile划分。通过坐标的基本操作，并使用核心GNA Mask 的精确定义，NATTENSim计算每个QTile访问的每个KVTile的坐标。如果作者考虑所有QTile的最坏情况（KVTile的最大数量），作者可以针对每个用例的自注意力计算比FLOP级加速更真实、更细粒度的加速上限。例如，一个完美的块Sparse Mask 和具有90%Sparse度的滑动窗口注意力 Mask 都具有FLOP级加速上限为
，但后者除非实现为向量矩阵乘法，否则永远不会恰好执行FLOPs的
。
使用NATT E NSim，作者可以实现以下功能：
-
- 比较不同的设计方案，并选择在实现难度和加速效果之间取得最佳平衡的那个方案。
-
- 寻找完全块Sparse的案例，其中无需进行细粒度 Mask ，并且模拟加速比与FLOP Level 的加速比相匹配，如图4所示。在动态KVtiling中，任何满足
均匀划分窗口大小，且
均匀划分步长的设置均可实现此目标。在静态KVtiling中，情况则不那么直接，可能需要一些模拟来找到这些点。 3. 3. 预测在特定参数和设计选择下，任何给定模型的端到端加速上限。
3.3. Blackwell架构的实现
作者从Blackwell [41]的CUTLASS FMHA核开始，该核在FP16精度下可达到1.2 petaFLOPs/s，在FP8精度（使用per-tensor缩放）下可达到1.7 petaFLOPs/s。尽管作者的实现特别针对Blackwell架构，但实现中的设计选择主要是架构无关的。该实现紧密遵循原始的FNA核[19]，唯一的区别是将多维tiling移出核外，转而依赖token置换。这也迫使使用静态KV tiling而不是原始FNA中的动态KV tiling。作者选择这种设计的原因如下：
-
- 若使用TMA进行数据移动，静态KVTile化是唯一选择。
-
- 将多维平铺融合到现有的FMHA Kernel 中可能会破坏太多关于序列模式的假设，即使CuTe能够促进布局转换并与TMA进行接口交互。
-
- 只要不受token置换和反向置换导致的内存传输时间限制，作者就没有太多性能提升空间。例如，考虑Cosmos-7B [1] 和 HunyuanVideo [27] 的用例，这分别仅为6.9GB和10.5GB。如果作者利用单块B200的8TB/s HBM带宽的一半，这将是FMHA时间的1.9%和1%，这仅会限制非常大的Sparse Level 。
-
- 在许多情况下，token置换和逆置换只需执行一次：在第一个transformer层之前进行置换，在最后一个层之后进行逆置换。大多数transformer架构对置换具有等变性质，ViT [15] 和 DiT [34]（在视觉领域具有代表性）都符合这一特性。
-
- 通过token置换，融合额外的上下文token非常简单。某些模型，如Hunyuan [27] 和 Flux [28]，通过跨注意力机制将视觉token与文本token进行交互。
‰
通常通过启动额外的FMHA核来支持这些场景，并使用它们的logsumexp“合并”两个输出。然而，作者在此同一核内实现此功能，允许某些KV token采用完全不同的布局和 Mask 。
由于作者实现了多维分块，空间局部性使作者能够将已访问的KV分块定义为Q分块中最后 Query 所需的最后一个KV分块坐标与Q分块中第一个 Query 所需的第一个KV分块坐标之间的范围。 Kernel 的大部分仍然对多维分块保持中立，只有生产者分块和softmax分块会考虑这一点。生产者分块将迭代索引映射到多维分块坐标，然后映射到静态一维分块索引，该索引直接通过CuTe与TMA加载的索引接口。同样，softmax分块（s）必须将一维Q和KV坐标映射回原始多维坐标，并应用细粒度GNA Mask 。然而，考虑到 Mask 的开销，作者还为完全块Sparse情况以及任何额外的KV Token 实现了谓词，在某些情况下可以完全消除作者的实现开销。
在当前版本中，作者通过PyTorch直接将token置换实现为复制操作。形状不能均匀整除的问题会手动填充，在反向token置换后裁剪输出张量，核函数处理KV填充的预测。在未来版本中，作者希望也能改进token置换，因为当前解决方案通常使用大约1/8的内存带宽。核函数和额外的内存操作以及潜在的填充直接集成到
中，并通过典型的na{1, 2, 3} d API暴露。
- 实验
为了评估GNA的适用性和性能，作者精心选择了具有适合Sparse注意力机制的多维布局 Token 的应用。选择应用的主要标准是它们的端到端工作负载中是否至少有40%是专门的自注意力机制。作者的最终候选者是Cosmos [1]（世界基础模型）、HunyuanVideo [27]（视频生成）和FLUX [28]（图像生成）。作者额外指出，FLUX只有在生成分辨率高于2K时才满足标准。作者在表1中展示了这些模型的负载分布。
4.1. GNA性能
作者首先考虑这些应用在注意力机制中引入约90%Sparse性的情况，并运行它们的计算形状通过
，遍历所有可能的步长值。对于每种使用场景，作者根据token布局的形状（特征图大小）选择了Q和KVTile形状
和
）。作者选择了能被KVTile形状
）整除的窗口大小，因为这样可以实现完美块Sparse的GNA形式，不浪费FLOPs，并有可能实现完美加速。STA [54]中也观察到了类似的现象，不同之处在于STA不区分
和
，并且它们的步长参数是固定的。
此外，作者根据一个简单但实用的过滤器对模拟器的结果进行剪枝。作者使用步长值的乘积作为衡量不同步长分组的标准，并消除任何在较大步长与较小步长相比性能没有提升的配置。例如，如果任何步长
的情况在分析加速比上（相对于NA本身，步长为1）表现更差，作者不会报告它。这尤其有帮助，因为较大步长会牺牲平移等变性并可能牺牲质量以换取潜在的性能提升，如果后者在某种设置下没有实现，那么该配置不太可能比其他配置有优势。
作者在图4中展示了一个通过NATTENSim获得的实验结果实例，该结果基于HunyuanVideo用例，其Sparse度约为90%。与大多数其他方法不同，在这个特定的用例中，作者拥有多种选择，它们具有不同的上限加速比，其中任何一种都可能通过进一步的微调或训练提供最佳的精度-效率权衡。作者还报告了作者的Blackwell核在三种90%Sparse度模型上的实际性能，并在图5中将其与NATTENSim报告的分析上限进行了比较。可以立即观察到两点：1. 在完美的块Sparse情况下，作者的核要么完全实现，要么几乎完全实现了分析加速比，这也是按FLOP计算的加速比（最大可达到的）。这并不令人意外，因为softmax阶段的细粒度 Mask 的使用取决于是否需要它。2. 在Hunyuan的情况下，即使标准的邻域注意力也非常接近实现其全部分析加速比。然而，在非完美块Sparse的情况下，细粒度 Mask 的成本仍然可能带来一些不可忽略的开销。
作者报告了表2至表4中的端到端加速措施。此外，在接下来的两个小节中，作者还将展示这些用例的定性结果以及一些定量基准测试。
4.2. 视频生成
作者对上述视频生成模型进行了实验：Cosmos-7B [1] 和 HunyuanVideo [27]。两者均采用各向同性Transformer架构（分别为DiT [34]和MMDiT [17]），并使用全自注意力机制。在两种情况下，作者都生成5秒的720p视频。对于Cosmos-7B，其token布局（DiT特征图形状）为
，作者使用窗口大小为
的GNA，其Sparse度约为
。
作者发现，在不进行进一步训练/微调的情况下，超过这一Sparse度水平无法保持视觉质量。作者尝试了性能最佳的GNA配置（步长
）和标准自注意力（步长
。由于初始扩散步骤对生成视频的整体结构有显著影响，作者保留自注意力机制用于前12步扩散，并使用GNA进行剩余的23步扩散。因此，GNA在整个端到端工作负载中的占比降至约
，进而进一步限制了可实现的端到端加速比至约
。作者在图6中展示了部分生成视频的样本帧，观察到GNA生成的视频质量与自注意力 Baseline 相当。
在HunyuanVideo中，token布局（MMDiT特征图形状）为
，作者使用窗口大小为
的GNA，其Sparse度约为
，步长为
。与Cosmos类似，作者保留自注意力机制用于部分初始扩散步骤，此处设置为15步，并在剩余的35步中运行GNA。这降低了GNA在整个端到端工作负载中的占比至约
，进而限制了可实现的端到端加速比至约
。作者再次尝试性能最佳的GNA配置（步长
）和标准NA（步长
）。输出视频的样本帧展示在图7中。
作者此外在VBench [22]上使用HunyuanVideo评估了GNA。在这种情况下，作者遵循STA [54]，将Sparse性应用于所有50步。然而，作者仅报告了无需任何额外训练或微调的结果。在表5中，NA和GNA在5.8.3%Sparse性下分别实现了VBench分数83.24%和83.40%，两者均与自注意力 Baseline 83.08%相当。然而，它们分别只能实现1.11倍和1.23倍的端到端加速。当使用91%的更高Sparse性时，VBench分数适度下降至82.36%和82.04%，但端到端加速分别显著提升至1.73倍和2.27倍。
4.3. 图像生成
作者在FLUX-1.dev上进行了图像生成实验，该模型与视频模型类似，采用各向同性Transformer架构（MMDiT [17]）并使用全自注意力机制。作者仅研究生成4K图像，这会产生
的特征图，因为分辨率较低的工作负载受自注意力成本的限制较小。为生成4K图像，作者使用URAE [51]中的 Adapter ，因为FLUX-1.dev原生不支持4K图像生成。
作者实验的窗口大小为
，步长为
、
、
，这些是从
中选出的最优选择。步长
是标准邻域注意力，步长
是完全块Sparse的，可以实现
的分析级速度提升。为保持质量和全局结构，作者遵循视频生成实验中的方法，在最初的扩散步骤中保留自注意力机制，具体为28个扩散步骤中的前9步。这使GNA在端到端工作负载中的占比缩小到约
，进而将端到端改进的上限提升到
。然而，作者成功实现了其中的
速度提升，超过自注意力 Baseline 。
在质量方面，即使使用高步长
，使用GNA时也很少观察到显著差异，如图8所示。作者还为这些配置提供了定量指标。遵循URAE [51]，作者在MANIQA [50]和QualiCLIP [2]上评估它们，通过使用来自HPDv2 [47]测试数据集的 Prompt 生成图像。作者还使用GenEval [18]评估文本到图像的对齐，这是一个以目标为中心的评估基准。该基准包含多个类别，每个类别负责测量不同的组合属性，如颜色、定位、属性绑定、目标数量等。作者在表6中报告了所有三个基准的结果。总之，作者观察到使用GNA的FLUX-1.dev在质量上与自注意力 Baseline 相当，同时在B200上提供高达
的端到端速度提升。
- 结论
滑动窗口和阻塞注意力是两种极端的以局部性为重点的Sparse注意力方法，前者具有平移等变性等归纳偏差[21,33,37,43]，并且可能具有更高的质量和表达能力，而后者则可能更高效且易于实现。
在本文中，作者扩展了邻域注意力模式家族，该家族在窗口大小、膨胀率和因果 Mask 的选择上已经具有高度的灵活性，并新增了一个“步长”参数。与卷积类似，该参数为滑动窗口效果添加了延迟步，使其能够实现许多现有的Sparse注意力方法。这些方法包括但不限于块状局部注意力[43]、块状/窗口自注意力[31]以及滑动Tile注意力[54]，此外还包括现有的滑动窗口[33, 37]和邻域注意力[21]方法。虽然滑动窗口注意力中的延迟步概念绝非新事物[43]，但与原始工作相比，作者出于不同原因重新审视了这一概念。Vaswani等人[43]关注的是显式内存操作的成本，而本文的重点是最大化这些方法的速度提升，使其与Sparse程度成正比（即从90%Sparse度实现10倍速度提升）。
作者创建了一个用于GNA的分析模型，名为NATTENSim，该模型能够模拟在不同设计选择下的tiling行为，并计算每个 Query tile访问的KV tiles数量，从而作者能够针对不同配置在GNA下计算更可靠的加速上限。这些指标有助于在不受实现差异影响的情况下，定量比较不同方法的加速上限性能。此外，作者发现存在许多GNA配置是完全块Sparse的，在这些配置下可以实现与FLOPs成比例的加速，并且它们不一定符合块注意力或滑动tile注意力的定义。
作者在Blackwell架构上进一步实现了GNA，基于CUTLASS FMHA，该实现能够在FP16下达到1.2 petaFLOPs/s的有效性能，在FP8下达到1.7 petaFLOPs/s，并证明在完全块Sparse配置的情况下，能够完全实现NATTENSim计算出的分析加速比。作者还重点介绍了GNA的3个潜在应用，所有这些应用作者都确认其大部分工作负载都集中在自注意力机制上，并报告了接近或达到预期FLOP加速比端到端加速比的结果。在Cosmos-7B [1]上，将23个35个扩散步骤中的56%Sparse性引入其中，作者实现了1.26倍的加速比，FLOP加速比为1.28倍。在HunyuanVideo [27]上，将35个50个扩散步骤中的91%Sparse性引入其中，作者完全实现了1.63倍的FLOP加速比。在FLUX-1.dev [28]上，将19个28个扩散步骤中的90%Sparse性引入其中，作者实现了1.45倍的加速比，FLOP加速比为1.46倍。上述三种配置仍然能够生成视觉上可接受的输出，无需任何进一步训练或微调。
作者希望泛化邻域注意力机制能够成为超越本文所研究的用例和硬件的超越光速局部注意力的解决方案。
- 未来工作
作者的新实现可以通过在存在全密集Tile且设置并非完全块Sparse的情况下进一步预测细粒度 Mask 来进一步优化。由于多维Tile保留了空间局部性，如果第一个 Query 和最后一个 Query 在QTile中的邻域交集跨越整个KVTile，则可以跳过细粒度 Mask 。作者当前的实现针对完全块Sparse配置执行此操作。除此之外，还可以对细粒度 Mask 逻辑本身进行性能优化。这些优化可以进一步缩小预期分析加速与部分块Sparse或完全不块Sparse情况下实际加速之间的差距。
Token permutation和reverse permutation也通过PyTorch以朴素方式实现，仅利用了B200峰值内存带宽的1/8，并且可以通过专门的拷贝 Kernel 和激活内存管理进一步优化。然而，如论文前文所述，在某些架构中，这些操作的调用次数可以大幅减少，并在特定假设下（即各向同性Transformer架构，以及GNA中无膨胀）实现。
参考
[1]. Generalized Neighborhood Attention: Multi-dimensional Sparse Attention at the Speed of Light
扫码加入👉「集智书童」交流群
（备注：方向+学校/公司+昵称）