Attention 优化全景：FlashAttention 的加速魔法与 MHA/GQA/MLA 的结构演进

副标题： Attention 的优化有两个正交方向------FlashAttention 用 Tiling + Online Softmax 省掉了 O(N²) 的中间矩阵；MHA→GQA→MLA 通过减少 KV Head 数量压缩了 KV Cache。两者叠加，才是现代推理引擎的全貌。

一、一个不合理的事实

Transformer 的 self-attention 公式，写过的人都认识：

复制代码

Attention(Q, K, V) = softmax(Q @ K^T / √d) @ V

假设序列长度 N=4096，头维度 d=64：

步骤	操作	输出大小
1	S = Q @ K^T	N × N = 4096 × 4096 = 16M 个元素
2	P = softmax(S)	N × N = 16M 个元素
3	O = P @ V	N × d = 4096 × 64 = 256K 个元素

最终输出只有 256K 个元素，但中间产出了一个 16M 的矩阵。 更极端一点：N=128K 时，attention 矩阵有 160 亿个元素，按 FP16 算就是 32GB------光存这一个中间结果就够爆掉一张 H100。

但真的需要存它吗？

注意一个事实：softmax 是按行独立的------第 i 行的结果只取决于 S[i,:] 和 V，和第 j 行无关。如果你能一行一行地算，算完就扔，就不需要同时持有整个 N×N 矩阵。

问题是：你做不到"一行一行地算"。因为 Q、K、V 本身太大，放在 HBM（显存）里。每次计算前要把数据从 HBM 搬到 SRAM（片上缓存），算完再写回 HBM。如果逐行算：

复制代码

逐行计算（天真方案）:
  for i in range(N):
      q_i = Q[i]        // 从 HBM 读 1 行 Q
      for j in range(N):
          k_j = K[j]    // 从 HBM 读 1 行 K
          s = q_i @ k_j
          P[i,j] = exp(s)
      O[i] = normalize(P[i]) @ V
      // 但每次内层循环都从 HBM 读，带宽浪费严重

每次只算一个元素，却要反复搬整个 K 和 V。带宽利用率极低。

FlashAttention 的答案是：一次搬一小块（tile），所有计算在 SRAM 里完成，中间结果不写回 HBM。

二、先理解战场：SRAM 和 HBM 的"带宽鸿沟"

GPU 的内存分为两级：

类型	容量	带宽	速度
HBM（显存）	大（40-80GB）	~2 TB/s	慢
SRAM（片上缓存）	小（~20MB/SM）	~19 TB/s	快 10×

SRAM 快 10 倍，但只有 20MB------连一个 4096×4096 的 FP16 矩阵（32MB）都装不下。

这就像一个工人（计算单元），工具就在手边（SRAM，快但小），仓库在 100 米外（HBM，大但慢）。如果每拧一个螺丝都要跑仓库拿一个零件，一天拧不了几个螺丝。聪明的做法是一趟搬一箱零件过来，拧完一箱再去搬下一箱。

标准 attention 的问题就是"一趟只搬一个零件"------写完 S 回 HBM，再从 HBM 读回来做 softmax，再写回去，再从 HBM 读回来乘 V。HBM 成了中转站，而中转的东西（S 矩阵）大得离谱。

复制代码

标准 attention 的数据流:

  HBM → SRAM: 加载 Q, K
  SRAM 计算: S = Q @ K^T
  SRAM → HBM: 写回 S (N×N, 巨大！)    ← 瓶颈
  HBM → SRAM: 加载 S
  SRAM 计算: P = softmax(S)
  SRAM → HBM: 写回 P (N×N, 巨大！)    ← 瓶颈
  HBM → SRAM: 加载 P, V
  SRAM 计算: O = P @ V
  SRAM → HBM: 写回 O (N×d, 正常大小)

FlashAttention 的目标：砍掉那两个"巨大"的 HBM 读写。S 和 P 从出生到死亡都在 SRAM 里。

三、第一个技巧：Tiling（分块）

