Training Free Exponential Context Extension via Cascading KV Cache

Authors: Jeffrey Willette, Heejun Lee, Youngwan Lee, Myeongjae Jeon, Sung Ju Hwang
Primary area: generative models
Keywords: transformer, efficiency, linear attention

Abstract:
The transformer's context window is vital for tasks such as few-shot learning and conditional generation as it preserves previous tokens for active memory. However, as the context lengths increase, the computational costs grow quadratically, hindering the deployment of large language models (LLMs) in real-world, long sequence scenarios. Although some recent key-value caching (KV Cache) methods offer linear inference complexity, they naively manage the stored context, prematurely evicting tokens and losing valuable information. Moreover, they lack an optimized prefill/prompt stage strategy, resulting in higher latency than even quadratic attention for realistic context sizes. In response, we introduce a novel mechanism that leverages cascading sub-cache buffers to selectively retain the most relevant tokens, enabling the model to maintain longer context histories without increasing the cache size. Our approach outperforms linear caching baselines across key benchmarks, including streaming perplexity, question answering, book summarization, and passkey retrieval, where it retains better retrieval accuracy at 1M tokens after four doublings of the cache size of 65K. Additionally, our method reduces prefill stage latency by a factor of 6.8 when compared to flash attention on 1M tokens. These innovations not only enhance the computational efficiency of LLMs but also pave the way for their effective deployment in resource-constrained environments, enabling large-scale, real-time applications with significantly reduced latency.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730444632264

thread_type:
official_review

<thread_object>
creator: reviewer_4zw7

<content>
summary:
In order to enable LLMs to preserve a long context without increasing the KV cache size, this paper proposes a training-free KV cache eviction policy with a cascading sub-cache algorithm. The algorithm allows one sub-cache window to accept a subset of tokens evicted from the previous sub-cache. This paper also introduces a linear prefill strategy that avoids the restrictive quadratic prompt complexity during the prefill stage by processing fixed-sized chunks (strides) of the prompt.

strengths:
1. The experiments and conclusions are sufficient, including the evaluation of latency and accuracy on different tasks and LLMs. Also, the display of the experimental figures are clear.
2. This work builds a new paradigm of KV cache with cascading sub-cache and effectively extends the context length without increasing complexity or requiring additional training. This work provides a promising solution for efficient edge-side LLM inference.

weaknesses:
1. I understand the proposed method is a generalization (more flexible version) of sliding window attention, but the main problem is that this paper does not clearly explain why the proposed cascading sub-cache is effective for KV cache eviction (intuitively). It is better to explain more about the background of cascading structure and the key motivations/observations/insights for introducing cascading KV cache.
2. Although experiments are sufficient, the baselines are kind of weak. For example, when evaluating on LongBench, there are more recent dynamic KV cache method such as PyramidKV [1], PyramidInfer[2], InfLLM [3], Quest [4], and FastGen [5]. It will be better to compare with one/two more recent dynamic KV cache method to further support your conclusions.
3. In Section 3.1, it is not clear that why use trunks can achieve linear attention complexity during the prefill stage, and what is the key difference between the proposed method and FlashAttention which also uses tiling to save the memory.
4. From Section 3.2, each sub-cache relies on the previous sub-cache. The sequential procedure seems to have an impact on the inference efficiency. How to avoid this problem and achieve much higer efficiency compared with StreamingLLM.
5. Token selection algorithm leverages exponential moving average (EMA) to track the historical attention score. This method seems similar to directly using accumulated attention score proposed by H2O. Can I understand this process as considering the different importance degrees of current and historical attention score based on the method of H2O, which is controlled by the hyper-parameter gamma.

Related references:
[1] PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling (arXiv 2024)
[2] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference (ACL 2024)
[3] InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory (arXiv 2024)
[4] Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference (ICML 2024)
[5] Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (ICLR 2024)

questions:
Please refer to the weakness part for questions.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730639848680

thread_type:
official_review

<thread_object>
creator: reviewer_ckne

<content>
summary:
The paper presents a new training-free approach which significantly extend the effective context length of LLM. The authors first introduce a strided prefill which reduces the prompt computing latency on 1M tokens by a factor of 6.8 compared to flash attention. Then illustrate their cascade KV cache maintaining mechanism. Given the same KV cache size, the cascade mechanism gains impressive improvements compared with the existing methods.

strengths:
1. Although the idea is simple, the experiment results are very good.
2. The idea is well implemented with careful code optimization.

weaknesses:
The experimental section would be better if the idea is evaluated on more data sets.

questions:
I am wondering the perplexity results on more evaluation datasets.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730662157588

thread_type:
official_review

<thread_object>
creator: reviewer_2xpr

<content>
summary:
The transformer’s context window is crucial for tasks like few-shot learning, but longer contexts increase computational costs, limiting large language models (LLMs) in long-sequence tasks. Recent key-value caching methods lower complexity but often discard tokens too soon and have inefficient prefill stages, leading to high latency. To solve this, the author propose a cascading sub-cache buffer system that selectively retains relevant tokens, supporting longer context without added cache size. This method improves accuracy and reduces latency across tasks, enhancing LLM efficiency for real-time applications in resource-limited environments.

strengths:
- The idea of cascade kv cache is simple and effective
- The design of this method carefully considered the implementation so as to get the actual speedup

weaknesses:
The baseline may not be sufficient, most of the experiment are comparing with streamingLLM

