1.1 注意力计算的计算与显存瓶颈

1.1 注意力计算的计算与显存瓶颈

读者读完这一节，应该能回答一句：标准注意力为什么"算力够、显存也够，却依然又慢又容易 OOM"——根子在 O(n²) 的显存物化和 HBM 的反复读写，而不是 GPU 乘法不够快。

上一节我们把注意力的数学形式摆出来了。这一节不喊口号，而是把"慢"和"费显存"拆成两个可以量化、可以归因的具体问题。很多同学第一次优化推理时会直觉地"加显存、换更大卡"，但如果定位错了瓶颈，加显存只是把问题往后推一步，并不能根治。我们先把数据流画清楚，再用量级和"屋顶线模型"把话说死，最后辨析几个最常见的误解。

1.1.1 朴素实现到底物化了多少东西

回头看那四行最朴素的计算：

Q = X · Wq      # [n, d]
K = X · Wk      # [n, d]
V = X · Wv      # [n, d]
S = Q · Kᵀ / √d  # [n, n]，注意力分数
P = softmax(S)   # [n, n]，注意力权重
O = P · V        # [n, d]
​

前三个矩阵乘（算 Q、K、V）的规模是 O(n·d²)，输入输出的张量都是 [n, d]，这个量级是"线性乘平方"，并不离谱。真正离谱的是中间两张矩阵 S 和 P：它们都是 [n, n] 的方阵。当序列长度 n 从 1K 涨到 8K、32K、64K，这个方阵的面积按 平方 膨胀。

关键点在于：朴素实现会把 S 和 P 完整写进 HBM（高带宽显存），算完 softmax 再读回来乘 V。换句话说，哪怕你只想算一个 token 的输出 O，它也要求 GPU 先把整张 n×n 的分数矩阵在显存里摊开。这就是"物化（materialize）"。物化本身没问题，但当 n 很大时，这张矩阵会直接顶爆显存，或者至少让显存被它占掉一大块，逼得 batch 上不去。

这里要强调一个工程直觉：显存里的"临时大矩阵"比模型权重更危险，因为权重是固定的、好预算的，而中间矩阵随输入序列长度浮动，你没法在部署前就拍胸脯说"我这张卡一定能跑 32K 上下文"。很多线上事故就是这么来的——压测时用的短样本好好的，一上长文档就 OOM，根子就在 S/P 这张随 n² 膨胀的矩阵。

1.1.2 显存随 n² 膨胀，数字会说话

我们把"S + P 两张矩阵"的半精度（fp16，每元素 2 字节）显存占用算一下：

单张 [n, n] 矩阵字节数 = n × n × 2
S 和 P 合计           = 2 × n² × 2 = 4 · n² 字节 ≈ 4·n² / 1024² MB
​

代入几个真实量级（fp16）：

注意这张图纵轴是显存、横轴是序列长度，曲线是二次方。当 n 从 8K 翻到 64K（8 倍），显存不是 8 倍，而是 64 倍。这就是为什么长文本、长上下文（例如 32K、128K 上下文的模型）对显存极其敏感——瓶颈不在模型参数，而在注意力中间矩阵。

一个常被忽略的真相：参数本身的显存是固定的，但注意力中间矩阵随序列长度平方增长。所以"小模型 + 超长上下文"反而可能比"大模型 + 短上下文"更吃显存。这也是为什么后来 KV Cache 和 PagedAttention 会成为推理优化的主战场：注意力相关的显存，才是长上下文时代真正的吞金兽。

1.1.3 更隐蔽的瓶颈：不是算力，是带宽

很多人以为"慢"是因为乘太多了。我们算一下注意力到底做了多少算术：

算 S = Q·Kᵀ   : n × n × d 次乘加  ≈ O(n²·d)
算 O = P·V     : n × n × d 次乘加  ≈ O(n²·d)
​

这部分计算量确实大，但现代 GPU 的 FLOPS 增长非常快。真正卡脖子的是 softmax 这一步对显存的读写。softmax 要对 S 的每一行做"减行最大值、取指数、求行和、归一化"。这意味着 GPU 必须：先把整张 S（[n, n]）从 HBM 读进片上 SRAM；做归约（max、sum）；把结果 写回 HBM；后面算 O = P·V 时，再把 P 从 HBM 读回来。

问题来了：过去十年，GPU 算力（FLOPS）大约每两年翻 2~3 倍，而 HBM 带宽每年只涨约 1.5 倍。结果就是"算术强度"（每搬一个字节能做多少运算）偏低的应用，会被带宽死死卡住。注意力恰恰是这种"低算术强度"操作：它要反复把大矩阵在 HBM 和算力单元之间搬来搬去，乘法单元常常在"等数据"。

