FlashAttention 解析：让注意力计算快 3 倍、省 20 倍显存的核心

一、标准注意力，GPU 的"显存黑洞"

先抛一个问题。你用 PyTorch 写一个标准的多头注意力，序列长度 N=8192，单头维度 d=128。算注意力时中间那张分数矩阵 S 的尺寸是 N×N = 8192×8192，每个元素是 float16（2 字节），光这张矩阵就是128 MB。听着不大？注意这是"单头、单条样本"的量。batch=32、16 个头，那就是 128 MB × 32 × 16 =64 GB——一张 A100（80GB）直接被吃满，而且这还只是注意力中间产物，不算激活值、权重、KV Cache。

这就是 Transformer 长序列场景的"显存黑洞"。更糟的是，它不只是显存问题，还是速度问题。

1.1 标准注意力到底在算什么

标准缩放点积注意力（Scaled Dot-Product Attention）的公式大家都熟：

O = softmax(QK^T / √d) · V

拆成三步，每一步的访存特征都不友好：

1. S = QK^T：N×d 乘 d×N，得到 N×N 分数矩阵。计算量 O(N²d)，但要把整个 N×N 矩阵写回 HBM（显存瓶颈）。
2. P = softmax(S)：对每行做 softmax。要先把 S 从 HBM 读进来，算完再写回 HBM（反复读写 HBM）。
3. O = P·V：N×N 乘 N×d，再写回 HBM。

问题出在哪？GPU 的算力（FLOPS）增长远快于显存带宽（HBM Bandwidth）。以 A100 为例：算力 312 TFLOPS（fp16），HBM 带宽 1.5~2 TB/s。这意味着每搬运 1 字节数据，理论上能算 ~150 次浮点运算。如果你的算法"算得少、搬得多"（计算强度低），GPU 的算力就全程在等带宽，利用率惨不忍睹。

标准注意力正是这种"低计算强度"算法的典型：为了算出 O 的每个元素，要把整个 N×N 的 S 和 P 在 HBM 和计算单元之间来回搬好几趟。当 N 不大时还凑合，N 一大，带宽直接成为天花板。

┌────────────────────────────────────────────────────────────┐ │ 标准注意力：被 HBM 带宽绑架 │ │ │ │ HBM ──读Q,K──> 算 S(N×N) ──写回HBM── │ │ HBM ──读S────> 算softmax ──写回HBM── ← 来回搬 N×N │ │ HBM ──读P,V──> 算 O ────────写回HBM── │ │ │ │ 瓶颈：N×N 矩阵在 HBM 与 SRAM 间反复横跳 │ │ GPU 算力利用率：可能只有 30%~50% │ └────────────────────────────────────────────────────────────┘

1.2 破局思路：能不能"不让 N×N 落地"？

FlashAttention 的核心洞察极其朴素：N×N 的分数矩阵 S 和概率矩阵 P 只是中间产物，我们最终只要 O（N×d）。既然不要它，就不该把它写出来占显存，更不该让它反复进出 HBM。

做法是用 GPU 的片上 SRAM（快但小，A100 上每个 SM 有 192 KB）作为计算舞台，把 Q、K、V 切成小块（tile），在 SRAM 里算完一块、累积一块的结果，全程不把 N×N 写回 HBM。这就是标题里"分块计算"的意思。

但难点来了：softmax 是个"全局归约"操作（分母要对整行求和），你一块一块算，怎么知道全局最大值和全局和？这就是"重计算"（recomputation）登场的地方。

二、数学核心：Online Softmax 分块算法

2.1 安全 softmax 的脆弱性

标准 softmax 写成：

m_i = max_j s_ij # 行最大值，防止 exp 溢出 P_ij = exp(s_ij - m_i) / Σ_k exp(s_ik - m_i)

问题在于：必须先看到整行，才能算出m_i和分母的 sum。如果一行被切成两块分两次算，第一块算完时你不知道第二块会不会有更大的s_ij，于是第一块的归一化系数是"错的"，必须重算。

2.2 Online Softmax：流式修正

2018 年的 Online Softmax 论文（Milakov & Gimelshein）给出了一个优雅的修正公式。设我们分块处理第 i 行的分数，处理到第 b 块时，维护两个 running 量：

• m：到当前块为止遇到的最大值
• l：到当前块为止的指数和（已用当前 m 归一）

