周师傅是城南酒楼最有名的掌勺，今晚他接到一单“百料佛跳墙”。他的灶台边有一块小案板，转身就能拿到调料，可惜一次只摆得下几碗料；十几米外的仓储区能装下几百种食材，可每次取货都要推车穿过半个后厨，来回一趟气喘吁吁。案板就是 GPU 里的 SRAM，仓储区就是 HBM（高带宽显存）。周师傅很快发现：一锅菜从备料到出锅，真正下锅翻炒不过十几分钟，大半时间全耗在“搬”上。 标准注意力算法正是那个“爱搬运”的厨师：它算出两个巨大的中间矩阵 S 和 P，用完就塞回仓库，下一步又搬回来。Flash Attention 则完全不同——它把食材切成小块，每次只搬一小批上案板，在案板上顺手完成所有调味（在线 Softmax），中间产物绝不回仓库。数学结果与标准算法分毫不差，但仓库到案板的搬运量少了一个数量级。 这篇文章会带你拆解这套“流水案板”的全套功夫：先看清 HBM 与 SRAM 的层次结构，再算算标准注意力浪费了多少搬运，接着学 Tiling 分块、在线 Softmax、复杂度量化，最后落到工程实现与调优经验。学完你就能理解：为什么大模型训练越来越依赖 Flash Attention 这类 IO 感知算子，以及 GPU 上“算得快”和“跑得快”之间的距离到底有多远。

岗位关联

对 AI 从业者来说，这是一堂性价比极高的性能课：搞懂“内存搬运”这个隐形瓶颈，你就能看懂大模型框架里 80% 的 attention 加速手段，无论是训练长上下文、优化推理延迟还是算子选型，都能一眼看穿底层原理。

目录
  1. 可视化演示
  2. 显存层次结构
  3. 标准注意力瓶颈
  4. 分块Tiling策略
  5. 在线Softmax算法
  6. IO复杂度量化
  7. 工程实现细节
  8. 性能调优经验
  9. AI 应用场景
  10. Python 演示
  11. 常见误区
  12. 面试题

可视化演示

本节配有交互动效（进入视口后自动加载，约 3KB）。

显存层次结构

周师傅的厨房里，最难的不是颠勺，而是“取料”。案板就在手边，转身能抓到葱姜蒜，可一次只摆得下几碗；仓库能装几百种食材，但每次往返都像穿越半个后厨。他掐表算过：一锅菜真正下锅翻炒不过十几分钟，运料却耗掉大半个钟头。所以周师傅立下第一条规矩：能一次拿齐的料，绝不跑第二趟。

GPU 的“两块地”和厨房如出一辙。HBM（High Bandwidth Memory）是显存，容量以 GB 计，A100 上就有 40—80GB，但带宽只有约 2TB/s；SRAM 是片上缓存，紧贴计算单元，每个 SM 只有一百多 KB，却能提供约 19TB/s 的带宽，快了近一个数量级。任何计算都要先把数据从 HBM 搬进 SRAM 才能动手，离计算单元越近，取数越快，地方也越小。

于是决定算子快慢的关键往往不是 FLOPs（每秒浮点运算次数），而是搬运时间：搬运时间 = 数据量 / 带宽。一个 N×N 矩阵要从 HBM 读一遍，耗时就是 N² 个元素除以 2TB/s。像 attention 这种“中间矩阵又多又大”的算子，天然是 memory-bound（受内存带宽限制）——计算单元常常饿着肚子等数据，空转的算力就是被浪费的锅气。

周师傅的结论是：少跑仓库比加快颠勺更值钱。这决定了后面所有策略的方向——不是把菜做得更复杂，而是把“搬运”降到最低。理解了这一点，你就握住了 Flash Attention 的灵魂。

