PaperReading:FlashAttention

原创 paperreading

1. 背景

在现代计算机架构中，内存层次结构是影响系统性能的核心因素之一。从PC到服务器，再到用于AI和高性能计算的GPU加速器，不同层次的内存（如SRAM、HBM、DRAM、Cache）扮演着不同的角色。在介绍 FlashAttention 之前，我们先来了解一下 CPU、GPU、SRAM、HBM、DRAM、L1/L2/L3 Cache 等概念，并解释它们之间的区别，因为 FlashAttention 的优化正是基于这些概念。

1. CPU 与 GPU：计算核心的分工

• CPU（中央处理器）

  • 功能：负责通用计算任务，包括系统控制、逻辑运算、指令调度。
  • 特点：单线程性能强、延迟低、适合复杂逻辑控制和低延迟任务。
  • 内存结构：配备多级缓存（L1/L2/L3 Cache）和主存（DRAM），用于提升数据访问速度和存储容量。

• GPU（图形处理器）

  • 功能：专为并行计算设计，适合图形渲染、深度学习等大规模数据处理。
  • 特点：多核心（数千线程）、高吞吐量、专注浮点计算和大规模数据并行。
  • 内存结构：拥有高速片上缓存（SRAM）和大容量显存（HBM），用于快速存取计算数据。

2. 内存层次结构：SRAM、HBM、DRAM 和 Cache

内存越靠近核心速度越快、容量越小

• SRAM（Static Random Access Memory）：GPU核心内部。速度极快、功耗低，但造价高、容量小（KB~MB级）。
• HBM（High Bandwidth Memory）：GPU核心外。大容量（几十GB）、高带宽（TB/s），比DRAM快但比SRAM慢。存放模型权重、激活、中间数据（Transformer中的Q/K/V矩阵等）。
• DRAM（Dynamic Random Access Memory）：计算机主板上的物理内存（内存条）。容量大（GB~TB），速度比SRAM和HBM慢，但造价低。存放操作系统、应用程序和数据，当GPU显存不足时也会使用主机DRAM。
• Cache（缓存）层次：L1/L2/L3：在CPU核心内（L1/L2）或共享（L3）。速度最快、容量最小（几十KB），每核独享。稍慢、容量稍大（几百KB~MB）。

3. SRAM 和 HBM 的角色

• SRAM主要用于计算高速缓存，用于存放当前正在计算的小数据块（例如Q/K/V的分块），负责「计算时手边处理」；
• HBM主要用于存储大量数据，用于存放大规模模型参数、输入数据、激活等，负责「存放数据备用」。

特性	SRAM（片上缓存）	HBM（高带宽内存）
位置	核心内部（紧邻计算单元）	核心外部（与计算单元连接）
容量	小（KB~MB）	大（几十GB）
速度	非常快（低延迟）	快但比SRAM慢
用途	缓存计算所需数据块	存储大量输入、中间结果
类比	计算工作台（临时用具）	仓库（大批物品）

2. FlashAttention v1

2.1 问题

以前（没有FlashAttention时）大多数深度学习框架和Transformer实现把计算交给HBM，但实际上，SRAM在GPU内部始终是存在的，但是未利用好。

• 深度学习框架的编程接口不够底层
  • PyTorch和TensorFlow等高层框架对开发者屏蔽了底层内存调度和管理（只要“声明Tensor”，自动放到GPU上，系统调度HBM和缓存）。
  • 很多常见的Attention实现（比如PyTorch标准Attention）并没有显式将数据分块调度到SRAM，而是默认HBM和计算单元直接传递大块数据。
• SRAM在传统Attention中：存在但“被动”
  • SRAM作为缓存始终在工作，存取的Q、K、V矩阵和中间结果未必能命中SRAM
  • 但是，这种缓存是自动的、透明的，不由开发者或算法显式控制（不像FlashAttention那样设计了分块计算和调度）；
  • 对于像Attention这样大规模、高维度、长序列的任务，SRAM容量太小，缓存命中率有限，大部分数据依然要回退到HBM访问，性能提升有限。

FlashAttention主动设计了“把分块数据放进SRAM计算”，让SRAM不再只是“被动缓存”，而是主动参与调度和计算。

论文中证明了 FlashAttention 在访问高带宽内存（HBM）时的复杂度是：
$$
O \left( \frac{N^2 d^2}{M} \right)
$$

这里：

• N：输入序列长度（token数）；
• d：注意力头的维度（例如64）；
• M：SRAM（片上缓存）的大小。

这意味着，序列 N 越长，IO 访问次数越多（但比传统方法少很多）；片上缓存（SRAM）越大 M，每次能加载更多数据块，访问次数就越少。相比之下，传统标准Attention的HBM访问复杂度是：
$$
\Omega(N d + N^2)
$$
也就是说，标准Attention需要：线性访问Q/K/V（$Nd$），存取完整 $N \times N$ 的注意力矩阵（$N^2$）。

$O$ 和 $\Omega$ 的含义

