ClusterKV: Manipulating LLM KV Cache in Semantic Space for Recallable Compression

发表时间: 2024-12 · arXiv:2412.03213 (DAC 2025)

原文: https://arxiv.org/abs/2412.03213

作者/机构：Guangda Liu, Chengwei Li, Jieru Zhao†, Chenqi Zhang, Minyi Guo. School of Computer Science, Shanghai Jiao Tong University.

速读

一句话结论 本文提出了 ClusterKV，通过在语义空间对 KV Cache 进行聚类来实现可召回的缓存压缩，在 32k 上下文下仅需 1k 到 2k 的缓存预算即可保持几乎无损的精度，并将解码吞吐量提升了 2.5 倍。

要解决什么问题 大语言模型在处理长上下文时，KV Cache 的显存占用和访存延迟会随上下文长度线性增长，成为自回归解码阶段的性能瓶颈。现有的 KV Cache 压缩方法主要卡在两个机制缺陷上。第一条路线是永久驱逐不重要的 token，但这忽略了 token 重要性的动态变化机制：在当前解码步被判定为权重低而丢弃的 token，可能在后续步骤中变得至关重要，永久驱逐会导致模型精度和生成质量不可逆地下降。第二条路线是支持历史 token 召回，但由于对所有历史 token 计算注意力权重的代价高达 $O(Ld)$，现有方法（如 Quest）只能退而求其次，按文本绝对位置将连续 token 划分为固定大小的页（page）进行召回。这种基于物理位置的划分机制会导致严重的内部碎片问题：被召回的一个页中往往只有一两个真正高权重的 token，其余全是无关 token，白白浪费了宝贵的缓存预算，使得真正重要的 token 无法被装入有限的显存中。

怎么做的 为了绕开物理位置划分带来的内部碎片卡点，ClusterKV 的核心思路是将召回粒度从“物理相邻的页”切换为“语义相近的簇（cluster）”。作者观察到，在语义空间（即 Key 向量空间）中距离相近的 token，对于同一个 Query 往往具有相似的注意力权重。因此，方法的第一步是基于 Key 向量对 token 进行 K-means 聚类。由于 Key 向量中存在数值极大的离群通道，使用欧氏距离或内积会导致聚类失效，因此作者定义了基于余弦相似度的语义距离公式：

$$\mathcal{D}(i, j) = 1 - \frac{\langle k_i, k_j \rangle}{|k_i| \cdot |k_j|}$$

在预填充阶段，除了保留前 16 个作为注意力沉淀（attention sinks）的初始 token 外，其余 prompt token 会被聚类成 $L/80$ 个簇（$L$ 为上下文长度），每个簇用其内部 token 的均值向量 $\mu$ 作为质心。在解码阶段，新生成的 token 每隔 320 步在内部进行一次增量聚类。方法的第二步是簇级别的选择与召回。对于当前解码步的查询向量 $q$，系统不再逐一计算所有历史 token 的权重，而是仅计算 $q$ 与各个质心 $\mu_i$ 的内积 $q \mu_i^T$。系统将质心按该内积值降序排列，依次选中权重最高的簇，并提取这些簇内包含的所有不连续 token 的 KV 值，直到选出的 token 总数达到预设的缓存预算 $B$ 为止。为了抹平聚类和动态索引带来的额外开销，ClusterKV 在系统层设计了三个关键部件：首先是异步聚类机制，将 GPU 上的聚类计算与当前层的注意力及前馈网络计算完全重叠；其次是定制的 CUDA 算子，通过在序列维度跨步分配线程并对通道维度进行分块，缓解了多头批量聚类时共享显存的原子加法（atomicAdd）写入冲突问题；最后是簇粒度的 GPU 缓存，保留上一步选中的 KV 数据，当前步只需从 CPU 内存中拉取未命中的簇，大幅减少了异构设备间的数据搬运。

