KV Cache 优化（三）：局部注意力、StreamingLLM 与 KV Pruning

前两篇分别减少了 KV head 的数量，或把每个 token 的 K/V 压缩成 latent。它们解决的是“每个 token 存得太宽”。这一篇换一个问题：

如果上下文已经长到几十万 token，是否真的需要保留每一个历史 token？

Sliding Window Attention、StreamingLLM 和 Pruning KV Cache 都在减少历史状态，但方式不同：

1
2
3
Sliding Window：从 Attention 规则上限制可见范围
StreamingLLM：保留开头的 attention sink，再保留最近窗口
KV Pruning：运行时根据重要性选择要留下的 token

1. 先看完整 Attention 的代价

在普通因果 Attention 中，位置 t 的 Query 会访问全部历史 K/V：

$o_t=\operatorname{softmax}\left(\dfrac{q_tK_{1:t}^T}{\sqrt{D_h}}\right)V_{1:t}$

如果上下文从 4K 增长到 128K，单个 Decode Query 需要读取的历史位置也增长了 32 倍。即使使用 KV Cache 避免了历史 token 的重复投影，当前 Query 仍然要扫描这些历史 K/V。

这就产生两类成本：

• 缓存容量：历史 token 越多，KV Cache 占用越大；
• 单步读取：每生成一个 token，都要访问越来越长的历史缓存。

2. Sliding Window Attention：只看最近窗口

设窗口大小为 w。位置 t 不再访问 K_{1:t}，而只访问最近的窗口：

$o_t=\operatorname{softmax}\left(\dfrac{q_tK_{\max(1,t-w+1):t}^T}{\sqrt{D_h}}\right)V_{\max(1,t-w+1):t}$

例如 w=4 时：

1
2
3
位置 8 的 Query：只能看 5、6、7、8
位置 9 的 Query：只能看 6、7、8、9
位置 10 的 Query：只能看 7、8、9、10

因此，在所有层都采用固定局部窗口、且不存在其他全局 Attention 通路时，Decode 只需保留最近 w 个 token 的 K/V，缓存上限不再随生成长度无限增长：

1
2
3
旧缓存：[4, 5, 6, 7]
处理 8：丢弃 4，保留 [5, 6, 7, 8]
处理 9：丢弃 5，保留 [6, 7, 8, 9]

如果窗口大小为 w，KV Cache 的序列维最多约为 w，而不是完整上下文长度 T。

2.1 为什么网络变深后，远处信息仍可能传过来？

假设每层只能看自己和前面两个位置。第一层中，位置 8 只能直接读取位置 6、7、8：

1
第 1 层：8 ← 6、7、8

但位置 6 的隐藏状态可能已经包含位置 4 的信息：

1
2
第 1 层：6 ← 4、5、6
第 2 层：8 ← 6、7、8

于是存在一条跨层信息路径：

$4\rightarrow6\rightarrow8$

再堆一层，位置 8 还可以通过位置 6 的上一层表示间接获得更远的信息。若每层窗口向左扩展 r 个位置，堆叠 L 层后，理论有效感受野大致可以扩展到 Lr 个位置。

这解释了“网络足够深时局部窗口的影响可能减弱”：远处信息可以经过中间 token 接力传递。但它并不等价于 Full Attention，因为 Full Attention 一层就能直接建立远距离连接，而局部 Attention 需要多层中转，中转过程也会压缩和混合信息。

这里跨层传递的不是原始 token，也不是上一层的 K/V，而是上一层完整 Transformer Block 的输出隐藏状态。以位置 8 为例：

1
2
3
4
5
第一层：Attention 汇总位置 6、7、8
       → 残差、Norm、FFN
       → 得到 h₈⁽¹⁾

第二层：用 h₈⁽¹⁾ 重新生成 q₈⁽²⁾、k₈⁽²⁾、v₈⁽²⁾

FFN 不直接混合不同 token，但会独立加工每个位置已经收集到的上下文信息。真正跨位置搬运信息的是 Attention；FFN 和残差连接则把这些信息加工、保留在传给下一层的隐藏状态中。

2.2 它改变了什么？

Sliding Window Attention 不是单纯的缓存管理技巧，它改变了 Attention 的可见性规则。原生局部模型会在训练时使用对应的 Mask；对全局模型进行推理时截断可以不更新权重，但不能保证质量不变：

1
2
训练时使用全局 Attention，推理时突然截断窗口
可能导致训练和推理分布不一致

有些模型从训练阶段就采用局部 Attention，有些模型则混合全局层、局部层或特殊全局 token。

3. StreamingLLM：为什么要保留开头 token？

直觉上，滑动窗口应该只保留最近 token。但实践中发现，序列开头的一些 token 即使语义上不重要，也可能收到很高的 Attention 权重。这类 token 被称为 attention sink，可以理解为注意力分布的“落脚点”。

如果简单删除这些开头 token，Softmax 的注意力分布可能变得不稳定，生成质量明显下降。

StreamingLLM 的基本策略是同时保留两部分：

1
[开头的 sink tokens] + [最近的滑动窗口]

例如总缓存预算为 8，保留 4 个 sink 和最近 4 个 token：

1
2
完整历史： [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]
Streaming：[1, 2, 3, 4,             9, 10, 11, 12]

当新 token 到来时，只淘汰中间那段旧窗口，不淘汰开头 sink：

$\text{Cache}_t=\text{Sink}_{1:s}\;\Vert\;\text{Recent}_{t-w+1:t}$

这里 s 是 sink token 数，w 是最近窗口大小。

3.1 Sink token 携带的是什么？

它们不一定是“最有语义价值的词”。更准确地说，模型可能把它们当作稳定注意力分布的锚点。它们承担的是一种计算角色，而不是简单的语义摘要角色。

