Flash Attention V1/V2/V3演进：主厨老火的“少搬多算”革命

可视化演示

展示Flash Attention通过分块Tiling将注意力计算从HBM移至SRAM的IO优化原理，以及V1/V2/V3三个版本的性能演进与内存优化对比。

背景与痛点

老火的招牌菜需要一张巨大的“食材关系表”：每一片食材要和桌上所有其他食材配对打分。过去，徒弟们从仓库把三箱原料（Q、K、V）全部搬到操作台，摊开算出整张表，再把表搬回仓库；要取用时又搬回来。操作台只有桌面大，食材稍微一多，桌子彻底堆满，这就是“内存溢出”。

这段故事映射的正是标准注意力。为了算出结果，必须先在显存（HBM）中构造并保存 N×N 的得分矩阵 S，再把 softmax 之后的概率矩阵 P 也存进 HBM。序列长度 N 从 512 涨到 4096、再到上万时，中间矩阵以 N² 的速度膨胀，显存很快爆掉。算一遍还要把 S 写一遍、读一遍、写 P、读 P，搬运量巨大。

老火盯着堆满桌面的食材意识到：真正该改的不是切菜手法，而是“搬运路线”。一次搬得再多，桌面就那么大；只有改变搬运策略，才能让这道菜真正做大。

核心思想

老火做了个实验：把一次搬一百种配料改成一次只搬十种，来回多走几趟，结果反而更快了。原来一次搬得多，路上堵车、堆在桌上也运转不开；只要一次搬“刚好算得完”的分量，勤搬反而高效。这就是“IO 感知”的直觉起点。

在 GPU 里，HBM 就是那个又大又远的仓库，容量几十 GB，但带宽有限；SRAM 是巴掌大的操作台，只有几十 KB，却带宽惊人。过去的算法把中间矩阵反复写回 HBM——等于在仓库门口排队。真正聪明的做法是：让计算紧贴着 SRAM 进行，把 HBM 只当作“源源不断的原料库”，而不是“临时中转站”。

IO 感知算法有三个具体抓手：第一是分块计算，只把需要的块搬上 SRAM；第二是算子融合，把 softmax 和矩阵乘合并成一个内核，中间结果不落 HBM；第三是重计算，反向传播时重新算一次得分，用计算换 IO。这三招后来成为 GPU 内核优化的通用方法论，也决定了 Flash Attention 每一代的设计走向。

老火在菜谱扉页上写下厨训：“最快的搬运，是少搬运。”这句训话成为三代改进的共同纲领。

V1分块原理

老火的新方案是“分盘备菜”：把整桌食材按一盘菜的份量分成小盘，每次只端一盘上操作台。做熟一小份后，马上和锅里已有的半成品精确拌匀；拌完再去端下一盘。这样操作台永远只放当下需要的一小盘，永远不会堆满。

在注意力计算里，这对应把 Q、K、V 按行切成小块：Q 切成若干 $B_r\times d$ 的小块，K、V 切成 $B_c\times d$ 的小块。每次只计算一小块得分 $S_{ij}$，再与之前累积的输出合并。只要块的大小能被 SRAM 装下，整张 N×N 的中间矩阵就永远不需要写进 HBM。

分块带来的最直接收益是显存：标准算法要保存完整的 N×N 中间矩阵，V1 只保存 N×d 的输出 O 和极小的运行统计量 m、l。HBM 访问量从 O(N²) 量级大幅下降，长序列训练终于有了喘息空间。

老火端出第一盘“V1 版炖菜”，发现桌面不再堆满，出菜速度还快了一倍多。但这只是开始——一个“看不见的调味难题”正等着他。

在线Softmax

分盘备菜撞上了一个新矛盾：传统炖菜要求“所有食材到齐才能调味”，因为归一化必须知道全局的最大鲜度值。可老火现在一盘一盘地做，怎么知道后面会不会出现更鲜的食材？如果急着调味，后面鲜度更高的食材一来，前面不就白调了吗？

