Dialogue Without Limits: Constant-Sized KV Caches for Extended Response in LLMs

We thank the reviewer for their feedback. We have addressed the points raised in the review below:

Q-1) Selection Bias in H$_2$O

R-1) H$_2$O retains tokens based on aggregated attention scores. This introduces a selection bias because it retains early tokens even when they do not significantly attend to newer tokens. Various prior works, like NACL [1], Keyformer, and EMS [2], have made similar observations. In NACL, the authors extensively study this problem using LongBench tasks on the Llama-2-7B model. While generating the last token, 96% of H$_2$O’s KV cache comprises entries from the first 200 tokens, while only 4% is derived from the subsequent 1199 tokens. Please refer to the NACL paper (particularly Figure 2 (a)) for details about this study.

[1] NACL: https://aclanthology.org/2024.acl-long.428.pdf

[2] EMS: https://arxiv.org/pdf/2412.08521

Q-2) SnapKV for Long-Form Generation

R-2) Indeed, SnapKV’s overwhelming memory requirements make it unsuitable for long-response tasks. While H$_2$O is more memory-efficient for these tasks, its performance is limited due to the problem of selection bias towards early tokens. MorphKV efficiently navigates this trade-off and achieves robust performance and high memory efficiency for both long-context and long-response tasks.

Q-3) H$_2$O for Long-Context Benchmarks

R-3) SnapKV is the state of the art for long-context tasks and already outperforms H$_2$O [1]. Hence, we only compare against SnapKV because outperforming it naturally implies outperforming H$_2$O (please see Section-5.3 of their paper).

[1] SnapKV: https://proceedings.neurips.cc/paper_files/paper/2024/hash/28ab418242603e0f7323e54185d19bde-Abstract-Conference.html

Q-4) Prompt for LM Judge

R-4) We use the publicly available LongWriter framework in our studies. LongWriter uses a single-answer grading approach, where an LLM judge is asked to assign a score to a single answer (please refer to Section 3.1 in [1] for further details). The prompt used for the LM judge can be found here: https://github.com/THUDM/LongWriter/blob/main/evaluation/judge.txt

[1] Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena: https://arxiv.org/abs/2306.05685

Q-5) Design Choices in H$_2$O Baseline

R-5) The performance of a KV cache pruning method depends on the quality of the retained KVs and how these are utilized for future tokens. In contrast, the memory savings depend on the number of KVs evicted. Naively reducing the memory capacity of H$_2$O also lowers its performance due to the limited accuracy of its token selection policy that exhibits a bias toward preserving early tokens (please see our earlier response).

Our studies using LongGenBench on the Mistral-24B-Instruct model show that MorphKV has 1.67$\times$ better accuracy when both H$_2$O and MorphKV operate within the same cache budget. To achieve comparable accuracy for the same task, H$_2$O must overcome the limitations of its token selection policy and retain a lot more KVs, increasing the memory footprint. MorphKV is 3.6$\times$ more memory-efficient compared to H$_2$O in this case.

Q-6) Local coherence and Distant Relative Heuristics in MorphKV

R-6) Local and distant heuristics enable MorphKV to capture useful information. Discarding recent tokens severely degrades performance, critical to maintaining the text’s local coherence. In contrast, distant heuristics are crucial to generating an apt response overall, as they capture information that establishes the global pre-text and relevance. MorphKV leverages the insight that due to transformers’ auto-regressive nature of output token generation, each token attends to all its past tokens and naturally holds awareness of the useful distant tokens. MorphKV retains only those distant tokens identified as relevant by preceding recent tokens, allowing the current output token to focus solely on this curated set of useful distant tokens.

Q-7) Adversarial Cases where Heuristics Overlook Important Tokens

R-7) The ability to retain useful or important distant tokens depends on the accuracy of the heuristics and the choice of the appropriate design hyperparameters. For example, the default implementation of MorphKV uses a window size of 32 because it works well across most representative tasks. However, for certain benchmarks that involve recalling information over very long time spans, such as LongGenBench, while this window size captures most of the useful distant tokens, the performance can be improved slightly by increasing the window size (which would allow capturing all useful distant tokens). Please see Table-3 and Table-4 in response to Reviewer bnHA for the detailed results. However, the probability of encountering such scenarios is very low because we carefully tune them, and the most appropriate window size can be well approximated based on the nature of the task and its information span.

We thank the reviewer for their feedback.

Q-1) Comparison with Prior Works

R-1) MorphKV’s superior performance mainly stems from its more accurate token selection policy.

StreamingLLM retains the KVs of the first few initial tokens called attention sinks and a sliding window of tokens, evicting all intermediate tokens between them. This creates a “context gap”, leading to partial capture of distant context as noted by prior works [1]. In contrast, MorphKV retains distant tokens (including important intermediate tokens) actively attended by recent tokens, enabling essential context capture throughout token generation.

