Efficient Long-Context LLM Inference via KV Cache Clustering

TL;DR

Chelsea introduces an online KV cache clustering framework for long-context LLMs that significantly reduces memory usage and accelerates inference without sacrificing model performance.

Contribution

The paper proposes a novel clustering method for KV caches in LLMs, enabling efficient memory use and faster inference through Chunked Soft Matching and cache merging.

Findings

Up to 80% reduction in KV cache memory usage.

Inference acceleration by up to 3.19 times.

Latency reduction of up to 2.72 times.

Abstract

Large language models (LLMs) with extended context windows have become increasingly prevalent for tackling complex tasks. However, the substantial Key-Value (KV) cache required for long-context LLMs poses significant deployment challenges. Existing approaches either discard potentially critical information needed for future generations or offer limited efficiency gains due to high computational overhead. In this paper, we introduce Chelsea, a simple yet effective framework for online KV cache clustering. Our approach is based on the observation that key states exhibit high similarity along the sequence dimension. To enable efficient clustering, we divide the sequence into chunks and propose Chunked Soft Matching, which employs an alternating partition strategy within each chunk and identifies clusters based on similarity. Chelsea then merges the KV cache within each cluster into a single centroid. Additionally, we provide a theoretical analysis of the computational complexity and the optimality of the intra-chunk partitioning strategy. Extensive experiments across various models and long-context benchmarks demonstrate that Chelsea achieves up to 80% reduction in KV cache memory usage while maintaining comparable model performance. Moreover, with minimal computational overhead, Chelsea accelerates the decoding stage of inference by up to 3.19$\times$ and reduces end-to-end latency by up to 2.72$\times$.