T=DBb\displaystyle T = \frac{D}{B_b}
搬运时间只由“搬多少”和“搬多快”决定。想提速，要么少搬，要么搬得快——Flash Attention 选择的是“少搬”。
  1. D（数据量）需要搬运的元素个数或字节数。比如读 Q、K 两个 N×d 矩阵，D=2Nd；而一个 N×N 的中间矩阵，D=N²，N=8192 时高达 6700 万个元素。
    D=2Nd或D=N2\displaystyle D = 2Nd \quad \text{或} \quad D = N^2
  2. B_b（搬移带宽）单位时间能搬多少数据。HBM 约 2TB/s，SRAM 约 19TB/s，差近一个数量级，这就是“仓库一趟顶案板十趟”。
    BbHBM=2.0 TB/s,BbSRAM=19 TB/s\displaystyle B_b^{HBM}=2.0\ \text{TB/s},\qquad B_b^{SRAM}=19\ \text{TB/s}
  3. T（搬运时间）数据量除以带宽，得到纯粹“在路上”的耗时。带宽不变时，数据量翻倍，耗时就翻倍；反过来，只要把 D 砍掉一个量级，时间就省下一个量级。
    T=DBb\displaystyle T = \frac{D}{B_b}

标准注意力瓶颈

老办法做佛跳墙是这样的：把几十种食材全部切好、焯好，分装成半成品送回仓库冷藏；要炖的时候，再一盒盒搬回案板。听上去很有条理，问题是一道菜来回要搬四趟——仓库门口排起长队，所有厨师都在等货。周师傅知道，这菜不是做不快，是被“搬运节奏”拖死的。

标准 attention 也这么干。它有三步：先算 S=QKᵀ，得到一个 N×N 的分数矩阵；再对 S 按行做 softmax 得到同样 N×N 的 P；最后算 O=PV。问题是 S 和 P 都是 N×N 的“半成品”。GPU 的 SRAM 根本装不下整个 N×N 矩阵（N=8192 时 S 有 6700 万个数），所以只能写回 HBM，下一步再读回来——半成品在仓库和案板之间反复折腾。

细数这笔账：S 写回一次、读回一次，P 写回一次、读回一次，光是这两个矩阵就产生 4N² 次元素搬运；再加上 Q、K、V、O 各一趟 Nd，总计约 4N²+4Nd。N 从 1K 涨到 8K，搬运量暴涨 64 倍。而 GPU 的算力余量远大于带宽余量，于是“数据在路上”的时间彻底盖过“数据在计算”的时间，SM 的算力大片大片刻在等数据。

周师傅一看账本：上菜太慢不是因为锅小，而是因为半成品在仓库和案板之间反复折腾。好，那就立新规矩——半成品一律不许回仓库。这就是 Flash Attention 的起点：让 S 和 P 永远不出案板。

IOstd=4N2+4Nd\displaystyle \text{IO}_{std} = 4N^2 + 4Nd
标准注意力的 HBM 搬运量被 N² 项支配：N 每翻一倍，搬运涨四倍。这就是大模型序列变长后训练越来越慢的根源。
  1. S 写回 HBM算完 S=QKᵀ 后，N×N 的 S 无法留在 SRAM，必须先存到显存。
    S→HBM: N2\displaystyle S \to HBM:\ N^2
  2. S 读回算 softmaxsoftmax 要按行处理 S，又把它从 HBM 读回来，同时把结果 P 写回显存。
    S←HBM: N2,P→HBM: N2\displaystyle S \gets HBM:\ N^2, \qquad P \to HBM:\ N^2
  3. P 再读回算 PV最后一步 O=PV 要从 HBM 读回 P 才能和 V 相乘。
    P←HBM: N2\displaystyle P \gets HBM:\ N^2
  4. Q、K、V、O 各一趟这四个 N×d 矩阵各读/写一次，贡献 4Nd，在 N 很大时可以忽略。
    IOstd=2N2+2N2+4Nd=4N2+4Nd\displaystyle \text{IO}_{std} = 2N^2 + 2N^2 + 4Nd = 4N^2 + 4Nd

分块Tiling策略

新规矩立下后，周师傅换了个思路：既然案板只放得下几碗料，那就把整锅菜拆成上百批，每批刚好摆满案板；切一批、炖一批、装一批，绝不让任何半成品离开案板。这样做菜慢了吗？不，反而快了——因为没人再整天推着小车往返仓库，所有人的时间都花在真正做菜上。

Flash Attention 的分块 Tiling 正是如此：把 Q、K、V 按行切成若干块，每次只取一个 Q 块和若干 K、V 块进 SRAM，在片上完成小块的 QKᵀ、softmax、PV，再更新输出 O 的对应行块。S、P 这两个 N×N 矩阵从头到尾只存在于 SRAM 里，压根不会出现在 HBM 中——半成品不出案板。

