一个 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 可以被选择性裁剪,但也把一个新的责任交给推理系统:缓存不再只是内存对象,而是影响模型后续可访问信息边界的有损状态。