效果如何 实验在单张 NVIDIA Ada 6000 GPU 上搭建，使用支持 128k 窗口的 GLM4-9B-Chat 评估模型精度，使用 Llama-3.1-8B 和 OPT-6.7B 评估推理效率。测试数据覆盖 LongBench 的 8 个长文本任务（最高 32k 上下文）以及 PG19 语言建模任务。对比基线包括代表“按页召回”路线的 Quest，以及代表“SVD 降维部分权重召回”路线的 InfiniGen。在量化结果方面，当缓存预算设定为极低的 1024 个 token 时，ClusterKV 在 LongBench 上的平均得分为 48.34，显著高于 Quest（43.23）和 InfiniGen（45.13），且极其逼近全量 KV Cache 的无损得分（49.01）。在 PG19 的困惑度测试中，ClusterKV 与全量 KV 的偏差仅为 0.5，而 Quest 和 InfiniGen 的偏差分别高达 4 和 2。在推理效率方面，在 32k 提示词和 1024 步解码的设置下，ClusterKV 相比全量 KV 实现了 2 倍的延迟加速和 2.5 倍的吞吐量提升。在与强基线的直接对比中，ClusterKV 的速度是 InfiniGen 的 2.3 倍（因为后者逐 token 选择的计算成本过高），并且在保持与 Quest 几乎相同推理延迟（偏差在 5% 以内）的前提下，提供了高得多的重要 token 召回率和模型精度。作者也指出了方法的一个局限与权衡：当初始聚类数量超过 400 时，召回率的提升会遭遇边际效应递减，因此将初始簇数量与上下文长度的比例固定为 1/80 是平衡精度与聚类开销的必要妥协。

主要贡献

大型语言模型（LLMs）在处理长文档问答和复杂逻辑推理等任务时，对长上下文的需求日益增加。然而，长上下文给推理效率带来了重大挑战，主要包括键值（KV）缓存的高昂内存成本以及因大量内存访问导致的延迟增加。为了缓解这一问题，近期的研究提出了压缩KV缓存以近似计算的方法。但现有方法存在明显缺陷：要么永久性地驱逐token（在后续推理中永远无法召回），要么以基于文本位置划分的“页面（pages）”为粒度来召回先前的token，这两种方式都会导致模型精度和输出质量的下降。

为了实现高效且准确的可召回（recallable）KV缓存压缩，本文提出了ClusterKV。该方法在“语义聚类（semantic clusters）”的粒度上召回token。作者针对聚类、选择、索引和缓存设计并实现了高效的算法和系统。实验结果表明，在32k上下文长度下，ClusterKV仅使用1k到2k的KV缓存预算，就在各项任务中实现了可忽略的精度损失，并获得了高达$2\times$的延迟加速和$2.5\times$的解码吞吐量提升。与目前最先进的可召回KV压缩方法相比，ClusterKV展示了更高的模型精度和输出质量，同时保持或超越了原有的推理效率。

KV压缩方法比较。绿框代表被选中用于注意力计算的token。
图1d中最后一步token的语义空间和注意力权重。较浅的框表示较大的权重。

背景知识与设计原则

LLM推理与KV缓存
LLM包含多个Transformer层，每层均包含一个多头注意力（MHA）模块和一个带有残差连接及归一化操作的前馈网络（FFN）在MHA中，输入张量被线性投影为每个头的查询、键和值张量（$Q, K, V \in \mathbb{R}^{N \times d}$），其中$N$是输入长度，$d$表示每个头的通道数或隐藏维度。MHA输出定义为$softmax(\frac{QK^T}{\sqrt{d}})V$，所有头的输出拼接后用于后续的FFN和归一化。对于生成式推理，LLM以自回归方式生成token，将每个生成的token附加到输入中以生成下一个token。为避免重新计算先前token的$K$和$V$，这些张量被存储在内存中以供重用，这被称为KV缓存。LLM推理包括预填充（prefill）和解码（decoding）两个阶段。预填充阶段处理整个输入序列，计算KV缓存并生成第一个输出token。在解码期间，最新生成token的查询向量$q$和先前token的$K, V$用于计算注意力以生成下一个token，公式化为$softmax(\frac{qK^T}{\sqrt{d}})V$，其中$q \in \mathbb{R}^{1 \times d}$，$K, V \in \mathbb{R}^{L \times d}$，$L$是先前token的上下文长度。

