大模型论文：FlashAttention Fast and Memory-Efficient Exact Attention with IO-Awareness(效率提升)

大模型论文：FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness(效率提升)

文章地址：2205.14135 (arxiv.org)

摘要

Transformer 在处理长序列时速度慢、内存开销大，其原因在于自注意力机制的时间和内存复杂度与序列长度呈二次关系。尽管已有一些近似注意力机制尝试通过牺牲模型精度来降低计算复杂度，但这些方法往往无法在实际运行时间上实现提速。我们认为，这些方法缺少的一个核心原则是 IO 感知（IO-aware） —— 即在算法设计中考虑 GPU 不同层级内存之间的数据读写开销。

我们提出了 FlashAttention，一种 IO 感知的精确注意力算法，通过 分块（tiling）机制 来减少 GPU 高带宽内存（HBM）与片上 SRAM 之间的数据读写操作。我们对 FlashAttention 的 IO 复杂度进行了分析，结果显示它相比标准注意力机制显著减少了对 HBM 的访问次数，并在一定范围的 SRAM 大小下达到了最优。

此外，我们还将 FlashAttention 拓展到 块稀疏注意力（block-sparse attention），得到了一个 近似注意力算法，该算法比现有所有近似方法都要快。

在实际应用中，FlashAttention 在 Transformer 训练中表现出显著优势：

• 在 BERT-large（序列长度 512）上，相比 MLPerf 1.1 的训练速度记录，整体训练时间加速 15%
• 在 GPT-2（序列长度 1000）上，加速达到 3 倍
• 在 Long Range Arena（序列长度 1000–4000）任务中，加速达到 2.4 倍

此外，FlashAttention 及其块稀疏版本还使 Transformer 能够处理更长的上下文，从而带来更高的模型质量：

• GPT-2 上的困惑度（perplexity）提升了 0.7
• 长文档分类任务上准确率提升 6.4 个百分点
• 成为第一个在 Path-X（序列长度 16K，准确率 61.4%） 和 Path-256（序列长度 64K，准确率 63.1%） 任务中取得优于随机猜测表现的 Transformer 模型

简而言之，FlashAttention 通过 IO 感知设计，真正做到了在实际训练时间上实现加速，同时提升了模型的能力与质量。

问题分析

• Transformers在长序列下自注意力机制复杂度为 O(n²)（n为序列长度），导致：

  • 内存爆炸
  • 训练时间剧增
  • 实际训练时间（wall-clock time）并无提升

• 自注意力模块的时间与内存复杂度随着序列长度呈二次增长（O(n²)）。因此一个关键问题是，如何提升注意力机制的速度和内存效率，从而帮助Transformer在长序列场景中更好运行

• 本文认为之前的方法没有考虑到的东西是：使注意力算法具备 IO 感知能力。即：仔细考虑不同类型GPU内存之间的数据读写成本，比如快的SRAM与慢的HBM之间的传输

