有很长一段时间，我都默认 Transformer 的二次复杂度是一个算力问题——序列一长，GPU 就会被自注意力矩阵的 FLOPs 烧穿。直到我在 A100 上估算了一遍 7B 模型的推理时间账，才发现自己把瓶颈看错了方向。
7B 模型用 FP16 保存权重，大约 14GB。A100 的 HBM 带宽约 2TB/s，把权重从显存里读一遍就需要 7ms。而生成一个 token 的前向计算量约为参数量两倍，14GFLOPs，按 312TFLOPS 的 FP16 峰值算，理论耗时只有 0.045ms。就算实际算力利用率只有十分之一，也不到 0.5ms。也就是说，模型每生成一个 token，等数据的时间是真正计算的十倍以上。它不是在算，它是在等。
二次复杂度，卡的其实不只是计算
注意力机制的每一次前向，都要算一个 n×n 的注意力矩阵——每个 query 和所有 key 做点积，n 是序列长度，于是复杂度是 O(n²)。这是所有讨论的起点，但很少有人继续追问第二步：这个 n×n 矩阵不只活在算式里，它还要被写进显存、被 softmax 读一遍、再被拿出来乘 V。换句话说，注意力层每次前向，都要往 HBM 里反复搬运 O(n²) 的数据。你日常觉得生成变慢，很多时候不是算子在饱和，而是这张 n² 矩阵的显存读写费用在爆炸。
The number of FLOPs is not a good proxy for runtime speed, which is instead dominated by memory access.
这句话写在 2022 年，三年过去，它依然是理解注意力加速最好的钥匙。翻译过来就一句话：FLOPs 不能代表运行速度，运行速度主要被内存访问锁死。
训练、预填充、解码，每个阶段卡的不是同一个项
人们总是笼统地说 Transformer 是 O(n²)，但把它拆开看，训练和推理是两个完全不同的瓶颈逻辑。
| 阶段 | 主要瓶颈 | 为什么 |
|---|---|---|
| 训练 | 算力 | batch 很大，权重被大量 token 复用，算术强度被拉高 |
| 预填充 | 算力与内存交替 | batch 小，n×n 注意力矩阵开始挤压 HBM 带宽 |
| 解码 | HBM 带宽 | batch 小，权重与 KV cache 的反复搬运占主导 |
训练阶段 batch 可以开到 1024×2048 tokens，同一个权重矩阵被海量 token 共享，计算强度极高，这时候加算力是真的有用。而推理解码阶段每步只生成一个 token，权重被复用的次数约等于 batch size，batch size 又被 KV cache 容量卡着脖子。结果就是：生成一个 token，要搬一遍 14GB 的权重，还要搬越来越长的 KV cache——内存带宽成了绝对瓶颈。
你日常和 ChatGPT 聊天的延迟，大部分就卡在解码这一步。
FlashAttention 的赌注：宁可多算，也要少搬
FlashAttention 做了一个看起来很反直觉的决策：它没有减少 FLOPs，反而因为 online softmax 多算了一点。它真正的贡献是把 QK^T 分块，让 n×n 注意力矩阵只活在 SRAM 里，不写回 HBM。配合 online softmax 的重归一化，一次前向对 HBM 的访问从 O(n²) 降到接近线性。
我第一次看到这个设计时很困惑：多算了怎么还能更快？后来想通了——HBM 的带宽只有 2TB/s，而芯片上的 SRAM 带宽高一个数量级，更关键的是它不用来回搬数据。GPU 的实际调度逻辑里，把数据从慢速显存搬到高速缓存，往往比在缓存里多做几十次算术运算还要贵。
所以 FlashAttention 优化的是 memory-bound，不是 compute-bound。这也是为什么很多模型启用它之后，FLOPs 一点没变，速度却能翻两三倍。
KV cache：二次复杂度的隐藏锁链
如果只有权重搬运，问题还不至于太难看。但 Transformer 每生成一个新 token，都要把历史上所有 key 和 value 重新读一遍才能算注意力。这些 K/V 被缓存在显存里，就是 KV cache。
KV cache 的大小随序列长度线性增长，而每生成一步，它都要被完整从头读一遍。生成的 n 步加起来，KV cache 的读取量就是 O(n²)。这才是二次复杂度在推理阶段最真实的显形——它甚至不是 FLOPs，是纯数据搬运。
这个机制有点像一个越拉越长的水管：序列越长，管子里存的水越多，每一步生成都要先把整根管子的水过一遍再输出。长上下文模型跑起来慢，很大程度上是被这根水管拖住了。
一张表看清四种优化思路
今天主流的长上下文优化，几乎没有一条是靠强行压低注意力 FLOPs 来提速的。它们的共同点是：让数据在 HBM 上少走几趟。
| 方案 | 减少什么 | 代表工作 |
|---|---|---|
| FlashAttention | 注意力矩阵的 HBM 读写 | FlashAttention |
| PagedAttention | KV cache 碎片化与显存浪费 | vLLM |
| GQA / MQA | KV cache 体积与带宽 | GQA 论文 |
| KV 量化 | KV cache 的存取字节数 | WKVQuant / ZipKV 等 |
vLLM 的 PagedAttention 甚至完全不减少 FLOPs，只是把显存碎片化问题治了，吞吐量就能成倍提升——因为它让更多请求同时在 GPU 上驻留，权重搬运被更大的 batch 摊薄了。2024 年苹果那篇 LLM in a Flash 更是把问题推到极端：当权重大到必须从闪存加载，I/O 直接取代 DRAM 成为新瓶颈。
那是不是把所有注意力改成线性就好了？
你可能会问：那上稀疏注意力或者线性注意力，把 O(n²) 变成 O(n) 不就得了？对，但要付代价。稀疏注意力要提前预设哪些 token 才需要关联，遇到长距离依赖很容易失手；线性注意力用核函数近似 softmax，通常在质量上打折。这就是为什么很多长上下文模型到今天还在坚持 full attention，同时靠更聪明的 KV cache 管理和并行方案硬扛。复杂度上的浪费，换来了架构的简单和可训练性，这笔账在工程上经常是划算的。
所以二次复杂度到底卡在哪？我的答案是分阶段的
回到标题的问题：是 HBM 带宽还是计算量？
我的回答是：理论上，当 n 趋向无穷，二次复杂度最终一定是算力先爆。但在你实际能买到的 GPU 上、在你日常能遇到的上下文长度里，先撞上来的几乎都是 HBM 带宽。训练阶段算力还会占上风，推理阶段的解码——尤其是长上下文解码——已经被内存带宽锁死了。
这对我理解模型架构演进有一个直接启发：看一个注意力变体厉不厉害，不要只盯着它 FLOPs 降了多少，要问它每生成一个 token 要搬多少 byte。下一代架构的竞争，表面上是算法竞争，本质上是数据搬运的物流竞争。
原创文章，作者：guanweilu，如若转载，请注明出处：https://guanweilu.cn/article/433.html