原论文：FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness，Tri Dao 等，Stanford / SUNY Buffalo，NeurIPS 2022，arXiv:2205.14135。
一句话总结#
FlashAttention 通过 IO-aware 的 tiling 和重计算，在不改变 attention 数学定义的前提下，大幅减少 HBM 读写，实现 2-4× 加速和 5-20× 内存节省，是 attention 优化的基础构件。
核心问题#
标准 attention 的内存瓶颈：
S = Q·K^T # 读写 N×N 矩阵到 HBM
P = softmax(S) # 读写 N×N 矩阵到 HBM
O = P·V # 读写 N×N 矩阵到 HBM对于序列长度 N，需要 O(N²) 的 HBM 读写。GPU 计算速度 » HBM 带宽，导致 memory-bound。
核心方法#
IO-Aware Tiling#
将 Q, K, V 分块（tile），在 SRAM（高速缓存）中完成计算：
for tile_Q in Q:
for tile_K, tile_V in K, V:
# 在 SRAM 中计算局部 attention
S_local = tile_Q · tile_K^T
P_local = softmax(S_local)
O_local += P_local · tile_V关键优化#
| 优化 | 说明 |
|---|---|
| Tiling | 分块加载到 SRAM，减少 HBM 访问 |
| Online softmax | 增量计算 softmax，避免存储完整 S |
| Recomputation | 反向传播时重计算 forward 中间值，不存储 |
内存复杂度#
| 方法 | 内存 | 说明 |
|---|---|---|
| 标准 Attention | O(N²) | 存储 S, P |
| FlashAttention | O(N) | 仅存储 O，中间值重计算 |
关键结果#
| 指标 | 效果 |
|---|---|
| 加速比 | 2-4×（vs 标准 PyTorch） |
| 内存节省 | 5-20× |
| 精度 | Exact：无近似，数学等价 |
| 序列长度 | 支持更长序列（内存不再是瓶颈） |
后续版本#
| 版本 | 改进 |
|---|---|
| FlashAttention-2 | 更好的并行化，减少 non-matmul FLOPs |
| FlashAttention-3 | 异步加载/计算，利用新硬件特性 |
| FlashAttention-3 (Hopper) | 针对 H100 的 Tensor Memory Accelerator |
与稀疏注意力的关系#
FlashAttention 优化密集 attention 的 IO，不改变 O(N²) 计算复杂度。稀疏注意力减少计算量到 O(N) 或 O(N log N)。两者正交：
- FlashAttention + 稀疏模式：稀疏的 tile 计算
- FlashAttention + 长序列：使 O(N²) 可接受的范围扩大
局限#
- 计算复杂度未变：仍是 O(N²)，只是内存优化
- 序列长度上限：SRAM 容量限制 tile size
- 硬件依赖：针对特定 GPU 架构优化
实现要点#
- Tile size 选择：平衡 SRAM 容量和并行度
- Softmax 稳定性：online softmax 的数值稳定性
- Kernel 融合：load/compute/store 流水线
个人理解#
FlashAttention 是 attention 优化的基础设施。它不改变 attention 的数学形式，只是更高效地实现。这使得它成为所有后续 attention 优化（包括稀疏注意力、量化 attention）的基础 kernel。理解 FlashAttention 是理解所有 attention 优化的前提。