DynamicKV: Task-Aware Adaptive KV Cache Compression for Long Context LLMs

TL;DR

DynamicKV introduces a task-aware, adaptive KV cache compression method for long-context LLMs, significantly reducing cache size while maintaining high performance across various tasks.

Contribution

It proposes a novel adaptive KV cache management strategy that dynamically adjusts token retention per layer based on task-specific activation patterns.

Findings

Retains only 1.7% of KV cache size with ~85% performance

Outperforms SOTA methods by 11% in Needle-in-a-Haystack test

Achieves efficient long-context processing with minimal cache usage

Abstract

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.