• 为什么仅优化FLOPs（浮点运算数）不够？

  • 许多近似注意力方法旨在降低注意力计算的计算量和内存需求。这些方法包括从稀疏近似sparse-approximation 到低秩近似low-rank approximation 。尽管这些方法将计算需求降低到与序列长度呈线性或接近线性的关系，但其中许多方法与标准注意力相比，并没有实现实际运行时间的加速，也没有得到广泛应用。一个主要原因是，它们专注于减少浮点运算次数(FLOP reduction)（这可能与实际运行速度并无关联），并且往往忽略了内存访问(memory access, IO）带来的开销

  • 🧠 举个例子：

    你在 Vim 编辑一个巨大的文件，如果每一次操作都要从慢速硬盘读写，哪怕你操作再快也没用。FlashAttention 就是尽量在缓存（SRAM）中完成计算，减少“去硬盘”（HBM）的次数
    

FlashAttention

• 本文提出 FlashAttention，一种在保证准确性的前提下，大幅减少内存访问的注意力机制。目标是避免对 HBM 中的注意力矩阵进行读写操作

• 我们采用两个关键技术(上图)：

  1. 分块计算（Tiling）：将输入分成多个小块（block），逐块处理注意力，从而避免一次性读写整个矩阵
  2. 重计算（Recompute）优化：在前向传播阶段只存储 Softmax 缩放因子，在反向传播时用它来快速重建注意力值，而不是从HBM读取完整矩阵

• 相比于原始Transformer的自注意力，FlashAttention具有以下特点：

  • 在不访问整个输入的情况下计算 softmax reduction(归约)
  • 在反向传播中不存储大型的中间注意力矩阵
  • 重新组织注意力计算过程，将输入分割成块，并对输入块进行多次遍历，从而逐步执行 softmax 归约（也称为平铺）
  • 存储前向传播中的 softmax 归一化因子，以便在反向传播中在芯片上快速重新计算注意力，这比从 HBM 读取中间注意力矩阵的标准方法更快

• 本文用 CUDA 实现 FlashAttention，将所有操作融合在一个 GPU kernel 中。这不仅加速了模型训练（在 GPT-2 上最高提升 7.6 倍），还降低了内存消耗（因 IO 减少而线性增长）

• 复杂度分析：FlashAttention 的HBM访问复杂度为 $O(N^2\cdot d^2\cdot M^{-1})$，其中：$N$ 是序列长度，$d $是头维度，$M $是SRAM大小。而传统注意力需要$ O(Nd + N²) $次 HBM 访问。FlashAttention最多能减少9倍的HBM访问

• 分块（Tiling）机制：将序列按块分割，每次只处理一个 block：

  • 比如：处理一个 1K 的序列时，我们将其分成 8 个 128 长度的 block，分别计算 attention
  • 每次操作只涉及小范围的局部缓存，极大减少了全矩阵访问

• 重计算技巧：FlashAttention 在反向传播时不再读取存储在HBM的注意力矩阵，而是快速用Softmax因子重建结果

• 原始的self-attention实现：给定输入序列 $Q,K,V\in {\mathbb{R}}^{N\times d}$，其中 $N$ 是序列长度，$d$ 是头维度，我们想要计算注意力输出 $O\in {\mathbb{R}}^{N\times d}$：
  $$S=Q{K}^{\top }\in {\mathbb{R}}^{N\times N},P=\mathrm{softmax}(S)\in {\mathbb{R}}^{N\times N},O=PV\in {\mathbb{R}}^{N\times d}$$

  其中 softmax 是按行应用的

  标准注意力机制的实现会将矩阵 $S$ 和 $P$ 存储到 HBM 中，这需要 $\mathcal{O}(N^2)$ 的内存。通常 $N\gg d$（例如，对于 GPT2，$N=1024$ 且 $d=64$）。我们在算法①中描述标准注意力机制的实现。由于部分或大多数操作是内存密集型的（例如 softmax），大量的内存访问导致实际运行时间较长。

• 分块 tiling：我们按块计算注意力。Softmax 会关联 $K$ 的列，所以我们通过缩放来分解大型的 softmax 计算。为了数值稳定性，向量 $x\in {\mathbb{R}}^B$ 的 softmax 计算如下：

$$m(x):=\underset{i}{\max }{x}_i,f(x):=\left[e^{x_1-m(x)},\dots ,e^{x_B-m(x)}\right],\ell (x):=\sum _{i}f(x{)}_i,\mathrm{softmax}(x):=\frac{f(x)}{\ell (x)}$$

• 对于向量 $x^{(1)},x^{(2)}\in {\mathbb{R}}^B$，我们可以将拼接后的向量 $x=\left[x^{(1)}{x}^{(2)}\right]\in {\mathbb{R}}^{2B}$ 的 softmax 分解为：

$$m(x)=m\left(\left[x^{(1)}{x}^{(2)}\right]\right)=\max \left(m\left(x^{(1)}\right),m\left(x^{(2)}\right)\right),f(x)=\left[e^{m(x^{(1)})-m(x)}f\left(x^{(1)}\right)e^{m(x^{(2)})-m(x)}f\left(x^{(2)}\right)\right]$$

$$\ell (x)=\ell \left(\left[x^{(1)}{x}^{(2)}\right]\right)=e^{m(x^{(1)})-m(x)}\ell \left(x^{(1)}\right)+e^{m(x^{(2)})-m(x)}\ell \left(x^{(2)}\right),\mathrm{softmax}(x)=\frac{f(x)}{\ell (x)}$$

• 因此，如果我们跟踪一些额外的统计信息（如 $(m(x),\ell (x))$），就可以一次计算一个块的 softmax。我们将输入 $Q,K,V$ 分割成块，计算 softmax 值以及额外的统计信息，然后合并结果。

• 重计算 Recomputation：我们的一目标之一是不为反向传播存储 $\mathcal{O}(N^2)$ 的中间值。反向传播通常需要矩阵 $S,P\in {\mathbb{R}}^{N\times N}$ 来计算关于 $Q,K,V$ 的梯度。然而，通过存储输出 O 和 softmax 归一化统计信息 $(m,\ell )$，我们可以在反向传播中从 SRAM 中的 $Q,K,V$ 快速轻松地重新计算注意力矩阵 $S$ 和 $P$。这可以看作是一种选择性梯度检查点技术。虽然梯度检查点技术已被提出用于减少所需的最大内存 ，但据我们所知，所有实现都必须以速度换取内存。相比之下，即使浮点运算次数增加，由于减少了 HBM 访问，重计算仍加快了反向传播

• 分块使我们能够在一个CUDA内核中实现我们的算法，从HBM加载输入，执行所有计算步骤（矩阵乘法、softmax、可选的掩码和随机失活、矩阵乘法），然后将结果写回HBM。这避免了反复从HBM读取输入和向HBM写入输出

• 具体算法过程如下

  

与原始的self-attention的对比

一、传统 Attention 是怎么做的？

以 Scaled Dot-Product Attention 为例，传统计算流程如下：

1. 给定输入向量 $X$，通过线性层变换成 $Q,K,V$：

   $Q=X{W}^Q,K=X{W}^K,V=X{W}^V$

2. 计算注意力分数：

   $A=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)$