MorphKV eliminates the token selection bias dominant in H$_2$O and Keyformer. MorphKV exploits the insight that due to the auto-regressive nature of token generation, each token attends to all its past tokens and thus, is naturally aware of useful distant tokens. MorphKV retains only those distant tokens identified as relevant by preceding recent tokens, allowing the current output token to focus only on a curated set of useful distant tokens.

SnapKV assumes that critical tokens from the input prompt remain critical throughout token generation. Based on this insight, it does not prune the KV cache during token generation. However, LLMs often suffer from degeneration while generating large responses [2] as each token attends to all past tokens, compounding noise from early decoding steps. As SnapKV retains KVs of all generated tokens, it is significantly more vulnerable to degeneration and performance degradation for long-response tasks. In contrast, MorphKV only retains important tokens in the KV cache enabling it to generate more coherent responses.

Degeneration is quantified using the rate at which phrases are repeated in the LLM response. Table-9 shows that SnapKV has 1.3$\times$ higher repetition rate than MorphKV on LongWriter.

Table 9: Text Degeneration via N-Gram Repetition

[1] Attention-Gate: https://arxiv.org/pdf/2410.12876

[2] Text Degeneration: https://arxiv.org/abs/1904.09751

Q-2) Why MorphKV outperforms Full-Attention (FA)

R-2) Although we intuitively expect FA to perform superior to KV cache pruning techniques, for the same reasons as described above (degeneration), it may actually perform worse for certain tasks. Recent works, like SnapKV, and PyramidKV also make similar observations. SnapKV argues that FA often introduces noise to the attention mechanism by retaining all tokens, hindering its ability to attend to only the most relevant tokens more strongly. This trend is heavily task-dependent and KV cache compression does not always guarantee outperforming full-attention. For example, while MorphKV outperforms FA for few long-response tasks, FA outperforms MorphKV by 1.8% on LongGenBench for Qwen2.5-32B-Instruct.

Q-3) Statistical testing

R-3) Our results do not exhibit any statistical variance because MorphKV performs greedy decoding, and therefore generated responses remain same across runs.

Q-4) Optimal Value of KV cache size

R-4) We can obtain the optimal values through Grid-Search with validation on the end task. For our evaluations, we perform a coarse-grained search and determine that a KV cache size of 4K tokens is a practical starting point for most tasks.

Q-5) Generalizability of MorphKV

R-5) Please refer to R-4 in response to Reviewer vs1e.

Q-6) Inference Time

R-6) Please refer to R-5 in response to Reviewer Jsuc.

Q-7) Comparison of MorphKV against combination of SnapKV and H$_2$O

R-7) MorphKV differs substantially from a naive combination of SnapKV and H$_2$O. This primarily stems from the token selection policy used in MorphKV. Although both H$_2$O and MorphKV prune the KV cache dynamically during the decode phase, they use different policies to identify the important KVs they must retain. Consequently, even for the same memory budget, the set of KVs retained by the two approaches is significantly different. Please refer to R-1 for a detailed discussion regarding how MorphKV compares with SnapKV.

Q-8) H$_2$O Illustration

R-8) We use Figure 3 only as an illustrative example to highlight that H$_2$O’s token retention policy introduces a bias towards early tokens due to aggregated attention scores, similar to Figure 2 (a) in NACL [1]. We will update the figure to highlight the dynamic token-eviction in H$_2$O.

[1] NACL: https://arxiv.org/pdf/2408.03675

Q-9) Compatibility with FlashAttention

R-9) Yes, MorphKV is already compatible with FlashAttention (please see Section 4 of the paper). While fused kernels do not provide attention scores directly, It is still possible to tweak the FlashAttention kernel to return partial attention matrices alongside the final attention output, which can be consumed by MorphKV. An in-depth discussion of this is beyond the focus of this work.

We thank the reviewer for their feedback. We have addressed the points raised in the review below:

Q-1) Analysis on Larger Models

R1) MorphKV remains effective for larger models, as demonstrated by our recent evaluations with larger models. We request the reviewer consult Table-1 and Table-2 in the response to Reviewer fibk.

Q-2) Comparison with Prior Works

R2) For a detailed comparison with prior works, including StreamingLLM and SnapKV, please refer to R-1 in response to Reviewer-FtVG.

Q-3) Choice of Benchmarks

R-3) Both long-context and long-response generation tasks cater to widely prevalent applications of LLMS and are key areas of focus [1-3]. Long-context generation not only represents common use cases of LLMs like text summarization and passage retrieval but can also be used for innovative training techniques such as multi-shot learning, where the model can learn from hundreds of training examples provided directly in the prompt.

