💻 computer science

FlashAttention for Scalable Vector Architectures

本文介绍了 FlashAttention-V，这是一种针对可扩展向量架构优化的分块式 FlashAttention 实现，它通过利用头间并行（inter-head parallelism）和高效的内存访问，显著加速了 CPU 上的 Transformer 推理，在预填充阶段实现了高达 42 倍的加速，在解码阶段实现了 11 倍的加速，同时也强调了当前量化格式在长向量执行中的结构性瓶颈。

原作者： Sonia Rani Gupta, Nikela Papadopoulou, Miquel Pericàs

发布于 2026-08-20

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Sonia Rani Gupta, Nikela Papadopoulou, Miquel Pericàs

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

从起草电子邮件的聊天机器人到调试软件的代码助手，现代人工智能系统都依赖于一种被称为“Transformer”的特定类型的计算机程序。这些程序通过观察句子中单词之间的相互关系来理解语言。为了实现这一点，它们使用了一种称为“注意力”（attention）的机制，这种机制就像一个聚光灯，允许系统专注于文本中最相关的部分，同时忽略其余部分。虽然这些系统通常在配备强大显卡的庞大数据中心中运行，但人们对在笔记本电脑、平板电脑甚至物联网中发现的专用微芯片等更小的日常设备上运行它们的需求正在日益增长。这些较小的设备通常使用另一种类型的处理器，这种处理器旨在通过处理被称为“向量架构”（vector architectures）的长行数据来同时处理大量计算。然而，一个主要的障碍出现了：注意力机制对内存的消耗极大，要求处理器不断地获取和存储大量数据，这降低了运行速度。

查尔姆斯理工大学（Chalmers University of Technology）和格拉斯哥大学（University of Glasgow）的研究人员通过重新设计向量处理器上注意力机制的工作方式解决了这个问题。他们创建了一种名为 FlashAttention-V 的新方法，该方法改变了计算机组织工作的方式，以适应这些处理器的独特形状。这种新方法不再试图强迫处理器一次只处理一小块数据，而是将多个独立的计算组合成一条单一的宽数据行。想象一下工厂的装配线，工人通常一次只能处理一件物品；这种新方法允许他们一次抓取一整托盘的物品并同时进行处理，从而大幅减少了往返于储物架之间的时间。通过重新排列操作顺序并更紧密地打包数据，研究人员发现他们可以使这些较小的设备运行得显著更快，尤其是在处理短文本爆发或设备逐个生成新词时。

团队在包括 TinyLlama、Llama 3.2 和 Qwen2.5 在内的多种语言模型上，利用真实硬件和详细的计算机模拟测试了他们的新方法。在一个名为 Banana Pi BPI-F3 的物理开发板上（该板使用 RISC-V 处理器），新方法证明了巨大的改进。当设备准备读取短输入时，新方法的速度是标准非优化版本的十二到十四倍。当设备逐词生成文本时，速度提高了四到五倍。研究人员还进行了广泛的模拟，以观察如果处理器采用更宽的数据线（能够同时处理更大块的信息）时，该方法表现如何。这些模拟显示，随着数据线变得更宽，速度增益持续攀升，在处理输入的最佳场景下，速度达到了基础版本的四十二倍。

然而，这项研究也揭示了一个关于这些设备能变得多快的明显限制。研究人员发现，虽然程序的注意力部分从这些更宽的数据线中获益匪浅，但系统的其他部分——特别是将数字转换为预测的层——则不然。这些层使用一种称为“量化”（quantization）的特定存储数字方式，通过压缩数据来节省空间。目前这种数据的打包方式与宽数据线之间存在结构性不匹配，迫使处理器进行额外且低效的工作来分离和重新组合这些数字。模拟显示，对于这些特定的层，通过同时处理更多数据所节省的时间完全被重新排列数据所耗费的时间抵消了。这意味着，虽然注意力机制可以变得极其快速，但整个系统的速度目前仍受到这些其他组件的制约，这表明未来的改进将取决于如何改变这些数字的存储方式，而不仅仅是处理方式。

这些发现为使人工智能在日常设备上更普及提供了一条清晰的路径。新方法 FlashAttention-V 成功地弥合了现代语言模型的设计与可扩展向量处理器能力之间的鸿沟。它证明了通过仅仅重新排列计算机处理任务的方式并更高效地打包数据，无需新硬件即可实现显著的性能提升。研究证实，对于短任务和实时文本生成，这些优化后的处理器可以非常有效。然而，它也作为一个警示，指出目前这些设备的压缩数据方式可能正在遭遇瓶颈，并且要释放未来更宽处理器在处理这些压缩数字时的全部潜力，可能取决于解决如何更智能地存储和移动这些数字的难题。

技术摘要：面向可扩展向量架构的 FlashAttention

问题陈述

在 CPU 上进行 Transformer 模型的推理正变得日益重要，特别是对于部署在边缘设备上的小语言模型（SLM）。虽然像 RISC-V 向量扩展（RVV）和 ARM 可伸缩向量扩展（SVE）这样的向量架构提供了极具前景的执行基座，但注意力机制模块仍然是主要的瓶颈，因为它对内存带宽的要求极高。

现有的向量 CPU FlashAttention 实现（例如在 llama.cpp 中）存在一个根本性的可扩展性限制：它们将向量长度（VL）与头维度（ $D$ ）绑定（通常为 64 或 128）。因此，当硬件支持显著大于 $D$ 的向量长度时（例如 512 位、2048 位或 8192 位），这些实现无法利用完整的 SIMD 宽度。它们会饱和在 $VL \le D$ 的状态，导致更宽的向量寄存器利用不足，无法发挥长向量架构的潜力。此外，标准的量化格式（特别是 Q8_0）在线性投影和前馈层中引入了结构上的不兼容性，造成了额外的瓶颈。

