随着上下文长度达到数万或数十万令牌,Transformer解码器中的键值缓存成为主要部署瓶颈。缓存存储每个层和头的键值,形状为(2, L, H, T, D)。对于Llama1-65B这样的标准Transformer,在128k令牌和bfloat16精度下,缓存约达335 GB,直接限制了批处理大小并增加了首令牌生成时间。
生产模型已沿多个轴压缩缓存。分组查询注意力在多个查询间共享键值,在Llama3中实现4倍压缩,GLM 4.5中12倍,Qwen3-235B-A22B中高达16倍,均沿头轴进行。DeepSeek V2通过多头潜在注意力压缩键值维度。混合模型将注意力与滑动窗口注意力或状态空间层结合,以减少维护完整缓存的层数。
这些变化未沿序列轴压缩。稀疏和检索式注意力在每个解码步骤仅检索缓存子集,但所有令牌仍占用内存。因此,实用的长上下文服务需要删除对后续令牌影响可忽略的缓存条目的技术。
英伟达的KVpress项目在一个代码库中收集了二十多种此类剪枝方法,并通过Hugging Face上的公共排行榜公开。H2O、Expected Attention、DuoAttention、Compactor和KVzip等方法均以一致方式评估。
KVzip是当前KVpress排行榜上最强的缓存剪枝基线。它使用复制粘贴预文本任务为每个缓存条目定义重要性分数。模型在扩展提示上运行,要求精确重复原始上下文。对于原始提示中的每个令牌位置,分数是重复段中任何位置分配回该令牌的最大注意力权重,在使用分组查询注意力时跨同一组中的头。低分条目被逐出,直到满足全局预算。
KVzip+改进了此分数。它将注意力权重乘以值贡献到残差流的范数,并通过接收隐藏状态的范数归一化。这更好地匹配了令牌在残差流中引起的实际变化,并提高了与下游准确性的相关性。
这些oracle分数有效但昂贵。KVzip需要在扩展提示上预填充,这使上下文长度加倍,对生产环境来说太慢。它也无法在解码期间运行,因为评分过程假设固定提示。
KVzap用小型代理模型替代oracle评分,直接操作隐藏状态。对于每个Transformer层和每个序列位置t,模块接收隐藏向量hₜ并输出每个键值头的预测对数分数。考虑两种架构:单线性层(KVzap Linear)和两层MLP,带GELU且隐藏宽度等于模型隐藏大小的八分之一(KVzap MLP)。
训练使用Nemotron预训练数据集样本中的提示。研究团队筛选27k个提示,长度在750到1,250令牌之间,每个子集采样最多500个提示,然后每个提示采样500个令牌位置。对于每个键值头,他们获得约120万训练对和23k验证对。代理模型学习从隐藏状态回归到对数KVzip+分数。