长上下文推理与KV缓存压缩
长上下文是LLM推理的新兴趋势，且LLM支持的上下文窗口正在迅速扩展（甚至高达1M tokens）。然而，长上下文推理会产生显著的内存和计算成本，解码期间的KV缓存大小和注意力计算复杂度随上下文长度线性增加，导致推理效率低下甚至失败。近期的研究揭示了注意力计算的稀疏性，即只有一小部分token对大部分注意力输出有贡献（【10，H2O: Heavy-hitter oracle for efficient generative inference of large language models+2023+NeurIPS】【13，Alisa: Accelerating large language model inference via sparsity-aware kv caching+2024+arXiv】）。这一观察使得通过选择token子集来近似注意力计算成为可能，公式化为$softmax(\frac{qK_S^T}{\sqrt{d}})V_S$，其中$K_S, V_S \in \mathbb{R}^{B \times d}$代表选中token的键和值，$B$是压缩后的KV缓存预算大小。通过设置固定预算，无论上下文长度如何，KV缓存大小和解码成本都能保持稳定。通常，$K_S$和$V_S$是基于注意力权重（$qK^T$）来选择的。

现有方法缺陷与动机
不可召回压缩的缺陷：KV缓存压缩应该是可召回的。虽然选择KV缓存子集进行计算能确保效率，但为选择而计算所有注意力权重会引入巨大开销。因此，现有工作通常仅对已被选中的token计算注意力权重，未在某一解码步被选中的token会被永久驱逐，永远不会在后续推理中被召回（【10，H2O: Heavy-hitter oracle for efficient generative inference of large language models+2023+NeurIPS】【11，Snapkv: Llm knows what you are looking for before generation+2024+arXiv】【12，Keyformer: Kv cache reduction through key tokens selection for efficient generative inference+2024+MLSys】）。然而，作者观察到token重要性在推理过程中是动态变化的。如图3a所示，在Llama3-8B的64个解码步骤中，最初不重要的token可能在后续步骤变得重要，反之亦然。不可召回的压缩无法捕捉这种动态特征，导致模型精度和输出质量下降。
现有可召回压缩方法的缺陷：通过与所有历史token计算注意力权重来实现可召回压缩会产生$O(Ld)$的不可接受成本。InfiniGen（【18，InfiniGen: Efficient generative inference of large language models with dynamic KV cache management+2024+OSDI+https://www.usenix.org/conference/osdi24/presentation/lee】）通过离线生成部分权重降低维度，但仍需存储部分键，且选择成本仍随$L$线性扩展。Quest（【15 ，Quest: Query-aware sparsity for efficient long-context llm inference+2024+ICML】）以页面粒度选择token，使用页面内所有token的每通道最大键来计算权重，将成本降至$O(Ld/page\_size)$。但是，简单按文本位置划分页面会导致重要token的内部碎片化。如图3b的热力图所示，选中的页面可能仅包含一两个重要token，却将不重要的token也纳入计算，浪费了KV缓存预算。

(a) 上下文长度为8192时，token重要性在解码步骤中的变化。(b) 页面粒度（page_size = 16）下重要token的内部碎片化。
(a) 上下文长度为8192时，token重要性在解码步骤中的变化。(b) 页面粒度（page_size = 16）下重要token的内部碎片化。

核心方法细节

问题建模与设计原则
对于上述近似注意力计算，令$K_S$为$(k_{i_1}, k_{i_2}, ..., k_{i_B})^T$，其中$I_T = \{i_1, i_2, ..., i_B\}$表示选中token的索引。目标是选择对注意力权重贡献最大的token，以尽可能接近原始计算。具体而言，$I_T$应当是$\arg\max \sum_{i \in I_T} q k_i^T$，即选择具有前$B$个最大注意力权重的token。基于此，作者观察到在语义空间中接近的token对于给定的$q$往往具有相似的注意力权重。因此，设计了在语义聚类粒度上进行KV选择的方法：首先在语义空间对token应用聚类；然后仅计算聚类表示（而非单个token）的注意力权重，并选择权重最大的聚类。由于聚类数量通常比token数量小一个数量级，基于聚类的选择显著降低了召回开销。