questions:
- How's your peformance compared with Minference?
- How's your performance compared with others on RULER task, which is a benchmark for evaluating context size
</content>
</thread_object>

</thread>
<thread>
creation date:
1730663614393

thread_type:
official_review

<thread_object>
creator: reviewer_hjx9

<content>
summary:
The paper proposes an extension to attention sink by (1) blockwise prefilling to optimize the latency of this period, (2) cascading caching management. The resulting framework shows quality and latency improvement over StreamingLLM.

strengths:
1. Prefilling observation is sound. It's good to bulk processing this stage, which StreamingLLM and H2O failed to do.
2. It's intuitionally beneficial to keep the middle context in addition to sliding window and attention sink, which the paper proves by passkey performance.

weaknesses:
1. It's pretty tricky to manage memory in this cascading way. I would say it's challenging to manage small amount of memory, e.g., two bytes, at arbitrary locations. It's very likely to cause memory fragmentation and hurt the throughput as a result, which a practical serving system should avoid.
2. There are a few earlier and more comprehensive works (compared to StreamingLLM) that tackles the prefilling stage with sparsity, e.g., MInference[1], it would be compare against [1].
[1] MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention

questions:
1. Related to weakness 1, is this method compatible with continuous batching? How well does the method perform compared to, or compatible with actual serving systems, e.g., vllm/sglang? 
2. Could you benchmark the latency and quality wrt to Minference? 
3. For FlashAttention, which version are you comparing to? FlashAttention-1, FlashAttention-2, or FlashAttention-3? Triton or cuda kernel?
4. What would be the memory overhead, and IO overhead for computing, storing and access the EMA for each token(and manipulate the resulting memory layout change)? Also, can this operation be fused into online softmax? It would be good to add a pseudo code section to illustrate how this part is done, e.g., based on FlashAttention logic.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732686724607

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Integration with MInference Kernel

comment:
Dear Reviewers,

Thank you again for your thoughtful and constructive feedback on our work. As multiple reviewers requested comparisons with MInference, we have conducted additional experiments and provided results in our previous responses.

To address these comments even more comprehensively, we have integrated the MInference attention kernel directly into our linear strided prefill strategy. This integration allows us to replace the dense leading attention window in our original method with MInference, while retaining the core striding and cascading cache strategies of our method.

Specifically, the attention operation uses the following overall flow to integrate with MInference:
  1. Retrieve tokens from cascading cache
  2. Perform attention for cached tokens + MInference window
  3. Add tokens from the leading MInference window to the cascading cache.
  4. Repeat 1-3 until entire sequence is prefilled.
  5. Perform generation using the keys and values in our cascading cache. 
  
## MInference Integration Results

By incorporating MInference, we observe significant performance improvements across most LongBench datasets, as summarized below:

|  | Narrative QA | HotPotQA | Qasper | Multi News | 2WikiMQA | Gov. Report | Avg. |
|---|---|---|---|---|---|---|---|
| Streaming LLM | 22.57 | 40.78 | 23.89 | 20.69 | 23.46 | 26.82 | 26.37 |
| Cascade (Ours) | 26.43 | 47.26 | 33.12 | **23.33** | 32.33 | **28.32** | 31.8 |
| Cascade (Ours) + MInference | **29.84** | **55.22** | **36.51** | 23.14 | **48.85** | 28.2 | **36.96** | 

This boost in performance is attributed to processing longer strides, enabling more tokens from prior steps to influence the current attention computation. For these experiments, the stride for MInference + Cascade was set to 100K, which is close to the maximum for the model (131K) minus an amount to account for the cascading cache size.

---

## Latency Comparison:

We also evaluated the per-layer latency of MInference, Cascade, and the combined method. The results are as follows:

|                             | 16K  | 32K  | 64K  | 131K | 262K | 524K  | 1M    |
|-----------------------------|------|------|------|------|------|-------|-------|
| Minference                  | 0.26 | **0.28** | **0.78** | 2.26 | 6.39 | 16.84 | OOM   |
| Minference + Cascade (Ours) | 0.34 | 0.68 | 1.3  | 3.04 | 5.77 | 11.88 | 24.13 |
| Cascade (Ours)              | **0.22** | 0.46 | 0.94 | **1.93** | **3.81** | **7.66**  | **15.34** |

These results demonstrate that while integrating MInference incurs slightly higher latency compared to our original method, it still scales linearly with context length and achieves lower latency than plain MInference above 131K—a property of interest for long-context streaming settings. Note that these latency results are for one individual self attention layer, as opposed to the full model forward which is part of the response to individual reviewers.

--- 

## Invitation for Further Discussion

We hope these additional results address your questions and concerns. We remain open to further discussion and would greatly value your thoughts on this updated integration during the remaining discussion period.

Thank you again for your time and insightful feedback.

Sincerely,

Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1733114084588

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Discussion period will end soon

comment:
Dear Reviewers,

As a gentle reminder, **we only have about 30 hours left for the author/reviewer discussion period.**

We politely ask that you review and respond to our updates after your initial feedback. We look forward to addressing any remaining questions which you might have and engaging in further discussion.

We sincerely thank you for your reviews which have improved our work, 
Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a cascading sub-cache buffer system that selectively retains relevant tokens, supporting longer context without added cache size. This method improves accuracy and reduces latency across tasks, enhancing LLM efficiency for real-time applications in resource-limited environments.

The paper presentation is clear. Extensive evaluations show the effectiveness of the proposed approach.