块大小怎么定？SRAM 里同时要放四样东西：Q 块、K 块、V 块、O 块，每块形状是 B×d，四份加起来 4Bd 个元素，不能超过 SRAM 容量 M。于是得到经典结论 B=M/(4d)。外层循环遍历 Q 的行块，内层循环遍历 K、V 的列块；内层每算完一块，就把这一批的“味道贡献”累加到 O 上，等外层全部跑完，O 就是完整的注意力输出。

周师傅把这套流程叫作“流水案板”——案板永远是满的，仓库永远是被冷落的。唯一的问题是：分批调味，怎么保证最后拼出来的味道和一次性下锅一模一样？这就轮到在线 Softmax 登场了。

B=M4d\displaystyle B = \frac{M}{4d}
块大小不是拍脑袋定的：SRAM 要同时装下 Q、K、V、O 四块，每块 B×d，四份加起来必须小于等于 M。案板有多大，批次就切多大。
  1. M（SRAM 容量）A100 上每个 SM 的 SRAM 约 192KB，以 fp16 计算大约能装 96000 个元素。
    M≈192 KB≈96000 个元素\displaystyle M \approx 192\ \text{KB} \approx 96000\ \text{个元素}
  2. 四块矩阵瓜分 M案板上要同时摆 Q 块、K 块、V 块、O 块，每块 B×d，总占用 4Bd 不能超限。
    4Bd≤M\displaystyle 4Bd \le M
  3. 解出块大小不等式两边同时除以 4d，得到块大小的上限，实际取整为向上取整。
    B=⌈M4d⌉\displaystyle B = \left\lceil \frac{M}{4d} \right\rceil

在线Softmax算法

难题来了。佛跳墙的汤汁讲究“整体平衡”：要先知道整锅最咸的底味，才能决定每味调料的比例。可周师傅一次只炖一批，怎么知道后面批次的咸淡？他的办法是在围裙口袋里揣一本小本子，记两个数字：见过的最咸的一勺（m），以及当前所有调料的比例累计（l）。每批新汤端上案板，先尝一口：若更咸，就把本子上旧的比例整体打个折扣，再混入新批次的贡献。

这正是在线 Softmax 的核心。行 softmax 需要全行的最大值和指数和，但分块后每块只有局部最大值。算法的巧妙之处在于：维护“运行中的”统计量（m 是见过的最大的分数，l 是指数归一和），每处理一个新块做一次合并更新，输出 O 也随之重缩放——小本子越记越厚，但永远只需要记两笔。

具体来说，假设已处理若干块，当前统计量是 (m, l, Õ)。新块分数 S_ij 的局部最大值为 m_ij：先比大小得到 m_new；旧输出和旧统计都建立在旧最大值 m 上，需要乘以折扣 r=e^{m-m_new}（若新最大值更大，折扣小于 1，相当于把整锅味道按新基准“兑淡”）；再累加新块的指数贡献与加权 V，最后统一除以 l_new 得到精确的 softmax 输出。整个过程数学上与标准 softmax 完全等价，同时天然继承了“减最大值防溢出”的数值稳定性技巧。

周师傅的小本子越记越顺：见过最咸的一勺，记下来；总比例，记下来。最后出锅时哪怕分了几百批，勾兑出来的汤汁与一次猛火炖出的味道分毫不差。分批没有牺牲味道，只是省掉了搬运。

