FlashAttention-4 算法解析：Blackwell 架构下的性能优化

FA4 在 Blackwell 上进一步升级了这一机制：

• 双向重叠 (Bidirectional Overlap)：不仅 Q、K、V 的加载是异步的，中间结果的回写（O）以及 Mask 的生成都进入了全异步流水线。
• WGMMA 指令升级：利用 Blackwell 新增的指令集，使得矩阵乘法（GEMM）和 Softmax 操作之间的等待时间几乎被压缩为零。

3.3 针对长序列的各种显存层级优化

FA4 针对超长 Context（例如 1M+ tokens）进行了 L2 Cache 驻留优化。由于 Blackwell 拥有更大的 L2 Cache，FA4 采用了更激进的 'Cache-Persistence' 策略，尽可能让 KV Cache 的热点块常驻 L2，而非反复刷回 HBM。

4. 深入浅出：算法原理与伪代码

FlashAttention 的核心思想一直是 Tiling（分块）。我们在计算 O = Softmax(QK^T)V 时，不生成巨大的 N×N 矩阵，而是切成小块计算。

在 FA4 中，这个过程加入了 量化 (Quantize) 和 反量化 (Dequantize) 的步骤。

伪代码 (FlashAttention-4 风格)

def flash_attention_4_fwd(Q, K, V):
    """
    Q, K, V 位于 HBM (显存)
    Block_Size_M, Block_Size_N: 分块大小 (针对 L2/SRAM 优化)
    """
    # 1. 初始化输出 O 和 统计量 l, m
    O = torch.zeros_like(Q)
    l = torch.zeros(Q.shape[0])
    m = torch.ones(Q.shape[0]) * -inf

    # 2. 将 Q, K, V 切分为块 (Tiles)
    # 外层循环：遍历 Q 的分块 (行方向)
    for i in range(0, N, Block_Size_M):
        # 异步加载 Qi 到 SRAM (使用 TMA)
        Qi = load_async(Q[i : i+Block_Size_M])

        # 内层循环：遍历 K, V 的分块 (列方向)
        for j in range(0, N, Block_Size_N):
            # FA4 特性：异步加载并动态反量化 (FP4 -> FP16/BF16 用于计算)
            # Blackwell 硬件会在加载时自动处理格式转换
            Kj = load_async_dequant(K[j : j+Block_Size_N])
            Vj = load_async_dequant(V[j : j+Block_Size_N])

            # --- 计算阶段 (Tensor Cores 满载) ---
            # 1. S_ij = Qi * Kj^T
            # 这里利用 Blackwell 的特定指令，甚至可以直接在低精度下做部分乘法
            S_ij = gemm(Qi, Kj.transpose())

            # 2. 在线 Softmax 更新 (Online Softmax)
            # 这一步需要高精度 (FP32) 以防止数值崩塌
            m_ij = row_max(S_ij)
            P_ij = exp(S_ij - m_ij)
            l_ij = row_sum(P_ij)

            # 3. 更新全局统计量 mi, li 并重新缩放旧的 Oi
            # 著名的 "Rescaling" 步骤
            mi_new = max(m[i], m_ij)
            alpha = exp(m[i] - mi_new)
            beta = exp(m_ij - mi_new)

            # 4. Oi = (alpha * Oi + beta * (P_ij * Vj)) / li_new
            # P_ij * Vj 这一步也是矩阵乘法
            Oi = scale_and_add(Oi, P_ij, Vj, alpha, beta)

            # 更新统计量
            m[i] = mi_new
            l[i] = alpha * l[i] + beta * l_ij

    # 将结果 Oi 写回 HBM
    store_async(O[i : i+Block_Size_M], Oi)
    return O

代码解读：

1. load_async_dequant: 这是 FA4 的精髓。在数据从 HBM 搬到 SRAM 的过程中，利用 TMA 硬件单元直接完成 FP4 到高精度的转换，或者直接加载 FP4 并在 Tensor Core 内部处理。
2. 流水线掩盖: 上述代码在 GPU 上执行时，gemm 计算和 load_async 是并行的。FA4 保证了计算单元永远不会因为等待数据而空转。

5. 性能对比：FA4 vs FA3

为了直观展示 FA4 的威力，我们假设在 Blackwell B200 平台上运行，序列长度为 8K。

扩充知识点：为什么 Softmax 这么难搞？
Softmax 包含指数运算 exp() 和求和 sum()。exp 是超越函数，计算昂贵，不像矩阵乘法那样容易并行。它需要读取一整行的数据才能算出最大值 max，这打破了数据的局部性。
FA4 通过特殊的硬件指令，将 Softmax 的部分计算融合进矩阵乘法的尾部，进一步减少了开销。

6. 总结与展望

FlashAttention-4 再次证明了一个观点：在后摩尔定律时代，算法必须顺应硬件的'脾气'。

• FA1/2 告诉我们要顺应 GPU 显存 (HBM) 的脾气。
• FA3 告诉我们要顺应 Tensor Core 和异步拷贝 的脾气。
• FA4 则告诉我们，如果要追求极致速度，必须学会接受 低精度 (FP4) 的不完美，并通过算法技巧（分块量化、在线校准）来弥补它。

对于开发者而言，FA4 意味着我们可以在同样的硬件预算下，训练 Context Window 更长的模型，或者让推理成本进一步降低。随着 Blackwell 硬件的普及，FA4 将成为未来一两年内大模型基础设施的标配。