💬 NLP

CentroidKV: Efficient Long-Context LLM Inference via KV Cache Clustering

CentroidKV 是一个简单且有效的框架，它通过一种使用分块软匹配和质心合并的在线 KV 缓存聚类方法，将长文本大语言模型（LLM）推理的内存使用量降低高达 75%，并将解码速度提升高达 1.92 倍。

原作者： Jie Hu, Shengnan Wang, Yutong He, Ping Gong, Jiawei Yi, Juncheng Zhang, Youhui Bai, Renhai Chen, Gong Zhang, Cheng Li, Kun Yuan

发布于 2026-06-16

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Jie Hu, Shengnan Wang, Yutong He, Ping Gong, Jiawei Yi, Juncheng Zhang, Youhui Bai, Renhai Chen, Gong Zhang, Cheng Li, Kun Yuan

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图阅读一本长达 10 万页的巨著，只为了回答关于第一句话的一个问题。随着阅读的深入，你的大脑会自然地试图记住每一个角色、每一个场景和每一个情节。在人工智能的世界里，这种“记忆”被称为 KV Cache（键值缓存）。

问题在于？随着故事变得越来越长，这种记忆变得如此庞大，以至于会撑爆计算机的大脑（GPU），让运行速度变得极其缓慢。这就像是背着一整个图书馆的藏书去跑马拉松。

现有的解决方案试图解决这个问题，但它们要么：

丢弃页面： 它们删除它们认为不重要的故事部分。但有时，50 页前看似“无聊”的一页可能正是解开结局的关键，这会导致 AI 产生困惑。
缩小字号： 它们压缩文本，但这往往会让阅读变得困难，并降低阅读速度。

CentroidKV 是一种更聪明、更高效的处理记忆的方法。以下是它的工作原理，我们使用简单的类比来解释：

1. “抱团取暖”策略 (聚类)

与其删除页面或缩小文本，CentroidKV 会寻找重复项。

想象一下，你正在组织一场拥有 10,000 名宾客的大型派对。许多宾客都穿着一模一样的红衬衫，留着同样的理发样式。与其记住每一个人的细节，CentroidKV 会说：“嘿，这 50 个人基本上是一样的。让我们把他们归为一组，并创建一个‘超级宾客’（质心/centroid）来代表他们所有人。”

工作原理： AI 扫描故事，并注意到某些单词或短语以非常相似的方式出现。它将这些相似的“标记”（tokens/单词）组合在一起，并用一个单一的、平均化的版本来替换整个组。
结果： 你从记住 10,000 名个体宾客，变成了只需记住几百个“超级宾客”。这在不丢失主线剧情的情况下，将内存占用量减少了高达 75%。

2. “分块”处理法 (分块软匹配)

你可能会问：“如果我有 100,000 页内容，你如何能在不耗费太长时间阅读的情况下找到重复项？”

如果你尝试将每一页与其它所有页面进行对比，那将耗时极长。CentroidKV 使用了一个名为 分块软匹配 (Chunked Soft Matching) 的巧妙技巧。

类比： 想象你正在整理一大堆脏衣服。你不需要把每一只袜子都与全屋所有的袜子进行对比，而是将脏衣服分成一个个小篮子（块/chunks）。
策略： 在每个篮子内部，AI 会寻找匹配的袜子。它使用一种特殊的“交替”方法来快速配对。这就像是在说：“在这个篮子里，让我们把红袜子和蓝袜子配对，但前提是它们必须非常相似。”
为什么快： 通过将问题分解成小的、易于管理的块，AI 即使面对极长的故事，也能瞬间完成这种分组。

3. “质量控制”过滤器

论文指出，你不能仅仅因为两个东西看起来像就将它们合并，否则你会丢失重要的细节。

类比： 想象你在合并一群人。你不会仅仅因为厨师和飞行员都戴着帽子就把他们合并在一起。你只会合并那些真正相似的人。
过程： CentroidKV 非常挑剔。它只合并那些非常、非常相似的组（高置信度）。如果两个事物只是“有点”相似，它就会保持原样。它在处理过程中也会变得越来越严格，以确保最终生成的“超级宾客”是原始组的准确代表。

