287 lines
7.4 KiB
Markdown
287 lines
7.4 KiB
Markdown
# Flash Attention
|
||
|
||
## 背景与动机
|
||
|
||
### 标准注意力的实现问题
|
||
|
||
传统Attention实现需要两步：
|
||
1. 计算注意力分数 S = Q @ K^T
|
||
2. 计算 softmax(S) @ V
|
||
|
||
中间结果 S 需要 O(N²) 显存存储（N=sequence length），成为大模型训练/推理的瓶颈。
|
||
|
||
### 问题分析
|
||
|
||
```python
|
||
# 标准注意力
|
||
def attention(Q, K, V):
|
||
# Q: (batch, heads, seq_len, head_dim)
|
||
# K, V: (batch, heads, seq_len, head_dim)
|
||
|
||
S = torch.matmul(Q, K.transpose(-2, -1)) # (batch, heads, seq_len, seq_len)
|
||
# S 需要 O(N²) 显存
|
||
|
||
A = F.softmax(S / sqrt(d), dim=-1) # 注意力权重
|
||
O = torch.matmul(A, V) # 输出
|
||
|
||
return O
|
||
```
|
||
|
||
对于 LLaMA-7B，seq_len=2048，注意力矩阵约 2×32×2048×2048×4bytes ≈ 1GB 每层，12层约12GB。
|
||
|
||
## Flash Attention核心思想
|
||
|
||
### Tiling（分块）技术
|
||
|
||
Flash Attention将N×N的注意力矩阵分块处理，每次只计算一小块，避免一次性存储完整矩阵：
|
||
|
||
```
|
||
分块策略：
|
||
将Q, K, V沿seq_len维度分为T个块
|
||
每块大小: block_size = 256 或 512
|
||
|
||
逐块计算：
|
||
For each block_i of Q:
|
||
For each block_j of K:
|
||
计算 block_ij = Q[i] @ K[j]^T
|
||
更新输出 O[i] = O[i] + block_ij @ V[j]
|
||
```
|
||
|
||
### 递归计算softmax
|
||
|
||
分块计算的关键是正确处理softmax的归一化：
|
||
|
||
```python
|
||
def flash_attention_block(Q_block, K_block, V_block, m_i, M):
|
||
"""
|
||
Q_block: (block_size, head_dim)
|
||
K_block, V_block: (block_size, head_dim)
|
||
m_i: 当前block的最大值（用于数值稳定softmax）
|
||
M: 累计的最大值
|
||
"""
|
||
# 计算当前block的注意力分数
|
||
S_ij = Q_block @ K_block.T / sqrt(d)
|
||
|
||
# 数值稳定的softmax: 减去每行最大值
|
||
row_max = S_ij.max(dim=-1, keepdim=True)[0]
|
||
S_ij = S_ij - row_max # 稳定性
|
||
|
||
# 计算exp(S_ij)
|
||
P_ij = torch.exp(S_ij)
|
||
|
||
# 更新累计状态
|
||
# m_new = max(m_old, row_max)
|
||
# d_new = d_old * exp(m_old - m_new) + exp(row_max - m_new) * sum(P_ij)
|
||
|
||
# 计算局部输出
|
||
# O_i = O_{i-1} * factor + P_ij @ V_block
|
||
```
|
||
|
||
### 在线softmax算法
|
||
|
||
```python
|
||
def online_softmax(x):
|
||
"""
|
||
在线计算softmax，利用数值稳定性技巧
|
||
"""
|
||
m = x.max(dim=-1, keepdim=True)[0] # 行最大值
|
||
x_safe = x - m # 稳定处理
|
||
|
||
# 分块累加
|
||
cumsum_exp = 0
|
||
cumsum_factor = 1
|
||
|
||
for block in split(x_safe):
|
||
exp_block = torch.exp(block)
|
||
|
||
# 关键：保持累积状态
|
||
new_cumsum_factor = cumsum_factor * exp(-block.max()) # 调整因子
|
||
cumsum_exp = cumsum_exp * new_cumsum_factor + exp_block.sum()
|
||
cumsum_factor = new_cumsum_factor
|
||
|
||
# 最终归一化
|
||
softmax_x = torch.exp(x_safe) / cumsum_exp
|
||
|
||
return softmax_x
|
||
```
|
||
|
||
## IO复杂度分析
|
||
|
||
### GPU内存层次
|
||
|
||
```
|
||
GPU HBM (High Bandwidth Memory):
|
||
- 带宽: ~1-2 TB/s
|
||
- 容量: ~80GB
|
||
- 特点: 高带宽但访问成本高
|
||
|
||
GPU SRAM (Shared Memory / L1 Cache):
|
||
- 带宽: ~10-20 TB/s
|
||
- 容量: ~100-200MB (per SM)
|
||
- 特点: 极低延迟但极有限容量
|
||
```
|
||
|
||
### Flash Attention IO优化
|
||
|
||
```python
|
||
def flash_attention_io_optimized(Q, K, V):
|
||
"""
|
||
IO复杂度分析：
|
||
- 标准Attention: O(N²) HBM访问
|
||
- Flash Attention: O(N²/d) HBM访问，其中d=block_size
|
||
|
||
通过将数据保留在SRAM中计算，减少HBM访问
|
||
"""
|
||
# SRAM容量限制：每个SM约192KB用于Tensor Core
|
||
# 假设block_size=256, head_dim=64
|
||
# 每块数据: 256 × 64 × 2 (FP16) × 2 (K, V) = 64KB
|
||
|
||
block_size = 256
|
||
num_blocks = seq_len // block_size
|
||
|
||
# 迭代加载K, V块到SRAM
|
||
for i in range(num_blocks):
|
||
K_block = K[:, :, i*block_size:(i+1)*block_size, :] # 加载到SRAM
|
||
V_block = V[:, :, i*block_size:(i+1)*block_size, :] # 加载到SRAM
|
||
|
||
# 在SRAM中计算
|
||
# ...
|
||
```
|
||
|
||
### 计算复杂度对比
|
||
|
||
| 方法 | 计算量 | HBM访问量 | 显存需求 |
|
||
|------|--------|----------|---------|
|
||
| 标准Attention | O(N²d) | O(N²) | O(N²) |
|
||
| Flash Attention | O(N²d) | O(N²d/block_size) | O(Nd) |
|
||
| 节省比例 | 相同 | ~block_size倍 | ~N倍 |
|
||
|
||
block_size通常为64或128，意味着HBM访问减少64-128倍。
|
||
|
||
## 与近似注意力的对比
|
||
|
||
### 近似注意力方法
|
||
|
||
| 方法 | 复杂度 | 精度 | 适用场景 |
|
||
|------|--------|------|---------|
|
||
| Flash Attention | O(N²d) | 无损 | 通用 |
|
||
| Sparse Attention | O(N²/ log N) | 近似 | 长序列 |
|
||
| Linformer | O(N) | 近似 | 长序列 |
|
||
| Reformer | O(N log N) | 近似 | 长序列 |
|
||
| Performer | O(N d²) | 近似 | 通用 |
|
||
|
||
### Flash Attention vs 近似方法
|
||
|
||
```
|
||
Flash Attention核心价值：
|
||
1. 无精度损失（exact computation）
|
||
2. IO优化带来实际加速
|
||
3. 数学上等价于标准Attention
|
||
|
||
vs 近似方法：
|
||
- Sparse Attention: 有精度损失，Flash Attention更快
|
||
- Linformer: 需要预设稀疏度，适用场景受限
|
||
- 结论: Flash Attention是exact方法中的最优实现
|
||
```
|
||
|
||
## 实现细节
|
||
|
||
### CUDA Kernel设计
|
||
|
||
```cuda
|
||
// Flash Attention CUDA Kernel伪代码
|
||
__global__ void flash_attention_kernel(
|
||
const float* Q, const float* K, const float* V,
|
||
float* O, int seq_len, int head_dim)
|
||
{
|
||
// 1. 加载Q到SRAM（分block）
|
||
__shared__ float Q_sram[BLOCK_M][HEAD_DIM];
|
||
|
||
// 2. 外循环遍历K, V块
|
||
for (int j = 0; j < seq_len; j += BLOCK_N) {
|
||
// 加载K, V块到SRAM
|
||
__shared__ float K_sram[BLOCK_N][HEAD_DIM];
|
||
__shared__ float V_sram[BLOCK_N][HEAD_DIM];
|
||
|
||
// 3. 计算当前块的attention
|
||
// 使用Tensor Core加速矩阵乘法
|
||
// 迭代累加避免溢出
|
||
}
|
||
|
||
// 4. 写入输出
|
||
}
|
||
```
|
||
|
||
### Flash Attention V2/V3改进
|
||
|
||
```
|
||
Flash Attention V2:
|
||
- 分块策略优化：更细粒度的循环 tiling
|
||
- 减少shared memory访问
|
||
- 支持sequence parallel
|
||
|
||
Flash Attention V3:
|
||
- 针对 Hopper 架构优化
|
||
- Tensor Core Memory Access Patterns
|
||
- FP8 支持
|
||
```
|
||
|
||
### 框架集成
|
||
|
||
```python
|
||
# PyTorch 2.0+ 使用Flash Attention
|
||
import torch
|
||
|
||
# 方式1: torch.nn.functional.scaled_dot_product_attention (SDPA)
|
||
output = F.scaled_dot_product_attention(
|
||
Q, K, V,
|
||
attn_mask=None,
|
||
dropout_p=0.0,
|
||
is_causal=True # 自动选择最优backend（包括Flash Attention）
|
||
)
|
||
|
||
# 方式2: 显式使用Flash Attention
|
||
from flash_attn import flash_attn_func
|
||
output = flash_attn_func(
|
||
Q, K, V,
|
||
causal=True,
|
||
softmax_scale=1.0 / sqrt(head_dim)
|
||
)
|
||
|
||
# 方式3: HuggingFace集成
|
||
model = AutoModelForCausalLM.from_pretrained(
|
||
"meta-llama/Llama-2-7b-hf",
|
||
torch_dtype=torch.float16,
|
||
attn_implementation="flash_attention_2" # 自动使用Flash Attention
|
||
)
|
||
```
|
||
|
||
## 性能收益
|
||
|
||
### Benchmark
|
||
|
||
```
|
||
硬件: A100 80GB
|
||
模型: LLaMA-7B, seq_len=2048
|
||
|
||
| 方法 | 注意力计算时间 | 显存占用 |
|
||
|------|--------------|---------|
|
||
| 标准Attention (FP16) | 120ms | 24GB |
|
||
| Flash Attention (FP16) | 35ms | 4GB |
|
||
| 提升 | 3.4x | 6x |
|
||
|
||
硬件: H100
|
||
| 方法 | 吞吐量 (tokens/s) |
|
||
|------|-------------------|
|
||
| 标准Attention | 1500 |
|
||
| Flash Attention | 5500 |
|
||
| Flash Attention + FP8 | 8500 |
|
||
```
|
||
|
||
## 使用注意事项
|
||
|
||
1. **因果掩码（Causal Mask）**：生成任务需要，Flash Attention支持`causal=True`参数
|
||
2. **序列长度**：需为2的幂次（Flash Attention内部tile要求）
|
||
3. **dtype**：支持FP16、BF16、FP32
|
||
4. **head_dim限制**：需要为8的倍数（128或64）
|
||
5. **混合精度训练**：前向用FP16，反向需要FP32的softmax累计（使用fused kernel） |