他灵机一动：边做边调。先按当前所见的最大值临时调味；如果后面来了更鲜的食材，就把锅里已经调好的部分乘以一个缩放系数“等比缩小”，再融入新食材。这就是在线 Softmax——始终维护三个状态：运行最大值 m、归一化分母 l、运行输出 o。每块到来时，用新块信息更新这三个状态。

这套在线更新在数学上与标准 softmax 完全等价：每当最大值更新，旧部分都被精确地乘以一个小于 1 的系数，而不是丢弃重算。V1 还在反向传播时“重算”得分矩阵，不保存 N×N 的 S——多花一点计算，换回大量 HBM 访问的节省。

老火把这招称作“看菜下锅、随时回勺”。不管后面的食材多惊艳，锅里的味道永远能精确对齐，绝不翻车。

V2并行优化

V1 的炖菜虽然不占桌面了，老火却看出新问题：主厨一次只盯一个锅，旁边灶台全空着；而且每端一盘菜回来，都要停下来把锅里的酱汁重新搅一遍（归一化），搅汁时所有人都在等。这是对 GPU 并行能力的巨大浪费。

V2 的解法是“多灶同开”：不再让一个线程块串行处理所有 K/V 块，而是让每个线程块独立负责一个 Q 块，各锅各炒、互不干扰。同时把“搅汁”从每盘都做的精细活改成“出锅时统一调味”——只在最终写回前做一次缩放，而不是在每个分块内反复做 rescaling。这一下就砍掉了大量非矩阵乘运算。

更进一步，V2 把 Q 的缩放从 kernel 内部挪到 kernel 外部预计算；把 mask 和 bias 合并进得分计算阶段。核心里剩下的几乎全是 GPU 张量核心最擅长的矩阵乘法。论文数据显示，在 A100 上 V2 相比 V1 实现了约 1.7~3.8 倍的加速。

老火看着一排灶台同时开火，总结道：以前是一锅一锅等，现在是一排一排炒。真正快的厨房，不是把动作做得更快，而是让所有灶台都在干活。

V3增强

到了第三代，老火换了新厨房——Hopper 架构 GPU。新厨房里添了一套“高精度快速秤”（FP8 张量核心）：称料速度快一大截，刻度稍微粗一点点，但对做菜来说足够。更妙的是，他把整个备菜流程改成流水线：一个帮厨专管端菜，一个专管炒菜，一个专管把成品送出去，三人各忙各的，互不等候。

技术上，V3 用 FP8 做矩阵乘法输入，在 FP32 中累加，并配合精细的缩放因子管理数值范围——这是“快速秤加校验秤”的组合。同时，V3 用 warp specialization 实现异步流水线：一部分 warp（线程束）专门负责从 HBM 加载数据、用 TMA 异步写回结果，另一部分 warp 专心做矩阵计算，加载与计算完全重叠，GPU 等待内存的时间被压到最低。

V3 还把前向和反向统一成单一内核，减少了内核启动次数；针对 Hopper 的硬件特性重新分配了共享内存的读写节奏。在 H100 上，V3 相比 V2 的 BF16 版本能再提升约 1.5~2 倍的吞吐，刷新了注意力计算的速度纪录。

老火的新厨房里三个人同时开工，成品源源不断送出。他终于明白：当搬运无法避免时，让它和做饭同时发生，就是最好的等待。

性能对比

三代菜谱摊开看：标准做法是整桌食材一次搬上桌，桌面爆炸；V1 是分盘端菜，桌面清爽但厨师反复回勺；V2 是多灶同开、出锅统一调味，后厨火力全开；V3 是快速秤加流水线，产量再创新高。三代的共同点：菜谱（输出结果）几乎没变，变的只是搬运策略。

用数据说话：Flash Attention 几乎没有减少计算量——FLOPs 与标准注意力相同；真正的收益全部来自 HBM 访问量的下降。标准注意力要读写 O(N²) 大小的中间矩阵，而分块之后，每次只搬一个能装进 SRAM 的小块。