思路很直接：Q 太大？切成小块。K 太大？也切成小块。每次只加载一小块 Q 和一小块 K 进 SRAM，算一小块 S，立刻用完，不存。

复制代码

把 Q 沿行切成 Tr 块，每块 Br 行:  Q₁, Q₂, ..., Q_Tr
把 K 沿行切成 Tc 块，每块 Bc 行:  K₁, K₂, ..., K_Tc
把 V 沿行切成 Tc 块，每块 Bc 行:  V₁, V₂, ..., K_Tc

外层遍历 Q 的块:
  for Q_block in Q_blocks:         // 每次加载 Br × d 的 Q
      for K_block, V_block in zip(K_blocks, V_blocks):
          // 加载 Bc × d 的 K 和 Bc × d 的 V 到 SRAM
          S_block = Q_block @ K_block^T    // Br × Bc，在 SRAM 里
          P_block = softmax(S_block)       // Br × Bc，还在 SRAM 里
          O_block += P_block @ V_block     // 累加到 O_block
      // 写完 O_block 到 HBM

Br 和 Bc 是精心选的值，让 Q_block、K_block、V_block、S_block、P_block 同时塞在 SRAM 里。

问题来了：softmax 不能分块做。

四、第二个技巧：Online Softmax

4.1 为什么 softmax 不能分块？

标准 softmax 分三步，必须看到整行的值：

复制代码

给定一行分数 [x₁, x₂, ..., x_n]:
  (1) m = max(x₁, ..., x_n)           // 找整行的最大值
  (2) y_i = exp(x_i - m)              // 减去 max 防止溢出
  (3) softmax_i = y_i / sum(y₁..y_n)  // 归一化

如果你只看了这行的前一半 [x₁..x_{n/2}]，你不知道最大值是不是在后一半。所以分块直接做 softmax 会出错。

但有一种做法叫 Online Softmax------可以在分块扫描时动态维护"已知的最大值"和"已知的指数和"，每当遇到更大的值，就修正之前的结果。

4.2 Online Softmax 是怎么工作的？

核心思想：扫描过程中，维护三个值：

m：到目前为止见过的最大值
l：以当前 m 为基准的指数和
O：以当前 m 为基准的 softmax 乘 V 的累加结果

每当遇到一个更大的 m ，就对已有的 l 和 O 做一次"缩放修正"。

复制代码

具体步骤（对一个 Q 行，逐 K 块处理）:

初始化: m = -∞, l = 0, O = 0

处理 K 块 j:
  S_j = Q_row @ K_j^T                    // 本块的分数
  m_new = max(m, max(S_j))               // 新的最大值
  
  如果 m_new > m（发现了更大的值）:
      修正因子 = exp(m - m_new)          // < 1，缩小之前的累加
      l = l × 修正因子                    // 旧指数和按比例缩小
      O = O × 修正因子                    // 旧输出按比例缩小
      m = m_new
  
  // 现在以新的 m 为基准，加入本块的贡献
  P_j = exp(S_j - m)                      // 分子
  l_j = sum(P_j)                          // 本块的指数和
  l = l + l_j                             // 更新总指数和
  O = O + P_j @ V_j                       // 累加本块的 softmax×V

当你把所有 K 块都处理完，此时的 O / l 就是精确的 softmax 结果。

4.3 一个具体数值例子

假设一行有三个分数 [2, 5, 3]，分两块处理：[2, 5] 和 [3]。

复制代码

标准 softmax:
  m = max(2,5,3) = 5
  exp: [e^{-3}, e^{0}, e^{-2}] ≈ [0.05, 1, 0.14]
  sum = 1.19
  softmax ≈ [0.042, 0.84, 0.118]

Online Softmax (分块):

处理第 1 块 [2, 5]:
  m = -∞ → 遇到 2,5 → m_new = 5
  修正: m 从 -∞ 变到 5，但 l 和 O 都是 0，无所谓
  P₁ = exp([2,5] - 5) = [e^{-3}, e^{0}] ≈ [0.05, 1]
  l₁ = 0.05 + 1 = 1.05
  O₁ = 1.05（假设 V 对应的累加）
  m = 5, l = 1.05