语义空间中的聚类
语义距离：由于对于给定的$q$，注意力权重仅与键张量相关，因此通过计算对应键向量之间的距离来衡量token间的语义距离。作者发现余弦相似度比L2或内积距离更合适，因为键向量中存在具有大幅值的异常通道（【19，Kivi: A tuning-free asymmetric 2bit quantization for kv cache+2024+ICML】），这会导致L2或内积距离发生剧烈变化。因此，token $i$和$j$在语义空间中的距离定义为$\mathcal{D}(i, j) = 1 - \frac{\langle k_i, k_j \rangle}{|k_i| \cdot |k_j|}$，余弦相似度越大的向量距离越小。
聚类过程：应用简单的K-means算法对键向量进行聚类（【20，K-means clustering algorithms: A comprehensive review, variants analysis, and advances in the era of big data+2023+Information Sciences+https://www.sciencedirect.com/science/article/pii/S0020025522014633】），如图4所示。首先随机采样键向量作为初始质心。随后，交替执行分配和更新步骤直到收敛。在分配步骤中，每个键向量基于距离$\mathcal{D}$被分配给最近的质心（即具有最大余弦相似度的质心），并获得相应的聚类标签。在更新步骤中，将分配给同一质心的键的平均值作为新质心。当分配不再改变时算法收敛，分配给同一质心的键形成一个语义聚类，质心作为聚类表示 。
推理阶段的聚类策略：在LLM推理期间，预填充阶段结束后，首先对提示（prompt）token的键向量应用聚类。例外情况是初始token（被称为attention sinks【9，Efficient streaming language models with attention sinks+2024+ICLR】），它们通常在聚类过程中表现为离群点。因此，始终保留前16个token，并对后续token应用聚类，生成$C_0$个质心。实验表明对于32k上下文，使用400个聚类能平衡效率和精度，因此设置$C_0 = \frac{L}{80}$。对于解码阶段生成的token，每$m$个解码步对生成的$m$个token的键向量应用聚类，创建$C_+$个新质心。由于将生成的键与预填充阶段的键一起聚类会产生巨大开销，因此仅在生成的token内应用聚类。为了分摊成本，将$C_+$和$m$分别设置为4和320。

聚类和选择过程。绿点代表键向量，紫点代表聚类质心。
聚类和选择过程。绿点代表键向量，紫点代表聚类质心。

语义聚类粒度的选择
选择机制：将聚类质心表示为$\mu_1, \mu_2, ..., \mu_C \in \mathbb{R}^d$。为了给给定查询$q$选择重要token，基于它们的注意力权重（即$q\mu_i^T$）按降序对这些质心进行排序。需要注意的是，虽然键使用余弦相似度距离进行聚类，但查询向量与质心之间的距离使用内积来衡量，因为它更好地对齐了注意力权重的计算。直观上，分配给具有较大注意力权重质心的键，对于给定的$q$往往也具有较大的注意力权重。因此，检索排序后的质心，并收集相应聚类中token的KV，直到选出前$B$个最重要的token，如图4所示。

效率考量
基于聚类的选择避免了重要token的内部碎片化，相比基于页面的选择能实现更高的精度。但它也引发了效率担忧。
担忧1：对于聚类，计算成本为$O(n_i C L d)$（其中$n_i$是收敛前的迭代次数，$C$是聚类数），这高于Quest（【15，Quest: Query-aware sparsity for efficient long-context llm inference+2024+ICML】）中获取页面表示（如每通道最大键向量）的成本$O(Ld)$。
担忧2：对于基于页面的方法中的选择，由于页面大小固定且页面内的token是连续的，所需页面数可以轻松计算为$B/page\_size$，选中token的索引可直接从选中页面的索引推导出来。然而，对于语义聚类，大小可能变化，且聚类内的token位置是动态且不连续的。因此，所需聚类数量和选中token的索引无法像页面那样轻易确定。
ClusterKV通过系统设计解决了这两个担忧。

