D ynamic KV : Task-Aware Adaptive KV Cache Compression for Long Context LLM s (2025)

Abstract: Efficiently managing the KV cache in Large Language Models (LLMs) is a critical challenge for long-context processing tasks such as retrieval-augmented generation (RAG), long text summarization, and multi-document analysis. Extending the context length substantially increases the KV cache size, leading to excessive memory consumption. Existing KV cache compression methods enforce a fixed pattern, neglecting task-specific characteristics, which hampers the effective retention of essential information while discarding less important tokens. In this paper, we introduce a novel Task-Aware KV cache mechanism that dynamically adjusts the KV cache size across different layers based on the characteristics of the tasks. Our approach builds on the significant observation of distinct activation patterns across layers in various tasks, which highlights 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, adapting 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 demonstrates exceptional performance on the LongBench dataset, retaining only 1.7% of the KV cache while preserving 90%, 87%, 78%, and 83% of the original accuracy for LlaMA-3-8B-Instruct, Mistral-7B-Instruct-v0.2, Qwen2-7B-Instruct, and InternLM-2.5-7B-Chat-1M, respectively. When the retained KV cache size is increased to 6.9%, the performance becomes nearly indistinguishable from that without any KV cache compression. 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 is available at repository https://github.com/DreamMr/DynamicK.