3. 加权求和：

   $\text{Attention}(Q,K,V)=A\cdot V$

存在的问题：

• 内存开销大：计算 $Q{K}^T$ 会产生一个 $(L,L)$ 的矩阵（L为序列长度），特别是当 L 很大时，内存需求巨大。
• 速度慢：现代 GPU 的 SRAM/寄存器非常快，但传统 attention 使用大量 HBM（高带宽内存）进行读写，频繁的内存交换导致效率低。

二、FlashAttention 是怎么做的？

FlashAttention 的核心思想是：将 attention 的所有操作融合在一个 GPU kernel 中，并利用块状计算（block-wise computation）减少内存读写，提高缓存命中率。

关键优化：

1. Block-wise 计算（分块 attention）

将序列拆分成小块，比如 $16\times 16$ 的子块，逐块计算 Attention，而不是一次性构建整个 $L\times L$ 的注意力矩阵。

2. 避免中间结果写入显存

传统 attention 会先计算出 $Q{K}^T$ 的完整矩阵，然后再做 softmax、再乘以 $V$。FlashAttention 直接在寄存器中计算每一步，并将 softmax 和乘 V 合并在同一个 kernel 内完成，省去了中间写入显存。

3. 精确 softmax 计算（numerical stable）

采用 在线算法（online algorithm） 来维护 softmax 所需的最大值和归一化因子，避免出现数值不稳定。

三、与原始 Attention 的区别（对比）

四、举个例子：对比原始与 FlashAttention 的执行流程

假设输入序列长度为 8（简化版）：

原始 Attention：

Q = X @ Wq
K = X @ Wk
V = X @ Wv
scores = Q @ K.T / sqrt(dk)
weights = softmax(scores)
output = weights @ V

这其中 scores 是完整的 $(8,8)$ 矩阵，写入显存后再进行下一步。

FlashAttention（伪代码概念）：

for q_block in blocks(Q):     # 取出 Q 的小块
    for k_block, v_block in blocks(K), blocks(V):
        scores = q_block @ k_block.T / sqrt(dk)
        max_score = max(scores)
        exp_score = exp(scores - max_score)      # 减 max 防止溢出
        norm = sum(exp_score)
        softmax_score = exp_score / norm
        output_block += softmax_score @ v_block

• 所有中间值都保存在 GPU 的寄存器或 L1/L2 cache 中，不写回显存
• softmax 的最大值和归一化分母是 在线累积 的，不需要完整的 attention map

总结

FlashAttention 是对 Attention 的底层实现进行优化，并未改变公式本身，而是改变了执行方式，实现高效的内存利用和更快的执行效率。尤其适用于大模型（如 GPT）在训练或推理长序列时的应用场景。

为什么要分块和重计算

为什么要分块？

HBM 是大容量但速度慢的内存，SRAM 是容量小但速度快。FlashAttention 的目标是：在SRAM里局部处理 attention，从而大幅减少 HBM 读写开销。

例子：假设序列长 N=4096，d=64。我们将其拆为：

• Q 分成 Tr = 64 块（每块64条向量）
• K, V 分成 Tc = 64 块
• 每次只需将 1 个 Q 块和 1 个 K/V 块从HBM加载到SRAM进行计算

这样就把原来的 O(N²) 读写复杂度优化为 O(N²·d²/M)，M 为SRAM大小

为什么要重计算？

在标准 attention 中，反向传播时需要使用前向的 S=QKᵀ，需要保存整个 S∈ℝⁿˣⁿ，占用巨大内存。 FlashAttention 中，只保存 Softmax 的归一化因子与部分最大值信息，在反向时快速重建，不需存整张矩阵