mnew=max⁡(m, mij),r=em−mnew\displaystyle m_{new} = \max(m,\ m_{ij}), \qquad r = e^{m - m_{new}}
lnew=r l+∑jesij−mnew,O~new=r O~+∑jesij−mnew vj\displaystyle l_{new} = r\,l + \sum_{j} e^{s_{ij} - m_{new}}, \qquad\tilde{O}_{new} = r\,\tilde{O} + \sum_{j} e^{s_{ij} - m_{new}}\,v_j
O=O~newlnew\displaystyle O = \frac{\tilde{O}_{new}}{l_{new}}
在线 Softmax 用“运行最大值 m + 运行指数和 l”两笔统计量，把分块算出的局部结果无缝拼接成全局结果，数学上分毫不差。
  1. 比大小得到新基准 m_new本块的局部最大值和运行中的最大值比大小，大者担任新的行最大值，保证后续指数运算不会溢出。
    mnew=max⁡(m, mij)\displaystyle m_{new} = \max(m,\ m_{ij})
  2. 折扣因子 r旧统计量建立在旧最大值 m 上，换了新基准后要整体乘以 r=e^{m-m_new}；若新最大值更大，r 小于 1，相当于把旧汤汁按新咸度“兑淡”。
    r=em−mnew\displaystyle r = e^{m - m_{new}}
  3. 更新指数归一和 l旧和乘上折扣，加上新块所有分数的指数贡献，得到合并后的分母。
    lnew=r l+∑jesij−mnew\displaystyle l_{new} = r\,l + \sum_j e^{s_{ij} - m_{new}}
  4. 重缩放输出累加器 Õ未归一化的输出同样先乘折扣，再叠加新块的加权和（softmax 分数乘 V 向量）。
    O~new=r O~+∑jesij−mnew vj\displaystyle \tilde{O}_{new} = r\,\tilde{O} + \sum_j e^{s_{ij} - m_{new}}\,v_j
  5. 归一化出锅全部块处理完后，把累加器除以 l_new，得到与标准 softmax 逐元素一致的最终输出。
    O=O~newlnew\displaystyle O = \frac{\tilde{O}_{new}}{l_{new}}

IO复杂度量化

光说“少跑仓库”不够，周师傅有本账：老办法一道菜在仓库和案板之间来回折腾，新办法案板流水作业，总搬运量到底省了多少？他专门把每次搬运的件数记下来，算了三天三夜。账算明白了，才知道这新流程值不值得铺开。

Flash Attention 论文也算了这笔账。标准注意力搬运量是 4N²+4Nd，与 N² 成正比，N 每翻一倍流量涨四倍，这是 attention 在大模型训练中越来越吃时间的原因。Flash Attention 的搬运量则取决于 SRAM 容量 M：块大小 B=M/(4d)，内层循环要遍历 N/B 轮，每一轮 K 和 V 各被读一块，所以 K、V 各被重复读取 N/B 次。

算下来，Flash Attention 的 HBM 访问量约 8N²d²/M（外加读一次 Q 和写一次 O 的 2Nd，可忽略）。两者之比约 M/(2d²)。以 A100 为例：每 SM 的 SRAM 约 96000 个 fp16 元素，取 d=64，则 M/(2d²)=96000/8192≈11.7——HBM 流量直接少一个数量级。注意这里只省“搬运”，计算量一分没少，Flash Attention 是精确算法，不是近似算法。

周师傅把账本一拍：省下的不是调味料，而是搬运工的工钱。案板越大（M 越大）、一批能放的料越多，省得越多；头维度 d 越大，一批能放的“行”越少，折扣就小些。这笔账清清楚楚——少搬一个数量级，速度自然上一个台阶。

IOflash≈8N2d2M,加速比≈IOstdIOflash≈M2d2\displaystyle \text{IO}_{flash} \approx \frac{8N^2d^2}{M}, \qquad\text{加速比} \approx \frac{\text{IO}_{std}}{\text{IO}_{flash}} \approx \frac{M}{2d^2}
HBM 流量的加速比约等于 M/(2d²)。SRAM 越大、头维度 d 越小，省得越多；典型 A100 配置下省约 10 倍。
  1. 内层循环轮数 T每个 Q 块要遍历 K、V 的所有列块，轮数等于 N/B，代入 B=M/(4d) 得 4Nd/M。
    T=NB=4NdM\displaystyle T = \frac{N}{B} = \frac{4Nd}{M}
  2. K、V 被反复读取每一轮都要读 K、V 各一块，总共读 2Nd·T 个元素。
    2Nd×T=8N2d2M\displaystyle 2Nd \times T = \frac{8N^2d^2}{M}
  3. 与标准流量作比标准的 4N² 除以 Flash 的 8N²d²/M，化简得到 M/(2d²)。
    4N28N2d2/M=M2d2\displaystyle \frac{4N^2}{8N^2d^2/M} = \frac{M}{2d^2}
  4. 代入典型数值A100 上 M≈96000（fp16 元素）、d=64，加速比约 11.7 倍。
    960002×642≈11.7\displaystyle \frac{96000}{2 \times 64^2} \approx 11.7

工程实现细节

