注意力为什么受限于显存带宽：用算术强度与分块把 FlashAttention 推导清楚 - 技术文章

2022 年之前，让 Transformer 处理长序列几乎只有一条路：想办法把自注意力的计算量从平方降到近线性。稀疏注意力、低秩近似、各种「近似注意力」层出不穷——它们在数学上都把 FLOPs 砍下来了，可真正跑起来却常常没有墙钟加速，也就没被广泛采用。FlashAttention 论文（arXiv:2205.14135，2022）给出的诊断一针见血：这些方法盯着 FLOPs，却忽略了一件更朴素的事——数据在显存里搬来搬去的成本。

本文要做的不是复述「FlashAttention 更快」这个结论，而是把它为什么快、快到什么程度、代价是什么从头推一遍。主线只有一条：一个算子的瓶颈到底在算力还是在带宽？判据是「算术强度」与 Roofline 模型；一旦判出注意力其实是带宽受限，FlashAttention 的全部设计——分块（tiling）、在线 softmax、反向重算、精细的线程划分——就都顺理成章。最后我们会看到，它甚至增加了总 FLOPs，却让墙钟时间下降——只要你接受「瓶颈在搬运而不是计算」这个前提，这笔账立刻就算得过来。

按本站惯例：凡引用他人结论处标注出处（论文名 + 章节或摘要）；凡本文自己算出来的数字（例如交叉点、显存流量估计），一律在文末「本文自行推导」一节单列，与官方数字分列，且只作量级一致性检查，不得作为官方数值引用。

一、先给直觉：GPU 是一座「算得快、喂得慢」的工厂

把 GPU 想象成一座工厂。**算力（FLOPS）**是工厂的产能；**显存带宽（bandwidth）**是把原料运进工厂的传送带。过去十几年，产能翻倍的速度一直快于传送带的翻倍速度——于是越来越多的算子不是「算不动」，而是「喂不饱」：工厂空转，等原料。

Horace He 在《Making Deep Learning Go Brrrr From First Principles》（2022）里把这件事归结为三类瓶颈：compute（真的在做浮点运算）、bandwidth（在显存里搬张量）、overhead（其余一切启动与调度开销）。判断你处在哪一类，能瞬间排除一大批无效优化：如果你在 bandwidth 受限区，给 GPU 堆更多 FLOPs 一点用都没有；反过来，如果你在 compute 受限区，把逻辑改写成更底层的语言来省 overhead 也帮不上忙。

关键在于：现代加速器上，「搬运」比「计算」贵得多。这不是一句修辞，而是一条可量化的存储层次事实。下面这张表把 GPU 上几个常见的存储层次按「谁快谁慢」排开（数量级仅供建立直觉，不同型号有出入）。

存储层次	位置	容量量级	带宽量级	术语
寄存器	每线程私有	几十至几百 B / 线程	最高	register file
SRAM（共享内存 / L1）	片内，每 SM	约百 KB / SM	约 19 TB/s（A100 全芯片聚合）	on-chip SRAM
HBM（显存）	片外	数十 GB	约 1.5–2.0 TB/s（全芯片；A100 40GB ≈1.55、80GB ≈1.94–2.04）	High Bandwidth Memory
NVLink（卡间）	机内 GPU 之间	—	数百 GB/s	inter-GPU
PCIe / 网络	机间	—	数十至数百 GB/s	inter-node

一句话读法：越靠近计算单元，越快、越小；越远，越慢、越大。这就是「非对称的存储层次」。FlashAttention 标题里的 IO-Awareness，说的就是别把数据在慢的那层（HBM）里反复倒腾。

以 A100 为例，它的片上 SRAM 聚合带宽约 19 TB/s，而片外 HBM 只有约 1.5–2.0 TB/s（FlashAttention 论文的硬件背景段落给出这两个量级）。也就是说，从 SRAM 取数比从 HBM 取数快大约十倍。这个约十倍的落差，是全文一切推导的起点。

再补一记更反直觉的：在 GPU 上不是所有 FLOP 都等价。现代 GPU 把绝大部分算力押在矩阵乘（GEMM，通用矩阵乘）的专用单元（Nvidia 称之为 Tensor Core）上。A100 的 FP16 与 BF16 矩阵乘峰值约 312 TFLOPs/s（dense），而非矩阵乘的 FP32 只有约 19.5 TFLOPs/s——每个非矩阵乘 FLOP 比矩阵乘 FLOP 贵约 16 倍（FlashAttention-2 博客原文口径，2023）。这条事实在第九节还会回来咬人：它意味着「顺手多做几次加法或乘方」在 Tensor Core 时代可能比想象中贵。

二、符号约定与术语表

先把后文要用的符号、英文名与含义固定下来，避免读岔。

符号 / 术语	含义	备注
N	序列长度（sequence length）	注意力的「行 / 列」规模
d	头维度（head dimension）	单个注意力头的通道数，常见 64 / 128 / 256
Q, K, V	查询、键、值矩阵	形状均约为 N 乘 d（单头）
S = QKᵀ	注意力分数矩阵	形状 N 乘 N，问题就出在这里
P = softmax(S)	归一化后的注意力权重	同样 N 乘 N
O = PV	注意力输出	形状 N 乘 d
B_q, B_k	分块后的块大小（行数）	块大小由 SRAM 容量 M 与头维度 d 共同约束
HBM	High Bandwidth Memory	GPU 片外显存，慢、大（数十 GB）
SRAM	片上静态随机存储器	每 SM 的共享内存 / L1，快、小（百 KB 级）
M	每 SM 的 SRAM 容量	FlashAttention 复杂度里的关键分母（≠ 块大小 B）
SM	流多处理器（Streaming Multiprocessor）	GPU 的调度/执行单元，A100 有 108 个
warp	线程束	GPU 调度的基本单位，一组 32 个线程
GEMM	通用矩阵乘	现代 GPU 上由 Tensor Core 加速
state space model	状态空间模型	与注意力并列的另一类序列建模结构
tiling（分块）	把大矩阵切成能塞进 SRAM 的小块分次处理	FlashAttention 的核心手法
online softmax	在线 / 增量式 softmax 归一化	不看到整行也能算对
recomputation（重算）	反向传播时不存中间量、用重算代替读取	用 FLOPs 换 HBM 访问
arithmetic intensity（算术强度）	每搬运一字节所执行的 FLOPs	判瓶颈的核心量
roofline model	用算术强度刻画算子性能上界的模型	见第三节
I/O complexity（IO 复杂度）	算法在存储层次间搬运数据的量级	FlashAttention 论文的落脚点

