FlashAttention分块计算优化大模型训练效率

1. FlashAttention的分块计算原理剖析

FlashAttention作为当前大模型训练中的关键技术突破，其核心创新点在于通过分块计算（Tiling）和内存访问优化，将传统Attention计算中O(N²)的HBM（高带宽内存）访问复杂度降低到O(N²d²/M)。这个看似简单的改进背后，蕴含着对GPU内存体系结构的深刻理解。

1.1 GPU内存层级与计算瓶颈

现代GPU通常包含多级内存结构：

• HBM：高带宽但高延迟的显存（16-80GB）
• SRAM：片上缓存（每SM约192KB）
• 寄存器：最快但容量最小

传统Attention计算存在两个致命问题：

1. Softmax操作需要全局归一化，导致必须将整个QK^T矩阵存入HBM
2. 反向传播时需要重新计算QK^T矩阵，造成重复内存访问

python复制# 传统Attention计算流程（伪代码）
def attention(Q, K, V):
    S = Q @ K.T  # [N,N]矩阵
    P = softmax(S / sqrt(d))
    O = P @ V    # [N,d]输出
    return O

1.2 分块计算实现方案

FlashAttention的解决方案是将计算分解为多个小块：

1. 将Q、K、V矩阵划分为大小为B×d的子块（B≈256-512）
2. 对每个子块计算局部Attention时：
   • 先将子块从HBM加载到SRAM
   • 在SRAM内完成QK^T、Softmax、PV计算
   • 通过在线重计算（Recomputation）避免存储中间矩阵

python复制# FlashAttention分块计算示例
def flash_attention(Q, K, V, B=256):
    O = torch.zeros_like(Q)
    for i in range(0, Q.size(0), B):
        Qi = Q[i:i+B]  # 加载到SRAM
        Oi = 0
        for j in range(0, K.size(0), B):
            Kj, Vj = K[j:j+B], V[j:j+B]