想把流水案板真正落地，光有想法不够。周师傅把三大工序——切配、焯水、下锅——合成同一个灶台流程，中间不再换人、换车、换库房；灶台上炖着的同时，徒弟们已在备下一批料，这叫“边算边搬”。后厨的动线彻底重组，半成品永远不落地。

工程上对应三件事。第一是 Kernel 融合：传统方案要启动多个 kernel（QKᵀ、softmax、PV），每个 kernel 都要从 HBM 读中间结果；Flash Attention 把三步写进同一个 CUDA kernel，一次启动、全程在 SRAM 内完成。第二是反向传播优化：标准反向需要 S、P 这两个 N×N 矩阵，Flash Attention 不存它们，而是从已保存的输出统计量 l 和 O 现场重计算，用额外计算换掉 N² 量级的显存。第三是并行度设计：每个 head 独立计算，batch×head 的并行度天然很高，多头之间互不干扰。

由此带来的显存收益是质的飞跃：前向不再需要为 S、P 申请 N×N 的显存块。标准注意力显存占用 O(N²)，Flash Attention 降到 O(N)——只有 Q、K、V、O 和每行的统计量 m、l。论文报告在 GPT-2 训练上端到端提速约 15%，注意力算子本身快约 3 倍；序列越长、头越多，省下的显存越可观。

周师傅的灶台改造完成：一个灶眼从头做到尾，中间不落地；徒弟流水线备料，灶火不停。仓库空了，锅气足了。显存不再是长序列的天花板，算力才真正花在了刀刃上。

Memstd=O(N2),Memflash=O(N)\displaystyle \text{Mem}_{std} = O(N^2), \qquad \text{Mem}_{flash} = O(N)
显存占用从 N² 降到 N：S、P 不再落显存，长序列模型终于不用靠“切段”硬撑了。
  1. S、P 是显存大户标准算法要为两个 N×N 矩阵各开一块显存，N 稍大的中型模型就会吃进数 GB。
    Memstd=2N2+O(Nd)\displaystyle \text{Mem}_{std} = 2N^2 + O(Nd)
  2. Flash 只留输出与统计量S、P 不落 HBM，显存只剩 Q、K、V、O 和每行的 m、l。
    Memflash=4Nd+O(N)\displaystyle \text{Mem}_{flash} = 4Nd + O(N)
  3. 量级对比N 很大时 4Nd 远小于 2N²，所以一个是 O(N²)，一个是 O(N)。
    O(N)≪O(N2)\displaystyle O(N) \ll O(N^2)

性能调优经验

新流程跑通了，周师傅没闲着：案板如果大了一圈，每批就能多切几刀；菜场送货的车多了，他调整批次让案板永远不空。他把这套“按硬件调批次”的经验记成厨艺手册，谁家灶台什么尺寸，一看便知怎么配菜。

调优的核心是“匹配硬件尺寸”。SRAM 越大，B=M/(4d) 越大，内层循环轮数越少，K、V 被重读的次数越少——H100 的 SRAM 更大、HBM 带宽翻倍，Flash Attention 的收益自然更明显。序列越长，N² 项越突出，省流量的优势越放大，这正是长上下文模型把 FlashAttention 当作标配的原因：8K、32K 甚至 128K 上下文，没有一个能离开它。

实践中的调优点包括：块大小不是越大越好，过大会挤占 SRAM 导致线程占用率下降；kernel 内部要组织共享内存的读写避免 bank conflict；用双缓冲把 HBM 搬运与 SRAM 计算重叠，让“搬运”和“计算”并行进行；反向传播时 m、l 的保存可以压缩存储。FlashAttention-2 进一步改进了线程块并行策略，让长序列上的吞吐再上一个台阶；配合 MQA/GQA 等变体减少 KV 读取，推理场景还能再省一笔带宽。

周师傅的经验是：算法是骨架，硬件是灶台，调优就是按灶台的火力来配批次。知道瓶颈在“搬”不在“算”，你的每一分调优力气都会花在刀刃上——这也是 Flash Attention 系列算子给整个 AI 工程界上的一课。

AI 应用场景

大模型长上下文训练加速

阶段：模型训练