一句话记住三个最容易混的量：N 是序列长度、d 是头维度、M 是 SRAM 容量。后文所有复杂度都以它们表达——请特别注意：N、d 是张量形状，M 是硬件容量参数，三者都不是模型规模，也与训练数据量无关。（另请留意 M 与块大小 B 的区别：M 是「片上能装多少」，B 是「一次处理多少行」，B 由 M 与 d 共同决定。）

三、形式化：算术强度、Roofline 与两类瓶颈

要给「算力受限还是带宽受限」一个可计算的判据，先定义算术强度：

I = \frac{\text{FLOPs performed}}{\text{Bytes transferred}}

它衡量「每搬运一字节，顺手做了多少计算」。搬到 HBM 的字节越少、算的 FLOPs 越多， $I$ 越大。

Roofline 模型给出算子的性能上界——实际吞吐会被「算力屋顶」和「带宽屋顶」中较低的那个卡住：

P_{\text{attainable}} = \min\big(\pi_{\text{peak}},\; I \cdot \beta\big)

其中 $\pi_{\text{peak}}$ 是算力峰值（FLOPs/s）， $\beta$ 是显存带宽（Bytes/s）。两条「屋顶」的交点，就是瓶颈的分水岭：

I^{*} = \frac{\pi_{\text{peak}}}{\beta}

当 $I < I^{*}$ ：算子的每个字节只换来很少的计算 → memory-bound（带宽受限），瓶颈是搬数据；
当 $I > I^{*}$ ：计算密度足够高 → compute-bound（算力受限），瓶颈是算。

一句话总结运行机制：光看 FLOPs 判断不了快慢，把 FLOPs 除以搬运的字节数、再和硬件的交叉点比，才知道该往哪儿使劲。

为什么这个判据比「看一眼 GPU 利用率高不高」这类问法靠谱？因为利用率是结果、算术强度是原因：一个带宽受限的算子在 A100 上跑出很低的算力利用率，并不是「没优化好」，而是它本就不可能更高——每搬运一字节只算几个 FLOP，天花板是由带宽定的，不是由 kernel 技巧定的。先算强度，再定方向，能省掉大量无效调参。

用一个具体的对照来体会这个判据。考虑两类极端算子：

矩阵乘：两个 N 乘 N 矩阵相乘，做约 $2N^{3}$ 次 FLOP，却只读 $2N^{2}$ 个元素、写 $N^{2}$ 个元素。算术强度随 N 线性增长，属于天生 compute-bound。
逐元素算子：比如一个激活函数，每个元素做 1 次 FLOP，却要把这个元素读进来再写回去（FP32 下往返 8 字节）。其算术强度只有约 $1 / 8 = 0.125$ FLOPs/byte。逐元素算子跑在 FP32 非矩阵乘单元上，其对应的交叉点约 $19.5 / 1.55 \approx 12.6$ FLOPs/byte——两者仍差约两个数量级，属于天生 memory-bound，再怎么调 kernel 也快不到哪去。

用 A100 的两个数估一下交叉点（这是本文的量级估算，见第十三节）： $\pi_{\text{peak}} \approx 312 \times 10^{12}$ FLOPs/s， $\beta \approx 1.55 \times 10^{12}$ B/s，于是

I^{*} = \frac{312 \times 10^{12}}{1.55 \times 10^{12}} \approx 2.0 \times 10^{2}\ \text{FLOPs/byte}

也就是说，每搬运一字节，得顺手做约 150 到 200 次浮点运算，才算「喂饱」了 A100（按 40GB 规格的带宽估算得约 200；若按 80GB 规格的约 2.0 TB/s，交叉点下移到约 156）。逐元素激活、归一化这类算子的算术强度远低于这个门槛——它们天生带宽受限。

下面是判据的流程图（图 1），它在后文会被反复用到。

flowchart TD
  X["给定算子 统计 FLOPs 与访存字节数"] --> I["算算术强度 I 等于 FLOPs 除以 Bytes"]
  I --> Q{"I 与交叉点 I星 比较"}
  Q -->|"I 小于 I星"| MB["带宽受限 memory bound"]
  Q -->|"I 大于 I星"| CB["算力受限 compute bound"]
  MB --> O1["优化方向 减少访存 分块 融合 提高复用"]
  CB --> O2["优化方向 张量核 降精度 提高利用率"]

图 1 说明了什么：先算算术强度，再和硬件交叉点比较，就能把优化方向一分为二——这正是后文判断「注意力该省 IO 还是省算力」的准则。

四、为什么早年的近似注意力没赢

在正式推导之前，先回答一个历史问题：既然注意力的成本随序列长度平方增长，为什么把 FLOPs 降到线性的近似方法，没能在工程上胜出？

原因有三层，值得逐层说清：

第一层：省的是算力，可瓶颈在带宽。 近似方法（稀疏、低秩、二者的组合）主要削减的是矩阵乘的 FLOPs。但按第三节的判据，注意力在长序列下越来越偏向带宽受限——它真正的成本在于把那个 N 乘 N 的中间矩阵在 HBM 上写来读去。只减 FLOPs、不动 IO，等于给一座「等原料」的工厂多买了几台机床，产能再高也闲置。

第二层：平方项的减少不等于墙钟的减少。 《FlashAttention》论文明确指出，很多近似方法虽然把计算复杂度降到线性或近线性，却没有显示出相对标准注意力的墙钟加速，因而没有被广泛采用。复杂度里的大 O 是一种渐近标尺，它既不告诉你常数项，也不告诉你实现里哪一步在搬数据。

第三层：近似是有代价的——换来了精度损失。 稀疏化、低秩化本质上都在改变模型的表达能力。与之相对，FlashAttention 反复强调自己是 exact attention（精确注意力，无近似）：它不改变注意力的数学结果，只是换一种算与被搬的顺序。对训练下游任务而言，「不损失精度就能加速」的吸引力，远大于「用一种近似换一点速度」。

三条路线的坐标轴对照

面对同一个平方成本，工程上有几条几乎正交的路线，别把它们混为一谈：

路线	改的是什么	是否损失精度	省的是	代表
朴素标准实现（会 materialize S、P 的那种）	基准	无	—	早期 PyTorch 式逐步实现
近似注意力	计算结构（稀疏 / 低秩）	有（换近似）	算力（FLOPs）	稀疏、低秩、线性注意力
FlashAttention	数据搬运次序（分块 / 重算）	无（精确）	显存带宽（IO）	FlashAttention 系列
长上下文工程	有效序列处理（窗口 / 检索 / 状态）	视方案	有效 N	滑动窗口、KV 压缩、状态空间模型（state space model）

