Training Free Exponential Context Extension via Cascading KV Cache

[W3] 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.

[Response] We are unsure what is meant by "use trunks." We assume this question relates to how we can achieve linear complexity during the prefill stage compared to Streaming LLM. There are actually two models in figure 6b which achieve linear complexity during prefill (Streaming LLM, and ours). The main difference is that the single token per step prefill of Streaming LLM is slow enough that it cannot beat quadratic attention implementations on realistic context lengths. Asymptotically, if the x-axis of figure 6b were extended far enough, Streaming LLM would eventually show lower latency than quadratic attention. Our method becomes faster that streaming LLM due to the fact that in one forward pass of the prefill stage, we process a chunk of tokens as opposed to processing a single token per forward pass as done in the original Streaming LLM prefill. We have added an illustration of this difference to the appendix in Figure 18. This change allows for a drastic speedup, improved performance (Figure 11), and maintains a fixed memory overhead when processing any arbitrary context length.

The key difference between our method and flash attention is that flash attention still has quadratic latency ($\mathcal{O}(n^2)$) while ours has a fixed cache size which caps the second dimension at a constant $k$, achieving an overall complexity of ($\mathcal{O}(nk)$), which can be expressed as $\mathcal{O}(n)$ after dropping constants by convention.

The two aforementioned aspects combined are what allows our method to show superior latency in figure 6b.

---

[W4] From Section 3.2, each sub-cache relies on the previous sub-cache. [...] seems to have an impact on the inference efficiency. How to avoid this problem and achieve much higer efficiency compared with StreamingLLM.

[Response] You are correct to note that this process has an effect on inference efficiency. This can be directly observed in Figure 6a where the zoomed in subfigure shows that our cascading cache with one cascade (equivalent to Streaming LLM) has lower latency than 4 cascades. However, we note that the difference is on the scale of 4 orders of magnitude smaller than the difference between the original Streaming LLM and our cascading cache due to the efficiency of our implementation which is why we needed to create the zoomed in subgraph of figure 6a to even visualize it.

So we believe our method in fact does achieve higher efficiency as compared with Streaming LLM, with a negligible overhead that increases with the number of cascades. We leave any further improvements on top of this as interesting lines of future work.

---

[W5] Token selection [...] (EMA). [...] similar to 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.

[Response] The motivation for using an EMA comes from the fact that a sum accumulation would provide a slight bias towards keeping older tokens which have had more strictly positive terms in the sum, while an EMA would not have this bias. To compare the sum and the EMA, we evaluated a sum accumulation on the LongBench experiment which shows that the EMA has a slight advantage over a sum, although a sum still performs much better than Streaming LLM.

---

We thank you for your time and effort in reviewing our work. We believe that that added experiments have improved the quality and thoroughness of our evaluation. If you have any remaining concerns after our response, we remain open to further discussion throughout the discussion period.

Thank you for your time in reviewing our work and highlighting the strengths of our method in increasing the context length while lowering the complexity without training. Please see our responses to your additional comments below:

---

[W1] [...] 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 [...] background [...] motivations/observations/insights for introducing cascading KV cache.

[Response] We apologize if this was not adequately highlighted in the original submission. We have revised the text of the intro to include more detailed motivation. To summarize, figure 1 visualizes the forgotten context tokens which are a side effect of sliding windows. Figure 2 shows the effect of this forgotten context on the passkey retrieval task, which suffers almost immediately as the context exceeds the sliding window. The motivation for the cascade pattern was derived as a simple and effective way to allow multiple exit paths for tokens such that unimportant tokens may be conditionally evicted before reaching the end of the cache buffer.

---

[W2] 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.

[Response]

Based on your feedback as well as other reviewers, we have added two new baselines, Minference and Pyramid KV, on LongBench. As Minference relies on fully dense decoding and Pyramid KV requires quadratic prompts, we truncate the input according to the standard in LongBench evaluation. All results below utilize Llama 3.1 8B Instruct.

For LongBench, we find that our method outperforms both Minference and PyramidKV on all but the Gov. Report Dataset which shows a thin margin between all models as seen in Figure 9. The average performance increase over all datasets is still heavily in favor of our cascading cache.

For InfiniteBench we allow a total cache size of 32K. we observe the same pattern as LongBench, whereby our model outperforms Minference by a large margin (21.5% and 18.2%) on the first two datasets, and only underperforms Minference by 2.3% on the last dataset where the overall differences between models is much smaller.

---

Continued in next comment...

Thank you for your time and care in reviewing our work and highlighting the simplicity, optimization, and strong experimental results we provide. Please see our responses to your additional comments below with 4 new evaluations. Perplexity results utilize Llama 3.1 8B, while every other result utilizes Llama 3.1 8B Instruct.

