第 1 章 基础：从注意力瓶颈说起

第 1 章 基础：从注意力瓶颈说起

本章导读：如果你曾经跑过一个稍微大一点的模型，大概率经历过这样的时刻——显存明明还剩一大块，推理速度却卡得让人怀疑人生；或者明明 batch 调大了一点，程序就报 OOM。这一章我们不急着上 FlashAttention 或 PagedAttention 的代码，而是先把"为什么慢、为什么费显存"这件事讲透。只有把瓶颈定位准确，后面所有的加速技术才有意义。读完本章，你应该能用一句话解释：标准注意力计算的复杂度到底卡在哪里，以及为什么"算力够、显存也够"却依然跑不快。

1. 注意力的数学形式与朴素实现

Transformer 的核心是自注意力。对于长度为 n、隐藏维度为 d 的输入，注意力计算可以写成：

Q = X · Wq
K = X · Wk
V = X · Wv
S = Q · Kᵀ / √d        # 注意力分数矩阵，形状 [n, n]
P = softmax(S)         # 概率分布
O = P · V              # 输出，形状 [n, d]
​

朴素实现会把中间的 S 和 P 两张 [n, n] 矩阵完整物化（materialize）到显存里。问题就出在这里：

• 计算复杂度 O(n²·d)：序列越长，分数矩阵平方膨胀。
• 显存复杂度 O(n²)：S 和 P 必须常驻 HBM（高带宽显存）。当 n=8192、半精度时，单这两张矩阵就要约 268 MB；n=65536 时超过 17 GB，远超单卡可用空间。
• 更隐蔽的瓶颈——显存带宽：现代 GPU 的计算能力（FLOPS）增长远快于显存带宽。softmax 需要对整行做归约（max、sum），意味着 S 要反复从 HBM 读进来、写回去。真正拖慢速度的往往不是乘法，而是"反复搬数据"。

2. KV Cache：推理特有的显存压力

训练时 n 是固定的序列长度；但自回归推理是逐个 token 生成。每生成一个新 token，都要用到前面所有 token 的 K、V。于是工程上引入 KV Cache：把历史 K、V 缓存下来，避免重复计算。

KV Cache 的显存占用是 2 · n · layers · hidden · batch · precision。以一个 7B 模型、序列 4K、batch 32 为例，KV Cache 轻松突破十几 GB。而且它是"按最大可能长度预分配"的——哪怕你实际只生成了 200 个 token，预分配的槽位也占着。

这里埋下了 PagedAttention 要解决的伏笔：连续显存分配导致的碎片与浪费。

3. 两个互补的优化方向

把上面的瓶颈拆开，其实是两个相对独立的问题：

1. 算子层面（计算 + 带宽）：如何让单步注意力又快又省显存？→ FlashAttention-2 的答案。
2. 系统层面（显存管理）：如何让 KV Cache 不再浪费、支持更大 batch？→ PagedAttention（vLLM）的答案。

它们不互斥，反而经常一起用：vLLM 既用 PagedAttention 管理显存，底层也调用 FlashAttention-2 做算子加速。这也是为什么本教程把两者放在同一本书里讲——它们是同一目标下的"双引擎"。

4. 本章小结与路线预告

• 标准注意力的真正瓶颈是 O(n²) 显存物化 + HBM 反复读写，而非单纯的算力不足。
• 推理阶段的 KV Cache 带来连续分配造成的碎片与预留浪费。
• FlashAttention-2 解决"算子效率"，PagedAttention 解决"显存管理"，二者互补。

下一章我们直接进入核心模块，先啃下 FlashAttention-2 的 tiling 与流水线思想，再拆解 PagedAttention 的分页式 KV Cache 设计。建议带着一个问题读下去：为什么"不把整个分数矩阵写出来"就能算对 softmax？