（本文的「标准实现」一律指会 materialize N 乘 N 矩阵的朴素实现；现代框架的 fused / SDPA 默认已不走这条路。）读这张表的关键是**「是否损失精度」那一列**：近似路线用「换精度」换速度，FlashAttention 用「换次序」换速度，二者不在同一坐标轴上，也因此可以叠加（FlashAttention 本身就能扩到块稀疏注意力）。

把这条脉络收敛成一句话：同样面对 N 平方的注意力成本，近似派选择「少算」，FlashAttention 选择「少搬」。后来的工程实践证明，在 GPU 的非对称存储层次上，后者是更根本的那一刀。这也解释了为什么它能同时拿下「更快」和「更省显存」——因为这两者都源自同一件事：不再生成那个巨大的中间矩阵。

五、推导一：标准注意力的 HBM 访问量为什么是 Ω(N·d + N 平方)

这是全景图里最关键的一步。先把「标准实现」每一句在存储层次之间搬了多少数据点清楚。

标准注意力（不含 FlashAttention）的计算步骤是：

S = Q K^{\top},\qquad P = \mathrm{softmax}(S),\qquad O = P V

标准实现的代价，可以逐句拆成下面这张存取清单来看。一句话概括：它必须把 N 乘 N 的中间矩阵落到 HBM——先写 S、读回算 softmax、写 P、再读 P 算 PV；这四步在标准实现里都要经过 HBM，而 FlashAttention（第七节会给出分块数据流图）把它们全部留在片上。这些读写就是「多余的」IO。

步骤	读	写
读 Q, K, V	N 乘 d 各一份	—
算 S = QKᵀ	Q, K	写 S（N 乘 N）
读 S、算 softmax	读 S（N 乘 N）	写 P（N 乘 N）
算 O = PV	读 P（N 乘 N）、V	写 O（N 乘 d）
反向（若存 P）	读 P（N 乘 N）	—

把访问量拆项写出来（单头，元素数以「读写次数」计、忽略常数因子——这一口径贯穿全文），得到标准注意力的 HBM 访问量：

\Omega\big(N d + N^{2}\big)

$N d$ 项：读 Q、K、V 三个 N 乘 d 矩阵，以及写回输出 O（N 乘 d）；
$N^{2}$ 项：那个 N 乘 N 的注意力矩阵 S 与 P 至少各写一次、读一次。

关键在于 N 平方项：当 N 翻一倍，N·d 项只翻一倍，N 平方项却翻四倍（例如 N 从 2K 到 4K）。以 2K 为基准：N 到 8K（翻 4 倍）时 N 平方项翻 16 倍，到 64K（翻 32 倍）时翻 1024 倍。（这里的 N 始终是序列长度，即一次前向要看多少个 token。）在长序列区间，N 平方项逐渐压倒一切——这就是「标准注意力带宽受限于 N 乘 N 矩阵搬运」的量化表述。

顺带一提，这也是显存占用的来源：那个 N 乘 N 矩阵要真金白银地占显存。N 取 8192、单精度、单头，S 就是 $8192^2 \times 4 \approx 268$ MB（本文推算）；多头与多批次叠加，很快就撑爆显存。所以标准注意力的显存是 O(N 平方)，而 FlashAttention 把它降到 O(N)。

这里要把一个常见误读就地按住：N 是序列长度，不是模型参数量，也不是训练数据量。后文出现的 2K、8K、64K 说的都是「一次前向要看多少个 token」。

六、推导二：在线 softmax —— 不看整行也能把归一化算对

要避免把整个 N 乘 N 的 S 写进 HBM，就得边算边归一化。但 softmax 的传统写法有个「先看全局」的硬依赖——你得先知道整行的最大值，才能数值稳定地求指数。

先看为什么必须先减最大值。若直接算指数：

\mathrm{softmax}(x)_i = \frac{e^{\,x_i}}{\sum_{k} e^{\,x_k}}

当某个 $x_i$ 较大（比如 1000）时， $e^{1000}$ 会直接溢出成 inf——FP16 下更是早就爆了。稳妥的做法是减去整行最大值再求指数：

\mathrm{softmax}(x)_i = \frac{e^{\,x_i - \max_j x_j}}{\sum_{k} e^{\,x_k - \max_j x_j}}

分子分母同乘 $e^{-\max_j x_j}$ 不改变比值，却把指数都压到 0 及以下，永不溢出。

问题来了：分母里的 $\max_j x_j$ 是整行的最大值。而分块计算时，我们一次只看到一块 $S^{(j)}$ ，看不到整行——这似乎要求「要么先扫一遍求最大值，要么把整行存下来」。在线 softmax（online normalizer） 解决了这个矛盾：它把「求最大值」和「累加分母」合并进一趟遍历，用一个可增量更新的三元组 $(\mathbf{m}, \boldsymbol{\ell}, O)$ 维护，遍历结束后自然得到正确结果。

该手法出自 Milakov 与 Gimelshein（NVIDIA）在 2018 年的论文《Online normalizer calculation for softmax》（arXiv:1805.02867），其动机正是减少内存访问：论文报告经典 softmax 加速最高约 1.3 倍，把 Softmax 与 TopK 融合后最高约 5 倍。FlashAttention 把它搬进了注意力核。

设我们按列块（block）依次处理 K、V，第 j 块得到分数 $S^{(j)} = Q_{\text{blk}} \, K^{(j)\top}$ 。维护量定义为： $\mathbf{m}$ 是到目前为止的行最大， $\boldsymbol{\ell}$ 是到目前为止去掉最大后的指数和（即归一化分母）， $O$ 是目前的未归一化输出累加。每处理一块，按下式增量更新：

\mathbf{m}^{(j)} = \max\!\big(\mathbf{m}^{(j-1)},\; \mathrm{rowmax}(S^{(j)})\big)

\boldsymbol{\ell}^{(j)} = e^{\,\mathbf{m}^{(j-1)} - \mathbf{m}^{(j)}} \odot \boldsymbol{\ell}^{(j-1)} + \mathrm{rowsum}\!\big(e^{\,S^{(j)} - \mathbf{m}^{(j)}}\big)

O^{(j)} = \mathrm{diag}\!\big(e^{\,\mathbf{m}^{(j-1)} - \mathbf{m}^{(j)}}\big)\, O^{(j-1)} + e^{\,S^{(j)} - \mathbf{m}^{(j)}}\, V^{(j)}