---

[W1/Q1] The experimental section would be better if the idea is evaluated on more data sets. I am wondering the perplexity results on more evaluation datasets.

We evaluated our model vs Streaming LLM on Wikitext2 test set which was concatenated into a single input consisting of 288K tokens. We use a prefill stride of 512 tokens with cache sizes of 128, 256, 512, 1024, and 2048. We find that the results are consistent with the PG19 perplexity results which show that our model performs consistently better than Streaming LLM. Please see the results in the table below:

We have also evaluated our model in terms of latency and performance against Minference [1], which is a recently published benchmark raised by other reviewers. Minference does not have a fixed memory overhead, and therefore goes out of memory for sequences over 131K on an A6000 (49GB) although the latency is competitive. Our model, however, maintains a fixed memory overhead for any context size and can easily handle 1M tokens on a single A6000.

We also performed an experiment evaluating against Minference and Pyramid KV on LongBench using the same truncation setting as we used for the quadratic model in Figure 9 due to the fact that Minference utilizes dense quadratic attention for every decoding step and Pyramid KV requires a quadratic prompt. Our Cascading KV Cache model maintains superior performance on every dataset except Gov. Report which shows the smallest margin between all models and methods in Figure 9.

We also extended our experiments to InfiniteBench [2] using a total cache size of 32K and compared to the Minference, finding the same pattern that was shown in LongBench where our model shows large improvements of (21.5% and 18.2%) on the first two datasets, and underperforms Minference by only 2.3% on the last dataset where the overall differences between models is much smaller.

---

Thank you for your time and effort in reviewing our work. Your suggestions have improved the thoroughness of our experiments. We hope we have adequately responded to your discussion points. If you have any remaining discussion points, please let us know. Thank you.

---

References

[1] Minference: https://arxiv.org/pdf/2407.02490

[2] Infinite Bench: https://arxiv.org/abs/2402.13718

Thank you for your time in reviewing our work and highlighting the strengths of our method being simple, effective, and fast. Please see our comments to your further responses below:

---

[W1] The baseline may not be sufficient, most of the experiment are comparing with streamingLLM

[Response] Streaming LLM is the key baseline for our method, as it is the only method from the relevant baselines which is truly linear in complexity with a fixed memory overhead like ours.

• In our experiments, we have included Streaming LLM, BigBird, full quadratic attention, SnapKV, H2O.
• The only experiment which compares to Streaming LLM without other baselines is the 1M token passkey retrieval which is due to the fact the the other methods have quadratic latency as shown in figure 6 which makes them infeasible for this task, and/or they require access to all tokens in the sequence (BigBird) which incurs a large overhead memory cost.

Streaming LLM and our Cascading KV cache are the only two methods which offer linear complexity with a fixed memory overhead for any sequence length.

---

[Q1] How's your peformance compared with Minference?

[Response] In terms of latency, our cascading cache is competitive with Minference at different context lengths, However Minference exceeds the memory of an A6000 (49GB) after 131K context length which is the environment we ran our experiments in. The full model latency for all layers (seconds) can be seen in the table below. All results posted below utilize Llama 3.1 8B Instruct.

We also performed an experiment evaluating against Minference and Pyramid KV on LongBench using the same truncation setting as we used for the quadratic models in Figure 9 due to the fact that Minference utilizes dense quadratic attention for every decoding step and Pyramid KV requires a quadratic prompt.

We find that our cascading KV cache outperforms Minference and Pyramid KV on every dataset except for Gov. Report which shows the tightest margin across all models and methods.

---

[Q2] How's your performance compared with others on RULER task, which is a benchmark for evaluating context size?

[Response] Compared to RULER, we believe that InfiniteBench [1] offers a more comprehensive and realistic assessment of model performance in long contexts. InfiniteBench evaluates tasks like multiple-choice reasoning, question answering, and summarization with average context lengths of 182K, 192K, and 171K, respectively. We used a total cache size of 32K. These tasks go beyond string recall, emphasizing a model's ability to extract and utilize meaningful information across large contexts—a critical capability for practical applications. We evaluate 3 datasets from InfiniteBench below:

We observe the same trend as seen in LongBench, whereby our model outperforms by a large margin (21.5% and 18.2%) on the first two datasets, and underperforms Minference by only 2.3% on the last dataset where the overall differences between models is much smaller.

---

Thank you for your time and effort in reviewing our work. Your thoughtful suggestions have improved the quality and completeness of our work. We hope we have adequately responded to your comments. If you have any remaining discussion points, please let us know. Thank you.

---

References

