"$O(n^2d)$"这个记号把太多东西压进了一个上标。注意力在序列变长时到底是计算瓶颈还是显存瓶颈？FFN 和注意力谁先顶不住？不同结论会导向完全不同的优化路线（FlashAttention、线性注意力、稀疏化各管一摊）。这篇把三笔账分开算。
FLOPs 账：两个矩阵乘各出 2n²d
忽略多头拆分（它只是把运算切成 $h$ 个并行小块，总数不变），前向注意力的矩阵乘有两个：
$$ S = QK^\top/\sqrt d \;(n\times d \cdot d\times n),\qquad O = PV \;(n\times n \cdot n\times d). $$
每个矩阵乘 $A_{m\times k}B_{k\times p}$ 的乘加次数是 $mkp$，按 2 FLOPs（一乘一加）计：
$$ \text{FLOPs}(QK^\top) = 2n^2d,\qquad \text{FLOPs}(PV) = 2n^2d,\qquad \text{合计 } 4n^2d. $$
对比同层的 FFN（标准配比 $d_{\text{ff}}=4d$）：$\text{FLOPs}=2\cdot n\cdot d\cdot 4d + 2\cdot n\cdot 4d\cdot d = 16nd^2$。两者之比
$$ \frac{\text{注意力}}{\text{FFN}} = \frac{4n^2d}{16nd^2} = \frac{n}{4d}. $$
注意力在 FLOPs 上超过 FFN 的临界点是 $n > 4d$。 以 LLaMA-7B（$d=4096$）算，$n>16\text{k}$ 后注意力开始占大头——这解释了为什么 4k 上下文时代没人着急优化注意力，128k 时代 FlashAttention 和线性注意力成了标配。注意这个 $4d$ 来自 FFN 配比和系数约定，量级结论（$n$ 超过 $d$ 的几倍后注意力主导）比精确系数可靠。
显存账：P 矩阵的 n² 是硬件杀手
训练的显存大头不是参数而是中间激活，注意力的中间产物按序列长度平方膨胀：
$$ S \;(n^2),\quad P \;(n^2),\quad P\odot\bar P \text{ 等反向中间量} \;(n^2). $$
以 fp16 存储、$h$ 个头、$L$ 层计，注意力激活的显存是 $L\cdot h\cdot n^2\cdot 2\text{B} \times \text{若干}$ 的级别。$n=32\text{k}$、$h=32$、$L=32$ 时仅 $P$ 一族就是几十 GB——这就是为什么朴素注意力在长序列训练时先 OOM，而不是先变慢。
对比参数和 optimizer state：7B 模型训练要存参数副本、梯度、Adam 二阶矩（混合精度下约 16 倍参数量），也是几十 GB。有趣的对照：短序列时参数是显存主角，长序列时激活是主角，ZeRO、梯度检查点、FlashAttention 分别从不同角度拆这两座山。
带宽账：decode 阶段的真正瓶颈
推理分两段。Prefill（处理 prompt）是矩阵乘密集阶段，GPU 算力利用率高。Decode（逐 token 生成）每步只算一个新 query 对全部历史 key：
$$ s_{n+1,j} = q_{n+1}^\top k_j\ (j\le n),\quad o_{n+1} = \textstyle\sum_j p_j v_j. $$
这一步的 FLOPs 只有 $O(nd)$，但需要把全部 $K,V$ 从显存搬一遍：访存量 $O(nd)$。算术强度（FLOPs/字节）极低，解码完全被 HBM 带宽卡死，GPU 算力利用率常常不到 5%。这就是 KV cache 存在的原因——把 $K,V$ 留在显存里避免重算，同时它也把显存变成了推理的第一资源约束。KV cache 的显存账：$2\cdot L\cdot h\cdot n\cdot d_h = 2Lnd_{\text{model}}$ 字节级（乘精度系数），$n$ 线性增长，长上下文下逼近参数本身的显存占用，GQA/MLA 和 KV 量化两条压缩路线都从这里出发。
三笔账汇成一张表
| 资源 | 注意力开销 | 增长阶 | 主导阶段 |
|---|---|---|---|
| FLOPs | $4n^2d$ | $n^2$ | prefill、长序列训练 |
| 激活显存 | $O(Lhn^2)$ | $n^2$ | 训练、长 prompt |
| 带宽（decode） | $O(Lnd)$ 访存 / 步 | $n$ | 自回归解码 |
这张表解释了优化方法的分工：FlashAttention 针对激活显存（不物化 $P$），线性注意力针对 FLOPs（去掉 $n^2$），GQA/MLA 针对 KV cache 显存，投机采样针对 decode 的低利用率。没有一种方法通吃，因为它们各自拆的是这本账的不同页。
小结
注意力的 $n^2$ 不是一个笼统的负担，而是拆成三笔各管一段的账：FLOPs $4n^2d$ 在 $n>4d$ 后压过 FFN；激活显存 $O(n^2)$ 是训练 OOM 的来源；decode 带宽 $O(nd)$/步决定推理延迟。理解账本结构才能看懂优化图谱——FlashAttention、线性注意力、GQA、投机采样各自削的是不同的项。下一篇打这本账里最硬的一项：能不能把 $n^2$ 直接降到 $n$，线性注意力的核函数改写。