Model Tells You Where to Merge: Adaptive KV Cache Merging for LLMs on Long-Context Tasks

TL;DR

This paper introduces KVMerger, an adaptive KV cache merging method that compresses cache data for large language models, enabling efficient long-context processing with limited memory while maintaining high performance.

Contribution

The paper presents a novel, dataset-independent KV cache merging algorithm that improves long-context LLM performance under memory constraints without significant accuracy loss.

Findings

KVMerger outperforms existing methods like H2O and CaM in long-context tasks.

It maintains high performance at 50% and 35% KV cache budgets.

Effective for models including Llama2-7B-chat and Llama2-13B-chat.

Abstract

How to efficiently serve Large Language Models (LLMs) has become a pressing issue because of their huge computational cost in their autoregressive generation process. To mitigate computational costs, LLMs often employ the KV Cache technique to improve the generation speed. While improving the computational efficiency, the storage requirements of the KV cache are substantial, particularly in long-context scenarios, leading to significant memory consumption. Existing KV cache eviction methods often degrade the performance of LLMs in long-context scenarios due to the information loss introduced by eviction. In this paper, we propose a novel KV cache merging approach, called KVMerger, to achieve adaptive KV cache compression for long-context tasks without significant performance degradation under constrained memory budgets. Our approach is inspired by the intriguing observation that key states exhibit high similarity at the token level within a single sequence. To facilitate merging, we develop an effective yet straightforward merging set identification algorithm to identify suitable KV states for merging. Our merging set identification algorithm stimulates the second observation that KV cache sparsity, from similarity perspective, is independent of the dataset and remains persistent at the model level. Subsequently, we propose a Gaussian kernel weighted merging algorithm to selectively merge all states within each merging set. We conduct extensive experiments to demonstrate the effectiveness of KVMerger for long-context tasks under constrained memory budgets, applying it to models including Llama2-7B-chat and Llama2-13B-chat. Using the LongBench and ZeroScroll benchmarks, we compare our method with other KV cache compression techniques, including H2O and CaM, showing that our method achieves superior performance across tasks with both 50% and 35% KV cache budgets.