FlashAttention精读：IO感知的精确注意力算法

一、精确注意力的性能瓶颈

标准Transformer的注意力计算在长序列场景面临严重的性能瓶颈。当序列长度n=4096、注意力头数h=32、维度d=128时，注意力矩阵的大小是 n×n=16M，FP16精度下占32MB显存。这还只是单个注意力头；多层多头叠加后，显存占用轻松突破数十GB。问题的根源不是计算本身（O(n²)的FLOPs现代GPU能承受），而是内存IO——attention矩阵的中间结果需要在HBM（高带宽内存）和SRAM（片上缓存）之间反复搬运，IO次数远大于FLOPs。

1.1 GPU内存层次结构

GPU内存层次（以A100为例）

SRAM（片上缓存）：20MB，~19TB/s带宽，访问延迟~1ns
HBM（高带宽内存）：40-80GB，~1.5TB/s带宽，访问延迟~400ns
DRAM（系统内存）：~1TB，~25GB/s带宽，访问延迟~100ns

性能差距：
├── SRAM vs HBM：带宽差13倍，延迟差400倍
├── HBM vs DRAM：带宽差60倍，延迟差4倍
└── 标准attention的瓶颈：HBM IO次数（O(n²)）
      

1.2 标准注意力的IO复杂度

标准注意力需要存储中间矩阵S=QK^T和P=softmax(S)，这两个矩阵各占O(n²)显存。在前向传播中，Q、K从HBM加载→计算S写入HBM→softmax从HBM读S→计算P写入HBM→计算O=PV写入HBM——IO总次数约O(n²)。反向传播需要再存储一次S和P，IO总次数翻倍。这与现代GPU计算能力严重不匹配——A100的SRAM计算能力是HBM IO的数十倍。

二、FlashAttention的核心思想：分块与重计算

2.1 分块计算（Tiling）

FlashAttention的核心思想是把Q、K、V分成固定大小的小块（如64×64），每次只把一块Q加载到SRAM，与对应块的K计算得到局部的S分块，立即在SRAM中完成softmax，再与对应块的V计算得到O分块。SRAM的大小是有限的（A100约20MB），分块策略让所有中间结果都留在SRAM中，避免HBM的频繁IO。

分块计算架构

Q矩阵（n×d）分块：Q_1, Q_2, ..., Q_Br
K矩阵（n×d）分块：K_1, K_2, ..., K_Bc
V矩阵（n×d）分块：V_1, V_2, ..., V_Bc

外层循环：Br（行块数）次
内层循环：Bc（列块数）次

for i = 1 to Br:
    加载Q_i到SRAM
    for j = 1 to Bc:
        加载K_j, V_j到SRAM
        S_ij = Q_i × K_j^T  (在SRAM中)
        P_ij = softmax(S_ij) (在SRAM中)
        O_ij = P_ij × V_j   (在SRAM中)
        累加到O_i
    写回O_i到HBM
      

2.2 在线Softmax（Online Softmax）

分块softmax的难点是：分块之间需要共享最大值和归一化分母。FlashAttention使用在线softmax算法，逐块更新全局最大值m和归一化分母l，再重新缩放已计算的O分块。这种"增量softmax"保证最终结果与全局softmax完全一致（数学等价），但分块实现避免了O(n²)显存占用。

在线softmax公式

全局状态：m（当前最大值）、l（当前归一化分母）、O（输出矩阵）

处理新块S_ij时：
m_new = max(m_old, rowmax(S_ij))  # 新最大值
P_ij = exp(S_ij - m_new)         # 局部未归一化概率
l_new = exp(m_old - m_new) * l_old + rowsum(P_ij)  # 新分母
O_new = exp(m_old - m_new) * O_old + P_ij × V_ij  # 重新缩放累加

关键性质：最终 O_final = softmax(QK^T) × V（数学等价）
      

2.3 重计算（Recomputation）反向传播

反向传播需要S和P矩阵，但FlashAttention前向传播时只存储了O和softmax的统计量（m和l）。为了不增加显存占用，FlashAttention反向传播时从Q、K、O重新计算S和P——这增加了30%的FLOPs但避免了O(n²)的显存占用。对于IO密集型操作，重计算的算力代价远小于IO收益，整体仍然更快。

三、FlashAttention-2的工程优化

3.1 减少非矩阵乘法操作

FlashAttention-1在外层循环Br次、内层循环Bc次中，每次都要做softmax的重缩放操作。FlashAttention-2改进了循环结构，让softmax只在最后做一次，中间不重缩放——把更多时间花在矩阵乘法上（GPU最擅长的操作）。实际加速比从理论值进一步提升2倍。

3.2 更好的并行化

FlashAttention-2的并行维度从（batch, num_heads, Br）扩展为（batch, num_heads, Br, Bc）。同时改进了线程块的工作分配，让每个SM（流式多处理器）的计算负载更均衡。在A100上达到理论峰值FLOPs的50-70%，远高于标准注意力的30-40%。

3.3 新增特性：因果掩码、Dropout

FlashAttention-2支持因果掩码（causal mask，GPT类模型需要）和注意力dropout，把这些常用操作也融合到分块计算中，进一步减少HBM IO次数。

四、性能对比与工程收益

4.1 端到端训练加速

4.2 显存节省

对于batch_size=4、num_heads=32、head_dim=128、seq_len=4096的典型配置，标准注意力需要约4GB的S和P矩阵显存，FlashAttention只需约200MB的统计量——显存占用降低20倍。这让长序列训练从"超算任务"变成"普通GPU可完成"。

4.3 精度一致性

FlashAttention是精确注意力——数学上与标准注意力完全等价（无近似），仅改变了计算顺序。但实际数值结果可能有1e-3量级的差异（因浮点数累加顺序不同），这通常不影响下游任务表现。

五、对长上下文LLM的深远影响

FlashAttention是大模型长上下文时代的关键基础设施。基于它的发展，2023-2024年涌现了多个变体：

5.1 FlashAttention-3（2024）

针对Hopper架构（H100）优化，使用TMA（Tensor Memory Accelerator）和WGMMA（Warp-Group Matrix-Multiply-Accumulate）硬件特性，在H100上达到理论峰值的75%。

5.2 FlashDecoding（2023）

针对推理时逐token生成的KV cache场景优化的FlashAttention变体，把每个token的注意力计算并行化到多个SM，解决"自回归生成时每步只用一个SM"的瓶颈，推理速度提升2-4倍。

5.3 与PagedAttention的协同

vLLM的PagedAttention借鉴了OS虚拟内存的"分页"思想管理KV cache，与FlashAttention结合实现高效的attention serving。这是现代LLM推理引擎（如vLLM、TensorRT-LLM）的核心架构。

六、经验教训：6个生产级实战启示