解决的问题：标准 attention 的 HBM 流量是 O(N²)，序列从 1K 拉到 8K 时显存与带宽双双告急，训练几乎被搬运拖垮。Flash Attention 用 Tiling 加在线 Softmax 把中间矩阵留在 SRAM，流量少一个数量级，长序列才跑得动。周师傅的道理：半成品不出案板，仓库就不再是瓶颈。

具体案例：在 GPT-2 中等规模训练中，把 attention 替换为 FlashAttention kernel，论文实测端到端训练提速约 15%，注意力算子本身快约 3 倍；当上下文从 1K 扩展到 8K 时，显存占用依然平稳，训练不必再频繁做序列截断。

推理服务降本增效

阶段：推理部署

解决的问题：自回归解码每个 token 都要读取全部 KV cache，推理延迟被 HBM 带宽卡死。减少 KV 读取次数与搬运量，直接决定单卡吞吐。周师傅的小本子哲学在这里变成：能少搬一趟就少搬一趟，延迟立刻降下来。

具体案例：在推理引擎（如 vLLM、TGI）中把 attention kernel 替换为 FlashAttention 实现，7B 级别模型的单卡并发数可提升 20%—50%，首 token 延迟明显下降；配合 MQA/GQA 减小 KV cache 体积，带宽压力进一步缓解。

超长文档问答服务

阶段：应用服务 / 数据接入

解决的问题：十万 token 的合同、论文要做摘要与问答，标准 attention 的显存 O(N²) 直接爆显存；Flash Attention 的显存 O(N)，消费级显卡也能处理超长上下文，不再需要把文档暴力切段导致语义断裂。

具体案例：在 24GB 显存的单卡上，用内置 FlashAttention 的推理框架加载 32K—128K 上下文模型，对整份年报或论文直接提问，文档无需预切分，回答质量显著优于分段检索方案。

Python 演示

python47 行
import numpy as np def softmax(x):    m = x.max(axis=-1, keepdims=True)    e = np.exp(x - m)    return e / e.sum(axis=-1, keepdims=True) def standard_attention(Q, K, V):    S = Q @ K.T    P = softmax(S)    O = P @ V    return O def flash_attention(Q, K, V, B=2):    N, d = Q.shape    O = np.zeros((N, d))        # 未归一化输出累加器    m = np.full(N, -np.inf)     # 每行运行最大值    l = np.zeros(N)             # 每行运行指数和    for i in range(0, N, B):        Qi = Q[i:i + B]        for j in range(0, N, B):            Kj = K[j:j + B]            Vj = V[j:j + B]            Sij = Qi @ Kj.T            m_ij = Sij.max(axis=1)            new_m = np.maximum(m[i:i + B], m_ij)            r = np.exp(m[i:i + B] - new_m)            pij = np.exp(Sij - new_m[:, None])            l_ij = pij.sum(axis=1)            O[i:i + B] = O[i:i + B] * r[:, None] + pij @ Vj            l[i:i + B] = l[i:i + B] * r + l_ij            m[i:i + B] = new_m        O[i:i + B] = O[i:i + B] / l[i:i + B, None]    return O if __name__ == "__main__":    np.random.seed(0)    N, d = 4, 3    Q = np.random.randn(N, d)    K = np.random.randn(N, d)    V = np.random.randn(N, d)    O1 = standard_attention(Q, K, V)    O2 = flash_attention(Q, K, V, B=2)    err = np.abs(O1 - O2).max()    print("max error:", err)    assert err < 1e-10, "flash attention is not equivalent!"    print("flash attention equals standard attention")
目的：用 NumPy 从零实现 Flash Attention 的分块 Tiling 与在线 Softmax，并验证它与标准注意力逐元素数学等价。运行时会把 B 设为 2，模拟“小案板分批炖汤”的过程。
  1. m = np.full(N, -np.inf)每行维护一个“见过的最咸分数”，初始负无穷保证第一块直接接管。
  2. new_m = np.maximum(m[i:i+B], m_ij)新老最大值比大小，谁大谁当新的行最大值 m。
  3. r = np.exp(m[i:i+B] - new_m)重缩放折扣：旧统计建立在旧最大值上，换基准要打这个折扣。
  4. O[i:i+B] = O[i:i+B] * r[:, None] + pij @ Vj旧输出按折扣重缩放后，叠加新块的 P·V 贡献。
  5. l[i:i+B] = l[i:i+B] * r + l_ij指数归一和同步更新，这是在线 softmax 的“勾兑比例”。
  6. O[i:i+B] / l[i:i+B, None]全部块算完后统一归一，得到与标准 softmax 一致的输出。

