面向可扩展向量架构的FlashAttention

# 面向可扩展向量架构的 FlashAttention  
来源：https://arxiv.org/html/2608.18656  
作者：Sonia Rani Gupta、Nikela Papadopoulou（英国格拉斯哥大学），邮件：[[email protected]](mailto:[email protected])  
Miquel Pericàs（瑞典查尔姆斯理工大学），邮件：[[email protected]](mailto:[email protected])  

###### 摘要  
在 CPU 上运行 Transformer 模型推理正变得日益重要，尤其是对于小型语言模型（SLMs），其中向量架构作为有前景的执行基础正崭露头角。注意力模块由于高内存带宽需求成为主要瓶颈；FlashAttention 通过算子融合来改善数据局部性并减少中间内存访问，从而缓解这一问题。本文提出了 FlashAttention-V，一种面向可扩展向量架构的分块 FlashAttention，它通过利用注意力头间的并行性、头间打包以高效利用超出头维度的向量长度，以及提升向量寄存器利用率与内存访问局部性，实现了从短到极长向量的高效适配。我们将 FlashAttention-V 集成到 llama.cpp 框架内的 ggml 中，并使用 gem5 和 Banana Pi BPI-F3 开发板在 TinyLlama、Llama 3.2、Qwen2.5 及 Pythia-410M 模型上进行评估。在 Banana Pi BPI-F3 上，我们验证了循环重排与跨注意力头循环展开的有效优化原则，其性能提升随模型规模增大而增强，在短上下文及解码阶段尤为显著。基于仿真的分析表明，在 512 位向量长度（VL）下，FlashAttention-V 在预填充阶段相比标量 FlashAttention 实现了 22 倍至 42 倍的加速；当扩展到 64 通道和 4096 位 VL 时，可进一步获得 2 倍至 2.5 倍的增益。在解码阶段，使用 512 位向量长度时，FlashAttention-V 相比标量 FlashAttention 实现了 8 倍至 11 倍的加速；由于单令牌、内存受限的执行特性，性能对向量宽度与通道数的敏感性逐渐降低。我们进一步发现 Q8_0 量化线性层在长向量执行下存在结构瓶颈，限制了算术摊销，该现象在 RISC-V RVV 和 Arm SVE 上一致出现，表明当前量化格式对长向量可扩展性构成了根本挑战。  