精度方面，在线 softmax 数学等价，实测误差通常在 1e-6 量级；V3 的 FP8 模式用 FP32 累加加缩放因子，也基本保持 BF16 质量。如今 Flash Attention 已内置于 PyTorch SDPA、HuggingFace Transformers、vLLM 等主流框架，普通开发者几乎无感地享受着这份加速。

老火把三年厨房账本挂上墙，写下结论：“算得少不如搬得少，搬得少不如搬得巧。”这就是 Flash Attention 的全部秘密。

AI 应用场景

Python 演示

import numpy as np  def flash_attention_v1(Q, K, V, block_size=32):    """Flash Attention V1 简化教学版(CPU numpy 模拟)"""    n, d = Q.shape    O = np.zeros((n, d))         # 未归一化输出累积    m = np.full(n, -np.inf)      # 每行运行最大值    l = np.zeros(n)              # 每行运行分母     # 外层循环: 遍历 Q 的块    for i in range(0, n, block_size):        qi = Q[i:i + block_size]        m_i = m[i:i + block_size].copy()        l_i = l[i:i + block_size].copy()        o_i = O[i:i + block_size].copy()         # 内层循环: 遍历 K/V 的块        for j in range(0, n, block_size):            kj = K[j:j + block_size]            vj = V[j:j + block_size]             # 当前块的得分矩阵 (Br, Bc)            s = qi @ kj.T / np.sqrt(d)            # 更新运行最大值            m_new = np.maximum(m_i, s.max(axis=1))            # 当前块的 exp 值(用新最大值做平移)            p = np.exp(s - m_new[:, None])            # 更新运行分母: 旧值按指数缩放 + 新块和            l_i = l_i * np.exp(m_i - m_new) + p.sum(axis=1)            # 更新运行输出: 旧值按指数缩放 + 新块贡献            o_i = o_i * np.exp(m_i - m_new)[:, None] + p @ vj            m_i = m_new         # 写回该块: 未归一化输出除以最终分母        O[i:i + block_size] = o_i / l_i[:, None]        m[i:i + block_size] = m_i        l[i:i + block_size] = l_i     return O  def softmax_attention(Q, K, V):    """标准注意力(对照实现)"""    d = Q.shape[1]    s = Q @ K.T / np.sqrt(d)    p = np.exp(s - s.max(axis=1, keepdims=True))    return (p / p.sum(axis=1, keepdims=True)) @ V  if __name__ == "__main__":    np.random.seed(42)    n, d = 128, 16    Q = np.random.randn(n, d)    K = np.random.randn(n, d)    V = np.random.randn(n, d)     out = flash_attention_v1(Q, K, V)    ref = softmax_attention(Q, K, V)    err = np.abs(out - ref).max()    print(f"max error: {err:.2e}")    assert err < 1e-6, "分块实现与标准实现不一致"    print("OK: Flash Attention V1 matches standard attention")

1. O = np.zeros((n, d))初始化未归一化输出累积，这是最终结果 O 的“锅”。
2. m = np.full(n, -np.inf)每行运行最大值初始化为负无穷，表示“还没见到任何食材”。
3. for i in range(0, n, block_size)外层循环：遍历 Q 的块，相当于一桌一桌翻台。
4. for j in range(0, n, block_size)内层循环：遍历 K/V 的块，把当前 Q 块和所有 K/V 块配对。
5. s = qi @ kj.T / np.sqrt(d)计算当前块的得分矩阵，对应公式 $S_j=Q_iK_j^\top/\sqrt{d}$。
6. m_new = np.maximum(m_i, s.max(axis=1))更新运行最大值，判断新块是不是新的“鲜度上限”。
7. p = np.exp(s - m_new[:, None])用新最大值做平移，得到数值稳定的指数值。
8. l_i = l_i * np.exp(m_i - m_new) + p.sum(axis=1)更新运行分母：旧分母按高度差缩放，加上新块和。
9. o_i = o_i * np.exp(m_i - m_new)[:, None] + p @ vj更新运行输出：旧输出缩放后叠加新块对 V 的加权贡献。
10. O[i:i + block_size] = o_i / l_i[:, None]该块全部处理完后，未归一化输出除以最终分母，写出最终结果。

常见误区

面试题