说得更直白一点：你买了一辆顶级跑车（算力），但修的是一条乡间小路（带宽），车再快也跑不出去。注意力在朴素实现下就是被"小路"卡死的那辆车。理解了这一点，你就不会再把所有希望寄托在"换更大显卡"上——你得修路，而"修路"在工程上就是把数据留在更快的 SRAM、减少 HBM 往返，这正是 FlashAttention 做的事。

1.1.4 用"屋顶线模型"把话说死

为了不空口说"带宽重要"，我们借用体系结构里常用的 屋顶线模型（Roofline Model） 来定量。它说：一个算子到底受限于"算力"还是"带宽"，取决于它的算术强度（arithmetic intensity = 总浮点运算数 ÷ 总数据搬运字节数），与机器的"平衡点"比较。

以一块公开的 A100 80GB 规格为例（具体数字请以你的硬件官方规格为准，这里只做量级演示）：

fp16 算力 ≈ 312 TFLOPS
HBM 带宽 ≈ 2.0 TB/s
机器平衡点 ≈ 312e12 / 2.0e12 ≈ 156 FLOPs/字节
​

• 如果某算子的算术强度 > 156 FLOPs/字节，它受算力限制（再快也快不过乘法单元）；
• 如果 < 156 FLOPs/字节，它受带宽限制（乘法单元在等数据，加再多算力也没用）。

现在看注意力里的 softmax。它要读整张 S（n² 个元素），每个元素约做 3 次运算（减 max、指数、除 sum），数据量约 n² × 2 字节（fp16）：

softmax 算术强度 ≈ 3·n² / (n² × 2 × 2) = 3/4 ≈ 0.75 FLOPs/字节
​

0.75 远远低于平衡点 156，所以 softmax 是极度带宽受限的。反观 Q·Kᵀ 这个矩阵乘，对很大的 n 算术强度会高一些（约 n/4 量级，n=8192 时约 2048，反而偏算力受限）。但整个注意力之所以慢，正是因为在标准实现里，为了算 softmax，我们必须把 [n,n] 的 S、P 反复在 HBM 上落地——这块 HBM 流量才是 Roofline 图上把注意力压在"带宽坡"下的真凶。

作者判断（踩坑经验）：90% 的初学者优化注意力时第一反应是"换更强的卡"，但如果瓶颈是带宽而不是算力，换卡收益有限（因为带宽没同比例涨）。真正的解法在后面两章：FlashAttention-2 的核心是"不把整张 S 写出来"，把数据留在更快的 SRAM 里分块算；PagedAttention 是把 KV Cache 的显存管得更聪明。一句话记住本节：注意力慢，主要慢在"反复搬数据"，而不是"算得慢"。

1.1.5 一个具体算例：n = 4096 时的账怎么算

把上面的公式落到一组真实数字上，印象会深很多。取 n = 4096、d = 4096、fp16：

Q·Kᵀ 的 FLOPs      ≈ 2 · n² · d = 2 × 4096² × 4096 ≈ 1.4 × 10¹¹
Q·Kᵀ 读写的字节     ≈ (读 Q,K 各 n·d·2) + (写 S  n²·2) ≈ 2×4096×4096×2 + 4096²×2 ≈ 1.3×10⁸
Q·Kᵀ 算术强度       ≈ 1.4e11 / 1.3e8 ≈ 1077 FLOPs/字节   （偏高，偏算力受限）

softmax 的 FLOPs   ≈ 3 · n² = 3 × 4096² ≈ 5.0 × 10⁷
softmax 读写字节    ≈ 读 S n²·2 + 写 P n²·2 ≈ 2×4096²×2 ≈ 6.7×10⁷
softmax 算术强度    ≈ 5.0e7 / 6.7e7 ≈ 0.75 FLOPs/字节     （极低，严重带宽受限）
​

这两个数对比很说明问题：Q·Kᵀ 自己是算力型的，但一旦配上 softmax 的"读 S → 写 P → 读 P"这一串 HBM 往返，整个注意力就被带宽拖下水。朴素实现里，S 和 P 要各在 HBM 上落一次地，这一来一回的带宽成本，比矩阵乘本身的算力还贵。FlashAttention 的聪明之处，就是让 S、P 全程待在 SRAM，根本不经历这两次 HBM 落地。

1.1.6 辨析：那"重计算"能不能救

有人会问：训练里常用的梯度检查点（gradient checkpointing）不是用"重算换显存"吗？推理能不能也把 S、P 不存、用的时候重算？答案是：推理每生成一个 token 都要用前面的 K、V，而注意力中间矩阵 S、P 是"用过即弃"的，但重算的代价是再跑一遍 Q·Kᵀ，这又把算力吃回来了。重计算是"用算力换显存"，对训练（显存比算力金贵）划算，对推理（我们要的是又快又省）不划算——它要么多花算力，要么多花带宽，没解决"带宽瓶颈"这个根本矛盾。FlashAttention 的思路更彻底：既不多物化、也不重算整张矩阵，而是分块在 SRAM 里"算一块、归一一块、累加一块"，一次过就出正确结果。

