Flash Attention解决了什么问题？
标准注意力计算需要将整个N×N的注意力矩阵存储在GPU高带宽内存（HBM）中。对于长序列，这个矩阵非常大——128K序列长度的注意力矩阵需要约64GB HBM。GPU核心（SM）与HBM之间的数据搬运成为瓶颈。
Flash Attention的核心创新：不在HBM中实例化完整注意力矩阵，而是在SRAM中分块计算。
GPU内存层次
理解Flash Attention需要先理解GPU的内存层次：
SRAM (片上共享内存)
├── 延迟: ~20 cycles
├── 带宽: ~19 TB/s (A100)
└── 容量: ~192KB per SM
HBM (高带宽内存)
├── 延迟: ~200+ cycles
├── 带宽: ~2 TB/s (A100)
└── 容量: 80GB (A100)
标准注意力的问题：在HBM中读写O(n²)大小的矩阵，受限于2TB/s的HBM带宽。
Flash Attention的解决思路：将计算分块，每块在SRAM中完成，减少HBM访问次数。
算法原理
标准注意力计算
S = Q @ K^T / √d # [N, N] 注意力分数
P = softmax(S) # [N, N] 归一化
O = P @ V # [N, d] 输出
问题：S和P是N×N矩阵，需要完整存储在HBM中。
Flash Attention的分块算法
将Q、K、V分成块，逐块计算：
输入: Q [N,d], K [N,d], V [N,d]
参数: block_size = B
初始化: O = 0, l = 0 (行和累积)
for j in range(0, N, B): # 遍历K, V的块
K_j = K[j:j+B] # 加载到SRAM
V_j = V[j:j+B]
for i in range(0, N, B): # 遍历Q的块
Q_i = Q[i:i+B] # 加载到SRAM
# 在SRAM中计算
S_ij = Q_i @ K_j^T / √d # [B, B]
P_ij = softmax(S_ij)
# 在线更新softmax（关键创新）
O_i += P_ij @ V_j # 累加到输出
在线Softmax（核心难点）
标准softmax需要看到整行才能归一化。但Flash Attention逐块计算，不能看到整行。
解决方法——在线Softmax（Online Softmax）：
# 维护两个状态: m (最大值), l (求和)
# 处理新块时:
m_new = max(m_old, max(S_ij)) # 更新最大值
l_new = l_old * exp(m_old - m_new) + sum(exp(S_ij - m_new)) # 更新归一化因子
O_new = O_old * (l_old * exp(m_old - m_new) / l_new) + P_ij @ V_j / l_new # 更新输出
这个数学技巧让我们可以逐块计算softmax而无需存储完整矩阵。
反向传播
标准注意力在反向传播时需要中间矩阵P。Flash Attention不存储P，如何反向传播？
方案：在前向传播中存储O和softmax归一化统计量（m和l），在反向传播中重计算P：
# 前向：存储O, m, l（O(N*d)而非O(N²)）
# 反向：用Q, K重计算P，再计算梯度
# 额外计算开销约2倍前向，但减少了O(N²)的存储
性能分析
时间复杂度
- 标准注意力: O(N²) HBM读写 + O(N²) 计算
- Flash Attention: O(N²/B) HBM读写 + O(N²) 计算
HBM读写减少B倍，计算量不变。
空间复杂度
- 标准注意力: O(N²) HBM存储
- Flash Attention: O(N*d) HBM存储（只存O和统计量）
实测性能
以A100 GPU为例，Llama-7B模型：
| 序列长度 | 标准注意力 | Flash Attention | 加速比 |
|---|---|---|---|
| 2K | 5ms | 2ms | 2.5x |
| 8K | 45ms | 8ms | 5.6x |
| 32K | OOM | 120ms | ∞ |
| 128K | OOM | 950ms | ∞ |
关键发现：序列越长，Flash Attention优势越大。长上下文场景下，标准注意力直接OOM，而Flash Attention可以运行。
Flash Attention 2/3的改进
Flash Attention 2
减少非矩阵乘法运算：
- 优化了online softmax的计算顺序
- 减少GPU线程同步开销
- 优化了并行度（沿序列维度并行）
效果：比Flash Attention 1快约2倍。
Flash Attention 3 (Flash Attention 3)
针对Hopper架构（H100）的优化：
- 异步化：使用TMA（Tensor Memory Accelerator）异步加载数据
- 低精度：支持FP8 Tensor Core
- 流水线：前一个块计算与下一个块加载重叠
效果：在H100上比Flash Attention 2快1.5-2倍。
实现要点
分块大小选择
block_size选择需要考虑:
- SRAM容量: 块太大溢出SRAM
- 计算效率: 块太小GPU利用率低
- 经验值: 128-256 tokens per block
多GPU并行
Flash Attention可以沿Q维度并行：
- 每个GPU处理不同的Q块
- 不需要GPU间通信（K和V共享）
- 适合Tensor Parallel
使用建议
# PyTorch 2.0+ 内置Flash Attention
import torch
attn_output = torch.nn.functional.scaled_dot_product_attention(
Q, K, V,
is_causal=True # causal mask
)
# 自动使用Flash Attention内核
# vLLM等推理框架默认使用Flash Attention
# 无需手动配置
总结
Flash Attention是近年来大模型领域最具影响力的算法创新之一。它不是改变注意力的数学含义，而是改变计算方式——通过减少内存搬运而非减少计算量来加速。这个"IO优化"思路解决了长上下文的核心瓶颈，使得128K甚至1M的上下文窗口成为可能。所有主流推理框架都已集成Flash Attention，它是长上下文推理的"隐形基础设施"。