处理第 2 块 [3]:
  本块最大 = 3 < m=5 → m 不变
  P₂ = exp([3] - 5) = [e^{-2}] ≈ [0.14]
  l₂ = 0.14
  l = 1.05 + 0.14 = 1.19
  O = O₁ + 0.14 × V₂_block

最终:
  softmax = (O₁ + 0.14 × V₂) / 1.19 = 标准 softmax 结果 ✓

关键：当遇到更大的 m 时，旧的累加通过 exp(m_old - m_new) 被"打折"，保证最终结果等价于一次性看到全部数据的标准 softmax。

五、两个技巧合体：FlashAttention 的完整算法

把 Tiling 和 Online Softmax 组装在一起：

复制代码

FlashAttention 前向传播:

输入: Q, K, V ∈ R^{N×d}（在 HBM 中）
输出: O ∈ R^{N×d}（写回 HBM）

把 Q 沿行切成 Tr 块，每块 Br 行
把 K, V 沿行切成 Tc 块，每块 Bc 行
（Br 和 Bc 的选择让 Q_block, K_block, V_block, S_block, P_block 全塞进 SRAM）

for i in range(Tr):
    O_i = zeros(Br, d)         // 输出块
    m_i = -inf                 // 本 Q 块的 per-row max
    l_i = 0                    // 本 Q 块的 per-row sum
    
    加载 Q_i 从 HBM 到 SRAM     // Br × d
    
    for j in range(Tc):
        加载 K_j, V_j 从 HBM 到 SRAM   // Bc × d
        
        S_ij = Q_i @ K_j^T             // Br × Bc, 在 SRAM 中
        
        m_new = max(m_i, row_max(S_ij))     // 逐行更新 max
        
        // 修正旧的 l 和 O（如果有更大的 max 出现）
        correction = exp(m_i - m_new)
        l_i = l_i * correction
        O_i = O_i * correction
        
        // 加入本块的贡献
        P_ij = exp(S_ij - m_new)            // Br × Bc, 仍在 SRAM
        l_i = l_i + row_sum(P_ij)
        O_i = O_i + P_ij @ V_j              // Br × d
        
        m_i = m_new
        
        // S_ij 和 P_ij 从未离开 SRAM ← 关键！
    
    O_i = O_i / l_i                         // 最终归一化
    写 O_i 到 HBM

整个算法中，S 和 P（两个 N×N 矩阵）没有一次被写入 HBM。 它们生存在 SRAM 里，被一块一块地消耗掉。

HBM 读写量对比

	标准 Attention	FlashAttention
读 Q	N×d	N×d
读 K	N×d	N×d × Tc（每块 Q 都要扫一遍 K）
读 V	N×d	N×d × Tc
写 S	N×N	0
读 S	N×N	0
写 P	N×N	0
读 P	N×N	0
写 O	N×d	N×d

FlashAttention 多了 K 和 V 的重复加载（外层 Q 循环每轮都要扫一遍 K），但省掉了 3 次 N×N 的读写。当 N 很大时，这个节省碾压重复加载的成本。

六、反向传播：重算代替存储

训练时除了前向，还需要反向传播算梯度。反向需要用到前向的中间结果------标准做法是把 S 和 P 存下来，反向时从 HBM 读。这又是 N×N 的存储和读取。

FlashAttention 的反向传播用了另一个经典技巧：不存中间结果，反向时重算。

复制代码

标准 attention 训练:
  前向: 算 S → 写 HBM → 算 P → 写 HBM → 算 O
  反向: 从 HBM 读 P → 从 HBM 读 S → 算梯度
  HBM 额外占用: N×N × 2 (S 和 P)