1️⃣ $O\bigl(\frac{N^2 d^2}{M}\bigr)$，表示上界，用于描述算法在最坏情况下的增长速度，意为“最多是这个量级”：
- 这里表示 FlashAttention的HBM访问复杂度的渐进上界。
- 意味着当序列长度 $N$、head维度 $d$、SRAM大小 $M$ 越来越大时，HBM的访问次数最多是 $\frac{N^2 d^2}{M}$ 级别。

2️⃣ $\Omega(Nd + N^2)$，表示下界，用于描述算法在最好的情况下的增长速度，意为“至少是这个量级”：
- 这里描述 标准Attention的HBM访问复杂度下界。
- 意味着无论你怎么优化，HBM访问次数至少是 $Nd + N^2$ 级别。

2.2 标准注意力实现

给定输入序列 $Q, K, V \in \mathbb{R}^{N \times d}$，其中 $N$ 是序列长度，$d$ 是head维度，我们希望计算注意力输出 $O \in \mathbb{R}^{N \times d}$：
$$
S = QK^\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}
$$
标准的注意力实现会将中间矩阵 $S$ 和 $P$ 存储在 HBM（高带宽内存）中，这会占用 ${O}(N^2)$ 的内存。通常 $N \gg d$（例如在 GPT-2 中，$N=1024$, $d=64$）。

由于大部分操作是内存受限（例如 softmax），大量的内存访问会导致运行时间变慢。这个问题在对注意力矩阵应用其他逐元素操作（比如对 $S$ 进行掩码或对 $P$ 进行dropout）时更加严重。

2.3 解决方案

在Transformer的注意力中，Softmax操作对中间矩阵 $S=QK^\top$ 进行归一化。但直接计算完整 $S$ 会占用大量内存。FlashAttention采用 分块（tiling） 和 重计算（recomputation） 技术，避免存储整个矩阵

1. 软最大值（Softmax）计算公式
对于向量 $x \in \mathbb{R}^B$：
$$
m = \max_i x_i ~~~~~ f_i = e^{x_i - m} ~~~~~ l = \sum_i f_i ~~~~~ \mathrm{softmax}(x) = \frac{f}{l}
$$

为了防止数值溢出，先减去最大值 $m$，再进行指数和归一化。

2. 块分解计算

如果将向量分为两块 $x^{(1)}, x^{(2)}$，拼接后 $x = [x^{(1)}, x^{(2)}]$：

• 计算块内最大值 $m_1 = \max x^{(1)}, m_2 = \max x^{(2)}$；
• 全局最大值 $m = \max(m_1, m_2)$；
• 分别计算块内指数和 $l_1, l_2$；
• 全局和：$l = e^{m_1 - m} l_1 + e^{m_2 - m} l_2$

合并后的Softmax：
$$
\mathrm{softmax}(x) = {\bigl[e^{x{(1)} - m}, e^{x{(2)} - m} \bigr]}/{l}
$$

这里的 $x{(1)}$ 指的就是 $x^{(1)}$，$x{(2)}$ 指的就是 $x^{(2)}$ （渲染问题）

📚 例子：四维向量Softmax分块计算

假设我们有一个长度为4的向量：
$$
x = [2, 5, 3, 1]
$$

将它分成两块：$x^{(1)} = [2, 5]$，$x^{(2)} = [3, 1]$

1. 计算第一块 $x^{(1)}$

• $m_1 = \max(2, 5) = 5$
• $f_1 = [e^{2-5}, e^{5-5}] = [e^{-3}, e^{0}] \approx [0.0498, 1]$
• $l_1 = 0.0498 + 1 = 1.0498$

2. 计算第二块 $x^{(2)}$

• $m_2 = \max(3, 1) = 3$
• $f_2 = [e^{3-3}, e^{1-3}] = [1, e^{-2}] \approx [1, 0.1353]$
• $l_2 = 1 + 0.1353 = 1.1353$

3. 合并两块

• 全局最大值 $m = \max(m_1, m_2) = \max(5, 3) = 5$
• 对每块指数值重新缩放：
  • 对块1：$f_1 = e^{m_1 - m} \cdot f_1 = e^{5-5} \cdot [0.0498, 1] = [0.0498, 1]$（不变）
  • 对块2：$f_2 = e^{3-5} \cdot [1, 0.1353] = e^{-2} \cdot [1, 0.1353] \approx [0.1353, 0.0183]$
• 合并和：
  $$
  l = l_1 \cdot e^{m_1 - m} + l_2 \cdot e^{m_2 - m} = 1.0498 \cdot 1 + 1.1353 \cdot e^{-2} \approx 1.0498 + 1.1353 \cdot 0.1353 \approx 1.0498 + 0.1537 = 1.2035
  $$
• 最终Softmax：
  $$
  \mathrm{softmax}(x) = \frac{[0.0498, 1, 0.1353, 0.0183]}{1.2035} \approx [0.041, 0.831, 0.112, 0.015]
  $$

整个 FlashAttention 的算法过程如下：

本文由 Yonghui Wang 创作，采用 知识共享署名4.0 国际许可协议进行许可
本站文章除注明转载/出处外，均为本站原创或翻译，转载前请务必署名
最后编辑时间为: May 31, 2025 03:23 pm