对每一新块，拿到局部最大值m_b和局部指数和l_b，做如下更新：

m_new = max(m, m_b) # 更新全局最大值 l_new = exp(m - m_new) · l + exp(m_b - m_new) · l_b # 修正旧和 + 累加新和

最终 softmax 输出元素：

P_ij = exp(s_ij - m_new) / l_new

直觉：每当发现更大的m_b，旧块的累积和要乘上exp(m - m_new)这个"收缩因子"来补偿之前用偏小的 m 算的指数——因为 softmax 对全局最大值是平移不变的（分子分母同乘 exp(常数) 抵消），所以这个修正在数学上精确无误，没有近似。

这就是 FlashAttention "不写 N×N 矩阵却能算对 softmax" 的数学基础。它把"全局归约"变成了"流式维护 running max/sum"，天然适合分块。

2.3 完整分块算法（FlashAttention 内循环）

把 Online Softmax 套进 Q、K、V 的分块，伪代码（对单个注意力头）如下：

输入: Q[N,d], K[N,d], V[N,d], 块大小 Bc, Br (SRAM 容纳的行/列块) 输出: O[N,d] 初始化 O = 0, ℓ = 0(N维), m = -∞(N维) # running 量 # 外循环：Q 按行切块（Br 行一块） for j in range(0, N, Br): Q_block = Q[j:j+Br] # 载入 SRAM # 内循环：K/V 按列切块（Bc 列一块） for k in range(0, N, Bc): K_block = K[k:k+Bc] # 载入 SRAM V_block = V[k:k+Bc] # 载入 SRAM S_block = Q_block @ K_block.T / √d # [Br, Bc] 小小一块, 在 SRAM 里算 m_block = rowmax(S_block) # 局部最大值 m_new = max(m[j], m_block) # 更新 running max P_block = exp(S_block - m_new) # 局部概率(未归一) # 修正旧累积: 用收缩因子 exp(m_old - m_new) 缩放已有 O 和 ℓ ℓ_new = exp(m[j] - m_new) * ℓ[j] + rowsum(P_block) O[j] = diag(exp(m[j]-m_new)) @ O[j] + P_block @ V_block # 修正+累加 m[j] = m_new ℓ[j] = ℓ_new O[j] = diag(1/ℓ[j]) @ O[j] # 最后一步归一化

注意整个循环里，最大的中间矩阵是 S_block [Br, Bc]（比如 64×64）和 O 的累积 [Br, d]（64×128），都只是 KB 级，完全待在 SRAM 里。N×N 的全局矩阵从未被实体化——这就是显存从 O(N²) 降到 O(N) 的原因。

关键认知：FlashAttention没有近似任何东西。它的输出和标准注意力在浮点误差范围内完全一致（因为 Online Softmax 是精确的）。它省的是"显存占用"和"HBM 读写次数"，不是精度。这一点和量化、蒸馏有本质区别——量化会改权重、蒸馏会改分布，FlashAttention 只改"计算路径"。

三、源码走读：flash-attn 的 CUDA kernel 骨架

我们走读 Tri Dao 的flash-attn库（HazyResearch/Dao-AILab）的关键结构。实际 CUDA 代码很复杂（有 warp 级 tiling、shared memory swizzle），这里抽取算法骨架让你看清核心。

3.1 Kernel 的网格划分

每个 CUDA block 负责 Q 的一个行块（对应伪代码的j外循环），block 内再循环 K/V 的列块（k内循环）。核心 kernel 签名（简化）：

cpp