可以类比为：

1
2
3
Attention 像一群人分配注意力
Sink token 像固定的公告栏
即使公告内容不重要，大家仍需要一个稳定的落点

StreamingLLM 因此不是“智能挑出最重要的语义 token”，而是保留固定前缀锚点和最新上下文。

原始 StreamingLLM 方法可应用于已有模型而不微调，但需要配套的缓存和位置编码管理。尤其使用 RoPE 时，不能只把若干已经旋转的 K 随意拼在一起，再忽略它们对应的位置；应遵循实现中的缓存位置重映射规则。

这里的 Streaming Attention 指这种面向持续输入流的有限缓存 Attention，不是 FlashAttention 内部逐块计算 Softmax 的“流式计算”。StreamingLLM 可以持续处理很长的输入流，却没有获得无限记忆：如果很早之前的订单号所在位置已被淘汰，模型以后不一定能准确找回它。保留 sink 也不等于自动保存所有历史事实。

4. Pruning KV Cache：按重要性淘汰 token

KV pruning 更主动：它给历史 token 打分，只保留预算内更重要的 token。

一个简化的评分方式是统计历史 Key 被当前 Query 关注的程度。例如在若干最近步骤中，位置 j 的累计注意力权重为：

$s_j=\sum_{t\in\mathcal{W}}\alpha_{t,j}$

其中 α_{t,j} 是第 t 步 Query 对历史位置 j 的 Softmax 权重，$\mathcal{W}$ 是观察窗口。这只是说明打分思想的简化公式，并非所有方法都采用同一分数。

例如 H2O 保留累计注意力较高的 heavy hitters 和近期 token；SnapKV 则利用 Prompt 末尾的观察窗口，为各个 head 选择重要的 Prompt KV，并考虑相邻位置的聚集。前者强调生成过程中的淘汰，后者强调生成前对 Prompt 缓存的筛选。

保留分数最高的 token：

1
2
3
历史 token：  1   2   3   4   5   6   7   8
重要性分数： .1  .8  .2  .1  .7  .1  .6  .2
预算=4：      保留 2、5、7，以及最近位置 8

实际方法会考虑更多因素，例如最近 token、局部连续片段、不同层和不同 head 的重要性，以及是否保留特殊 token。

4.1 为什么不能只保留“注意力最高”的一个 token？

因为一次注意力分数只是当前 Query 的观点：

• 当前 Query 关注的位置，不代表下一个 Query 仍然关注；
• 某个 token 当前分数低，未来可能变得重要；
• 不同 head 关注的 token 可能完全不同；
• 删除后，后续 Query 的 Softmax 分布会发生变化。

因此 KV pruning 是有损策略，需要在显存节省和模型质量之间取舍。它还需要额外的打分、索引和缓存管理开销。

4.2 删除 K/V 后会发生什么？

被保留的缓存可以写成一个子集：

$K'=[k_{i_1},k_{i_2},\ldots,k_{i_m}],\qquad V'=[v_{i_1},v_{i_2},\ldots,v_{i_m}]$

当前 Query 改为：

$o_t'=\operatorname{softmax}\left(\dfrac{q_t(K')^T}{\sqrt{D_h}}\right)V'$

注意力不再看到被删除的位置。这里不是把它们的 Value 置零这么简单，因为 Softmax 的分母也会变化，剩余 token 的权重会重新归一化。

5. 三种方法放在一起比较

它们和 GQA/MLA 的区别在于优化维度不同：

1
2
3
GQA：减少每个 token 的 KV head 数
MLA：减少每个 token 的 KV 表示宽度
局部/淘汰方法：减少保留的 token 数

6. 一个共同的缓存预算例子

假设原始上下文有 4096 个 token，模型每层每 token 的 KV 状态需要 4 KiB：

1
2
3
完整 Cache：4096 × 4 KiB = 16 MiB/layer
窗口大小 512：512 × 4 KiB = 2 MiB/layer
压缩比例：1/8

如果采用 StreamingLLM，保留 32 个 sink 和 480 个最近 token，总预算仍为 512：

1
32 个固定前缀 token + 480 个最近 token = 512 个 token

如果采用 pruning，则可以动态选择 512 个 token，但需要付出重要性评估和索引管理成本。

7. 这些方法和 KV Cache 的关系

它们都没有改变 KV Cache 的基本对象：缓存的仍然是各层历史 token 的 K/V，而不是文本、Attention 权重或最终输出。

改变的是缓存集合：

$\{1,2,\ldots,T\}\rightarrow\mathcal{S},\qquad |\mathcal{S}|\ll T$

未来 Query 只访问集合 S 中的 K/V。于是显存和单步读取量都下降，但被删除位置的信息不再能直接参与后续 Attention。

8. 总结

1
2
3
4
完整 Attention：保留并读取所有历史 token
Sliding Window：只保留最近 token
StreamingLLM：保留开头 sink + 最近 token
KV Pruning：按运行时重要性选择 token

最重要的区别是：

Sliding Window 是固定规则，StreamingLLM 是带注意力锚点的固定规则，KV Pruning 是动态选择规则。它们都通过减少 token 数量节省 KV Cache，但都会牺牲一部分全局信息通路。

系列导航

• 第一篇：MQA 与 GQA
• 第二篇：MLA
• 本篇：局部注意力与缓存淘汰
• 第四篇：投机解码

参考资料

• Mistral 7B, 2023.
• Efficient Streaming Language Models with Attention Sinks, 2023.
• H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models, 2023.
• SnapKV: LLM Knows What You Are Looking for Before Generation, 2024.

• Previous
  KV Cache 优化（二）：深入理解 Multi-head Latent Attention
• Next
  KV Cache 优化（四）：Speculative Decoding 投机解码