FlashAttention 训练:
  前向: 算 O（不存 S 和 P）
  ← 只把 O 和 softmax 的归一化因子 (m, l) 写回 HBM  ← 只有 O(N×d + N)
  反向: 
    从 HBM 读 O, m, l
    重新跑一遍前向 tiling 流程 → 当场重算 S 和 P
    用重算出的 S 和 P 算梯度

重算的成本和 Tiling 一起被分摊：反向过程复用前向的 tiling 逻辑，每个 Q 块重算对应的 S 和 P 块，算完梯度就扔。额外存储从 O(N²) 降到了 O(N)。

这就是 FlashAttention 能支持超长序列训练的核心原因------训练状态下显存占用由没有 S/P 决定的，不存就省掉了巨大的开销。

七、FlashAttention → FlashAttention-2 → FlashAttention-3

7.1 FlashAttention（2022，Dao et al.）

奠定了两个核心思路：Tiling + Online Softmax + 反向重算。在当时的 A100 上，训练加速 3-7×，让 GPT-4 级别的长上下文训练成为可能。

但初版的并行策略有浪费。 初版沿 batch 和 head 维度做并行------每个 thread block 处理一个 head 的一个 Q 块。当 batch 和 head 数不够大时，thread block 不够多，GPU 的 SM 闲置。

7.2 FlashAttention-2（2023，Dao）

核心改进：把并行从 batch/head 维度挪到序列长度维度。

复制代码

FlashAttention-1:
  外层 for Q_blocks (串行)
    内层 for K_blocks (串行)
  ← Q 块之间不能并行
  
FlashAttention-2:
  for Q_blocks in parallel:        ← Q 块可以并行！
    内层 for K_blocks (串行)

Q 块之间没有依赖（每个 Q 块独立产出自己的 O 块），所以可以全部分配给不同的 thread block 并行处理。这样序列越长、GPU 利用率越高。实测训练加速再提升 2×，总加速达到标准 attention 的 5-9×。

7.3 FlashAttention-3（2024，Dao et al. for Hopper）

针对 H100/H200 的 Hopper 架构做了专门优化：

利用 FP8 的异步数据搬运（TMA，Tensor Memory Accelerator）
将 GEMM 和 softmax 重叠执行（warp specialization：不同 warp 各司其职）
H100 上速度达到 740 TFLOPS（理论峰值的 75%），比 FlashAttention-2 再快 1.5-2×

三版对比

	年份	核心改进	相对上版本提升	总加速（vs 标准）
FlashAttention	2022	Tiling + Online Softmax	---	3-7×
FlashAttention-2	2023	序列维并行 + 减少非乘操作	2×	5-9×
FlashAttention-3	2024	H100 FP8 + 异步 + warp 分工	1.5-2×	~10×+

八、Attention 的结构变体：MHA → MQA → GQA → MLA

前面七节讲的都是"同样的 attention，怎么算得更快"。但 attention 本身的结构也可以变------目的不是加速计算，而是减少 KV Cache 的显存占用。

回顾一下 KV Cache 的构成。每个 transformer 层，对于每个 attention head：

复制代码

KV Cache 大小 = 2 × batch × seq_len × num_layers × num_heads × head_dim × dtype_bytes

对于 Llama 2 7B：32 层、32 个 head、head_dim=128、batch=1、seq_len=4096、FP16：

复制代码

KV Cache = 2 × 1 × 4096 × 32 × 32 × 128 × 2 bytes
         = 2,147,483,648 bytes ≈ 2 GB

2GB 的 KV Cache，仅用于一个请求的 4096 个 token。 高并发场景下，KV Cache 往往比模型权重还占显存。这就是为什么 MQA/GQA/MLA 这些变体会出现------它们都在回答同一个问题：能不能少存几个 K 和 V？

8.1 Multi-Head Attention（MHA）：标准方案，最费显存

标准的多头注意力：Q、K、V 各有 H 个头，独立计算。

复制代码

MHA: Q 有 8 个头，K 有 8 个头，V 有 8 个头

  Q₁ → K₁,V₁ → head₁
  Q₂ → K₂,V₂ → head₂
  Q₃ → K₃,V₃ → head₃
  ...
  Q₈ → K₈,V₈ → head₈

