6.2 FlashAttention V2详解

本文深入剖析 FlashAttention V2 相比 V1 的核心改进：调换内外循环顺序、优化线程块内的工作分配、减少非矩阵乘运算，最终在 A100 上达到理论 FLOPS 的 50-73%。

📑 目录

1. V1 的性能瓶颈回顾
2. 核心改进一：调换循环顺序
3. 核心改进二：减少非矩阵乘运算
4. 核心改进三：Warp 级工作分配优化
5. Causal Mask 优化
6. 前向传播完整算法
7. 反向传播优化
8. 性能分析与对比
总结
自我检验清单
参考资料

1. V1 的性能瓶颈回顾

1.1 V1 的实际利用率

FlashAttention V1 在 A100 上只达到了理论 FLOPS 的 25-40%。对于一个理应是 Compute-Bound（通过减少 IO 后）的算子来说，这个利用率并不理想。

瓶颈主要来自三个方面（按对性能影响从大到小排序）：

📊 问题	📝 原因	💡 影响
循环顺序不佳	外循环遍历 K/V，导致 O 块被反复读写、Q 块被重复加载	大量额外的 HBM 流量，且无法跨 Q 并行
非 GEMM 运算占比高	Softmax rescaling 中有大量逐元素操作	Tensor Core 利用率低
Warp 切分维度不当	V1 沿 K/V 维度切分给 Warp，多个 Warp 共同更新同一行 O，需要合并部分 Softmax 结果	频繁的 Warp 间通信与同步，并行效率低

1.2 硬件特性分析

A100 GPU 的关键参数：

FP16 Tensor Core 峰值：312 TFLOPS
HBM 带宽：2 TB/s
SRAM 容量：192 KB/SM（Shared Memory，可配置最大 164 KB 供用户使用）
Warp 数量：每个 Thread Block 可包含多个 Warp（通常 4-8 个）

要充分利用 Tensor Core，需要确保计算流水线不被非 GEMM 操作打断。

2. 核心改进一：调换循环顺序

2.1 V1 的循环结构问题

V1 中外循环遍历 K/V 块，内循环遍历 Q 块。这带来两个问题：

Q 和 O 的重复读写：每处理一个新的 K/V 块，所有 Q 块和对应的输出 O 都需要重新加载和写回 HBM
并行度受限：外循环 K/V 块之间因 online softmax 的 running max/sum 累积存在串行依赖，V1 只能在 batch × heads 维度上并行。当 batch 和 heads 较小、序列较长时（长上下文场景），无法利用序列维度的并行性，SM 利用率不足

2.2 V2 的新循环结构

V2 将循环顺序调换为外循环遍历 Q 块，内循环遍历 K/V 块：

外循环: for i = 1 to T_r (遍历 Q 块) — 分配给不同 Thread Block
    从 HBM 加载 Q_i 到 SRAM（只加载一次）
    初始化 O_i = 0, m_i = -inf, l_i = 0
    内循环: for j = 1 to T_c (遍历 K/V 块)
        从 HBM 加载 K_j, V_j 到 SRAM
        计算 S_ij，更新统计量和 O_i（全在 SRAM 中）
    将最终 O_i 写回 HBM（只写一次）

2.3 收益分析

这个看似简单的调换带来了三重收益：

收益 1：O 只在最后写回一次

在 V1 中，每处理一个 K/V 块，O 都要被读取再写回，共 $2T_c$ 次 HBM 访问（ $T_c$ 次读 + $T_c$ 次写）。V2 中 O 只在内循环结束后写回一次。

收益 2：外循环天然可并行

不同的 Q 块之间完全独立（每个 Q 块计算自己的输出行），可以直接映射到不同的 Thread Block，实现 GPU 级别的并行。

收益 3：Q 只加载一次

每个 Thread Block 只负责一个 Q 块，Q 在 Shared Memory 中驻留整个内循环的过程。

📌 关键点：V1 选择 K/V 外循环源自 online softmax 的累积逻辑（每扫完一列 K/V 块就更新一次 O 的归一化），代价是 O 必须反复读+写。V2 发现：让 Q 只加载一次更优，因为减少 O 的读写比减少 K/V 的读收益更大。

3. 核心改进二：减少非矩阵乘运算

3.1 问题：非 GEMM 指令的瓶颈

在 Attention 的 Tiling 计算中，除了两次矩阵乘法（ $QK^\top$ 和 $PV$ ），还有大量逐元素操作：