[1] Infinite Bench: https://arxiv.org/pdf/2402.13718

[Q1] Related to weakness 1, is [...] compatible with continuous batching? [...] compatible with actual serving systems, e.g., vllm/sglang?

[Response] Yes, our method is compatible with continuous batching because it can be adopted into PagedAttention frameworks (vLLM, SGlang). We intend to provide these integrations in a future release of our code. If you foresee any way in which our method may be problematic with continuous batching, then please let us know for further discussion.

---

[Q3] Which version are you comparing to? FlashAttention-1,2,3? Triton or cuda kernel?

[Response] The models which state vanilla quadratic Flash Attention use the official FlashAttention-2 cuda kernels. For our model which utilizes FlashAttention, we use a modified FlashAttention2 Triton kernel. We have added this distinction to the text whenever referencing flash attention, and also added the distinction between Triton and cuda kernels at the end of section 4 paragraph 1.

---

[Q4] What would be the memory/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? [...] add a pseudo code section to illustrate how this part is done.

[Response] Since the EMA is a single scalar for every item in the cache, the memory overhead would be 2 bytes per vector which is stored in the cache (the total cache size is a constant). The computation of the EMA in the inner loop consists of a simple instantiation of the vector of coefficients, and a sum over the columns of the current block of $QK^\top$ in the inner loop. The final EMA update is performed using an atomic sum.

This operation can be fused into the online softmax. The main obstacle to overcome when fusing into the online softmax is the normalization constant $c$, which is not known at each block of computation in the flash attention kernel. We overcome this with a simple approximation which assumes that the remaining mass in the normalization constant is proportional to the currently accumulated mass. Concretely, if there are are 10 total KV blocks $\text{TB} = 10$ in the flash attention computation and we have already computed 5 blocks $\text{CB} = 5$, then we adjust the current normalization constant as:

$$c^\prime = c + c \frac{\text{TB} - \text{CB}}{\text{CB}}$$

which is equivalent to $c^\prime = c + c ((10-5)/5) = 2c$ when computing the EMA after computing block 5 in the example above. We have added pseudocode for this process in algorithm 3 in our updated paper. Note that figure 3 shows the above approximation is effective compared with no token selection. Additionally, we performed a further experiment on LongBench to evaluate whether a sum over attention scores is better than the EMA. We find that the EMA performs best.

---

Thank you for your time and effort in evaluating our work. Your suggestions have improved the quality and completeness of our work. We hope we have adequately addressed the points you raised in your review. If you have any remaining discussion points, please let us know. Thank you.

---

References

[1] Infinite Bench: https://arxiv.org/abs/2402.13718

Thank you for your time and care in reviewing our work, and highlighting the strengths of our prefill method, as the prefill stage is not considered in many prior works. Please see our response to your other comments below.

---

[W1] 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.

[Response] Our cache allocates a contiguous memory buffer at initialization. Therefore, All sub-cache buffers are actually part of the same contiguous buffer, with boundaries being set analytically. By doing this, we mitigate any possible memory fragmentation issues. However, due to the update procedure, the stored tokens will not necessarily be in sorted order. Therefore, we also track the integer positional indices as well in order to apply the correct positional encodings.

---

[W2] 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

[Response] Our cascading cache is optimized for scaling to extremely large contexts on a fixed memory budget, a scenario where Minference struggles due to its memory overhead. In terms of latency, our cascading cache is comparable with Minference at different context lengths, However minference exceeds the memory of an A6000 (49GB) after 131K context length which is the environment we ran our experiments in. The full model latency for all layers (seconds) can be seen in the table below. All additional experiments below utilize Llama 3.1 8B Instruct models.

We compare Minference and Pyramid KV with our cascading method on LongBench, we truncate the inputs in the same manner as for the quadratic model in Figure 9 due to the fact that Minference utilizes full dense attention for each decoding step and Pyramid KV utilizes a quadratic prompt. We find the following results:

The only dataset where Minference outperformed our model was on Gov. Report which by far shows the narrowest gap between models across all experiments shown in figure 9 in our paper.

Additionally, we added a comparison to Minference on 3 datasets from InfiniteBench [1] which evaluates tasks like multiple-choice reasoning, question answering, and summarization with average context lengths of 182K, 192K, and 171K, respectively. WE used a total cache size of 32K. Like LongBench, these tasks evaluate a model's ability to extract and utilize meaningful information across large contexts—a critical capability for practical applications.

We observe the same trend as seen in LongBench, whereby our model outperforms by a large margin (21.5% and 18.2%) on the first two datasets, and underperforms Minference by only 2.3% on the last dataset where the overall differences between models is much smaller.

---

Continued in next comment...