（记号说明： $\mathrm{rowmax}(\cdot)$ 取每行最大值、 $\mathrm{rowsum}(\cdot)$ 对每行求和、 $\mathrm{diag}(\cdot)$ 把向量张成对角矩阵、 $\odot$ 表示逐元素相乘。）

这三个式子里藏着全部聪明劲儿，逐项说清：

第一式：新的行最大，是「旧最大」和「本块行最大」二者取大。它只会单调不减。
第二式：如果最大值变了（ $\mathbf{m}^{(j)} > \mathbf{m}^{(j-1)}$ ），就必须把之前累加的分母按指数因子 $e^{\mathbf{m}^{(j-1)} - \mathbf{m}^{(j)}}$ 缩放回来——因为之前那些项是以旧最大值为基准做指数的，现在基准换了，得补偿。本块新贡献的指数和则用新最大值稳定地算出来。
第三式：对输出累加做同样的缩放补偿，再把本块的 $e^{S^{(j)} - \mathbf{m}^{(j)}} V^{(j)}$ 加进来。注意 O 的缩放与分母的缩放用的是同一个因子，这正是它能「边累加边修正」而最终结果精确的原因。

遍历完所有块后， $\mathbf{m}^{(\text{final})}$ 就是真正的整行最大值， $\boldsymbol{\ell}^{(\text{final})}$ 就是真正的分母，最终输出一次性归一化：

O = \mathrm{diag}\!\big(\boldsymbol{\ell}^{(\text{final})}\big)^{-1} O^{(\text{final})}

正确性直觉：把整个流程想成「先按一个临时基准近似地累加，一旦发现更大的基准就把旧的按比例换算过去」。因为始终在做「以当前最大值为基准」的指数，数值上永远稳定（不会 $e^{\text{大数}}$ 溢出）；又因为每次缩放都成比例地同时作用于分母和输出，最终比值不变，结果依然是精确 softmax，不是近似。这也是 FlashAttention 反复强调「exact attention, no approximation」的技术底座。

用一个极小的例子验证这条直觉。设某行只有两块分数，第一块的值是 3 与 1，第二块的值是 5 与 2。

处理第一块： $\mathbf{m} = 3$ ， $\boldsymbol{\ell} = e^{3-3} + e^{1-3} = 1 + e^{-2} \approx 1.135$ 。
处理第二块前，本块行最大是 5： $\mathbf{m}_{\text{new}} = \max(3, 5) = 5$ 。
旧的 $\boldsymbol{\ell}$ 按 $e^{3-5} = e^{-2}$ 缩放： $1.135 \times e^{-2} \approx 0.1536$ ；再加本块指数和 $e^{5-5} + e^{2-5} \approx 1.0498$ ，得到 $\boldsymbol{\ell} \approx 1.2034$ 。
而整行真值为 3、1、5、2，最大值 5，分母等于 $e^{-2} + e^{-4} + 1 + e^{-3} \approx 0.1353 + 0.0183 + 1 + 0.0498 \approx 1.2034$ ——完全一致。

这个小算例说明：在线 softmax 的「修正」不是拍脑袋的近似技巧，而是分块求和加基准换算这一代数恒等式的直接实现。

图 2 给出这个迭代的数据流。

flowchart TD
  S0["初始化 行最大 m 为负无穷 分母 l 为 0 输出 O 为 0"] --> L1["取第 j 块 K 与 V"]
  L1 --> S1["算本块分数 S_j 等于 Q块 乘 K_j 转置"]
  S1 --> M["更新行最大 m_new 等于 max 旧m 与 本块行最大"]
  M --> L["更新分母 l_new 等于 e 的 旧m减m_new 次幂 乘 l 加 本块指数和"]
  L --> O["更新输出 O_new 等于 缩放后的旧 O 加 本块贡献"]
  O --> C{"还有 K V 块吗"}
  C -->|是| L1
  C -->|否| Z["输出 O 除以 l 得到精确结果"]

图 2 说明了什么：整条注意力前向就是「取一块、更新三元组、并入输出」的循环，全程只需在 SRAM 里维护每块的小块累加量（约块行数乘 d 的量级），而全序列的统计量 $\mathbf{m}$ 、 $\boldsymbol{\ell}$ 只是 O(N)——两者不是一回事，前者是「一次处理多少行」，后者是「整条序列的元信息」。

七、推导三：分块之后，FlashAttention 的 HBM 访问量 O(N 平方 d 平方除 M)

有了在线 softmax，就能名正言顺地分块了。FlashAttention 的做法是：

把 Q 切成 $N / B_q$ 个块、把 K、V 切成 $N / B_k$ 个块，块大小的选择要让每个块连同累加量能塞进 SRAM（容量 M）；
外层循环过 K、V 块，内层循环过 Q 块，在 SRAM 里做局部矩阵乘加在线 softmax；
只把最终的 O（N 乘 d）写回 HBM，那个 N 乘 N 的中间矩阵全程不落 HBM。

图 3 就是这条分块数据流：数据在 HBM 与片上 SRAM 之间按块来回，但从不落 N 乘 N 的中间矩阵。

flowchart LR
  A["HBM 里的 Q K V"] -->|按块加载 K V| B["片上 SRAM"]
  A -->|按块加载 Q| B
  B --> C["分块算 S 等于 Q 乘 K 转置"]
  C --> D["在线 softmax 归一化"]
  D --> E["累积 P 乘 V"]
  E -->|只把最终 O 写回| F["HBM 里的输出 O"]
  D -.->|保存归一化统计量| G["SRAM 寄存器"]

图 3 说明了什么：与第五节的「标准实现」正好相反——标准实现要把 S、P 两个 N 乘 N 矩阵落到 HBM，FlashAttention 则把它们全部留在片上，只把最终的 O 写回。

FlashAttention 论文证明：其 HBM 访问量是

O\!\left(\frac{N^{2} d^{2}}{M}\right)

其中 M 是 SRAM 容量。把它和标准注意力的 $\Omega(Nd + N^2)$ 对比（论文引言写作 $\Omega(Nd+N^2)$ 、定理 2 写作 $\Theta(Nd+N^2)$ ；本文统一用更弱的 $\Omega$ 口径，方向一致），能读出三件事：

M 在分母上：SRAM 越大，能一次塞下的块越大，需要重复搬运的次数越少，HBM 访问越少。这就是「利用非对称存储层次」的定量含义。
没有独立的 N 平方显存项：那个平方项被「除以 M」压下去了。对典型的 d（64 或 128）与 M，FlashAttention 的 HBM 访问量比标准实现少最多约 9 倍（论文引言与其图 2 口径；摘要只说「HBM 访问更少」）。
论文还给出下界：任何精确注意力算法都无法在所有 SRAM 容量上渐近优于这个访问量（论文证明了一个下界，Proposition 3）。也就是说，FlashAttention 在这个意义上是渐近最优的——它省的是 IO，不是靠牺牲精度。