结果：更快、更轻量

由于 AI 现在需要携带的“背包”（压缩后的记忆）变小了：

阅读速度更快： 其“解码”（生成下一个词）的速度提升了高达 1.92 倍。
能承载更多用户： 系统可以同时为多达 4 倍的用户 提供服务，因为它不再容易耗尽内存。
不会遗忘： 尽管缩小了记忆，AI 回答问题的表现几乎与拥有完整、未压缩记忆时一样出色。

它做不到的事（局限性）

论文诚实地说明了这种方法无法做到的事情：

并非对所有情况都有效： 如果故事依赖于非常具体且随机的代码（例如仅出现一次的唯一 ID 编号），AI 可能会难以保留那个精确的细节，因为它倾向于将相似的事物进行分组。它擅长处理故事和含义，但在寻找精确、随机的字符串方面并不完美。
仍留在 GPU 上： 目前，这种分组操作发生在计算机的主处理器上。作者建议，未来我们可以将这种分组操作放在更慢、更便宜的处理器（CPU）上进行，然后只将结果发送给主处理器，但目前尚未实现这一点。

总结： CentroidKV 就像一位聪明的图书管理员，他意识到大型图书馆中的许多书籍其实只是同一故事的重印本。与其保留 1,000 本副本，他只保留一份“母本”，并附上一张便条说：“这本代表了 1,000 本书。” 这节省了空间，加快了搜索速度，并保持了故事的完整性。

技术摘要：CentroidKV

问题陈述

具有超长上下文窗口（高达 1M token）的大语言模型（LLMs）在部署时面临着显著的键值（KV）缓存挑战。KV 缓存用于存储先前 token 的键（Key）和值（Value）以避免重复计算，其规模随上下文长度线性增长。这造成了推理延迟和吞后吞吐量的瓶颈，因为生成每个 token 都需要访问整个缓存。此外，庞大的内存占用也增加了部署难度。

现有的减少 KV 缓存大小的方案包括：

驱逐（Eviction）： 根据注意力分数丢弃不太重要的 token。这通常会导致性能下降，因为最初被认为不重要的 token 在未来的生成过程中可能会变得至关重要。
合并（Merging）： 将待驱逐的 token 与保留的 token 进行合并。然而，现有方法通常依赖于初步的驱逐来定义合并集，导致语义相距甚远的 token 被强行分组，从而产生质心偏差（centroid bias）并造成信息丢失。
量化与剪枝（Quantization and Pruning）： 这些方法提供的效率提升有限，或者会引入高昂的计算开销。

因此，需要一种能够基于内在相似性进行压缩、且不会丢弃关键信息或产生过高计算成本的方法。

方法论：CentroidKV

本文介绍了一种名为 CentroidKV 的在线 KV 缓存聚类框架。其核心创新是 分块软匹配（Chunked Soft Matching, CSM） 算法，旨在基于相似性进行聚类，同时保持较低的计算复杂度。

核心组件

