KV 缓存、稀疏注意力与 FlashAttention

第 21 课结尾我们记了一笔账：注意力要让每个 token 和所有 token 两两比较，$n$ 个 token 就有 $n\times n=n^2$ 个格子。当时一带而过，这一课来还这笔账。

把数字摆具体：GPT-4 上下文 128k token，若全用标准注意力，一层的注意力表就有 128000² ≈ 164 亿个格子，乘上约 120 层——一次推理光注意力就得算两万亿次，FFN 还没算进去。 这道平方的账单，是长上下文路上最硬的一堵墙。

最直觉的省法：每个 token 别再看全部 n 个位置，只看左右 W 个邻居。 这叫滑动窗口注意力，计算量从 $O(n^2)$ 一下降到 $O(n\cdot W)$（W 固定，相当于随 n 线性增长）：

可只看邻居有个明显代价：看不到远处——开头那句关键信息，到结尾就盲了。 所以实战中会打几个补丁：留几个「全局 token」（比如开头的指令）让所有位置都能看到它、它也能看到全文； 或者让窗口「跳着看」（膨胀窗口），用同样的格子数覆盖更远。Longformer（2020）、Mistral（2023）就靠这类组合，把上下文撑到几十万 token。

在动用聪明算法之前，生成阶段有一笔纯属浪费的重复账，白送的优化。回想第 23 课的自回归：模型一个词一个词地吐， 每吐一个，就把它接到末尾再跑一遍。问题来了——每跑一遍，它都把前面所有词的 K、V 从头又算了一次。

可那些词没变啊！「I」的 K、V，在第 2 步、第 3 步、第 100 步都是同一个值。 于是有了 KV 缓存：算过的 K、V 存起来，下一步只算新来的那一个词，其余直接从缓存取：

代价是拿显存换速度：缓存要一直占着内存，序列越长、占得越多—— 这也是为什么超长对话特别吃显存。后来的 MQA / GQA（第 20 课提过的多头变体）就是为了缩小这个缓存而生的： 让多个头共享同一份 K、V，缓存立刻小一大截。省内存这件事，从注意力一路追到了缓存。

这条路最反直觉：一格都不少算，结果和标准注意力分毫不差，却能快好几倍。秘密不在「算」，在「搬」。GPU 有两层存储：

• HBM（显存）：几十 GB，容量大，但读写慢；
• SRAM（片上缓存）：只有几十 MB，但读写快十几倍。

标准注意力的真正瓶颈，是那张 $n\times n$ 的大矩阵：它太大，塞不进 SRAM， 只能写进慢吞吞的 HBM，再读回来做 softmax，再写回去……来回搬运几趟。瓶颈根本不是计算，是搬运。

可这里有个拦路虎：softmax 要先看到一整行所有分数（求最大值、求总和）才能归一化。 要是把行切成小块、一块一块进 SRAM，看第一块时还不知道后面有没有更大的数，怎么算？FlashAttention（2022，斯坦福）的核心绝活，就是一招「在线 softmax」：边读边修。

诀窍就在那句「发现更大的数，就把已经算的部分等比缩小一下」。靠一个running 最大值和一个running 分母， softmax 不用一次看到整行，也能流着算、且结果精确。于是注意力可以分块塞进 SRAM 算完就走，那张 n² 大矩阵压根不用落地到 HBM：

它不改变结果、不丢信息，所以成了无脑就该开的默认项：今天 PyTorch、各大训练框架都内置了它。 稀疏注意力（省算）和 FlashAttention（省搬）还能叠加使用，长上下文这才真正跑得起来。