SRAM 容量 M 对块大小的约束（一次数量级估算）。 A100 每个 SM 的共享内存上限约 192 KB（这是论文口径的片上 SRAM 合计；NVIDIA 规格表里每 SM 可配置 shared memory 与 L1 的划分略有不同，实际用于存储的块容量还要扣掉寄存器等占用，论文取用的 M 比这个上限小）。若取头维度 d 等于 128、FP16（每元素 2 字节），一个块的一行需要约 256 字节；单看容量上限，能放下数百行（本文推算）。这给出的定性结论是：块大小既不能大到超出容量、也不宜小到让 IO 重复次数上升。FlashAttention-2 论文最终取的块大小是 64 与 128 的组合（出自其正文 §3 末段的 Tuning block sizes）——这是这两头之间的折中（注意：块大小是 B_q、B_k，与 SRAM 容量 M 不是一回事）。本文的粗估只是容量上限，与上述取值属同一量级，不用于解释论文的具体取值。

需要特别强调一个常被误读的点：FlashAttention 并没有减少矩阵乘的 FLOPs（它甚至因为反向要重算而增加了 FLOPs）。它减少的是HBM 访问。当算子本是带宽受限时，「多算一点、少搬很多」反而更快——这正是第一、三节判据的直接推论。换句话说，FlashAttention 是一笔**「用便宜的计算换昂贵的搬运」**的交易，而它成立的前提，是那 16 倍的算力单价差。

八、反向传播：为什么「重算」比「存下来」更省

训练时的反向传播需要 S 或 P（去算 dS、dQ、dK、dV）。标准实现会把前向的中间矩阵存下来供反向读取——但那意味着把 N 乘 N 的 P 写进 HBM 再读回来，正是我们千方百计要避免的 IO。

FlashAttention 的选择是不存 P，反向时重算：前向只保留每块的归一化统计量（行最大 $\mathbf{m}$ 和分母 $\boldsymbol{\ell}$ ，都是 O(N) 的量），反向时用它们配合 Q、K、V 在 SRAM 里重新算出 S 与 P 的块，再立即用于梯度计算。这就是标题里的 recomputation（重算）。

代价与收益一句话说清：

代价：多花一些 FLOPs（重算前向的一部分）；但这些 FLOPs 是矩阵乘，在 M 足够时相对便宜。
收益：HBM 访问量从「读写一个 N 平方矩阵」降到「与标准前向同量级」，显存占用从 O(N 平方) 降到 O(N)。

因为反向和前向一样是带宽受限，用便宜的重算 FLOPs 换掉昂贵的 HBM 读写是划算的。这也解释了一个看似矛盾的事实：FlashAttention 在增加总 FLOPs 的前提下，墙钟时间反而下降——只要瓶颈在搬运，账就算得过来。

还需要点明一个配套细节：重算之所以成立，靠的是前向把 $\mathbf{m}$ 与 $\boldsymbol{\ell}$ 当作检查点（checkpoint）留下来。它们是 O(N) 的量，留着几乎不占地方，却让反向不必知道整个 P 矩阵。这其实是「梯度检查点」（gradient checkpointing：为省显存，前向只存少量中间量、反向时按需重算）思想在注意力内部的落地——一切为了不让中间矩阵落到慢速存储上。

重算与「存下来」：一笔可以算清的账

把两种反向策略的代价并排放一张表，账就一目了然：

策略	前向额外存储	反向额外读取	反向额外计算	主要代价
存下 P	写 N 乘 N，再读 N 乘 N	读 P（N 乘 N）	少	HBM 带宽
重算	只存 O(N) 统计量	无（就地重算）	重算 S 与 P 的 FLOPs	少量矩阵乘 FLOPs

判据只有一条：把「多算的 FLOPs」按矩阵乘单价、把「少搬的字节」按带宽单价折算，谁便宜听谁的。在 A100 这类带宽与算力落差悬殊的硬件上，天平明显压在「重算」一侧。

需要提醒的是，这个结论并不普适：当 SRAM 极小（重算要反复把 Q、K、V 搬回片上）或算力本就紧张时，重算的账可能反过来。所以第八节的结论必须连同第七节的 M 一起读——分块容量决定了重算到底划不划算。

从检查点公式看「为什么可以不存」

把上述直觉写成式子会更清楚。设前向按块处理，第 j 块保留的就是最终的行最大与分母，那么反向要重建第 j 块的归一化权重，只需一式计算：

P^{(j)} = \mathrm{diag}\!\big(\boldsymbol{\ell}^{(\text{final})}\big)^{-1} \exp\!\big(S^{(j)} - \mathbf{m}^{(\text{final})}\big)

也就是说，反向需要的全部信息都可以被 O(N) 的统计量「压缩」下来：只要留住最终的 $\mathbf{m}$ 与 $\boldsymbol{\ell}$ ，第 j 块的权重就能在片上重算，而不必把整个 P 矩阵存进 HBM。这正是「省掉 N 平方显存」在数学上成立的原因——不是近似，而是把可重算的东西用紧凑统计量代替。

反过来，如果没有这套统计量，反向就必须依赖存下来的 P，N 平方的显存与 IO 一步也省不掉。所以第一到第八节的推导其实一环扣一环：在线 softmax 给出统计量，统计量让重算成立，重算让 N 平方的存储与 IO 归零。

九、从 FlashAttention-2 到 4：并行划分与被忽略的「非矩阵乘 FLOPs」

FlashAttention（v1）已经把 N 平方的显存 IO 干掉了，为什么还需要 v2？答案不在 IO 复杂度，而在GPU 利用率：v1 的注意力 kernel 整体只达到理论峰值 FLOPs/s 的 25–40%（对比对象是成熟的 GEMM），还有距离。FlashAttention-2（arXiv:2307.08691，2023-07-17）针对三处下手：

减少非矩阵乘 FLOPs：现代 GPU 上矩阵乘走 Tensor Core、其余 FP32 走普通通路。A100 上 312 TFLOPs/s 的矩阵乘对比 19.5 TFLOPs/s 的非矩阵乘（约 16 倍），意味着每一个非矩阵乘操作都贵得离谱。v2 重写了在线 softmax 的缩放、边界检查、因果掩码等操作，在输出不变的前提下减少这类「贵的 FLOPs」。
更好的并行：v1 只在 batch 维和头维并行（一个头一个线程块），而 A100 有 108 个 SM（流多处理器）——当批大小或头数很少（长序列常见情形）时，线程块不够多，SM 用不满。v2 额外沿序列长度维并行，让长序列这个区间也能吃满 SM。
更好的 warp 间划分：把「切分 K 与 V」的 split-K 方案改成「切分 Q」，避免各 warp（线程束，一组 32 个线程）把中间结果写回共享内存再同步相加，减少 SRAM 读写。

