AI 技术AI 推理 · 18/18
#SnapKV#KV Cache Compression#长上下文#Attention Pattern#推理加速#LLM Serving

AI 推理系列(十八):SnapKV——为什么模型真正关注的历史 Token,可能不到 10%?

围绕长上下文 Decode 中持续增长的 KV Cache,解释 SnapKV 如何利用 Prompt 末尾观察窗口估计各 Attention Head 的关键历史位置,并分析按 Head 选择、局部聚合、显存与带宽收益、多轮 Agent 失效以及指令遵循风险。

一个 32K Token 的检索增强请求完成 Prefill 后,模型已经为全部输入生成了 KV Cache。接下来的答案也许只有 300 Token,但每生成一个新 Token,Attention Kernel 都要从显存读取一份不断增长的历史状态。即使模型实际只反复关注标题、问题相关段落和少量近期上下文,冷门位置的 KV 仍会占据显存,并参与后续访存。

SnapKV 针对的不是 KV 的数值精度、物理布局或跨请求复用,而是一个更激进的问题:Prefill 已经生成的历史 KV,是否都必须进入 Decode 阶段。它观察 Prompt 末尾一小段查询位置的 Attention 分布,从中估计每个 Head 后续可能持续关注的历史位置,然后只保留有限预算内的关键 KV 与近期窗口。这里的“不到 10%”应理解为部分实验配置中的高压缩现象,而不是所有模型、任务和请求都具有固定的 10% 有效 Token 比例。

KV Cache 中为什么会出现大量冷数据

自回归 Decode 的每一步只新增一个 Query,但需要读取历史 Key 和 Value。对于完整 Attention,历史长度为 N 时,每一步的 KV 读取量大体随 N 增长。长上下文请求的 Decode 因而经常受到显存容量和内存带宽共同限制:上下文越长,单请求占用越大,每轮生成需要搬运的数据也越多。

“理论上可以关注”与“实际持续关注”并不相同。在长文档问答中,某些 Head 会偏向最近 Token,某些 Head 会长期关注少量分散位置,还有一些 Head 会在标题、分隔符或问题相关片段附近形成稳定热点。大量历史位置虽然保留在缓存中,但在后续生成中的 Attention 权重长期很低。

这种稀疏性不代表被忽略位置在语义上毫无价值。它只说明在当前 Prompt、当前问题和当前生成轨迹下,它们没有被这些 Head 高频读取。若问题、角色指令或生成目标改变,原来的冷位置可能重新变成关键位置。SnapKV 的收益和风险都来自这种上下文依赖性。

观察窗口为什么能够提供预测信号

SnapKV 把 Prompt 末尾的一段 Query 位置作为 Observation Window。假设输入由文档、指令和最终问题组成,Prompt 末尾通常包含用户当前任务、输出要求和靠近生成起点的上下文。模型在计算这些位置时,已经需要决定应从前文读取哪些内容。

如果多个观察位置反复把较高 Attention 分配给同一组历史位置,可以把这些位置看作当前任务下的候选工作集。SnapKV 利用这一现象,在 Decode 开始前汇总观察窗口对历史 KV 的访问模式,形成每个 Head 的重要性分数。

这里不是在准确预言未来每一步 Attention,而是在做带有误差的工作集估计。它依赖两个经验条件:当前问题已经出现在 Prompt 中,以及 Prefill 尾部的访问模式与短期 Decode 具有一定持续性。单轮长文档问答通常更符合这两个条件;后续任务会改变的多轮 Agent 则更容易破坏它们。

从完整 Prompt 到压缩缓存的执行流程

SnapKV 的压缩发生在 Prefill 完成之后、正式 Decode 之前。完整 Prompt 仍需要经过模型,因此它不会直接减少 Tokenization、Prefill FLOPs 或长 Prompt 的首轮 Attention 计算。它节省的是之后需要驻留和读取的 KV。

flowchart TD
  A[长 Prompt Prefill] --> B[生成完整 KV]
  B --> C[读取末尾观察窗口 Attention]
  C --> D[按 Head 汇总历史位置分数]
  D --> E[局部聚合相邻热点]
  E --> F[选择关键历史位置]
  F --> G[保留近期窗口]
  G --> H[构建压缩 KV Cache]
  H --> I[后续 Decode]
  I --> J{任务关注点改变?}
  J -->|否| K[继续使用压缩缓存]
  J -->|是| L[存在遗漏关键 KV 风险]