减去最大值： $S_{ij} - m_i$
求指数： $\exp(\cdot)$
行求和： $\text{rowsum}(\cdot)$
Rescaling：乘以修正因子

这些操作无法使用 Tensor Core，只能用 CUDA Core，成为计算流水线中的瓶颈。

3.2 V2 的优化：延迟 Rescaling

V2 将 Softmax 的 rescaling 步骤延迟到内循环结束后统一执行：

V1 的做法（每个 K/V 块都 rescale）：

O_i^{(j)} = \frac{l_i^{(j-1)} \cdot e^{m_i^{(j-1)} - m_i^{(j)}}}{l_i^{(j)}} \cdot O_i^{(j-1)} + \frac{e^{S_{ij} - m_i^{(j)}}}{l_i^{(j)}} \cdot V_j

V2 的做法（只维护未归一化的累积）：

\tilde{O}_i^{(j)} = e^{m_i^{(j-1)} - m_i^{(j)}} \cdot \tilde{O}_i^{(j-1)} + e^{S_{ij} - m_i^{(j)}} \cdot V_j

最终一次性归一化：

O_i = \tilde{O}_i^{(T_c)} / l_i^{(T_c)}

3.3 收益量化

这个优化减少了内循环中每次迭代的三次逐元素操作：

旧 O 的反归一化：乘以 $l^{(j-1)}$
旧 O 的重新归一化：除以 $l^{(j)}$
新贡献的归一化：除以 $l^{(j)}$

V2 把这「2 次除法 + 1 次乘法」全部从内循环移除，只在最后做一次除法。当内循环迭代次数 $T_c$ 较大时，节省相当可观。

4. 核心改进三：Warp 级工作分配优化

4.1 V1 的 Warp 分配方案：沿 K/V 维度划分

V1 中，Thread Block 内的多个 Warp 沿 K/V 维度切分同一个 $S_{ij}$ 块。

假设 Thread Block 有 4 个 Warp，K/V 块大小为 $B_c = 128$ ：

Warp 0 负责 K/V 的第 0-31 列
Warp 1 负责 K/V 的第 32-63 列
Warp 2 负责 K/V 的第 64-95 列
Warp 3 负责 K/V 的第 96-127 列

每个 Warp 只能算出 $S_{ij}$ 的一部分列，因此：

计算 row-max / row-sum 时，必须跨 Warp 合并各自的局部结果
计算 $P \cdot V$ 时，多个 Warp 都会向同一行 O 累加贡献，需要再次同步

这种「split-K」式分配导致频繁的 Warp 间同步与 shared memory 通信。

4.2 V2 的 Warp 分配：沿 Q 维度划分

V2 将 Q 块进一步细分给不同的 Warp：

假设 Thread Block 有 4 个 Warp，Q 块大小为 $B_r = 128$ ：

Warp 0 负责 Q 的第 0-31 行
Warp 1 负责 Q 的第 32-63 行
Warp 2 负责 Q 的第 64-95 行
Warp 3 负责 Q 的第 96-127 行

每个 Warp 独立计算自己负责的 Q 行与完整 K/V 块的注意力，互不干扰。

4.3 对比两种分配策略

❌ 沿 K/V 分 Warp（V1 方案）	✅ 沿 Q 分 Warp（V2 方案）
多个 Warp 需要合并部分 Softmax 结果	每个 Warp 独立完成 Softmax
需要 shared memory 或 shuffle 通信	无需 Warp 间同步
输出 O 的同一行被多个 Warp 共同更新	输出 O 的不同行由不同 Warp 独占

💡 提示：如果沿 K/V 维度分割给不同 Warp，每个 Warp 只看到 Softmax 输入的一部分，必须在 Warp 间通信来合并最大值和求和，这正是 V1 性能低下的原因之一。

5. Causal Mask 优化

5.1 Causal Mask 的作用

在自回归模型（如 GPT）中，第 $i$ 个 token 在生成时只能”看到”自己以及前面的 token，不能窥探未来——即位置 $i$ 只能 attend 到位置 $j \leq i$ 。这是自回归语言模型训练时保持「因果性（causality）」的关键约束：模型必须基于历史预测下一个 token，而不能直接抄答案。

在实现上，这通过一个下三角掩码矩阵 来实现，掩码在 Softmax 之前加到注意力分数上： $\tilde{S} = QK^\top / \sqrt{d} + M$ 。被屏蔽的位置加上 $-\infty$ 后，经 $\exp$ 变为 0，从而对 Softmax 输出无贡献。

