DynamicKV: Task-Aware Adaptive KV Cache Compression for Long Context LLMs
This addresses the computational bottleneck of KV cache compression for users of long-context LLMs in applications like RAG and summarization, representing a strong incremental improvement over fixed-pattern methods.
The paper tackles the problem of inefficient KV cache management in large language models for long-context tasks by proposing DynamicKV, a method that dynamically adjusts token retention per layer based on task-specific characteristics, achieving ~85% of full KV cache performance while retaining only 1.7% of the cache size on LongBench.
Efficient KV cache management in LLMs is crucial for long-context tasks like RAG and summarization. Existing KV cache compression methods enforce a fixed pattern, neglecting task-specific characteristics and reducing the retention of essential information. However, we observe distinct activation patterns across layers in various tasks, highlighting the need for adaptive strategies tailored to each task's unique demands. Based on this insight, we propose DynamicKV, a method that dynamically optimizes token retention by adjusting the number of tokens retained at each layer to adapt to the specific task. DynamicKV establishes global and per-layer maximum KV cache budgets, temporarily retaining the maximum budget for the current layer, and periodically updating the KV cache sizes of all preceding layers during inference. Our method retains only 1.7% of the KV cache size while achieving ~85% of the Full KV cache performance on LongBench. Notably, even under extreme compression (0.9%), DynamicKV surpasses state-of-the-art (SOTA) methods by 11% in the Needle-in-a-Haystack test using Mistral-7B-Instruct-v0.2. The code will be released.