工程实现需要把选中的逻辑 Token 位置重新组织成 Attention Kernel 可读取的 KV 张量或索引结构。压缩完成后,Decode 只能访问保留下来的历史位置与近期窗口。被移除的 KV 不会因为后续 Attention 突然需要它们而自动回来,除非系统重新执行相关 Prompt 的 Prefill 或保存了另一份可恢复状态。

为什么必须按 Head 选择

不同 Attention Head 学到的访问模式可能不同。一个 Head 偏向局部语法和最近位置,另一个 Head 可能跟踪实体、分隔符或文档中特定区域。若使用一套全局 Token 排名,强势 Head 的热点可能占满预算,掩盖其他 Head 所需的历史位置。

SnapKV 因而在 Head 粒度上统计重要性并分配保留位置。对标准多头注意力,可以理解为每个 Query Head 形成自己的候选集合;在 GQA 或 MQA 中,多个 Query Head 共享较少的 KV Head,实际实现需要把 Query Head 的分数汇总到对应 KV 组,避免选择结果与物理缓存布局不一致。

按 Head 选择会引入新的不规则性。不同 Head 可能保留不同位置,如果底层 Kernel 只擅长读取统一、连续的 KV 序列,理论上的压缩比例未必会等比例转化为端到端速度。高效实现需要控制索引、Gather、张量重排和 Kernel 启动的额外成本。

局部聚合解决了什么问题

直接选择 Attention 分数最高的位置,容易把预算集中在同一段连续区域。例如实体名称、代码语句或标题附近的多个相邻 Token 可能同时获得高分。逐个保留它们未必错误,但会降低缓存对整段长上下文的覆盖范围。

SnapKV 使用局部池化或聚合思想平滑历史位置的重要性,使相邻热点形成更稳定的区域分数,再从中选择代表性位置。它不等同于语义聚类,也不会把一段文本压缩成摘要;底层保存的仍是原 Token 对应的 Key 和 Value,只是选择过程不再完全依赖单点峰值。

聚合窗口过小,选择结果可能被噪声和单个极端分数支配;窗口过大,又可能把相邻但语义不同的区域混在一起。观察窗口长度、聚合核大小、每个 Head 的保留预算以及近期窗口大小,都会共同影响压缩质量。

SnapKV 节省什么,又不节省什么

成本或能力SnapKV 的影响仍然存在的边界
Prompt Prefill 计算基本不减少完整 Prompt 仍需生成初始 KV 和观察窗口 Attention
Decode KV 显存按保留预算下降每层、每个 KV Head 仍需保存选中位置的 K/V
Decode 内存带宽历史读取量可明显下降索引、Gather 和不规则访问可能抵消部分收益
模型参数不修改目标模型权重和主要计算结构不变
训练成本可作为 training-free 方法使用不代表对所有模型无需适配或质量验证
历史可访问性只访问保留位置淘汰 KV 无法在当前缓存中精确恢复
输出质量在适合的任务和预算下可接近完整缓存高压缩、多指令和任务转移可能导致退化

论文报告的内存效率和 Decode 加速来自特定模型、输入长度、压缩预算、硬件与 Kernel 实现。内存压缩比例通常比端到端速度提升更显著,因为 Decode 仍包含投影、MLP、采样、调度和其他不能被 SnapKV 消除的成本。

它与其他 KV 优化处在不同层级

GQA 减少 KV Head 数量,MLA 改变历史状态的表示形式,KV 量化减少每个元素的位宽,PagedAttention 改善物理分配,Prefix Caching 复用已经计算过的公共前缀。SnapKV 则选择删除部分历史位置对应的 KV。

技术主要动作是否丢弃历史位置主要风险
GQA / MQA多个 Query Head 共享 KV Head模型结构与质量权衡
MLA缓存低维 Latent否,改变表示需要模型与专用 Kernel 支持
KV 量化降低 KV 数值精度量化误差累积
PagedAttention分页管理物理 KV Block映射和 Kernel 复杂度
Prefix Caching跨请求复用相同前缀命中率和缓存路由
Sliding Window只保留最近位置是,按时间淘汰远距离信息直接丢失
SnapKV根据观察窗口选择关键位置是,按注意力模式淘汰未来查询变化与选择误差

这些方法可以组合,但组合后的有效缓存预算、数据布局和 Kernel 支持必须共同设计。例如先用 GQA 缩减 KV Head,再对每个 KV 组应用 SnapKV,可以同时降低 Head 维度和序列维度的缓存量;继续叠加低比特量化,还会改变带宽与误差分布。每增加一层压缩,都应重新测量质量,而不是把单项论文收益简单相乘。