In contrast, long-response generation is crucial for tasks such as story generation, paragraph completion, comprehensive question-answering, content creation, etc. While prior works, such as SnapKV, are optimized explicitly for long-context tasks, they are inefficient for long-response tasks. MorphKV addresses this critical bottleneck.

[1] Google: https://cloud.google.com/transform/the-prompt-what-are-long-context-windows-and-why-do-they-matter

[2] IBM: https://research.ibm.com/blog/larger-context-window

[3] Amazon: https://aws.amazon.com/blogs/security/context-window-overflow-breaking-the-barrier/

Q-4) Effect of KV cache budget on MorphKV

R-4) Thanks for the detailed feedback. We have conducted additional studies using LongGenBench tasks for the Mistral-24B-Instruct [1] model to evaluate the impact of varying cache capacities on performance. We observe that the performance of both MorphKV and H$_2$O increases with the number of cached tokens. At the same KV cache capacity, MorphKV remains more effective than H$_2$O. We will revise the Appendix to include these results.

Table 7: Performance with increasing KV cache capacity on LongGenBench

[1] Mistral-24B-Instruct: https://mistral.ai/news/mistral-small-3

Q-5) Is dynamic token selection for caching computationally expensive? If so, how does the runtime of MorphKV compare to SnapKV and H$_2$O?

R-5) KV cache compression methods must navigate complex trade-offs encompassing accuracy, inference time, throughput, and memory footprint. Optimizing for a single metric alone is insufficient for practical adoption.

MorphKV prunes the cache at every token processing step, which increases the computational overhead. Our implementation employs several optimizations to reduce these overheads, such as CPU offloading and prefetching with dedicated CUDA streams that minimize the time overhead associated with loading attention weights.

Compared to a similar dynamic token eviction policy, such as H$_2$O, MorphKV achieves 8% faster inference time (while remaining more accurate). MorphKV’s runtime is higher than SnapKV’s, which is expected because SnapKV is tailored for long-context tasks, not long-response ones.

However, the reduced memory footprint of the KV caches enables MorphKV to deliver a much higher throughput (up to 4.68$\times$) due to a larger batch size, effectively compensating for the degradation observed in the runtime of each request. In conclusion, MorphKV outperforms existing KV cache compression techniques by achieving the best overall balance: it reduces KV cache memory usage by up to 85%, increases throughput by 4.68$\times$, and achieves an inference time improvement of 8% over state-of-the-art methods (such as H$_2$O) – all while preserving or improving accuracy. Table 8 summarizes these results.

Table 8: Comparison of key metrics across KV cache compression methods for Mistral-7B on LongBench

Q-6) MorphKV on Retrieval Tasks

R-6) MorphKV dynamically evicts tokens based on attention weight scores of the recent window tokens. Hence, it is not biased towards retention or eviction of initial tokens during generation. We evaluate MorphKV on the LongBench suite, which includes many retrieval tasks (single-doc QA, multi-doc QA, Passage Retrieval etc.), and MorphKV shows robust performance across all these tasks.

We appreciate the reviewer's feedback on both the structure of our paper and the results. We have addressed the points raised in the review below:

Q-1) Analysis on Larger Models

R-1) MorphKV remains effective even for larger models. We evaluate MorphKV using long-response tasks from the LongGenBench suite on an Nvidia H100 GPU equipped with 80 GB of HBM-2e memory on two models (Mistral-24B-Instruct [1] and Qwen2.5-32B-Instruct [2]). Even for a single task, SnapKV consumes significant KV cache sizes: 6.16 GB (50.26 GB including weights) and 9.7 GB (71.14 GB including weights), respectively, nearly exhausting the GPU's available HBM capacity. In contrast, MorphKV requires only 0.80 GB and 1.06 GB, respectively. The tables below summarize these results. The evaluated model sizes are consistent with prior works, like SnapKV (7B), H$_2$O (30B), and Keyformer (7B).

Table-1: Mistral-24B-Instruct on LongGenBench

Table-2: Qwen2.5-32B-Instruct on LongGenBench

[1] Mistral-24B-Instruct: https://mistral.ai/news/mistral-small-3

[2] Qwen2.5-32B-Instruct: https://qwenlm.github.io/blog/qwen2.5/

We thank the reviewer for their appreciation of our work. We have addressed the point raised in the review below:

Q-1) Paper Could Explore PyramidKV, Ada-KV, HeadKV Advancements

R-1) Thank you for the excellent suggestion. MorphKV is orthogonal to methods like PyramidKV, which optimizes the KV cache across layers, and Ada-KV and Head-KV, which optimize the KV cache across attention heads. Thus, integrating MorphKV with these methods can improve its efficacy even further. In fact, we have already conducted some preliminary studies wherein we simulate a layer-wise optimization by naively deactivating MorphKV for the first three layers based on prior studies that describe the importance of information in the early layers [1,2,3]. Our studies show that this design is even more effective (with roughly 10% additional memory requirements) than the default implementation of MorphKV, which prunes the KV cache uniformly across all model layers. The results are summarized in Table 5 below.