系统概览
系统概览如图5所示。在预填充阶段，键张量在GPU上通过语义聚类（SC）进行处理，生成质心和相应的元数据。生成的KV张量被卸载到CPU内存中。在解码阶段，计算查询向量$q$和聚类质心的注意力权重以确定每个聚类的重要性。该结果连同聚类元数据用于生成选中token的索引（$I_T$），然后利用这些索引将选中的KV（$K_S, V_S$）从CPU内存加载到GPU内存以进行注意力计算。GPU上维护了一个选中KV的缓存，因此只需加载尚未缓存的KV（$K'_S, V'_S$）。每隔$m$个解码步，对生成的$m$个token执行聚类和KV卸载。

ClusterKV的系统概览。绿框代表在GPU上运行的组件。
ClusterKV的系统概览。绿框代表在GPU上运行的组件。

语义聚类优化
ClusterKV在系统和内核级别优化了聚类的效率。
系统级优化：如图6所示，ClusterKV异步应用聚类，在QKV投影和RoPE模块计算出键之后立即启动。这允许聚类与当前层的注意力及FFN计算，以及下一层的QKV投影和RoPE重叠执行。通过重叠这些过程，ClusterKV最小化了与聚类相关的开销。
内核级优化：由于聚类需对每个头单独应用，一个关键优化是实现跨头的批量（batched）聚类。对于分配步骤（主要涉及argmin操作以及键与质心之间的矩阵乘法），有高效的批量Torch内核可用（【21，Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation+2024+ASPLOS+https://doi.org/10.1145/3620665.3640366】）。因此，重点优化质心更新步骤，并实现了一个自定义CUDA内核，其中不同的头由独立的ThreadBlocks并行处理，如图7所示。该内核检索键，累加分配给同一聚类的键，在共享内存中记录相应的累加计数，并计算平均值作为新质心 。
缓解写入冲突：共享内存中存在潜在的写入冲突问题，因为同一聚类的累加键和计数需要使用atomicAdd进行计算。为了缓解这一问题，应用了多项优化。如图7所示，由于距离较远的token倾向于被分配到不同的聚类，沿序列维度以跨步（strided）方式排列线程，以最小化并发处理同一聚类内键的发生率。此外，将通道维度划分为$P$个分区，每个分区由一个线程处理。这使得$\frac{BlockSize}{P}$个键被并行处理。在选择$P$时存在权衡：使用较大的$P$会减少并行处理的键数并增加沿序列维度的迭代次数；而使用较小的$P$会增加线程内沿通道维度的迭代次数，并可能增加写入冲突的发生率。为了确定最佳的$P$值，对不同的$P$值进行了离线分析。发现在$BlockSize = 512$时，对128个通道使用$P = 16$或$P = 32$能获得最佳性能。

聚类与其他操作的重叠执行。
质心更新内核的执行。在此示例中，一个ThreadBlock包含8个线程（BlockSize = 8），通道维度分为两个分区。

选择与索引机制
细节如图8所示。聚类后，ClusterKV存储聚类质心和相应的元数据，包括聚类大小、前缀和以及排序后的索引。然后获取聚类大小。聚类标签和键索引按标签排序，并存储排序后的索引。
在解码期间，ClusterKV计算查询向量$q$和聚类质心$\mu$之间的注意力权重，按注意力权重降序对聚类进行排序，以确定最接近$q$的聚类。接下来，ClusterKV收集并重新排序它们对应的聚类大小，并计算前缀和。如果将KV缓存预算设置为3，匹配第二个前缀和，则将选择前2个最接近的聚类。然后收集选中聚类的标签、大小和结束位置，以确定选中token的索引$I_T$。注意，当选中聚类大小的总和超过预算时，ClusterKV会裁剪最后一个选中聚类的token以遵守预算限制。
与聚类类似，ClusterKV实现了高效的CUDA内核，使用独立的ThreadBlocks并行处理KV头的索引。频繁访问的与聚类大小相关的元数据被存储在共享内存中以提升性能。

