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

TL;DR

CTKVR introduces a two-stage KV retrieval method for long-context LLMs that balances efficiency and accuracy, significantly improving throughput with minimal accuracy loss.

Contribution

The paper proposes a novel centroid-then-token KV retrieval scheme that enhances long-context LLM inference efficiency by combining lightweight centroid indexing with token-level refinement.

Findings

Achieves less than 1% accuracy degradation.

Provides 3x and 4x throughput speedups on Llama-3-8B and Yi-9B models.

Effective across diverse GPU hardware.

Abstract

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.