常见误区

误区名字带 Flash，估计是用随机采样或稀疏化来近似注意力吧？不然怎么会快那么多？

正解它不是近似方法。Flash Attention 的计算量、参与的位置、softmax 的分数一个都不少，只是重排了执行顺序：分块后让中间矩阵 S、P 只留在 SRAM，不再写回 HBM。数学上与标准 softmax 逐元素等价，端到端精度无损失。就好比周师傅把佛跳墙分几十批炖，料一点没少，只是省掉了来回仓库的搬运。

为何常见：Flash 一词容易让人联想到之前的稀疏/近似注意力（如 Longformer、Linformer），加上论文主打“快”，读者先入为主地脑补成近似。

误区我这显卡显存几十 GB，中间矩阵写回显存有什么大不了的？容量够用就行。

正解瓶颈是带宽，不是容量。写一次 N×N 的 S 矩阵虽然“放得下”，但要把 HBM 带宽占满好几毫秒，而计算单元在几微秒内就把数据“饿死了”。算得快、等得久，正是 memory-bound 算子的典型症状。周师傅的仓库再大，运一趟要半个钟头，菜就凉了。

为何常见：人们习惯用容量衡量“够不够”，GPU 宣传页上最显眼的也是显存大小，带宽这个隐性指标常被忽略。

误区既然分块是精髓，那把块设得越大，一次处理得越多，不是越快吗？

正解块大小被 SRAM 容量硬约束：B=M/(4d)，因为 Q、K、V、O 四块要同时驻留。超过 M 就会溢出到 HBM，搬移反而变多；块太小则 K、V 重读轮次 T=N/B 变大，HBM 流量上升。真实调优是在“正好塞满 SRAM”附近找最优，而不是无限增大。周师傅的案板就那么大，硬塞只会把料挤到地上。

为何常见：矩阵乘法中块越大缓存利用率越高是常见经验，容易类比过头；没注意到 SRAM 是稀缺资源，四份矩阵要抢同一块地方。

误区在线 softmax 是增量更新的，结果肯定和一次性算的 softmax 不一样，是不是牺牲了精度？

正解它在数学上等价于标准 softmax。维护的 m（运行最大值）和 l（指数和）经过合并更新后，与直接对整个行算出的最大值、和完全一致；唯一差别是浮点累加顺序，误差约 1e-6 量级，而且它还继承了“减最大值”的数值稳定性技巧。论文实验确认端到端精度无损失。

为何常见：“在线/增量”在不同领域常代表近似（如在线学习、随机梯度），读者按惯性理解；加上看不见完整的 P 矩阵，心里不踏实。

面试题

Q1. Flash Attention 加速的核心手段是什么？

单选题 · easy

  • A.减少矩阵乘法的计算量
  • B.减少 HBM 与 SRAM 间的搬运← 正确答案
  • C.用近似分数替代精确分数
  • D.降低 softmax 的数值精度

答案：B。Flash Attention 的计算量和标准注意力完全一样，它省的是“搬运”：S、P 不再写回 HBM。周师傅的道理——菜的做法没变，变的是不再把半成品送回仓库。C、D 都是近似或降精度思路，与 Flash Attention 的精确性相悖。

Q2. 关于 GPU 中 HBM 与 SRAM 的说法，正确的是？

单选题 · easy

  • A.HBM 容量大但带宽较低← 正确答案
  • B.HBM 容量小但带宽很高
  • C.SRAM 容量大且带宽很高
  • D.SRAM 容量大但带宽较低

答案：A。HBM 是“大仓库”，几十 GB 容量、约 2TB/s 带宽；SRAM 是“小案板”，每 SM 一百多 KB、约 19TB/s 带宽。B、C、D 都颠倒了容量与带宽的关系。记住：容量看 HBM，速度看 SRAM。

Q3. 标准注意力中反复往返 HBM 的中间矩阵是？

单选题 · mid

  • A.Q 和 K 两个输入矩阵
  • B.V 和 O 两个输出矩阵
  • C.S 和 P 两个 N×N 矩阵← 正确答案
  • D.只有一个 Q 矩阵

