Technical Summary: CentroidKV

Problem Statement

Large Language Models (LLMs) with extended context windows (up to 1M tokens) face significant deployment challenges due to the Key-Value (KV) cache. The KV cache, which stores keys and values of preceding tokens to avoid re-computation, scales linearly with context length. This creates a bottleneck for inference latency and throughput, as every generated token requires accessing the entire cache. Furthermore, the large memory footprint complicates deployment.

Existing solutions to reduce KV cache size include:

• Eviction: Discarding less critical tokens based on attention scores. This often leads to performance degradation because tokens deemed unimportant initially may become crucial for future generations.
• Merging: Combining tokens to be evicted with retained tokens. However, existing methods often rely on preliminary eviction to define merging sets, forcing semantically distant tokens to be grouped, resulting in centroid bias and information loss.
• Quantization and Pruning: These offer limited efficiency gains or introduce high computational overhead.

There is a need for a method that compresses the KV cache based on intrinsic similarity without discarding critical information or incurring prohibitive computational costs.

Methodology: CentroidKV

The paper introduces CentroidKV, an online KV cache clustering framework. Its core innovation is the Chunked Soft Matching (CSM) algorithm, designed to cluster KV states based on similarity while maintaining low computational complexity.

Key Components

1. Observations:

   • Localized Similarity: Key states exhibit high cosine similarity along the sequence dimension, with similar tokens clustering in localized regions.
   • Distance Correlation: As token distance increases, the cosine similarity of key states generally decreases monotonically and follows a convex trend.

2. Inference Pipeline:

   • Pre-filling: The sequence is processed, and if the KV cache exceeds a budget, the first clustering round is executed.
   • Decoding: The cache grows with each token. Clustering is invoked periodically (every $g$ steps) to maintain the cache within the budget.
   • Preservation: Attention sinks and recent tokens are preserved before partitioning to ensure model stability.

3. Chunked Soft Matching (CSM):

   • Chunking: The sequence is divided into chunks to handle long contexts efficiently.
   • Alternating Partition: Within each chunk, tokens are partitioned into two sets ($A$ and $B$) in an alternating manner. Theoretical analysis proves this strategy is optimal given the observed monotonic and convex similarity trends, ensuring highly similar tokens are placed in different sets to facilitate matching.
   • Matching: Similar to Bipartite Soft Matching (BSM) used in Vision Transformers, CSM identifies highly similar token pairs across sets $A$ and $B$.
   • Compression: Instead of merging all pairs, the algorithm ranks candidate matches by similarity and merges only the top fraction controlled by a compression ratio $r$.
   • Progressive Scheduling: The compression ratio $r$ is linearly annealed over multiple clustering rounds ($r = r_{init} - j \cdot \delta_r$). This makes merging increasingly selective in later rounds, preserving high-confidence structures.

4. Aggregation:

   • For each cluster, the key and value states are merged into a single centroid using degree-weighted averaging. The "degree" ($n_t$) tracks the cumulative contribution of tokens merged into a centroid across rounds.
   • The attention mechanism is adapted to use these centroids, incorporating the cluster degree as a log-term in the softmax calculation to approximate the original attention output.

Key Contributions

• CentroidKV Framework: A simple, plug-and-play solution for online KV cache clustering that improves inference efficiency in long-context scenarios.
• Chunked Soft Matching Algorithm: A novel clustering algorithm that extends Bipartite Soft Matching to KV cache compression. It utilizes an alternating partition strategy within chunks, which is theoretically proven to be optimal for the observed similarity patterns.
• Theoretical Analysis: The paper provides a formal proof of the optimality of the intra-chunk partitioning strategy and analyzes the computational complexity, demonstrating that CSM achieves lower complexity ($O(ncd)$) compared to K-Means ($O(inkd)$) or full pairwise mesh methods ($O(n^2d)$).
• Empirical Validation: Extensive experiments on Llama-3.1-8B and Mistral-7B models across RULER and LongBench benchmarks.

Results

• Memory Efficiency: CentroidKV achieves up to a 75% reduction in KV cache memory usage (reducing cache size by ~4x) while maintaining comparable model performance.
• Inference Speed:
  • Decoding Latency: Accelerates the decoding stage by up to 1.92× (measured as Time Per Output Token reduction).
  • Throughput: Increases serving throughput by up to 4× when integrated into the vLLM framework.
• Accuracy:
  • On the RULER benchmark, CentroidKV achieves the best or near-best performance across various cache budgets (25%–75%), particularly on tasks requiring semantic understanding (e.g., S-NIAH-1/2, MK-NIAH-1, VT).
  • On LongBench, it performs consistently well on Single-Doc QA, Multi-Doc QA, and Summarization.
  • Limitations: Performance drops slightly on Synthetic tasks and fine-grained retrieval tasks (e.g., UUID-based keys) where exact token matching is critical. The authors note that content-aware methods generally struggle with semantically arbitrary random strings compared to position-based methods like StreamingLLM.

Significance and Claims

The paper claims that CentroidKV offers a favorable trade-off between accuracy and efficiency. By leveraging the intrinsic similarity of key states and employing a lightweight, single-pass clustering algorithm, it avoids the performance degradation associated with naive eviction and the high computational overhead of iterative clustering methods like K-Means.

The authors emphasize that CentroidKV is a lightweight, plug-and-play solution that does not require model retraining. It effectively addresses the memory bottleneck of long-context LLMs, enabling faster decoding and higher throughput without significantly compromising the model's ability to understand long-range dependencies.

Limitations noted by the authors:

• The current implementation focuses on GPU-based compression and does not explore memory offloading strategies (e.g., CPU clustering).
• The compression ratio schedule is manually specified, and the method restricts cache reduction to at most half at each step. Future work could explore adaptive compression strategies.