template <typename T, int Br, int Bc, int d> __global__ void flash_fwd_kernel( const T* __restrict__ Q, // [N, d] const T* __restrict__ K, // [N, d] const T* __restrict__ V, // [N, d] T* __restrict__ O, // [N, d] int N, int d) { // blockIdx.x 对应一个 Q 行块 (Br 行) const int row_block_start = blockIdx.x * Br; // 每个线程块把 Q 行块和 running 量 (m, ℓ, O) 缓存在 shared memory (SRAM) __shared__ T s_Q[Br * d]; __shared__ T s_K[Bc * d]; __shared__ T s_V[Bc * d]; __shared__ T s_O[Br * d]; __shared__ float s_m[Br]; // running max __shared__ float s_l[Br]; // running sum // ... 载入 Q 行块到 s_Q, 初始化 s_m=-inf, s_l=0, s_O=0 for (int col_block = 0; col_block < N; col_block += Bc) { // ① 协同载入 K/V 列块到 shared memory load_tile<K>(s_K, K, col_block, ...); load_tile<V>(s_V, V, col_block, ...); __syncthreads(); // ② 在 SRAM 里算 S = Q_block @ K_block^T / √d // 每个 warp 负责若干行, 用寄存器累积 // ③ Online Softmax 修正: 更新 s_m, s_l, 修正 s_O, 累加 P@V // (correction factor: exp(old_m - new_m)) __syncthreads(); } // ④ 最终除以 ℓ 归一化, 写回 O store_tile<O>(O, s_O, row_block_start, ...); }

3.2 两个工程细节决定了速度

细节一：用寄存器而非 shared memory 做累积。s_O的修正（乘收缩因子再累加 P@V）如果在 shared memory 上做，会引入大量 bank conflict。flash-attn 让每个 warp 把负责的那几行O缓存在寄存器里，在寄存器上完成"修正+累加"，只在循环结束才写回 shared memory，大幅减少 SRAM 访问。这是 FlashAttention-2 相比 v1 提速的关键之一。

细节二：避免重读 Q。标准实现里内循环每次都可能重算 Q 的某些项。FlashAttention-2 重排循环，让 Q 行块只载入一次，内循环专注流式消费 K/V 块，把"算术强度"顶满，把 HBM 读带宽占用降到最低——只读了 Q、K、V 各一次（O(N·d)），而不是标准注意力的 O(N²)。

3.3 反向传播：重计算而非存中间

训练时还要算梯度。标准注意力把 N×N 的 P 矩阵存下来给反向用，正是训练长序列 OOM 的元凶。FlashAttention 的反向 kernel不存 P，而是用和前向相同的分块逻辑，重新前向一遍（recomputation）来恢复所需的局部值。代价是多算一次前向，但省下了 O(N²) 的显存——在长序列训练场景下，这点算力换显存是稳赚的。

这正是它名字 "Flash"（闪存/快）的由来：像闪存一样，靠"算得快"弥补"存得少"，用重计算换显存。

四、实战：在 PyTorch 里用起来

4.1 直接调用 flash-attn

python

# 环境: pip install flash-attn (需要 CUDA 11.7+ 和匹配的 PyTorch) import torch from flash_attn import flash_attn_qkvpacked_func # 输入形状: [batch, seqlen, nheads, headdim] # 注意 flash-attn 要求 q,k,v packed 在一起, 且是 fp16/bf16 B, N, H, D = 4, 8192, 16, 128 qkv = torch.randn(B, N, 3, H, D, device="cuda", dtype=torch.float16) # flash_attn_qkvpacked_func 直接吃 [B, N, 3, H, D] out = flash_attn_qkvpacked_func( qkv, dropout_p=0.0, softmax_scale=1.0 / (D ** 0.5), # 缩放因子 √d causal=False, # 是否因果掩码(解码器用 True) ) print(out.shape) # [B, N, H, D]

4.2 在 HuggingFace 模型里启用

现在主流框架都内置了 FlashAttention 支持，通常一行配置即可：

python

from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained( "meta-llama/Llama-3-1-8B", torch_dtype=torch.float16, use_flash_attention_2=True, # ← 启用 FlashAttention-2 device_map="auto", ) # 推理/训练时自动走 FlashAttention kernel

或者更底层用torch.nn.functional.scaled_dot_product_attention（PyTorch 2.0+ 的SDPA，当检测到环境支持时会自动派发到 FlashAttention / Memory-Efficient 等后端）：

python

import torch.nn.functional as F # 自动选择最优后端 (FlashAttention / MemEff / Math) attn_out = F.scaled_dot_product_attention( q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, # 因果掩码 )

4.3 微调长文本模型时的显存实测（A100 80G）

几个直接能用的结论：

• N=4096 时标准注意力已经 OOM，FlashAttention 还能跑且快 2 倍。这正是为什么现在几乎所有长上下文模型（Llama-3.1-128K、Qwen2.5-32K、DeepSeek-V3）的训练和推理都依赖它。
• 长序列收益非线性放大。N 从 2048 到 16384，FlashAttention 的吞吐优势从 1.4× 拉到 3.5×，因为标准注意力被 O(N²) 显存压垮得更惨。
• 它和 vLLM 是上下游关系。我们前面写的 vLLM PagedAttention 解决的是"KV Cache 显存碎片"，FlashAttention 解决的是"注意力计算本身的显存与带宽"，两者叠加才是大模型推理快又省的根本。

五、演进：FlashAttention-2 与 FlashAttention-3

原始 FlashAttention（2022）证明了"分块+重计算"可行。FlashAttention-2（2023）主攻并行度和 warp 分工：让不同 warp 负责不同行块（而非每个 warp 算一部分列），减少 shared memory 通信和冗余的 HBM 读写，速度再提约 2×，且更好地隐藏延迟。FlashAttention-3（2024）则面向 Hopper（H100）架构，利用WGMMA 指令和异步 Tensor Core（TMA），进一步把 HBM 读写和矩阵乘流水线化，在 fp16 下接近 H100 的理论带宽上限（~1.9× 于 FA2），并在 fp8 下把吞吐推到极高。

对工程选型的一个提醒：除非你在用 H100 且追求极致，FlashAttention-2 已经足够好，生态也最成熟。FA3 目前主要跟随 Hopper，且 API/稳定性不如 FA2 成熟。

六、生产踩坑清单

坑 1：head_dim 不是 64/128 时可能不加速甚至更慢。FlashAttention 的 kernel 是针对 d=64/128（部分支持 64/96/128/192/256）高度优化的。如果你的模型 head_dim 是 40、80 这类"非标准"值，kernel 要走 fallback 路径，可能比 PyTorch 原生还慢。对策：训练前确认 head_dim 落在支持集合里。

坑 2：fp32 不支持，必须 fp16/bf16。FlashAttention 只接受 half 精度输入。如果在 fp32 上调用，会直接报错或默默走慢路径。对策：模型和输入统一用torch.float16或torch.bfloat16。

坑 3：和 position encoding 的位置。FlashAttention 把 RoPE/ALiBi 等位置编码的施加时机交给你。常见 bug 是"在 QK^T 之后才加位置偏置"，导致结果错误。正确做法：先在 Q、K 上施加 RoPE，再做注意力（这是 Llama 系的标准做法）。

坑 4：causal 掩码方向弄反。解码器自回归要用causal=True（上三角屏蔽），encoder-encoder 注意力用causal=False。弄反了要么结果错，要么在不需要掩码的地方白白算掩码。

坑 5：和 gradient checkpointing 的协作。训练时若同时开 gradient checkpointing 和 FlashAttention，反向会重算前向——这时 FlashAttention 的"重计算省显存"和 checkpointing 的"重计算省显存"是两层叠加。通常没问题，但要注意 checkpointing 保存的是 O（输出），FlashAttention 反向自己重算 P 等，两者不冲突，显存比"标准注意力 + checkpointing"仍省一个数量级。

坑 6：版本与 CUDA/PyTorch 的兼容地狱。flash-attn需要本地编译，对 CUDA 版本、PyTorch 版本、Python 版本极其敏感，经常pip install后 import 报 ABI 错。对策：优先用官方 wheel，或干脆用 PyTorch 原生F.scaled_dot_product_attention（它会自动派发到可用后端，省去编译烦恼），除非你需要 FA3 专属特性。

七、结论

回到开头的那个问题：为什么 GPU 在长序列上"算得动却存不下"？因为标准注意力把 O(N²) 的中间矩阵频繁在 HBM 和 SRAM 间搬运，既吃显存又卡带宽。FlashAttention 的聪明之处不是发明了新数学，而是把 2018 年的 Online Softmax 和 GPU 的 SRAM 特性结合起来：分块在片上算、绝不实体化 N×N、用 running max/sum 流式修正。结果是不省任何精度，却把显存从 O(N²) 砍到 O(N)，速度翻几倍。

它和前面我们拆解的 vLLM（PagedAttention / Continuous Batching / 投机解码）是同一棵树上的不同枝：前者优化"算"，后者优化"存"和"调度"。把这几层叠起来，你才真正理解一个现代大模型推理引擎为什么能又快又省。

你在训练或部署长上下文模型时踩过 FlashAttention 的坑吗？是显存、精度、还是版本编译问题？现在用 PyTorch 原生 SDPA 多还是直接调 flash-attn 多？

往期回顾：

llama.cpp Q4 量化原理拆解：10GB 显存跑 70B 模型的秘密