每个 head 独立存一份 KV Cache

优点：每个 head 学习不同的注意力模式（有的看语法、有的看语义、有的看位置），表达能力最强。

缺点：KV Cache 大小和 head 数成正比。32 个 head 就要存 32 组 K、V。

8.2 Multi-Query Attention（MQA）：极端省显存，但伤质量

2019 年 Google 提出：所有 Q head 共享同一组 K 和 V。

复制代码

MQA: Q 有 8 个头，K 和 V 只有 1 组

  Q₁ ─┐
  Q₂ ─┤
  Q₃ ─┤
  ... ├─→ K, V (只有 1 组)
  Q₈ ─┘

所有 Q head 都 query 同一组 KV

优点：KV Cache 缩减为原来的 1/H。32 个 head 的模型，KV Cache 直接降到 1/32。

缺点：表达能力下降。不同 head 本来应该关注不同的语义模式，但被迫共享同一组 K、V 意味着它们只能从"同一个钥匙串"上找不同的齿。PaLM 用了 MQA，但在一些 benchmark 上确实比同规模 MHA 模型差。

8.3 Grouped-Query Attention（GQA）：折中方案

2023 年 Meta 在 Llama 2 中引入：Q head 分成 G 个组，每个组共享一组 K、V。

复制代码

GQA（8 个 Q head，4 个组，每组 2 个 Q）:

  组1: Q₁,Q₂ → K₁,V₁
  组2: Q₃,Q₄ → K₂,V₂
  组3: Q₅,Q₆ → K₃,V₃
  组4: Q₇,Q₈ → K₄,V₄

KV Cache = 标准 MHA 的 G/H = 4/8 = 1/2

方案	KV 数量	KV Cache（相对 MHA）	质量
MHA	每个 head 独立	1×（基线）	最优
GQA	G 组共享	G/H	接近 MHA
MQA	全 head 共享	1/H	有损失

GQA 在实际模型中（Llama 2 70B，8 组）：KV Cache 降到 MHA 的 1/4，但生成质量几乎无损。 因为 8 个 Q head 中，确实有些 head 的注意力模式天然相似------与其让它们各自独立学出来相似的结果，不如直接共享。

目前主流模型大多选 GQA：

Llama 2 70B：8 组
Llama 3：8 组
Mistral：滑动窗口 + GQA
Gemma 2：GQA

8.4 Multi-head Latent Attention（MLA）：DeepSeek 的方案

DeepSeek-V2（2024）提出了 MLA，比 GQA 更进一步：不只是在 head 之间共享 KV，而是把 KV 压缩到一个低维 latent 空间里，推理时从 latent 向量解压出 K 和 V。

复制代码

MHA:  每层存 H 组 KV，每组 d 维 → H × d 维
GQA:  每层存 G 组 KV，每组 d 维 → G × d 维
MQA:  每层存 1 组 KV，d 维       → d 维
MLA:  每层存 1 组 latent，d_latent 维 → d_latent 维（远小于 d）

MLA 的 KV Cache 不是存 K 和 V 本身，而是存一个压缩后的潜在向量 c。推理时：

复制代码

c → W_UK → K（解压）
c → W_UV → V（解压）

d_latent 远小于 num_kv_heads × head_dim。 比如 DeepSeek-V2，MLA 的 KV Cache 只有 MHA 的约 1/7，而且训练时通过特殊的低秩分解保证解压质量。

8.5 结构变体和 FlashAttention 的关系

一个常见的误解是"用了 FlashAttention 就不需要 MQA/GQA 了"------它们是正交的，解决不同的问题。

复制代码

FlashAttention 解决的问题:  attention 的中间矩阵（S, P）太大
  → 解决方法: Tiling + Online Softmax，中间矩阵不写 HBM

MQA/GQA/MLA 解决的问题:  KV Cache 本身太大
  → 解决方法: 减少 KV head 数量 / 压缩 KV 维度