答案：C。标准注意力把 S=QK^T 和 P=softmax(S) 这两个 N×N 矩阵写回 HBM，下一步又读回，共 4N² 次搬运。Q、K、V、O 各只有一趟 Nd，不是主要矛盾。C 才是“来回折腾”的元凶。

Q4. Flash Attention 用到的关键技术包括？

多选题 · mid

  • A.分块 Tiling 把计算装入 SRAM← 正确答案
  • B.在线 Softmax 合并行统计量← 正确答案
  • C.随机掩码采样部分注意力
  • D.Kernel 融合避免中间写回← 正确答案

答案：ABD。A、B、D 正是 Flash Attention 的三大支柱。C 是近似注意力（如 BigBird）的思路，Flash Attention 不采样、不掩码，每个位置都精确计算，只是重排了计算顺序。

Q5. 在线 Softmax 中的符号 m 记录的是什么？

单选题 · mid

  • A.当前行已见过的最大分数← 正确答案
  • B.当前块的分数平均值
  • C.整个矩阵的分数之和
  • D.输出向量的模长平方

答案：A。m 是“见过的最咸一勺”——运行中的行最大值。新块局部最大值 m_ij 与它比大小，决定是否需要对旧输出打折扣重缩放。B、C、D 都不是 m 的含义；在线 softmax 的 m 只负责数值稳定性与重缩放基准。

Q6. 块大小公式 B=M/(4d) 中数字 4 表示？

单选题 · mid

  • A.每份数据要复制四份拷贝
  • B.SRAM 同时放 Q、K、V、O 四块← 正确答案
  • C.整个矩阵被切成四块处理
  • D.浮点数恰好占四个字节

答案：B。SRAM 里要同时驻留 Q 块、K 块、V 块和输出 O 块，每块 B×d，总元素 4Bd 不能超过 M，所以 B=M/(4d)。这是“案板刚够摆四碗料”的数学表达。A、C、D 都是想当然的猜测。

Q7. Flash Attention 中 K、V 被重复读取的轮次 T 是？

单选题 · hard

  • A.T=N/B，SRAM 越大轮次越少← 正确答案
  • B.T=N/B，SRAM 越大轮次越多
  • C.T=N/B，轮次与 SRAM 无关
  • D.T=d/B，与序列长度无关

答案：A。T=N/B=4Nd/M。B 越大每轮处理的 K 块越大，总轮数越少，K、V 重读次数越少，HBM 流量越低。这解释了为什么 SRAM 越大 Flash Attention 收益越大，也支撑了加速比 M/(2d²)。

Q8. 关于 Flash Attention 的 HBM 流量，正确的有？

多选题 · hard

  • A.标准注意力量级是 O(N²)← 正确答案
  • B.Flash 流量约 O(N²d²/M)← 正确答案
  • C.典型配置下少一个量级← 正确答案
  • D.Flash 的 FLOPs 也减半了

答案：ABC。A、B 是论文比较的两个复杂度；C 对应加速比 M/(2d²)≈10 倍的结果。D 错误：Flash Attention 不省任何计算，FLOPs 与标准相同，它只省 HBM 搬运。

Q9. Flash Attention 的结果与标准注意力的关系是？

单选题 · easy

  • A.精确等价，仅浮点舍入差异← 正确答案
  • B.有损近似，误差完全不可控
  • C.与标准的结果完全不同
  • D.只在短序列时近似成立

答案：A。在线 softmax 维护的 m、l 与一次算出 softmax 的统计量数学上一致，输出逐元素等价，只是浮点运算顺序不同带来约 1e-6 量级的舍入误差。B、C、D 都把它误解成了近似方法。

Q10. 已知 M=100K 元素、d=100，Flash 相比标准的 HBM 流量加速比约为？

单选题 · extreme

  • A.约 1 倍，流量完全没有减少
  • B.约 5 倍，与 M/(2d²) 估算吻合← 正确答案
  • C.约 10 倍，对应加速比 d²/M
  • D.约 1000 倍，因为访问归零

答案：B。M/(2d²)=100000/(2×100²)=5，所以选 B。C 把公式记反了（d²/M 远小于 1）；A 忽略了两者常数差异；D 错误，Flash 仍有 Q/K/V 搬入和 O 搬出，HBM 访问不会归零。周师傅会说：仓库省了九成人力，但不是零人力。