实测：v2 比 v1 快约 2 倍，注意力 kernel 前向可达约 230 TFLOPs/s（A100，FP16 与 BF16，约合理论峰值的 73%；反向约 63%），端到端训练 GPT 类模型达约 225 TFLOPs/s（模型 FLOPs 利用率约 72%）。它还支持头维度到 256（覆盖 GPT-J、CodeGen、Stable Diffusion 1.x）以及 MQA 与 GQA（多查询 / 分组查询注意力：多个查询头共享同一组 K/V 头，以缩小 KV 缓存）。块大小通常取 64 与 128 的组合，随头维度与设备 SRAM 容量调整。

后续版本继续沿硬件走：官方仓库（Dao-AILab/flash-attention）里，FlashAttention-3 面向 Hopper（H100 等，其博客发布于 2024 年，tridao.me/blog/2024/flash3/）；FlashAttention-4 用 CuTeDSL（NVIDIA CUTLASS 的 Python 领域专用语言）写成，面向 Hopper 与 Blackwell，其博客发布于 2026-03-05（tridao.me/blog/2026/flash4/），配套论文为 arXiv:2603.05451。仓库许可为 BSD-3-Clause（其提交历史显示曾由 Apache-2.0 改为 BSD）。

把 v1 到 v2 的演进翻成一句工程格言：第一阶段省 IO（降复杂度），第二阶段榨并行（提利用率）。两件事的判据不同——前者看 HBM 访问量，后者看 SM 占用率与非矩阵乘开销。

十、对照表：标准注意力 vs FlashAttention

把散在前文的结论收进一张表，便于横向对照。

维度	标准注意力	FlashAttention
显存占用	O(N 平方，要存 N 乘 N 的 S 与 P)	O(N，只存 O(N) 的统计量)
HBM 访问量	Ω(N·d + N 平方)	O(N 平方 d 平方除 M)
是否精确	精确	精确（no approximation）
中间矩阵落 HBM	是	否
反向传播	读存下来的 P	用统计量重算 P
FLOPs	基准	更多（含重算）
是否减少矩阵乘 FLOPs	—	否（省的是 IO）
提升空间	—	v2 靠并行与划分把利用率拉到 50–73%

至于墙钟收益，不同口径的数字必须分开读，别混成一句「快 N 倍」：

口径	数字	相对谁	出处（版本与年份）
FlashAttention v1 注意力层	约 2–4 倍	相对「优化基线」（非 v2 自身）	FA-2 摘要（2023）
GPT-2 注意力	最高约 7.6 倍	相对 PyTorch 注意力实现	FA-1 引言与其图 1（2022）
端到端训练 · BERT-large	约 +15%	相对 MLPerf 1.1 训练记录	FA-1 摘要（2022）
端到端训练 · GPT-2	约 3 倍	相对 HuggingFace 实现	FA-1 摘要（2022）
端到端训练 · Long-Range Arena	约 2.4 倍	相对标准注意力基线	FA-1 摘要（2022）

特别提醒：本表第一行「约 2–4 倍」是 v1 相对优化基线的数字，不是 v2 自身；v2 相对 v1 是约 2 倍（第九节），两者别混读。

表里最容易看错的是「FLOPs 更多、墙钟更快」这一点——它不是矛盾，而是带宽受限这一前提的直接后果。判据始终是第一节与第三节那两句话：先看瓶颈在哪，再谈优化。

十一、常见误读澄清

误读 1：「FlashAttention 减少了 FLOPs」。 错。它甚至增加了 FLOPs（因为反向要重算）。它减少的是 HBM 访问量。把「省 IO」说成「省算力」，会让整条推理链失焦。

误读 2：「它是近似注意力」。 错。它是精确注意力，结果与标准实现一致（在浮点误差范围内）。稀疏、低秩才是近似路线，二者不在同一坐标轴上。

误读 3：「N 是模型参数量」。 错。N 是序列长度。本文三个核心量里，N、d 是张量形状量，M 是硬件容量参数；把 N 误当模型规模，会把复杂度结论套到完全不同的对象上。

误读 4：「有了 FlashAttention，长上下文就免费了」。 不完全对。它把注意力的显存从 O(N 平方) 降到 O(N)、把 IO 压下去，但其他部分（如 KV 缓存随序列线性增长、前馈层、通信）仍会增长。它解决的是注意力那一环的瓶颈，不是长上下文的全部成本。

误读 5：「SRAM 越大越好，所以买共享内存最大的卡就行」。 只对一半。块大小受寄存器与共享内存双重约束，且与头维度耦合；太大反而触发寄存器溢出或直接跑不起来。这是第七节「容量换带宽」的边界所在。

误读 6：「换个更快的 GPU 或更高精度就能解决」。 不必然。带宽受限的算子在换卡后，只有当新卡的带宽 / 算力比值发生变化、或 SRAM 容量变化改变了分块效率时才会明显改善；单纯堆算力（例如换更高峰值算力的同代卡）可能收益很小。正确的动作是回到 Roofline：先看交叉点变没变，再看实现的分块是否吃满了新卡的 SRAM。

误读 7：「它只对训练有用，推理用不上」。 不对。推理的 prefill（预填充）阶段处理长 prompt 时，注意力同样是带宽受限，FlashAttention 类实现同样有效；而 decode 阶段每步只处理一个 token，瓶颈转移到 KV 缓存的读写与访存，这是另一个话题（KV 量化、分页、前缀共享）。分清 prefill 与 decode 两个阶段的瓶颈，是推理优化的前提。

十二、工程含义：什么时候它帮不上忙

把判据落到实际选型上，可以给出几条工程结论：

长序列、小批量或少头：收益最大。此时标准实现的 N 平方 IO 占比高、SM 又用不满，FlashAttention 的两项优势同时生效。
短序列：N 平方项不占主导，收益退回到「省显存加省一部分 IO」，别指望翻倍级加速。
本来就是 compute-bound 的场景：FlashAttention 省的是 IO，若某个算子已经算力受限，它帮不上忙——按 Roofline 判据，先确认瓶颈在哪，再决定要不要上。
反向重算的适用边界：重算省 IO、费 FLOPs；当 FLOPs 已经紧张或 SRAM 太小（M 小、重算代价高）时，收益会缩小。
数值精度：在线 softmax 的「边累加边修正」在低精度（FP16、BF16，乃至 FP8）下要求更小心地处理缩放与累加，这也是后续版本持续优化的方向。
选型建议：把「序列长度」当第一维来问——短序列上它主要帮你省显存，长序列上它才同时帮你省时间。

