论文

SelKV:逐token合并或丢弃的选择性KV缓存合并与注意力补偿

SelKV: Selective KV Cache Merging with Per-Token Merge-or-Drop and Attention Compensation

模型推理KV Cache

摘要

大语言模型(LLM)自回归生成文本,依赖随上下文长度线性增长的键值(KV)缓存,构成主要瓶颈。近期压缩方法经token合并缓解该成本,但常依赖不加区分的聚合,损害表示并引入“注意力下垂”——合并后的token尽管编码了多个输入,却获得与单个token相同的softmax质量。我们提出一个免训练的双组件KV缓存压缩框架。第一,软余弦门基于值向量相似度自适应调节合并决策,抑制或丢弃不相似token以保持语义保真。第二,我们引入注意力比例补偿机制,应用由prefill注意力统计导出的解码期logit偏置,纠正合并引起的softmax失衡。在LongBench(16个英文数据集)上仅保留25% KV缓存的评估中,该框架相对代表性一次性压缩基线取得强劲的压缩后表现;在被评估的分组查询注意力(GQA)模型上尤其稳健,保持近乎无损的生成质量。此外,该方法在复杂多文档QA任务上超过全缓存基线,并在100k token下提供3.3倍解码加速。

SelKV:逐token合并或丢弃的选择性KV缓存合并与注意力补偿:论文配图
图 1:选择性 KV 缓存压缩管道概述。预填充后,该方法对令牌重要性进行评分,选择保留的令牌,将逐出的令牌路由到合并目标,应用软余弦门,补偿合并位置的注意力,并在解码之前重新定位 RoPE 索引。