选择和索引的过程。绿框代表聚类存储的元数据。此处KV预算为3。
选择和索引的过程。绿框代表聚类存储的元数据。此处KV预算为3。

缓存选中Token的KV
ClusterKV在GPU上维护了一个聚类粒度的缓存来存储选中token的KV，减少了从CPU内存到GPU内存的不必要数据传输并提升了性能。在解码阶段，缓存保留过去$R$个解码步骤中选中token的KV，以及对应的选中聚类的标签。在当前解码步，将选中聚类的标签与缓存保留的聚类标签进行比较。只有不在缓存中的聚类的KV才从CPU内存加载。缓存在内存使用和缓存有效性之间引入了权衡。在实践中，发现设置$R = 1$（即仅保留上一个解码步的KV状态）能取得良好的平衡。

实验环境

  • 硬件配置：NVIDIA Ada 6000 GPU。

  • 数据集：

    • LongBench中的8个数据集（2WikiMQA, TriviaQA, HotpotQA, MultiFieldQA, MuSiQue, NarrativeQA, Qasper, GovReport），涵盖单文档QA、多文档QA、少样本学习和摘要等任务，上下文长度最高达32k。
    • PG19数据集（用于语言建模任务）。

  • 模型架构：

    • GLM4-9B-Chat（支持高达128k的上下文窗口，用于模型精度评估）。
    • Llama-3.1-8B（用于推理性能评估）。
    • OPT-6.7B（专门用于与InfiniGen进行延迟对比）。

  • 基线方法：Quest 和 InfiniGen。为了与Quest对齐，ClusterKV和InfiniGen在前两层禁用了选择功能，使用全量KV缓存。

实验结果

1. LongBench模型精度评估
* 实验内容：在256、512、1024和2048的KV缓存预算下，比较不同方法在8个LongBench数据集上的精度（GovReport使用ROUGE-L，其余使用F1分数）。
* 实验结果：ClusterKV在大多数设置下优于Quest和InfiniGen，并且在仅1k到2k tokens的预算下，就实现了与使用全量KV缓存相当的精度。表1的平均分数显示ClusterKV相较于基线有显著提升。
* 引用图表：图9（不同方法在LongBench上的结果）。

不同方法在LongBench上的结果。
不同方法在LongBench上的结果。

2. 语言建模任务（困惑度）
* 实验内容：在PG19测试集上评估语言建模的困惑度，输入长度从1到32000 tokens，统一设置KV预算为1024。
* 实验结果：ClusterKV的困惑度与全量KV高度一致（偏差最大仅为0.5），而Quest和InfiniGen的偏差分别达到了约4和2。
* 引用图表：图10（KV缓存预算为1024 tokens时的语言建模困惑度）。

KV缓存预算为1024 tokens时的语言建模困惑度。
KV缓存预算为1024 tokens时的语言建模困惑度。

3. 重要Token的召回率及消融实验
* 实验内容：在NarrativeQA（32k上下文）上提取样本，计算不同方法在推理过程中重要token的召回率（预算从256到2048）。并测试ClusterKV不同距离度量和聚类数$C_0$的影响。
* 实验结果：ClusterKV在所有预算下实现了最高的召回率。消融实验证实，余弦相似度优于L2和内积距离；增加$C_0$能提升召回率，但当$C_0 > 400$时提升不再显著，证实$C_0 = 400$是精度和效率的平衡选择。
* 引用图表：图11（重要token的召回率：(a) 不同方法之间的比较，(b) ClusterKV不同配置之间的比较）。

重要token的召回率：(a) 不同方法之间的比较，(b) ClusterKV不同配置之间的比较。
重要token的召回率：(a) 不同方法之间的比较，(b) ClusterKV不同配置之间的比较。

4. 推理效率比较（vs 全量KV缓存）
* 实验内容：在提示长度（$P$）8k到32k，解码长度（$D$）256到1024下，评估ClusterKV的推理延迟。
* 实验结果：对于$P = 32k$和$D = 1024$，在1024预算下，ClusterKV实现了$2\times$的延迟加速，解码吞吐量提升高达$2.5\times$。此外，ClusterKV的聚类开销极小，仅占预填充时间的6%到8%，占总推理时间不到2%。
* 引用图表：图12（ClusterKV与全量KV缓存配置的推理延迟比较）。