Integrating MorphKV with methods like PyramidKV, Ada-KV, and Head-KV as a generalized solution requires a deeper understanding of how the methods complement each other, careful tuning of the design parameters, and development of a robust policy that dynamically combines them for optimal performance. Such an extensive study is beyond the scope of our current paper, and we reserve it for future work.

Table 5: MorphKV with basic layer-wise optimization for Mistral-7B on LongGenBench

MorphKV configuration has a capacity of 2000 tokens

[1] PyramidKV: https://arxiv.org/pdf/2406.02069

[2] SqueezeAttention: https://openreview.net/forum?id=9HK2rHNAhd

[3] Layer-condensed KV cache: https://aclanthology.org/2024.acl-long.602.pdf

Q-2) Trade-Off between Computational Complexity and Memory Efficiency

R-2) By default, MorphKV prunes the KV cache after each token is processed to attain a constant-sized cache, which incurs additional computational overheads. The overheads can indeed be lowered by adopting a lazy token-selection policy, at the cost of slight performance degradation. The computational savings scale in the number of pruning steps skipped by the lazy token-selection policy. For example, if MorphKV prunes the cache after every 10 tokens processed, the computational overheads reduce to 0.1$\times$ of that required in the default MorphKV.

To study this further, we implement a variant of MorphKV, which allows the KV cache to exceed the pre-allocated capacity before pruning it back. Table 6 compares the performance of the default and lazy variant of MorphKV. We will include these results in the revised manuscript.

Table 6: MorphKV with lazy token-selection for Mistral-7B on LongGenBench

MorphKV configuration has a capacity of 4000 tokens

Q-3) Inference Time

R-3) We request the reviewer to consult R-5 in response to Reviewer Jsuc.

Q-4) Generalizability of MorphKV

R-4) We already evaluate MorphKV using the LongBench [1] suite that encompasses tasks such as question answering [2,3] based on academic, legal, government, literature and financial reports, reasoning [4-6] from articles and reports from wikipedia, and summarization [7] on documents from academia and product development (please see Section-5.3). MorphKV delivers robust performance across these tasks, showing its generalizability across various domains

[1] LongBench: https://arxiv.org/pdf/2308.14508

[2] NarrativeQA: https://ar5iv.labs.arxiv.org/html/1712.07040

[3] Qasper: https://paperswithcode.com/dataset/qasper

[4] HotpotQA: https://paperswithcode.com/dataset/hotpotqa

[5] 2wikiMQA: https://aclanthology.org/2020.coling-main.580/

[6] PassageCount: https://arxiv.org/pdf/2308.14508

[7] QMSum, VCsum: https://arxiv.org/pdf/2308.14508

We thank the reviewer for their feedback. We have addressed the points raised in the review below:

Q-1) Effect of Window Size on MorphKV

R-1) The window size indeed affects the performance of MorphKV. In practice, all KV cache pruning methods rely on hyperparameters: examples include total cache capacity in H$_2$O and Keyformer, window sizes in StreamingLLM and SnapKV, and the number of layers for KV cache projection in SwiftKV. Careful tuning of these parameters is essential for practical adoption and achieving substantial memory savings [1,2].

Similarly, MorphKV also benefits from tuning its window size parameter. Our experiments show that a window size of 32 consistently offers high performance across all LongBench suite tasks, including diverse scenarios such as code generation and literature review. Although increasing window sizes can slightly improve performance for tasks requiring information recall over longer spans, it typically yields diminishing returns, as described in Table 3 and Table 4. Hence, the window size can be coarsely approximated based on the nature of the task in terms of the information span to achieve good performance.

Table 3: Mistral-24B-Instruct performance on LongGenBench

Table 4: Qwen2.5-32B-Instruct performance on LongGenBench

NOTE: each scenario above has a total KV cache capacity of 4000 tokens

[1] Accelerating Enterprise LLM Workloads: https://www.snowflake.com/en/engineering-blog/swiftkv-llm-compute-reduction/

[2] StreamingLLM Integration into TensorRT-LLM: https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/llama#run-llama-with-streamingllm

Q-2) Ease of Usage in Real-World Systems

R-2) Currently, MorphKV is implemented in the widely used transformers library from HuggingFace (see setup in Section-4). MorphKV is integrated inside both standard Attention and FlashAttention modules and uses cache offloading to accommodate demanding memory-intensive use cases. Consequently, MorphKV is compatible with a vast spectrum of models hosted on HuggingFace. Our approach is consistent with prior works and facilitates seamless adoption across LLMs for real-world scenarios.