从Self-Attention到FlashAttention注意力计算加速的演进之路一、注意力机制的计算瓶颈不是模型设计问题而是硬件利用问题标准Self-Attention的FLOPs随序列长度平方增长——这是每个NLP研究者都熟知的事实。但优化这个问题的思路在2022年经历了一次范式转换在此之前主流努力集中在降低计算复杂度稀疏注意力、低秩近似、核方法FlashAttention的出现则揭示了一个被忽视的事实——注意力计算的实际瓶颈不在FLOPs而在GPU的显存带宽。具体来说标准Attention实现需要将中间矩阵QK^T, softmax结果写回HBMHigh Bandwidth Memory然后在下一次计算时再读出来。这些中间矩阵的大小也是O(N^2)的而HBM的读写带宽约1.5TB/s on A100远低于SRAM的计算带宽约19.5TB/s。FlashAttention 通过分块tiling和重计算recomputation将绝大部分中间结果留在SRAM中从而将显存访问复杂度从O(N^2)降为O(N)。flowchart TB subgraph Standard[标准 Attention: 多次读写 HBM] S1[Q, K 从 HBM 读入] -- S2[计算 SQK^T] S2 -- S3[S 写回 HBM (O(N²))] S3 -- S4[S 从 HBM 读回] S4 -- S5[计算 Psoftmax(S)] S5 -- S6[P 写回 HBM (O(N²))] S6 -- S7[P 从 HBM 读回] S7 -- S8[计算 OPV] S8 -- S9[O 写回 HBM] end subgraph Flash[FlashAttention: 单次遍历 HBM] F1[Q, K, V 从 HBM 读入] -- F2[分块加载到 SRAM] F2 -- F3[SRAM 内完成 SQK^T] F3 -- F4[SRAM 内完成 softmax(online)] F4 -- F5[SRAM 内完成 OPV] F5 -- F6[仅 O 写回 HBM (O(N))] F2 --|下一块| F2 end二、Online Softmax分块计算的关键前提标准Softmax需要三次遍历一次算max数值稳定一次算exp求和一次算除法。这在分块场景下意味着每块都需要等所有块算完才能继续——完全不兼容分块策略。Online Softmax 的解决思路是维护一个不断更新的running_max和running_sum每处理一个新块时用当前的running_max重新scale已经累积的running_sum然后加上当前块的贡献。这样每个块只需要访问一次且最终结果与标准Softmax数值等价不是近似。import torch import torch.nn.functional as F import math def online_softmax_kernel( Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, block_size: int 128 ) - torch.Tensor: 使用Online Softmax的分块注意力计算。 为什么使用Online Softmax 标准Softmax需要全局max进行数值稳定但分块计算时 无法提前知道全局max。Online Softmax通过维护 running_max 和 running_sum在每个块内动态修正 最终结果与标准softmax数值等价。 B, H, N, D Q.shape scale 1.0 / math.sqrt(D) # 初始化累积变量 O torch.zeros_like(Q) running_max torch.full( (B, H, N, 1), -float(inf), deviceQ.device, dtypeQ.dtype ) running_sum torch.zeros( (B, H, N, 1), deviceQ.device, dtypeQ.dtype ) n_blocks (N block_size - 1) // block_size for j in range(n_blocks): # 加载当前 K, V 块 j_start j * block_size j_end min(j_start block_size, N) K_j K[:, :, j_start:j_end, :] # [B, H, block, D] V_j V[:, :, j_start:j_end, :] # [B, H, block, D] # 计算当前块的注意力分数 S_j torch.matmul(Q, K_j.transpose(-2, -1)) * scale # Online Softmax: 更新running_max m_prev running_max.clone() m_curr torch.max(S_j, dim-1, keepdimTrue).values running_max torch.maximum(m_prev, m_curr) # 计算修正因子exp(m_prev - running_max) correction torch.exp(m_prev - running_max) # 更新running_sum和输出 P_j torch.exp(S_j - running_max) # [B, H, N, block] running_sum correction * running_sum P_j.sum(dim-1, keepdimTrue) # 修正之前累积的O加上当前块的贡献 O correction * O torch.matmul(P_j, V_j) # 最终归一化 O O / running_sum return O def verify_online_softmax(): 验证Online Softmax与标准实现数值一致。 torch.manual_seed(42) B, H, N, D 2, 8, 512, 64 Q torch.randn(B, H, N, D, devicecuda) K torch.randn(B, H, N, D, devicecuda) V torch.randn(B, H, N, D, devicecuda) # 标准实现 scale 1.0 / math.sqrt(D) attn torch.softmax(Q K.transpose(-2, -1) * scale, dim-1) O_standard attn V # Online实现 O_online online_softmax_kernel(Q, K, V, block_size128) max_diff (O_standard - O_online).abs().max().item() print(fMax difference: {max_diff:.2e}) assert max_diff 1e-5, f差异过大: {max_diff} if __name__ __main__: verify_online_softmax()三、FlashAttention-2 的改进减少非矩阵乘法操作FlashAttention-2 相较于 v1 的核心改进在于调整了循环嵌套的顺序。v1 的外层循环遍历 Q 的块内层循环遍历 K/V 的块v2 将外层循环改为遍历 K/V 的块内层遍历 Q 的块。这个看似微小的改动带来了两个实质性的收益减少线程块间的同步次数。v1 中每个 Q 块处理完后需要等待所有 K/V 块完成才能写回v2 中每个 Q 块可以在所有 K/V 块处理完之前就开始下一轮循环。更好地利用 Tensor Core。v2 的循环顺序使得矩阵乘法的维度更接近 Tensor Core 偏好的形状M128, N128, K128减少了 padding 浪费。在 A100 上的实际测试中v2 相比 v1 的理论加速比约为 2x实际加速比在 1.7x-2.0x 之间差距取决于序列长度。对于序列长度超过 4096 的场景加速效果最明显。四、FlashAttention 的适用边界FlashAttention 并非在所有场景下都是最优解短序列场景N 512分块策略的额外开销tiling的循环控制、寄存器分配占比超过显存带宽节省的收益。此时标准Attention实现反而更快。非标准Attention变体FlashAttention 的 CUDA kernel 针对标准 scaled dot-product attention 深度优化。如果你的 Attention 包含自定义的 mask 模式、相对位置编码注入或跨模态交互可能需要回到标准实现或等待社区适配。FP32精度要求FlashAttention 默认使用 FP16/BF16 计算以适配 Tensor Core。如果任务对精度敏感如科学计算场景需要验证 FP16 带来的精度损失是否可接受。非NVIDIA硬件FlashAttention 的 kernel 是 CUDA 实现的。在 AMD GPU、TPU 或 Apple Silicon 上不可用需要寻找对应平台的替代实现如 AMD 的 Composable Kernel 库或 Apple 的 MLX。五、总结FlashAttention 带来的核心洞察是在深度学习计算的优化中算术强度的提升往往比FLOPs的减少更有效标准 Attention 的瓶颈不在计算量而在显存带宽FlashAttention 通过分块计算将显存访问复杂度降低一个量级。Online Softmax 是实现分块计算的关键数学前提且结果与标准实现数值等价。FlashAttention-2 通过调整循环顺序进一步优化了GPU占用率。短序列和非标准变体场景下FlashAttention 的优势会减弱或消失。