## introduction large language models (llms) @xcite are exerting a considerable influence in the field of natural language processing (nlp), driving advancements in summarization, translation, code generation, etc. @xcite @xcite @xcite . recent developments in llms @xcite have been scaled up to handle long contexts, with llama3 @xcite processing up to 32k tokens and internlm @xcite handling 1m tokens. scaling llms to longer contexts introduces significant latency due to the quadratic complexity of attention. a common solution is to cache key and value (kv) status @xcite , reducing computation. however, this comes at a high memory cost -for example, caching 100k tokens in llama2-7b @xcite still requires over 50gb of memory. to address this issue, recent studies have explored the optimization of kv caching, including kv cache quantization @xcite , token dropping @xcite , architectural improvements to transformers @xcite , kv cache fusion @xcite , and hierarchical sharing and constraints @xcite . existing kv cache compression methods enforce a fixed pattern (as shown in figure 1 ), such as a hierarchical pyramid structure @xcite or a structure similar to fastgen's fixed internal pattern @xcite , or they fix the length of the kv cache to selectively retain tokens across different layers @xcite . however, llms require different numbers of layers when handling different types of tasks. for example, for knowledge-based question-answering tasks, only the first few layers are needed to achieve high accuracy, while for complex reasoning tasks (e.g., mathematics and code generation), more layers are often required to achieve higher accuracy @xcite . thus, we raise a question: do different types of tasks all follow a fixed pattern? to examine this question, we aim to systematically investigate the design principles of the kv cache compression across different tasks. inspired by @xcite , we first investigate how information flow is aggregated through attention mechanisms across different layers in four types of tasks, including single-and multi-document qa, summarization, synthetic tasks and code completion. we find that the attention distribution varies for different types of tasks. for example, in summarization tasks, the upper layers require a small kv cache size, while code completion tasks need larger kv cache sizes in the upper layers. this implies that for code completion tasks, upper layers require maintaining a larger kv cache size, in contrast to pyramidkv @xcite , where the kv cache size decreases as the layer depth increases. building on this insight, we propose a task-aware adaptive kv cache compression method, named dynamickv. specifically, we first calculate an attention score for the most recent few tokens and all other tokens, which in rag @xcite can be viewed as calculating the relevance of the most recent query to the retrieved text. then, we preset a temporary storage to hold the temporary kv cache states and gradually calculate the size of the final retained temporary storage at each k layer by calculating the size of the correlation mean. it should be noted that at each update, the value is gradually normalized, and the retained temporary storage at each layer is always smaller than the previous one. this temporary storage is determined by the number of tokens that need to be retained, and its size is much smaller than the original cache, thus imposing minimal memory overhead. experiments demonstrate that our dynamickv can retain full performance while utilizing only 6.9% of the tokens, and in extreme scenarios, it preserves 90% of the performance with just 1.7% of the tokens. furthermore, experiments on the needle in a haystack benchmark revealed that dynamickv significantly outperforms state-of-the-art (sota) methods. ## related work potential patterns of attention in llms. the transformer architecture @xcite has driven progress in nlp through layered refinement of inputs. bert @xcite reveals a hierarchical structure in intermediate layers via @xcite : surface features dominate lower layers, evolving into syntactic and semantic representations toward the top. this underscores the capability of llms to encode both lexical and complex linguistic information across layers. for decoder-only models, fan et al. (2024) demonstrate that intermediate layers suffice for simple tasks, challenging the necessity of full-depth inference. training strategies like @xcite further optimize efficiency by introducing layer-wise dropout, enabling early computation exit. concurrently, kv cache optimization has emerged as a critical direction. @xcite propose cross-layer attention (cla) to halve cache size via cross-layer attention sharing, while @xcite (ada-kv) dynamically optimize eviction policies by analyzing cross-layer attention patterns. these works highlight the interplay between attention dynamics @xcite and memory-efficient computation. token drop strategies in kv cache compression. vary in approach but share a focus on identifying influential tokens. attention-based methods like fastgen @xcite and scissorhands @xcite use attention patterns for pruning. memory-aware approaches include streamingllm @xcite , which prioritizes streaming via attention sinks, and h2o @xcite , which employs cumulative attention scoring for greedy eviction. hierarchical methods like pyra-midkv @xcite adapt by layer but lack generalizability. snapkv @xcite offers task-agnostic compression by selecting key positions per head. dynamic frameworks such as lazyllm @xcite enable flexible token revival, and ada-kv @xcite improves overall performance by optimizing eviction loss bounds over uniform strategies. ## preliminary studies to systematically investigate the attention mechanism across layers in llms for long-context inputs, we conduct a fine-grained analysis on four different types of tasks: single-and multi-document question answering (qa), summarization, synthetic tasks, and code completion. experimental setting. in particular, we focus our analysis on llama @xcite , visualizing the distribution and behavior of attention across layers to gain deeper insights into its internal mechanisms. inspired by @xcite , we calculate the average attention scores between the most recent tokens and all other tokens. based on these scores, we then identify the top-k (128 multiplied by the number of layers) tokens with the highest attention across all layers. observations. as shown in figure 2 (a), we use boxplot to visually present the distribution of four different types of tasks across different layers. we find that different tasks show significantly different token retention rates at a fixed layer. for example, at early layers, the spread is wide, indicating large task-specific variation. to further understand the distribution of token retention rates across different tasks, we visualize the token retention rates across all layers for each task, as shown in figure 1 (b) . we find that ❶ synthetic task shows higher retention rates in earlier layers, ❷ code completion shows higher retention rates in the earlier layers as well as the last three layers, and ❸ document qa and summarization exhibit different retention dynamics compared to others. insight. the tokens to retain at each layer should adapt dynamically based on the task type. ## dynamickv previous work on kv cache compression @xcite often allocaates a fixed kv cache size across llm layers. however, as our analysis in § 3 demonstrates, attention patterns are not identical across different layers with different types of tasks. therefore, using fixed kv cache size across layers on different tasks may lead to suboptimal performance. thus, we propose dy-namickva dynamic layer-adaptive kv cache compression method. dynamickv consists of two steps: (1) dynamic budget allocation and (2) progressive cache update. ## experiments we conduct comprehensive comparative and ablation experiments to verify the effectiveness of our dynamickv. in § 5.1, we introduce the models, datasets and baselines used in our experiments. ## conclusion we investigate task-specific attention patterns in llms processing long-context inputs and find distinct attention distributions across tasks. to address this, we propose dynamickv, a layer-adaptive kv cache compression framework that dynamically optimizes kv cache allocation per layer. we evaluate the effectiveness and generalizability of dynam-ickv through experiments on 16 datasets from the longbench benchmark, demonstrating its broad applicability and performance benefits. from the results, we mainly conclude that: (1) a wave-like pattern is followed in complex reasoning tasks (e.g., code completion tasks); (2) a pyramid-like pattern is followed in synthetic and summarization tasks; (3) the dynamic hierarchical adaptive dy-namickv approach is capable of formulating a relatively appropriate kv cache retention strategy in accordance with diverse tasks. particularly, in the circumstance of maintaining an extremely small kv cache size, the effect is significantly enhanced. in the future, we hope that there is a more suitable method to perform kv cache compression without increasing the computation.