## 1. 引言  
基于 Transformer 的模型支撑着现代 AI 系统，驱动着从聊天机器人、编程助手到机器翻译、AI 智能体等广泛应用，覆盖多种规模。大型语言模型（LLMs）通常具有数千亿参数，多部署于数据中心；而参数规模在 10 亿至 100 亿之间的小型语言模型（SLMs）则针对边缘部署进行优化。在这一谱系中，Transformer 推理均要求高吞吐、低延迟与强能效，无论是大规模服务器部署还是电池受限的边缘平台。注意力模块是所有 Transformer 架构的核心组件。多头注意力（MHA）通过并行计算多个注意力头扩展了原始设计，使模型能够捕获多样化的关系模式。为降低推理期间存储键值（KV）缓存的内存与带宽开销，多查询注意力（MQA）和分组查询注意力（GQA）等变体通过在查询头间共享键与值来减少 KV 缓存占用。这些注意力变体在多种硬件平台上已得到高效执行探索。注意力实现在 GPU 和专用 AI 加速器上均经过优化，以最大化训练与推理的吞吐量。同时，由于 CPU 具有高可用性、低延迟以及跨边缘设备与数据中心的可移植性，相关工作也在 CPU 上探索优化 Transformer 推理。尽管这些变体在查询、键、值的组织方式上有所不同，但它们均依赖于相同的底层自注意力计算。自注意力需要计算和存储大型中间矩阵，例如得分矩阵 $S=QK^{T}$（其中 $Q$ 和 $K$ 分别表示查询与键），以及生成的 softmax 注意力权重。这导致巨大的内存容量与带宽开销，并随序列长度增长而愈发显著。为解决此内存瓶颈，Rabe 和 Staats 提出了一种基于在线 softmax 的高效内存注意力算法，无需存储完整的中间得分矩阵 $S$ 即可实现精确的自注意力。在此基础上，FlashAttention 进一步通过分块与融合内核优化执行，将中间数据保留在 GPU 片上内存中。尽管最初为 GPU 优化，近期工作已探索在配备向量单元的 CPU 上运行 FlashAttention 以提升 Transformer 推理性能。现代向量处理器采用长度无关的指令集架构（ISA），如 RISC-V 向量扩展（RVV）和 Arm 可扩展向量扩展（Arm SVE），支持可变向量长度。这使得它们在高性能计算（HPC）和 AI 工作负载中颇具吸引力，可利用这些架构暴露的 SIMD 并行性。然而，FlashAttention 在头维度上暴露的向量级并行性有限。查询、键和值的维度为 $N \times D$，其中 $N$ 为序列长度，$D$ 为头维度。FlashAttention 输出具有相同形状 $N \times D$，其中 $N$ 变化范围大（令牌生成时坍缩为 1），而 $D$ 通常较小，为 64 或 128。现有的 RISC-V、ARM 和 x86 上的 FlashAttention 实现遵循一种受限设计：向量长度（VL）受限于 $D$。llama.cpp 中的向量化 FlashAttention 实现均限制 $VL \leq D$。Titopoulos 等人在 RVV 上针对 $VL \leq D$ 对 FlashAttention 进行了向量化。基于 NEON 的融合注意力实现 MEATTEN 在头维度 $d_k$ 上进行向量化，从而施加了相同限制。对于 16 位精度，这些实现仅能利用最多 1024–2048 位的向量长度，而量化更激进的模型利用长度更短。这限制了 FlashAttention 在长向量架构上的可扩展性，其中向量长度超过头维度，导致单指令多数据（SIMD）资源利用不足。本文通过优化 FlashAttention 以适应广泛向量长度的向量架构，弥合了这一差距，实现对传统与长向量寄存器的高效利用。我们提出 **FlashAttention-V**，一种重新设计的 FlashAttention 算法，它利用跨注意力头的并行性。通过循环重排、循环展开及头间打包，FlashAttention-V 将多个头映射到单个向量寄存器中，从而高效利用超出头维度 $D$ 的向量长度。循环展开进一步暴露指令级并行性，使 FlashAttention-V 在预填充与解码阶段均可高效利用向量寄存器，并在传统与长向量长度间实现扩展。我们的设计原生支持 MHA 和 GQA，避免了共享键值的冗余内存访问。我们在 llama.cpp 框架内的 ggml 中实现 FlashAttention-V，该框架常被用作研究平台。我们使用 RISC-V ISA 和 Arm SVE ISA 的向量内核函数来向量化内核。我们在仅解码器的大语言模型（LLMs）上评估 FlashAttention-V，这些模型运行在两个阶段：预填充（处理输入令牌并计算 KV 缓存）和解码（逐次生成一个令牌）。我们使用三个 GQA 模型：TinyLlama、Llama 3.2 和 Qwen2.5，以及 Pythia-410M 来评估 MHA。我们使用 QEMU 验证了基于 Arm SVE 的 FlashAttention-V 功能正确性，以证明其跨向量 ISA 的可移植性。我们在 Banana Pi BPI-F3 RISC-V 开发板上使用 gem5 对 FlashAttention-V 进行 RVV 评估，并评估其在高达 8192 位向量长度下的可扩展性。为反映真实硬件行为（固定数量的向量通道处理更长向量长度会增加每条指令的执行延迟），我们基于 Vitruvius+ 的通道微架构信息，扩展了 gem5 的 MinorCPU，为其添加了与向量长度成比例的 SIMD 功能单元延迟模型。最后，我们将分析范围扩展至注意力机制之外，量化了量化线性投影与前馈层中的瓶颈，表明向量打包与掩码归约开销显著限制了在较长 VL 下的可扩展性。本文的贡献如下：  
- • 我们提出了 FlashAttention-V，一种面向向量架构的 FlashAttention 算法，它利用跨注意力头的并行性并引入头间打包，实现了对超出头维度的向量长度（$VL > D$）的有效利用。结合循环重排、头间展开、最大化向量寄存器重用以及原生 MHA/GQA 支持，FlashAttention-V 在 gem5/RVV 上使用 512 位向量长度时，在预填充阶段相比标量（非向量化）FlashAttention 实现了 22 倍–42 倍的加速，在解码阶段实现了 8 倍–11 倍的加速。在 Banana Pi BPI-F3 上的验证确认了这些增益，相比标量 FlashAttention 可达 12 倍–14 倍。  
- • 我们使用 gem5 评估了 FlashAttention-V 的可扩展性，确立了在高达 8192 位向量长度下，预填充阶段 2 倍–3 倍、解码阶段 1.2 倍的加速上限。通过延迟感知的逐通道模型，我们表明配置多达 64 个向量通道和 4096 位向量时仍保持强可扩展性，同时指出了限制超过 4096 位向量性能的向量长度开销。这些发现为长向量架构上的硬件及向量长度感知的 FlashAttention 优化提供了可操作的洞见，此前尚未有量化分析。  
- • 我们分析了 Q8_0 量化与线性投影及前馈模块中长向量执行之间的结构不兼容性。通过在 RISC-V 上的微基准测试分析，我们确定打包与掩码归约是主要瓶颈，在 2048 位 VL 下消耗了 60% 的执行周期。在 2048 位 VL 下的 Arm SVE 验证确认这些开销与架构无关，源于 Q8_0 的交错权重-缩放布局，该布局引入了与 VL 相关的成本，阻碍了算术摊销。这突显了量化模型在长向量执行方面的基本限制。  

