Huajishe Notes

1-Flash-Attention

7 min read

Flash Attention

背景与动机

标准注意力的实现问题

传统Attention实现需要两步：

  1. 计算注意力分数 S = Q @ K
  2. 计算 softmax(S) @ V

中间结果 S 需要 O(N²) 显存存储（N=sequence length），成为大模型训练/推理的瓶颈。

问题分析

# 标准注意力
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的归一化：

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算法

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优化

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访问量显存需求
标准AttentionO(N²d)O(N²)O(N²)
Flash AttentionO(N²d)O(N²d/block_size)O(Nd)
节省比例相同~block_size倍~N倍

block_size通常为64或128，意味着HBM访问减少64-128倍。

与近似注意力的对比

近似注意力方法

方法复杂度精度适用场景
Flash AttentionO(N²d)无损通用
Sparse AttentionO(N²/ log N)近似长序列
LinformerO(N)近似长序列
ReformerO(N log N)近似长序列
PerformerO(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设计

// 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 支持

框架集成

# 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）

Graph View