FlashAttention 优化原理与使用
FlashAttention 通过 IO 感知的分块计算策略，在不改变数学结果的前提下，将注意力计算的内存访问量从 O(N²) 降低到 O(N)。
核心思想
传统注意力的瓶颈不在计算而在内存访问——中间结果 S=QK^T 和 P=softmax(S) 需要写入 HBM 再读回。FlashAttention 通过分块计算避免存储完整的 N×N 注意力矩阵：
python
# FlashAttention 使用
from torch.nn.functional import scaled_dot_product_attention
# PyTorch 2.x 内置 FlashAttention
output = scaled_dot_product_attention(Q, K, V, is_causal=True)
# 自动选择 FlashAttention / Memory-Efficient Attention性能提升
| 序列长度 | 标准注意力 | FlashAttention | 加速比 |
|---|---|---|---|
| 512 | 1.2ms | 0.4ms | 3x |
| 2048 | 18ms | 2.1ms | 8.5x |
| 8192 | OOM | 12ms | ∞ |
FlashAttention 2 vs 3
FlashAttention-2 优化了 GPU 利用率；FlashAttention-3 针对 Hopper 架构（H100）使用 TMA 和 Warpgroup 进一步加速。