ClusterKV与全量KV缓存配置的推理延迟比较。
ClusterKV与全量KV缓存配置的推理延迟比较。

5. 推理效率比较（vs SoTA压缩方法）
* 实验内容：分别与InfiniGen（基于OPT-6.7B，预算256）和Quest（基于Llama-3.1-8B，预算1k）进行延迟比较。
* 实验结果：相比InfiniGen，ClusterKV平均加速$2.3\times$（因InfiniGen逐token选择开销极大，其延迟与全量KV相当，而ClusterKV选择开销仅占解码延迟约5%）。相比Quest，ClusterKV的延迟表现极其接近（偏差不超过5%），但同时提供了显著更高的模型精度。
* 引用图表：图13（推理延迟比较：(a) ClusterKV与InfiniGen，(b) ClusterKV与Quest）。

推理延迟比较：(a) ClusterKV与InfiniGen（预算256 tokens），(b) ClusterKV与Quest（预算1k tokens）。
推理延迟比较：(a) ClusterKV与InfiniGen（预算256 tokens），(b) ClusterKV与Quest（预算1k tokens）。

6. 缓存有效性分析
* 实验内容：在NarrativeQA（32k）样本上分析聚类粒度缓存的命中率。
* 实验结果：当$R = 1$和$R = 2$时，平均命中率分别为63%和74%。与直接从CPU内存加载相比，缓存机制将解码吞吐量分别提高了$2.3\times$和$3\times$。

结论

本文引入了ClusterKV，通过在语义聚类粒度上召回token，实现了高效且准确的KV缓存压缩。ClusterKV在显著提升推理效率的同时，保持了模型精度。

参考文献引用总结（核心方法与背景部分）

  • 【9，Efficient streaming language models with attention sinks+2024+ICLR】：说明了初始token（attention sinks）在语义空间中常作为离群点存在，因此在聚类时需要被保留。
  • 【10，H2O: Heavy-hitter oracle for efficient generative inference of large language models+2023+NeurIPS】：指出了注意力计算的稀疏性，并指出传统方法因开销大而仅对已选token计算权重，导致永久驱逐。
  • 【11，Snapkv: Llm knows what you are looking for before generation+2024+arXiv】：同【10】，说明现有工作不可召回的缺陷。
  • 【12，Keyformer: Kv cache reduction through key tokens selection for efficient generative inference+2024+MLSys】：同【10】，说明现有工作不可召回的缺陷。
  • 【13，Alisa: Accelerating large language model inference via sparsity-aware kv caching+2024+arXiv】：指出了注意力计算的稀疏性，即小部分token贡献了大部分输出。
  • 【15，Quest: Query-aware sparsity for efficient long-context llm inference+2024+ICML】：介绍了基于页面粒度选择token的方法，其降低了选择开销，但本文指出其存在内部碎片化问题。
  • 【18，InfiniGen: Efficient generative inference of large language models with dynamic KV cache management+2024+OSDI+https://www.usenix.org/conference/osdi24/presentation/lee】：介绍了通过降维实现可召回压缩的方法，但指出其仍需存储部分键且成本随上下文线性扩展 。
  • 【19，Kivi: A tuning-free asymmetric 2bit quantization for kv cache+2024+ICML】：解释了为何选择余弦相似度而非L2或内积，因为键向量中存在大幅值的异常通道。
  • 【20，K-means clustering algorithms: A comprehensive review, variants analysis, and advances in the era of big data+2023+Information Sciences+https://www.sciencedirect.com/science/article/pii/S0020025522014633】：本文采用的K-means聚类算法的基础文献 。
  • 【21，Pytorch 2: Faster machine learning through dynamic python bytecode transformation and graph compilation+2024+ASPLOS+https://doi.org/10.1145/3620665.3640366】：在聚类的分配步骤中，使用了该文献提供的现有高效批量Torch内核 。