RAG 单轮问答适合,多轮 Agent 风险更高

在单轮 RAG 中,完整文档、检索片段、用户问题和输出要求通常在 Prefill 时已经确定。观察窗口能够看到当前问题,选择出的历史位置更可能覆盖本次答案所需证据。若答案较长,压缩后的缓存还会在多个 Decode 步骤中持续复用,因此收益更容易积累。

多轮 Agent 的条件不同。第一次任务可能要求总结架构,第二次改为查找某个 API 参数,第三次又要求遵循最前面的安全规则。若每轮只继承上轮压缩后的 KV,新问题需要的历史位置可能早已被删除。缓存压缩会从一次推理优化,变成不可逆的会话状态裁剪。

长期 Agent 更稳妥的设计是把可恢复的原始对话、文档或外部 Memory 保存在 KV Cache 之外。每轮任务变化明显时,可以重新构造 Prompt、重新检索证据或重新 Prefill,而不是假设第一次观察窗口选出的工作集永久有效。系统 Prompt、角色规则和工具约束还可以设置保护策略,避免被一般的注意力分数竞争淘汰。

生产落地应观察哪些指标

只统计“压缩了多少 Token”不足以判断是否值得部署。至少需要联合观察以下信号:

  • 各层、各 KV Head 的实际保留比例,而不是只看全局平均值;
  • 观察窗口内的 Attention 集中度,以及选择集合在相邻请求间的稳定性;
  • 压缩、重排和索引构建时间,占完整 Prefill 与 Decode 延迟的比例;
  • Decode 阶段实际读取的 KV 字节数和 Kernel 时间;
  • TTFT、TPOT、P95/P99 延迟和单位 GPU Goodput;
  • 长文档问答、精确复制、前部指令、多轮任务切换等专项质量指标;
  • 不同 Prompt 长度和压缩预算下的退化拐点。

还应把 System、Developer、Tool Schema、用户内容和检索文档分角色评测。某种压缩策略即使总体困惑度变化不大,也可能优先损伤位于 Prompt 前部的规则,表现为格式约束失效、工具调用越权或安全指令遗忘。平均准确率很难暴露这种风险。

从固定预算走向动态和两阶段压缩

后续工作沿两个方向扩展了 SnapKV。一类方法继续压缩 Decode 访问:先在 Prefill 后删除大部分历史 KV,再在生成过程中使用稀疏 Attention,避免每一步读取全部保留位置。RocketKV 属于这种两阶段思路。另一类方法尝试根据请求、层、Head 或当前状态动态分配不同压缩率,而不是给所有样本使用同一预算,VarRate 等工作反映了这一方向。

动态方法试图解决一个现实问题:不同输入的可压缩性差异很大。重复性强、证据集中、问题明确的 Prompt 可以使用更小预算;多指令、跨段推理、精确复制和未来查询不确定的 Prompt 需要更保守的缓存。系统应把压缩率视为质量与资源之间的在线决策变量,而不是部署时一次设定的常数。

研究也逐渐从平均任务分数转向更细的副作用。指令遵循风险尤其重要,因为被压缩掉的并不只是“背景文本”,也可能是 System Prompt、输出格式、拒答边界和工具规则。KV Cache Compression 若进入生产系统,需要像量化和模型蒸馏一样建立独立回归集,而不能只依赖通用长上下文基准。

结语

SnapKV 把长上下文缓存问题从“怎样保存全部历史”改写成“本次生成真正需要保存哪些历史”。它利用 Prompt 尾部观察窗口估计 Attention 工作集,按 Head 选择并聚合重要位置,同时保留近期窗口,从而缩小 Decode 阶段的驻留 KV 与读取带宽。

这种方法最有吸引力的场景,是问题在 Prefill 时已经明确、后续生成会长期复用同一批证据的单轮长上下文请求。它最危险的场景,则是任务不断变化、前部指令不能遗忘、被淘汰信息无法重新获取的长期会话。SnapKV 证明 KV Cache 可以被选择性裁剪,但也把一个新的责任交给推理系统:缓存不再只是内存对象,而是影响模型后续可访问信息边界的有损状态。

资料来源

  1. SnapKV: LLM Knows What You are Looking for Before Generation
  2. SnapKV 官方实现
  3. RocketKV: Accelerating Long-Context LLM Inference via Two-Stage KV Cache Compression
  4. VarRate: Training-Free Variable-Rate KV Cache Compression for Long-Context LLMs
  5. The Pitfalls of KV Cache Compression