它们可以而且应该叠加使用。 实际上，几乎所有现代模型都同时用了 FlashAttention（加速计算）+ GQA/MLA（节省 KV Cache）。两者不替代，只互补。

	FlashAttention	GQA	MLA
解决什么	中间矩阵读写太贵	KV Cache 存太多	KV Cache 存太多（更进一步）
怎么做的	Tiling + Online Softmax	Q head 分组共享 KV	KV 压缩到 latent 空间
省什么	S/P 矩阵的 HBM IO	num_kv_heads 的倍数	K、V 的维度
需要改模型结构？	不需要	需要（从训练开始）	需要（训练 + 特殊设计）
和 FA 叠加？	---	✅ 完全兼容	✅ 完全兼容

九、影响：FlashAttention 改变了什么

9.1 让长上下文从"显存装不下"变成"只影响计算时间"

标准 attention 训练时，显存中要保留 S 矩阵（N×N）。N=8K 时 S 占 256MB，N=32K 时 S 占 4GB，N=128K 时 S 占 64GB------存中间结果比存模型权重还大。

FlashAttention 把这个 O(N²) 存储降到了 O(N)（只存 O 和 softmax 的归一化因子）。这意味着显存不再限制序列长度，长上下文真正变成了"慢一点但能跑"而不是"直接 OOM"。

9.2 工业界全面采用

框架	支持方式
PyTorch	`torch.nn.functional.scaled_dot_product_attention` 默认调用 FlashAttention
HuggingFace Transformers	`model.to_bettertransformer()` 或自动启用
vLLM	内嵌 FlashAttention / FlashInfer
SGLang	基于 FlashInfer（FlashAttention 的 decode 优化变体）
GPT-NeoX / Megatron-LM	训练标配

9.3 "省掉中间矩阵"成为一类优化的范式

FlashAttention 的成功激发了一类"kernel fusion + tiling"的优化方法。同一思路后来被用于：

FlashInfer：针对 decode 阶段（batch 小、序列长）做了专门优化
FlashDecoding：把长序列的 attention 沿 K/V 维度分块并行，专攻 decode
FlashFFT：用同样的思路优化长序列的卷积计算

十、总结

FlashAttention 不是在优化"怎么算"------它是在优化"不存什么"。 标准 attention 把中间结果当成必须存在 HBM 里的东西；FlashAttention 证明了 S 和 P 矩阵可以当场消费、当场丢弃，从而避开了 O(N²) 的显存读写。

三个层面的理解：

复制代码

表层（现象）：attention 速度提升 5-10×，长序列不 OOM 了
中间层（机制）：Tiling 分块 + Online Softmax 数值技巧 = 全程 SRAM 计算
深层（物理）：用 SRAM 的高带宽做算术，不让中间结果碰 HBM

设计哲学：
  "最快的 IO 是不做 IO"
  标准 attention 的问题不是算力不够，是写了太多用一次就扔的数据
  FlashAttention 的答案：不写、不存、当场算完

FlashAttention 之所以被广泛视为 2020 年代以来 Transformer 基础设施最重要的突破之一，不是因为它加速了多少倍------是因为它把"显存放不下"变成了"只是算得慢一点"，从根本上改变了"能跑多大上下文"的物理约束。

附录：进一步阅读

Dao et al. (NeurIPS 2022): FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness --- 原始论文
Dao (2023): FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
Shah et al. (2024): FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision
Tri Dao 的 FlashAttention GitHub: github.com/Dao-AILab/flash-attention
Zihao Ye 的关于 Online Softmax 的博客（UW）

这个专栏专注于 LLM 推理与部署的底层技术------从算子开发、显存管理、调度策略到生成优化，用可运行的代码和可复现的实验，把论文里的"黑科技"变成可以理解、可以动手验证的知识。如果你对"GPU 到底在干什么"和"推理框架做了什么优化"这类问题感兴趣，欢迎来主页翻翻其他文章。