本文其余部分组织如下：第 2 节提供 Transformer 模型背景；第 3 节综述相关工作；第 4 节描述 FlashAttention 与自注意力的算法转换，并分析 Q8_0 量化下的线性投影与前馈层；第 5 节介绍实验平台与设置；第 6 节和第 7 节评估 FlashAttention-V 并分析量化模型与长向量执行间的结构不兼容性；第 8 节总结全文。  

## 2. 背景  
本文考虑仅解码器模型，使用多头注意力（MHA）或其高效变体分组查询注意力（GQA），后者通过在查询头间共享键与值来降低 KV 缓存开销。这些模型均基于相同的底层自注意力机制构建。自注意力使每个令牌能够关注所有其他令牌。输入嵌入被投影为查询（Q）、键（K）和值（V），并通过缩放点积与 softmax 计算注意力：  
$$\text{Attention}(Q,K,V)=\operatorname{softmax}\!\left(\frac{QK^{\top}}{\sqrt{D}}\right)V$$  
其中 $D$ 为头维度。多头注意力跨多个头并行应用此操作以捕获多样化关系。FlashAttention 是自注意力的优化实现，可减少内存使用并提高计算效率。它不在内存中物化完整的注意力矩阵，而是将计算分块，维护运行最大值与归一化因子，以实现数值稳定的在线 softmax，避免了中间内存开销。llama.cpp 是基于 ggml 构建的开源 CPU 推理框架，ggml 是支持量化操作与向量化指令的张量库，兼容 AVX、Arm NEON、Arm SVE 和 RVV。在 llama.cpp 中，注意力模块同时使用了自注意力与 FlashAttention 实现。  

## 3. 相关工作  
已有许多研究聚焦于优化注意力机制。我们首先讨论自注意力的算法改进，随后介绍 FlashAttention 变体。Rabe 和 Staats 提出了一种避免 $O(n^2)$ 内存使用的高效内存自注意力算法。FlashAttention 算法遵循类似原理，是自注意力的优化实现，旨在降低内存带宽。Dao 等人进一步提出了 FlashAttention-2 进行了优化。接下来我们讨论面向 CPU 的自注意力特定实现。