快速自查清单

遇到一个「注意力太慢」的现场，按下面的顺序问，能很快定位该不该上这类 IO 优化：

自查项	若答案是「是」	若答案是「否」
序列长度是否占主导（长上下文、大批 token）？	继续	收益有限，先做别的
标准实现是否把 N 乘 N 矩阵写回了 HBM？	有优化空间	已被框架处理
该算子是否为带宽受限？	是：值得省 IO	否：先解算力瓶颈
训练是否在反向存了中间注意力矩阵？	启用重算更划算	—
SRAM 容量是否够放下目标块大小？	分块效率高	调小块或换设备
是否有可观测的 HBM 流量与占用数据佐证？	以数据为准	先上性能分析工具

清单的最后一行最重要：任何「该不该优化」的判断都以实测的显存流量与占用为准，不要凭本文的复杂度结论直接下结论——复杂度告诉你「方向」，profiling 才告诉你「幅度」。

十三、本文自行推导（量级一致性检查）

本节只做量级核对，所有数字均为本文按公开硬件参数推算，不是官方数值，不得作为官方引用。

推导 1｜A100 的 Roofline 交叉点。 取矩阵乘峰值 $\pi_{\text{peak}} \approx 312 \times 10^{12}$ FLOPs/s（FP16 与 BF16，来自 FlashAttention-2 博客口径），显存带宽 $\beta \approx 1.55 \times 10^{12}$ B/s（A100 40GB HBM 规格量级；80GB 约 2.0 TB/s，故 $\beta$ 取值会使交叉点在约 150–200 FLOPs/byte 间浮动）。代入 $I^{*} = \pi_{\text{peak}} / \beta$ ：

I^{*} = \frac{312 \times 10^{12}}{1.55 \times 10^{12}} \approx 2.0 \times 10^{2}\ \text{FLOPs/byte}

结论：约 150–200 FLOPs/byte 区间。任何「每字节只算几个 FLOP」的算子都稳稳落在带宽受限区。仅作量级检查。

推导 2｜标准注意力把 N 乘 N 矩阵落 HBM 的代价（单层）。 设 N 取 8192、d 取 128、头数乘批等于 32（单精度 4 字节）。单头 S 有 $N^2 = 8192^2 \approx 6.7 \times 10^{7}$ 个元素，约 $6.7 \times 10^{7} \times 4 \approx 2.7 \times 10^{8}$ 字节，约 268 MB；32 个头合计一次写回约 8.6 GB。标准前向要写 S、读 S、写 P、读 P，若按「S 与 P 各一写一读」计，矩阵往返流量约 34 GB。以 $\beta \approx 1.55$ TB/s 计，约需 20 毫秒量级每层（单矩阵一写一读约 17 GB、约 11 毫秒，S 与 P 合计翻倍）。FlashAttention 不落这两个矩阵，这笔流量几乎归零。假设为：单精度、未计 kernel 融合带来的额外读写。仅作量级检查。

推导 3｜被「省掉」的是 IO 不是计算。 标准注意力核心矩阵乘的 FLOPs 约 $4 N^{2} d$ （两次 N 阶矩阵乘各约 $2N^{2}d$ ）。N 取 8192、d 取 128 时约 $4 \times 6.7 \times 10^{7} \times 128 \approx 3.4 \times 10^{10}$ FLOP 每头，32 头合计约 $1.1 \times 10^{12}$ FLOP，约 1.1 TFLOP 每层。仅作量级检查。

推导 4｜为什么重算划算。 重算多出的 FLOPs 是矩阵乘（单价便宜），省下的是 HBM 读写（单价贵）。只要「多算的 FLOPs 乘矩阵乘单价 < 少搬的字节乘带宽单价」，重算就是赚的。把第三节的 16 倍单价差代入，天平明显偏向重算。仅作方向性论证，非精确成本模型。

推导 5｜块大小的容量约束。 A100 每 SM 共享内存上限约 192 KB；头维度 d 取 128、FP16 每元素 2 字节时，一行 K 块占约 256 字节，按容量上限可容纳约 $192 \times 1024 / 256 = 768$ 行，扣除 Q、K 两块与累加量后降至数百行量级。仅作容量上限的粗估，实际受寄存器与具体布局影响。

推导 6｜标准注意力的算术强度为何随序列增长而下降。 只看矩阵乘，FLOPs 约 $4N^{2}d$ 。标准实现要在 HBM 上往返的字节，除读 Q、K、V（约 $3Nd$ 个元素）与写输出 O（约 $Nd$ 个元素）外，还必须多搬 S 与 P（各约 $N^{2}$ 个元素、各一写一读，共约 $4N^{2}$ 个元素）。以 FP16（2 字节）计：

I_{\text{std}} \approx \frac{4N^{2}d}{2\,(4Nd + 4N^{2})} \;\xrightarrow[N \to \infty]{}\; \frac{4N^{2}d}{8N^{2}} = \frac{d}{2}

d 取 128 时趋于约 64 FLOPs/byte，明显低于交叉点约 200。这正说明：序列越长，标准注意力的算术强度越是被那 4N² 个元素的搬运拉低——瓶颈在搬运，不在算。FlashAttention 把 N 平方项去掉后，算术强度回升，才可能靠 Tensor Core 吃满（口径：FP16、2 字节每元素、只计 HBM 往返、不计 L2 命中；本文粗估，仅作量级检查）。

十四、边界、未解与小结

逐条列清本文能证明与不能证明的边界：