形象地看，这个 mask 让 $N \times N$ 的注意力矩阵只有下三角（含对角线）有效，上三角全部置零。这意味着理论上约一半的计算是无用的，这正是 V2 块级跳过策略要利用的结构。

5.2 朴素实现的浪费

如果不做特殊处理，即使 mask 为 0 的位置也会执行完整的 GEMM 和 Softmax 计算，浪费约一半的算力。

5.3 V2 的块级跳过策略

V2 引入了块级（Block-level）的 Causal Mask 优化：

对于第 $i$ 个 Q 块（行范围 $[iB_r, (i+1)B_r)$ ）和第 $j$ 个 K 块（列范围 $[jB_c, (j+1)B_c)$ ）：

完全在 mask 下方（ $(i+1)B_r - 1 \geq (j+1)B_c - 1$ 且 $iB_r \geq (j+1)B_c - 1$ ，即 Q 块最小行号 ≥ K 块最大列号）：块内所有位置都满足 row ≥ col，完全可见，正常计算
完全在 mask 上方（ $(i+1)B_r - 1 < jB_c$ ，即 Q 块最大行号 < K 块最小列号）：块内所有位置都满足 row < col，完全不可见，直接跳过
跨越对角线：块内既有可见位置也有不可见位置，需要逐元素应用 mask

对角线以下（全可见）    对角线块（部分mask）    对角线以上（全跳过）
┌─────────────┐      ┌─────────────┐      ┌─────────────┐
│ ■ ■ ■ ■ ■ ■ │      │ ■ ■ ■ □ □ □ │      │ □ □ □ □ □ □ │
│ ■ ■ ■ ■ ■ ■ │      │ ■ ■ ■ ■ □ □ │      │ □ □ □ □ □ □ │
│ ■ ■ ■ ■ ■ ■ │      │ ■ ■ ■ ■ ■ □ │      │ □ □ □ □ □ □ │
└─────────────┘      └─────────────┘      └─────────────┘

def flash_attention_v2_forward(Q, K, V, B_r, B_c, causal=False):
    N, d = Q.shape
    O = zeros(N, d)
    L = zeros(N)  # logsumexp，用于反向传播

    # 外循环（并行化到不同 Thread Block）
    parallel for i in range(0, N, B_r):
        Q_i = Q[i:i+B_r]  # 加载到 SRAM，整个内循环驻留

        # 初始化累积量（SRAM/寄存器）
        m_i = full(B_r, -inf)
        l_i = zeros(B_r)
        O_tilde = zeros(B_r, d)  # 未归一化的累积输出

        # 内循环：遍历 K/V 块
        kv_end = min(N, (i + B_r)) if causal else N
        for j in range(0, kv_end, B_c):
            K_j = K[j:j+B_c]  # 从 HBM 加载
            V_j = V[j:j+B_c]  # 从 HBM 加载

            # 计算局部注意力分数
            S_ij = Q_i @ K_j.T / sqrt(d)

            # 应用 causal mask（仅对角线块需要）
            # i, j 已是绝对起始坐标；当 K 块最大列号(j+B_c-1) > Q 块最小行号(i) 时本块跨对角线
            if causal and i < j + B_c - 1:
                apply_causal_mask(S_ij, i, j, B_r, B_c)

            # 更新统计量
            m_new = max(m_i, rowmax(S_ij))
            P_ij = exp(S_ij - m_new[:, None])
            l_new = exp(m_i - m_new) * l_i + rowsum(P_ij)

            # 更新未归一化输出
            O_tilde = exp(m_i - m_new)[:, None] * O_tilde + P_ij @ V_j

            m_i = m_new
            l_i = l_new

        # 最终归一化并写回
        O[i:i+B_r] = O_tilde / l_i[:, None]
        L[i:i+B_r] = m_i + log(l_i)  # 保存 logsumexp 给反向

    return O, L

6.2 与 V1 的关键差异

📊 方面	V1	V2
外循环	K/V 块	Q 块（并行）
O 的 HBM 访问次数	$2T_c$ 次（ $T_c$ 读 + $T_c$ 写）	1 次（仅最终写回）
内循环归一化	每步反归一化+重新归一化（2 除 + 1 乘）	仅做指数缩放，最后统一除一次
Warp 切分维度	沿 K/V（split-K，Warp 间需通信）	沿 Q（每 Warp 独占若干行，无需通信）
Causal 优化	无	块级跳过
保存的统计量	$m$ 和 $l$	$L = m + \log(l)$ （合并为 logsumexp）