1.1.7 辨析：是不是只要堆 HBM 带宽就够

另一个常见误解是"既然瓶颈是带宽，那我买 HBM 带宽更高的卡（比如 H100 比 A100 带宽高约 1.7 倍）是不是就解决了"。部分正确，但远远不够。第一，带宽提升慢、且贵；第二，更关键的是，朴素实现里那两次 HBM 落地是"不必要的浪费"——带宽再高，你把一倍流量砍掉九成，都比单纯堆带宽划算得多。所以正确思路永远是"先减少不必要的数据搬运，再考虑堆硬件"。FlashAttention 把 S/P 的 HBM 落地几乎清零，等价于"凭空给你一张带宽高一个量级的卡"，而且不花一分钱买硬件。

1.1.8 给后面两章埋个伏笔：能不能不写出整张 S

如果把 softmax 的公式展开，你会发现它有一个很好的性质：softmax 是按行归一化的，每一行的输出只依赖这一行自己的 max 和 sum，而不需要看到其他行。

softmax(S)_i = exp(S_i − rowmax(S)_i) / Σⱼ exp(S_ij − rowmax(S)_i)
​

既然每一行独立，我们完全可以把 S 切成很多小块，每块在 SRAM 里算完局部 max/sum，再跨块"修正"得到全局正确的 softmax，全程不把完整的 [n, n] 矩阵物化到 HBM。这正是 FlashAttention 系列的核心思想——我们会在第 2 章、第 3 章详细拆，并用"online softmax"这种具体算法把"分块还不丢精度"这件事讲透。

1.1.8 动手验证：用一行代码测出显存随 n² 增长

光看公式还不够，建议你真的跑一下，把「平方增长」变成屏幕上的曲线。下面这段代码不依赖任何推理框架，只用 NumPy 估算朴素注意力 S/P 两张矩阵的 fp16 显存，你可以直接复制运行（其中具体数值请以你的真实环境为准）：

import numpy as np

for n in [1024, 2048, 4096, 8192, 16384, 32768]:
    bytes_s = 2 * n * n * 2          # S 和 P 两张 [n,n] 矩阵，fp16=2 字节
    mb = bytes_s / 1024 / 1024
    print(f"n={n:6d}  S+P 显存 ≈ {mb:8.1f} MB")
​

运行后你会看到，n 每翻倍，显存直接变成 4 倍。把打印出来的数字画成图，就是 1.1.2 那张二次曲线。这一步的意义不只是「验证」，更是帮你建立直觉：以后再有人问「能不能把上下文从 4K 开到 32K」，你脑子里立刻能反应出「显存要涨 64 倍，得先想清楚中间矩阵怎么处理」。

1.1.9 心智模型：把注意力想象成「查表 + 归一」

如果你想给这一节建立一个能长期记住的心智模型，我建议你把标准注意力看成两件事：第一是「算一张巨大的相关性表 S」（[n,n]），第二是「对这张表的每一行做归一得到权重 P，再拿权重去加权求和 V」。朴素实现为了做第二步，必须把整张表先落盘（HBM）再读回来——这就是所有麻烦的根源。FlashAttention 的突破，本质上就是把「算表」和「归一」合并到同一个 SRAM 小循环里一次完成，表格根本不需要落盘。记住这个「查表 + 归一」的拆分，第 2 章看到 tiling 时你会瞬间懂它为什么有效。

1.1.10 本节小结

• 标准注意力的显存瓶颈是 O(n²) 的两张中间矩阵物化，不是模型参数；长上下文下它比权重更吃显存，且随输入浮动、难预算。
• 真正的性能瓶颈常常是 HBM 带宽（反复读写），而非 GPU 算力；用 Roofline 模型可定量判断（softmax 算术强度仅 ~0.75 FLOPs/字节）。
• 重计算"用算力换显存"对推理不划算；单纯堆 HBM 带宽也治标不治本，根本解法是"减少不必要的数据搬运"。
• 因为 softmax 按行独立，存在"分块计算、不物化全矩阵"的可行路径——这是后续 FlashAttention 的引线。

带着这个问题进入 1.2：显存压力除了"中间矩阵"，还有推理特有的 KV Cache，它是另一个量级更大的显存消耗源。

读后自测（1.1）：合上这一节，试试不看原文回答三个问题——① S 和 P 两张矩阵的形状是什么、为什么随 n 平方膨胀？② 用 Roofline 模型说，softmax 为什么是带宽受限而非算力受限？③ 如果让你向老板解释「为什么换了更强的卡推理还是慢」，你会用哪句话？能答上来的，说明这一节真读到脑子里了；答不上来的，回到 1.1.3 和 1.1.4 再读一遍。