观察结果：
- 局部相似性： Key 状态在序列维度上表现出高度的余弦相似性，相似的 token 会聚集在局部区域。
- 距离相关性： 随着 token 距离的增加，key 状态的余弦相似度通常呈单调递减趋势，并遵循凸函数趋势。
推理流水线：
- 预填充（Pre-filling）： 处理序列，如果 KV 缓存超过预算，则执行第一轮聚类。
- 解码（Decoding）： 缓存随每个 token 增长。每隔 $g$ 步调用一次聚类，以将缓存维持在预算范围内。
- 保留（Preservation）： 在进行分区之前，先保留注意力汇点（attention sinks）和近期 token，以确保模型的稳定性。
分块软匹配 (CSM)：
- 分块（Chunking）： 将序列划分为若干块，以高效处理长上下文。
- 交替分区（Alternating Partition）： 在每个块内，通过交替方式将 token 分为两个集合（ $A$ 和 $B$ ）。理论分析证明，鉴于观察到的单调且凸的相似性趋势，该策略是最优的，它确保了高度相似的 token 被分配到不同的集合中，从而便于匹配。
- 匹配（Matching）： 类似于 Vision Transformer 中使用的二分软匹配（Bipartite Soft Matching, BSM），CSM 识别集合 $A$ 和 $B$ 之间高度相似的 token 对。
- 压缩（Compression）： 该算法并不合并所有对，而是根据相似度对候选匹配进行排序，并仅合并由压缩率 $r$ 控制的前一部分。
- 渐进式调度（Progressive Scheduling）： 压缩率 $r$ 在多个聚类轮次中进行线性退火（ $r = r_{init} - j \cdot \delta_r$ ）。这使得合并在后期变得越来越具有选择性，从而保留高置信度的结构。
聚合（Aggregation）：
- 对于每个聚类簇，使用度加权平均（degree-weighted averaging）将 key 状态和 value 状态合并为一个单一的质心（centroid）。“度”（ $n_t$ ）用于追踪跨轮次合并到某个质心的 token 的累积贡献。
- 注意力机制经过适配，通过在 softmax 计算中引入聚类度作为对数项，以近似原始的注意力输出。

主要贡献

CentroidKV 框架： 一种简单、即插即用的在线 KV 缓存聚类解决方案，可提高长上下文场景下的推理效率。
分块软匹配算法： 一种新型聚类算法，将二分软匹配扩展到 KV 缓存压缩。它利用了块内的交替分区策略，该策略在理论上被证明对于观察到的相似性模式是最优的。
理论分析： 本文提供了块内分区策略最优性的形式化证明，并分析了计算复杂度，证明了 CSM 的复杂度（$O(ncd)$）低于 K-Means（$O(inkd) $）或全成对网格方法（$ O(n^2d)$）。
实证验证： 在 RULER 和 LongBench 基准测试中，对 Llama-3.1-8B 和 Mistral-7B 模型进行了广泛实验。

结果

内存效率： CentroidKV 实现了高达 75% 的 KV 缓存内存减少（将缓存大小减少约 4 倍），同时保持了相当的模型性能。
推理速度：
- 解码延迟： 加速了解码阶段，最高可实现 1.92× 的加速（以每个输出 token 的时间减少量衡量）。
- 吞吐量： 当集成到 vLLM 框架中时，服务吞吐量最高可提升 4×。
准确性：
- 在 RULER 基准测试中，在各种缓存预算（25%–75%）下，CentroidKV 均取得了最佳或接近最佳的性能，特别是在需要语义理解的任务（如 S-NIAH-1/2, MK-NIAH-1, VT）上。
- 在 LongBench 上，它在单文档问答（Single-Doc QA）、多文档问答（Multi-Doc QA）和摘要任务中表现持续出色。
- 局限性： 在合成任务和细粒度检索任务（例如基于 UUID 的键）中，性能略有下降，因为这些任务对精确的 token 匹配至关重要。作者指出，与基于位置的方法（如 StreamingLLM）相比，内容感知方法在处理语义随机的字符串时普遍表现较差。

意义与主张

本文声称 CentroidKV 在准确性和效率之间提供了理想的权衡。通过利用 key 状态的内在相似性并采用轻量级的单次遍历聚类算法，它避免了由于朴素驱逐导致的性能下降，也避免了像 K-Means 这样迭代聚类方法带来的高昂计算开销。

作者强调，CentroidKV 是一个轻量级、即插即用的解决方案，不需要模型重训。它有效地解决了长上下文 LLM 的内存瓶颈问题，在不显著损害模型理解长程依赖能力的前提下，实现了更快的解码速度和更高的吞吐量。

作者注明的局限性：

当前实现侧重于基于 GPU 的压缩，尚未探索内存卸载策略（例如 CPU 聚类）。
压缩率调度是手动指定的，且该方法限制每一步的缓存减少量不超过一半。未来的工作可以探索自适应压缩策略。