发表时间: 2025-06 · arXiv:2506.11418 (TMLR 2026)
原文: https://arxiv.org/abs/2506.11418
Jie Hu, Shengnan Wang, Yutong He, Ping Gong, Jiawei Yi, Juncheng Zhang, Youhui Bai, Renhai Chen, Gong Zhang, Cheng Li, Kun Yuan
Peking University, Huawei Technologies, University of Science and Technology of China
一句话结论 本文提出了 CentroidKV,一种基于相似度对 KV Cache 进行在线分块聚类和合并的方法,在保持模型性能的同时将 KV Cache 显存占用降低了 75%,并将长上下文推理的吞吐量提升了最高 4 倍。
要解决什么问题 大语言模型在处理多轮对话等复杂任务时,上下文窗口已扩展至百万级别,这使得 KV Cache 的显存占用随序列长度线性暴涨。由于自回归生成的特性,每生成一个新 token 都需要读取完整的 KV Cache,这不仅导致了严重的显存容量瓶颈,更使得显存带宽成为拖慢推理延迟与吞吐量的核心卡点。现有的主流缓解方案是 KV Cache 驱逐(如基于累积注意力分数丢弃 token),但这种做法会永久丢失被剔除的 token 信息,导致模型在长上下文场景下性能显著下降,因为当前看似不重要的 token 在未来的生成中可能变得至关重要。为了弥补这一缺陷,后续一些基于合并的方法尝试将要驱逐的 token 与保留的 token 进行合并。然而,这些方法仍然依赖于初步的 token 驱逐来定义合并集合,这往往会强行把语义上相距甚远的 token 绑在一起,造成中心点偏移和严重的信息丢失,最终依然会损害模型性能。因此,真正的卡点在于:如何能在不永久丢弃上下文信息、不引入高昂计算开销的前提下,根据 token 的内在语义大幅压缩 KV Cache 的体积。
怎么做的 作者观察到,Key 状态在序列维度上表现出极高的局部相似性,且相似度随 token 距离增加呈单调递减的凸函数趋势。基于此,本文提出了 CentroidKV,核心思路是不丢弃 token,而是将语义相似的 KV Cache 在线聚类并合并为中心点(Centroid)。为了绕开全局聚类(如 K-Means 需多次迭代计算全局距离矩阵)极高的计算复杂度卡点,该方法设计了三个关键部件。首先是分块(Chunking),在保留初始的注意力沉淀(attention sink)和最近生成的 token 后,将长序列切分为多个局部数据块,把匹配范围限制在局部以大幅降低复杂度。其次是分块软匹配(Chunked Soft Matching),该算法受视觉 Transformer 中的二分软匹配启发,在每个数据块内,采用交替划分的策略将 token 分配到两个集合中(作者通过理论证明,在相似度递减规律下,这种交替划分能最大化保留高相似度边缘,是最优策略),然后在所有数据块的两个集合之间寻找余弦相似度最高的 token 对作为候选聚类。最后是 KV Cache 压缩,框架不会合并所有匹配对,而是按照相似度排序,只合并排在前面的比例 $r$。为了在多轮聚类中过滤掉低置信度的匹配,压缩率 $r$ 会随着聚类轮数 $j$ 线性衰减:
$$r = r_{\mathrm{init}} - j \cdot \delta_r$$确定要合并的聚类后,系统会记录每个聚类包含的原始 token 数量(即聚类度 $n_t$),并使用基于度的加权平均来计算新的 Key 中心点:
效果如何 实验在 Llama-3.1-8B-Instruct 和 Mistral-7B-Instruct-v0.2 模型上进行,上下文长度最高测试至 128K。对比基线包括基于固定位置规则驱逐的 StreamingLLM,以及基于内容感知的驱逐方法 SnapKV 和 PyramidKV。在 RULER 和 LongBench 评测基准上,当 KV Cache 预算被压缩到原大小的 25% 到 75% 时,CentroidKV 在单文档问答、多文档问答和摘要等绝大多数需要语义聚合和推理的任务上均取得了最优或次优的准确率,在极低预算下显著优于直接丢弃 token 的基线方法。在效率方面,在 128K 上下文长度下,CentroidKV 将 KV Cache 显存占用降低了 75%,同时将解码阶段的每个 token 生成时间(TPOT)加速了最高 1.92 倍,且首次 token 生成时间(TTFT)的额外开销极小。在集成到 vLLM 推理引擎后,得益于显存占用的减少,系统能够容纳更多并发请求,使服务吞吐量提升了最高 4 倍。该方法的局限性在于,由于它依赖语义相似度进行合并,在处理由无语义随机字符串组成的 UUID 检索任务(如 S-NIAH-3),以及需要精确 token 匹配的细粒度检索任务时,激进的合并会模糊精确的 token 级别信息,导致性能出现下降。此外,当前方法采用手动指定的压缩率调度,且每次聚类最多只能将缓存减半,尚未实现自适应压缩。
核心问题:随着大型语言模型(LLMs)在多轮对话等复杂实际应用中的需求增加,模型已经能够支持高达1M token的上下文窗口。然而,在长上下文场景中部署LLMs面临着与键值(KV)缓存相关的巨大挑战。KV缓存存储了所有先前token的键和值以避免重复计算,其内存需求随上下文长度线性增长。由于LLMs的自回归特性,生成每个token都需要访问整个KV缓存,这使其成为推理延迟和吞吐量的显著瓶颈。此外,巨大的KV缓存对内存容量提出了极高的要求。
现有方法的局限性:现有的KV缓存压缩方法(如驱逐、合并、量化和通道剪枝)存在明显缺陷。基于历史注意力分数丢弃不重要token的KV缓存驱逐方法,通常会导致显著的性能下降,因为当前上下文中被认为不重要的token可能在未来的生成中变得至关重要。虽然一些方法引入了补偿机制(将驱逐的token与保留的token合并),但它们依赖于初步的token驱逐来定义合并集,这往往迫使语义上相距较远的token被分组在一起,导致质心偏差和信息丢失。另一方面,现有的聚类方法(如K-means)计算开销过大,不适合对延迟敏感的在线推理。
研究目标与创新点:
为了应对这些挑战,本文提出了一种简单而有效的在线KV缓存聚类框架——CentroidKV。该方法旨在通过基于内在相似性将KV缓存聚类为质心来进行压缩,从而在不丢失关键未来信息的情况下减少内存占用。
1. 引入了CentroidKV,这是一个用于在线KV缓存聚类的轻量级即插即用框架。其核心创新是一种名为分块软匹配(Chunked Soft Matching)的新型聚类算法。
2. 在理论上分析了该框架的计算复杂度,并基于观察到的相似性模式,正式证明了块内划分策略的最优性。
3. 实验表明,CentroidKV在对模型性能影响极小的情况下,实现了高达 $75\%$ 的KV缓存内存使用量减少。此外,它将解码阶段的推理速度提升了高达 $1.92\times$,并将吞吐量提高了高达 $4\times$。
KV缓存聚类基础:为了简化说明,聚焦于特定层内的单个注意力头。LLM推理包含预填充(pre-filling)和解码(decoding)两个阶段。在预填充阶段,模型生成第一个token并初始化KV缓存,将提示词的键和值状态分别存储为 $K \in \mathbb{R}^{n \times d}$ 和 $V \in \mathbb{R}^{n \times d}$。在解码阶段,给定当前输入的查询、键和值状态 $q, k, v \in \mathbb{R}^{1 \times d}$,KV缓存更新为 $K = [K, k]$,$V = [V, v]$。原始的注意力输出计算如下:
基于余弦相似度的距离度量:将KV缓存分解为单个token,表示为 $K = [k_1, \ldots, k_n]$ 和 $V = [v_1, \ldots, v_n]$。方法中利用余弦相似度作为键状态之间的距离度量:
Observation 1:键状态在序列维度上表现出高度局部的相似性:使用Llama-2-7B-32K模型,从WikiText-2数据集中随机采样长度为4K的序列进行零样本推理。观察发现,跨不同层和头的token之间存在高余弦相似度,且具有高相似度的token倾向于聚集在局部区域内。这一观察表明,可以利用KV缓存聚类来实现高效推理而不影响准确性,同时也启发了在局部区域内(而非整个序列)识别相似token以提高聚类效率的设计。
Observation 2:随着token距离的增加,键状态的余弦相似度通常呈单调递减并呈现凸性趋势:在序列中随机采样多个token,将局部区域定义为256个token,并计算样本和注意力头的平均相似度。结果发现,随着token距离的增加,键状态之间的余弦相似度总体遵循单调递减趋势,且这种相关性相对于距离呈现凸性(convex)。这一观察为设计一种最小化识别相似token集计算复杂度的高效聚类算法提供了动机,并为后续证明框架中划分策略的最优性提供了经验支持。
整体推理流水线:集成CentroidKV的LLM推理整体流水线包含在算法1中。缓存预算 $B$ 由提示词长度 $n$ 和缓存比率 $R$ 决定。在预填充阶段,通过Flash Attention【Tri Dao 等人,Flashattention: Fast and memory-efficient exact attention with io-awareness,2022,NeurIPS】计算注意力输出后,如果KV缓存大小超过预算,则调用CentroidKV进行第一次聚类,将KV缓存压缩为质心并记录聚类度。在解码阶段,KV缓存每步增长一个token,导致持续超出内存预算。为了摊销压缩开销并保持推理效率,CentroidKV每隔 $g$ 个解码步被周期性地调用一次。由于二分软匹配在单次传递中最多只能将缓存大小减少一半,因此需要执行多轮聚类,以逐步压缩KV缓存直到满足预定义的预算。
# 算法1:集成CentroidKV的推理流水线
Require: 缓存比率 R, 压缩比率 r, 最大解码长度 T, 注意力sink n1, 最近预算 n2, 间隔步数 g, 块大小 c
1: 预填充阶段: Q, K, V ∈ R^{n×d}
2: 缓存长度 s = n, 聚类度 N = [1] · n, 缓存预算 B = R · n # 初始化
3: O = FlashAttn(Q, K, V)
4: while s > B do
5: K, V, N, s = CentroidKV(K, V, N, s, n1, n2, r, c) # 预填充后的第一次聚类
6: end while
7: 解码阶段: q, k, v ∈ R^{1×d}
8: for i = 1, ..., T − 1 do
9: K = [K, k], V = [V, v], N = [N, 1], s = s + 1 # 更新缓存质心
10: O = softmax(qK^T /√d + log N) · V # 公式3
11: if s ≥ B + g then
12: K, V, N, s = CentroidKV(K, V, N, s, n1, n2, r, c) # 每g个解码步进行一次聚类
13: end if
14: end for
分块软匹配(Chunked Soft Matching, CSM):基于序列维度上键状态的高局部相似性观察,CentroidKV首先将键状态划分为块。由于注意力sink和最近token对模型性能具有重要性【Guangxuan Xiao 等人,Efficient streaming language models with attention sinks,2023,arXiv】,在划分之前会保留它们。该算法受到视觉Transformer(ViT)中用于token合并的二分软匹配(BSM)算法【Daniel Bolya 等人,Token merging: Your vit but faster,2022,arXiv】的启发。BSM首先将输入token分为两个不同的集合 $A$ 和 $B$;接着,对于集合 $A$ 中的每个token,向集合 $B$ 中最相似的token绘制一条边;在这些边中,仅保留排名靠前的相似连接;通过这些边保持连接的token随后被合并为一个聚类,而其他的保持不变;最后拼接集合 $A$ 和 $B$ 形成输出。直接将BSM应用于KV缓存聚类面临两个挑战:长上下文场景下大序列维度导致的匹配过程低效,以及如何以最小化对模型准确性影响的方式将序列最优地划分为集合 $A$ 和 $B$。CSM通过前置的分块步骤解决了第一个挑战,直接提高了计算效率。与K-means等通常需要多次迭代更新以建立稳定聚类的方法相比,CSM通过执行单次传递的聚类过程,以最小的token交互实现了最低的复杂度(如Table 1所示,Mesh方法计算所有token的成对余弦相似度成本极高,而CSM的块大小 $c$ 相对较小)。为了解决第二个挑战,CSM以交替的方式将每个块划分为两个集合 $A$ 和 $B$。CSM的核心思想是将高度相似的状态分配到不同的集合中,确保具有高相似度的token对被放置在集合之间,而不是集合内部。在计算相似度之后,CSM聚合来自所有块的边,并生成用于KV缓存合并的候选匹配对。
KV缓存压缩策略:在确定候选聚类后,CentroidKV执行选择性合并以构建压缩的KV缓存。并非所有匹配对都被合并,而是根据相似度对所有候选匹配进行排序,并且仅合并由压缩比率 $r$ 控制的顶部部分。具体而言,$r < 1.0$ 时仅保留前 $r$ 比例最相似的匹配,丢弃置信度较低的匹配,从而避免了来自弱相似token的噪声聚合。为了进一步提高跨轮次的聚类质量,采用了一种逐渐减小的 $r$ 调度策略,其初始化为 $r_{\mathrm{init}}$,并以衰减率 $\delta_r$ 进行线性退火:
其中 $j$ 是聚类轮次索引。这种调度使得合并随着时间的推移变得越来越具有选择性,在后续轮次中保留了高置信度的结构。对于每个合并的聚类,维护一个度计数器 $n_t$ 以跟踪其跨轮次的累积贡献。键和值状态都使用基于度的加权平均进行聚合。具体来说,对于具有对应度 $n_1, n_2, ..., n_t$ 的聚类 $k_1, k_2, \ldots, k_t$,质心计算为:
数据集:
硬件配置:提供超过100 TFLOPS计算能力的高性能GPU(具体如在vLLM实验中使用单张A100 GPU)。
准确性评估(RULER与LongBench):
- 实验内容:在不同KV缓存预算(25%-75%)下,对比CentroidKV与StreamingLLM、SnapKV、PyramidKV在RULER和LongBench上的准确性。
- 实验结果与分析:
- 在RULER基准测试中(Fig 4, Table 2, Table 3),CentroidKV在两个模型上均取得了最强的整体性能。随着缓存预算的减少,性能差距扩大,突显了CentroidKV在激进压缩下的鲁棒性。在S-NIAH-1/2、MK-NIAH-1等任务上表现最佳。但在涉及UUID类型键值的S-NIAH-3和MK-NIAH-3任务上出现性能下降,因为CentroidKV基于语义相似度聚类,难以区分缺乏有意义语义结构的随机字符串。
- 在LongBench基准测试中(Fig 5, Table 4, Table 5, Table 10, Table 11),CentroidKV在单文档问答、多文档问答和摘要任务上表现持续优异。但在合成任务和某些细粒度检索式任务(特别是Mistral模型上的25%预算时)出现下降,因为激进的token合并会模糊检索查询所需的精确token匹配和位置信号。
- 结论:CentroidKV在需要语义聚合和推理的任务中表现最强,但在依赖精确token级信息的检索密集型场景中存在局限性。
效率评估(延迟与内存):
消融实验(分块策略与压缩比率调度):
vLLM服务框架集成评估(附录D):
与ClusterKV的对比(附录E):
- 实验内容:在统一2048预算下对比CentroidKV和ClusterKV(Table 13, Table 14)。
- 实验结果与分析:ClusterKV将聚类作为一种检索机制,保留完整的KV缓存,导致TTFT和TPOT增加,且未减少GPU内存使用。而CentroidKV直接用质心替换原始状态,显著提高了长上下文解码效率并实现了实际的内存节省(在64k时减少高达17.0%)。
统计显著性分析(附录F):
- 实验内容:使用非参数自举法(bootstrap hypothesis testing)估计聚合分数的采样分布。
- 实验结果与分析:在RULER上(Table 15, Table 16),在最激进的压缩(25%预算)下,CentroidKV的95%置信区间与第二好的方法完全不重叠,证实了其优势具有统计学意义。在LongBench上(Table 17, Table 18),CentroidKV在适度预算下表现相当,但在激进压缩下因检索密集型任务出现下降。
本文提出了CentroidKV,这是一种简单而有效的在线KV缓存聚类框架,旨在提高长上下文LLM推理的效率。CentroidKV在大多数任务上不降低准确性的前提下,将KV缓存内存使用量减少了高达 $75\%$,同时将解码阶段加速了高达 $1.92\times$,并将吞吐量提高了高达 $4\times$。
未来工作展望:本文侧重于在GPU上压缩KV缓存,未研究内存卸载策略。未来一个有希望的研究方向是在CPU上执行聚类,并将生成的质心传输到GPU。此外,CentroidKV目前采用手动指定的压缩比率调度,并在每个聚类步骤将缓存减少限制在最多一半,未来的工作可以通过设计自适应压缩策略来解决这一问题。
附录A:理论结果与证明
我们在理论上证明了以交替方式将每个块划分为两个集合 $A$ 和 $B$ 的划分策略的最优性。直观地说,划分策略应保留具有最高相似度的边以产生更好的聚类。基于第4节中的Observation 2,我们考虑一个凸且单调递减的分数函数 $f : \mathbb{N} \to \mathbb{R}$,它将距离映射到重要性分数。利用分数函数 $f$,可以通过求解以下优化问题来实现最优划分:
以下定理表明,对于任何凸且单调递减的分数函数 $f$,我们提出的交替划分策略始终是问题方程(6)的最优解。
定理 A.1。定义划分集 $\mathcal{P}_{2n} = \{(A, B) ~ | ~ |A| = |B| = n, \text{ 且 } A \cup B = [2n]\}$。如果函数 $f : [2n - 1] \to \mathbb{R}$ 满足 $f(1) - f(2) \ge f(2) - f(3) \ge \cdots \ge f(2n - 2) - f(2n - 1) \ge 0$,则有:
这里,我们使用符号 $[k]$ 表示不大于 $k$ 的正整数集,即 $[k] := [1, k] \cap \mathbb{Z}$。
证明。当 $n = 1$ 时,结果是平凡的。在下文中,我们假设 $n \ge 2$。考虑以下映射:
对于任意 $l \in [2n - 1]$,定义 $\mathcal{S}_{2n, l} := \{(\phi_{2n}(x), \phi_{2n}(x + l)) \mid x \in \mathbb{Z} \cap [-l + 2, 2n - 1]\}$。对于 $\forall (A, B) \in \mathcal{P}_{2n}$,定义