论文性质边界：本文的复杂度结论（ $\Omega(Nd+N^2)$ 、 $O(N^2d^2/M)$ 、下界、加速倍数）来自 FlashAttention 系列论文与官方博客、仓库的公开表述，本文只做转述与推导复现，未独立复跑其基准测试。
口径边界：文中「9 倍」「7.6 倍」「2–4 倍」「230 与 225 TFLOPs/s」「72% MFU」「108 SM」等数字均出自对应来源的特定硬件与设定（多为 A100、特定头维度与序列长度），不可外推到其他 GPU 或配置；v2 之后的 v3 与 v4 面向 Hopper 与 Blackwell，数值另计。
被外部推测包裹的部分：本文未展开后续版本（v3 与 v4）的具体加速比与实现细节，只引仓库与博客说明其面向的硬件与语言；如需精确数字须回一手论文或博客。
论文自认的未解项：FlashAttention-2 原文把「块大小的自动调优」列为未来工作（当时是手动按头维度调四种组合），说明工程上仍有自动化空间。
仍存在的质疑：FlashAttention 并非「万能加速」——它对短序列、compute-bound 场景收益有限，且重算会抬高 FLOPs；把它当银弹是误读。
本文自己的限度：第十三节全部数字为量级估算，含「单精度或 FP16、未计融合、按 40GB 规格带宽」等简化假设，只能用于建立直觉，不能用于容量规划或性能承诺。
事实能证明与不能证明的边界：能证明的是「标准实现的 HBM 访问含独立 N 平方项、FlashAttention 把它压掉」；不能据此断言某个具体工程一定提速多少——那取决于序列长度、头维度、SRAM 容量与实现质量。
适用范围边界：本文讨论的是精确注意力的 IO 优化；稀疏或线性注意力是另一条路线（换精度或换结构），二者可叠加（FlashAttention v1 亦扩展到块稀疏注意力），但不在本文推导范围。

小结。 一句话把全文收束：注意力的墙钟瓶颈常常不在「算得不够快」，而在「搬得太多」。判据是算术强度与 Roofline——A100 的交叉点在约 150–200 FLOPs/byte 区间；标准注意力因为必须搬运一个 N 乘 N 的中间矩阵，HBM 访问量含独立的 N 平方项，算术强度随序列增长反而被拉低，于是长序列下越来越带宽受限。FlashAttention 的三件事——分块（tiling）、在线 softmax、反向重算——合起来把 HBM 访问量从 $\Omega(Nd+N^2)$ 压到 $O(N^2d^2/M)$ ，并且不牺牲精度；v2 再通过更好的并行与线程划分，把矩阵乘利用率从 25–40% 拉到 50–73%。它省的是 IO，代价是更多（但廉价的）FLOPs——这正是第一性原理给出的、最简单的划算买卖。

数据截至 2026-10-09。 上游论文与官方材料均为公开可访问；本文未对任何硬件做独立基准复现，厂商披露的性能数字尚无本文侧独立复现。

十五、📋 来源与核验记录

一手论文与官方材料

Dao, T., Fu, D. Y., Ermon, S., Rudra, A., Ré, C. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135（v1 2022-05-27，v2 2022-06-23）。https://arxiv.org/abs/2205.14135 （另见 HTML 版 https://arxiv.org/html/2205.14135v2 ）｜访问 2026-10-09。关键数据出处：IO 复杂度 O(N 平方 d 平方除 M) 对比标准注意力 Ω(Nd+N 平方)（论文引言用 Ω、定理 2 用 Θ，本文统一用 Ω）、最多约 9 倍更少 HBM 访问（论文引言与其图 2，非摘要）、渐近最优的下界（Proposition 3）、GPT-2 注意力 7.6 倍（引言）、BERT-large 15%、GPT-2 3 倍、Long-Range Arena 2.4 倍（摘要）、SRAM 与 HBM 带宽量级（约 19 TB/s 对比约 1.5–2.0 TB/s）、192 KB / SM、A100 108 个 SM。
Dao, T. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691（2023-07-17）。https://arxiv.org/abs/2307.08691 ｜访问 2026-10-09。关键数据：比 v1 快约 2 倍、前向 230 TFLOPs/s（约 73% 峰值）、反向最高约 63%、端到端 225 TFLOPs/s、由 25–40% 提升至 50–73% 理论峰值、头维度上限 256、MQA 与 GQA 支持、块大小取 64 与 128 的组合（正文 §3 末段 Tuning block sizes）、块大小自动调优列为未来工作。
Dao, T. FlashAttention-2（Stanford CRFM 博客，2023-07-17）。https://crfm.stanford.edu/2023/07/17/flash2.html ｜访问 2026-10-09。关键数据：A100 矩阵乘 312 TFLOPs/s 与非矩阵乘 19.5 TFLOPs/s（每非矩阵乘 FLOP 贵约 16 倍）、108 个 SM、split-Q 划分、tiling 加 recomputation、端到端 225 TFLOPs/s 与 72% 模型 FLOPs 利用率。（注：块大小的具体取值出自来源 2 的论文正文，不在本篇博客。）
Dao-AILab/flash-attention 官方实现仓库。https://github.com/Dao-AILab/flash-attention ｜访问 2026-10-09。关键事实：FlashAttention-3 面向 Hopper（博客 2024 年，tridao.me/blog/2024/flash3/）、FlashAttention-4 用 CuTeDSL 面向 Hopper 与 Blackwell（博客 2026-03-05，tridao.me/blog/2026/flash4/；论文 arXiv:2603.05451）、仓库许可为 BSD-3-Clause（提交历史显示由 Apache-2.0 改来）。
Milakov, M., Gimelshein, N.（NVIDIA）Online normalizer calculation for softmax. arXiv:1805.02867（v1 2018-05-08，v2 2018-07-28）。https://arxiv.org/abs/1805.02867 ｜访问 2026-10-09。关键数据：减少内存访问、经典 softmax 加速最高约 1.3 倍、融合 Softmax 与 TopK 最高约 5 倍。
NVIDIA A100 Tensor Core GPU 产品页。https://www.nvidia.com/en-us/data-center/a100/ ｜访问 2026-10-09。用途：核对 A100（含 80GB HBM）产品定位与精度支持范围；A100 规格页含 FP16 Tensor Core 312 TFLOPS、FP32 19.5 TFLOPS、80GB 内存带宽约 1935 GB/s（PCIe）/ 2039 GB/s（SXM）。

第三方 / 他人分析

He, H. Making Deep Learning Go Brrrr From First Principles（2022）。https://horace.io/brrr_intro.html ｜访问 2026-10-09。用途：三类瓶颈（compute、bandwidth、overhead）框架、A100 312 TFLOPs/s 与「非矩阵乘只有 19.5 TFLOPs/s（远低于 312）」的对比论据。

本文自行推导

第十三节的交叉点 I 星约 150–200 FLOPs/byte、N 乘 N 矩阵搬运流量与耗时估计、4N 平方 d 的 FLOPs 估算、块大小容量上限（约 768 行起算）、以及「标准注意力算术强度随 N 增长趋于 d/2」的推演，还有正文中「单头 S 约 268 MB」等数字，均为本文按公开硬件参数推算的量级一致性检查，非官方数值，不得作为官方引用。