方法论

作者提出了 FlashAttention-V，一种专门针对可伸缩向量架构重新设计的 FlashAttention 算法。其核心创新在于利用了跨注意力头（inter-head）的并行性，而不仅仅是在单个头内部进行并行。

关键算法变换

跨头打包（Inter-Head Packing）： 当 $VL > D$ 时，FlashAttention-V 将多个注意力头映射到单个向量寄存器中。通过重排序循环并应用循环展开，该算法将来自多个头的 $Q$ 、 $K$ 和 $V$ 元素打包进连续的向量寄存器。这使得系统能够利用超过头维度的向量长度。
循环重排序与展开： 算法通过重排序循环来提高空间局部性并实现单位步长（unit-stride）的内存访问。它对注意力头进行循环展开，以增加指令级并行性（ILP）并最大化向量寄存器的占用率。
GQA/MHA 感知： 该设计原生支持分组查询注意力（GQA）和多查询注意力（MQA）。对于共享的键值头，算法仅加载一次数据并将其在向量寄存器中进行复制，从而避免冗余的内存访问。
分块执行（Blocked Execution）： 实现采用了受 FlashAttention-2 启发的缓存感知分块策略，以将中间数据保留在寄存器中并减少内存流量。

实验设置

实现： 集成到 llama.cpp 中的 ggml 框架内，使用 RISC-V 和 ARM SVE 指令集。
评估模型： TinyLlama, Llama 3.2, Qwen2.5 (GQA 模型) 以及 Pythia-410M (MHA 模型)。
硬件/仿真：
- 真实硬件： Banana Pi BPI-F3 (RVV 256-bit)。
- 仿真： 配置为 RISC-V 顺序执行 RiscvMinorCPU 的 gem5。该模拟器经过扩展，以模拟与向量长度和通道数成比例的现实向量延迟（基于 Vitruvius+ 微架构），向量长度从 512 位扩展到 8192 位，通道数从 8 个扩展到 64 个。
基准测试： 非向量化 FlashAttention (ggml-scalar) 和现有的向量化 FlashAttention (ggml-vec)。

核心贡献

FlashAttention-V 算法： 一种通过利用跨头并行性，实现有效利用超过头维度向量长度（$VL > D$）的新方法。在 gem5 仿真中，它在预填充（prefill）阶段相比标量 FlashAttention 实现了 22×–42× 的加速，在解码（decode）阶段实现了 8×–11× 的加速（基于 512 位向量长度）。
可扩展性分析： 论文提供了 FlashAttention-V 在高达 8192 位向量长度下的定量分析。研究指出，虽然理论加速比随向量宽度增加，但实际性能受限于向量功能单元的延迟。作者发现，具有 64 个向量通道和 4096 位向量 的配置实现了最优的可扩展性（预填充阶段获得 2×–2.5× 的增益），而超过此点后，由于延迟开销，收益会递减。
量化瓶颈的识别： 作者表征了 Q8_0 量化 与长向量执行在线性投影和前馈层中的结构性不兼容。他们证明了 Q8_0 中权重与缩放因子的交错布局迫使进行显式的打包和掩码归约操作，这在 2048 位 VL 下消耗了高达 60% 的执行周期。这种开销抵消了长向量带来的算术收益，这一发现通过 RVV 和 Arm SVE 均得到了验证。

结果

真实硬件 (Banana Pi BPI-F3)： FlashAttention-V 实现了相比标量 FlashAttention 12×–14× 的加速。它在短上下文（ $N \le 128$ ）下表现最为出色，这在边缘推理中非常常见。
预填充阶段 (仿真)：
- 在 512 位 VL 下，加速比在 22×（TinyLlama）到 42×（Qwen2.5）之间。
- 扩展到 8192 位 VL 为具有较小头维度的模型（可以打包更多头到每个寄存器）提供了额外收益（最高约 3×）。
- 在现实的延迟缩放（64 通道）下，4096 位配置维持了相比 512 位基准 2×–2.5× 的加速，而 8192 位配置由于延迟惩罚出现了收益递减。
解码阶段： FlashAttention-V 在 512 位 VL 下实现了 8×–11× 的加速。然而，由于解码过程的单 Token 特性，性能对向量宽度和通道数的敏感度在超过 512 位后逐渐降低，因为每次迭代的工作量不足以抵消打包/解包的开销。
量化层： 微基准测试显示，Q8_0 线性层中的打包和掩码归约开销在长向量执行时占据了主导地位，阻碍了算术摊销。

意义与主张

本文声称填补了 FlashAttention 与可伸缩向量架构之间的空白，通过实现远超头维度的向量长度的高效执行。作者断言，他们的优化（循环重排序、跨头打包和展开）是可迁移的原则，能够显著加速边缘设备上的 Transformer 推理。

至关重要的是，论文强调了当前量化策略在长向量执行中的一个根本局限性。作者认为，虽然 FlashAttention-V 成功扩展了注意力机制，但 Q8_0 量化线性层的结构布局为前馈网络中的类似扩展构成了障碍。作者得出结论，未来的工作必须探索替代的量化格式或内存布局，以消除显式的打包开销，因为目前的长向量算术收益正被量化模型中的数据移动成本所抵消。

这项工作为硬件设计者和软件工程师提供了具有指导意义的见解，即 64 通道、4096 位向量配置在延迟开销占据主导地位之前，是实现可扩展性的“甜点位”（sweet spot）；同时，仅优化注意力算子对于实现端到端性能提升是不够的，必须解决由量化引起的瓶颈问题。