CTkvr: KV Cache Retrieval for Long-Context LLMs via Centroid then Token Indexing

📄 arXiv: 2512.15550v1 📥 PDF

作者: Kuan Lu, Shuhang Lin, Sai Wu, Yichen Yao, Junhan Yang, Huan Li, Wei Chu, Xu Yinghui, Yuan Qi, Gang Chen

分类: cs.CL

发布日期: 2025-12-17

💡 一句话要点

提出CTKVR：一种基于质心和Token索引的长文本LLM KV缓存检索方法

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本LLM KV缓存 检索效率 质心索引 Token索引 RoPE CPU-GPU协同

📋 核心要点

  1. 现有动态KV选择方法在长文本LLM中面临准确性和效率的权衡，块级索引牺牲准确性，token级索引效率低。
  2. CTKVR利用RoPE后相邻查询向量的相似性，提出质心-token两阶段检索策略，平衡效率和准确性。
  3. 实验表明，CTKVR在多个基准测试中性能优越，精度损失小于1%，并在Llama-3-8B和Yi-9B上实现了显著的吞吐量提升。

📝 摘要（中文）

大型语言模型（LLM）越来越多地应用于多轮对话等长文本场景。然而，长文本对推理效率提出了重大挑战，包括来自键值（KV）缓存的高内存开销和因过度内存访问而增加的延迟。现有的动态KV选择方法难以权衡：块级索引通过检索不相关的KV条目来降低准确性，而token级索引因低效的检索机制而导致高延迟。本文提出了一种新颖的质心-token KV检索方案CTKVR，以解决这些限制。CTKVR利用了一个关键观察：位置相邻的查询向量在旋转位置嵌入（RoPE）后表现出高度相似性，并共享其大部分top-k KV缓存条目。基于此，CTKVR采用两阶段检索策略：在预填充期间预先计算轻量级质心以进行质心粒度索引，然后进行token级细化以进行精确的KV检索。这种方法平衡了检索效率和准确性。为了进一步提高性能，我们实现了一个优化的系统，用于使用CPU-GPU协同执行进行索引构建和搜索。实验表明，CTKVR在多个基准测试中实现了卓越的性能，且准确性下降不到1%。同时，在不同的GPU硬件上，CTKVR在96K上下文长度下，Llama-3-8B和Yi-9B的吞吐量分别提高了3倍和4倍。

🔬 方法详解

问题定义：论文旨在解决长文本LLM推理过程中，由于KV缓存过大和频繁访问导致的效率瓶颈问题。现有方法如块级索引牺牲了精度，而token级索引效率低下，无法满足长文本场景的需求。

核心思路：论文的核心思路是利用RoPE后相邻位置的查询向量具有高度相似性的特点，这意味着它们倾向于访问相似的KV缓存条目。因此，可以通过先进行粗粒度的质心索引，再进行细粒度的token索引，从而在效率和精度之间取得平衡。

技术框架：CTKVR采用两阶段检索框架。第一阶段是质心索引阶段，在预填充阶段预先计算每个token的质心，并构建质心索引。第二阶段是token检索阶段，首先根据查询向量找到最近的质心，然后仅在与该质心相关的token集合中进行精确的token级检索。整个过程利用CPU-GPU协同执行来加速索引构建和搜索。

关键创新：CTKVR的关键创新在于提出了质心-token两阶段检索策略，它避免了块级索引的粗糙性和token级索引的低效性。通过利用RoPE的特性，实现了高效且精确的KV缓存检索。

关键设计：CTKVR的关键设计包括：1) 质心的计算方法，例如可以使用k-means等聚类算法；2) 质心索引的构建方式，可以使用哈希表等数据结构；3) token级检索的优化，例如可以使用向量相似度搜索算法；4) CPU-GPU协同执行的策略，例如可以将索引构建放在CPU上进行，而将搜索放在GPU上进行。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

CTKVR在多个基准测试中表现出色，精度损失小于1%。在96K上下文长度下，Llama-3-8B和Yi-9B的吞吐量分别提高了3倍和4倍。这些结果表明，CTKVR能够显著提高长文本LLM的推理效率，同时保持较高的模型精度。

🎯 应用场景

CTKVR适用于需要处理长文本的LLM应用，例如多轮对话系统、长文档摘要、代码生成等。通过提高推理效率，CTKVR可以降低部署成本，并支持更大规模的上下文窗口，从而提升用户体验和模型性能。该研究对于推动LLM在实际场景中的应用具有重要意义。

📄 摘要（原文）

Large language models (LLMs) are increasingly applied in long-context scenarios such as multi-turn conversations. However, long contexts pose significant challenges for inference efficiency, including high memory overhead from Key-Value (KV) cache and increased latency due to excessive memory accesses. Recent methods for dynamic KV selection struggle with trade-offs: block-level indexing degrades accuracy by retrieving irrelevant KV entries, while token-level indexing incurs high latency from inefficient retrieval mechanisms. In this paper, we propose CTKVR, a novel centroid-then-token KV retrieval scheme that addresses these limitations. CTKVR leverages a key observation: query vectors adjacent in position exhibit high similarity after Rotary Position Embedding (RoPE) and share most of their top-k KV cache entries. Based on this insight, CTKVR employs a two-stage retrieval strategy: lightweight centroids are precomputed during prefilling for centroid-grained indexing, followed by token-level refinement for precise KV retrieval. This approach balances retrieval efficiency and accuracy. To further enhance performance, we implement an optimized system for indexing construction and search using CPU-GPU co-execution. Experimentally, CTKVR achieves superior performance across multiple benchmarks with less than 1% accuracy degradation. Meanwhile, CTKVR delivers 3 times and 4 times throughput speedups on Llama-3-8B and Yi-9B at 96K context length across diverse GPU hardware.