Training Free Exponential Context Extension via Cascading KV Cache

## Abstract

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.

## 1 Introduction

Figure: Figure 1: Attention matrices from Streaming LLM (Xiao et al., ) and Cascading KV Cache (Ours), both with the same total cache size.
Refer to caption: x1.png

Large language models (LLMs) have become indispensable in a wide range of applications, from natural language processing to AI-driven content generation. However, their deployment is often hindered by the significant computational resources required, particularly during the quadratic attention operation in inference. Despite recent advancements like Flash Attention, which reduce memory overhead, the quadratic growth of latency and compute costs with input size remains a bottleneck, especially when processing long input sequences. This challenge is exacerbated in streaming applications, where high latency directly impacts user experience and operational costs.

Existing methods, such as sliding window approaches (Beltagy et al., ; Jiang et al., ), attempt to manage long sequences but impose a static limit on the model’s ability to retain context due to the fixed window size. As a result, valuable tokens are naively discarded once they fall outside the window, leading to irreversible loss of context. While some recent methods have tried to stabilize this process by preserving a small number of initial tokens (Xiao et al., ) (Figure 3 top), they still operate within a rigid, static framework with a naive eviction policy that fails to account for the relative importance of tokens within the sequence. Additionally, the linear inference procedure from Xiao et al. () processes a single token per step which makes it impractical to apply the method during the prompting stage when many tokens need to be processed together in parallel.

Figure: Figure 2: Passkey retrieval accuracy up to 1M tokens given a cache size of 65K. Our Cascading cache maintains higher accuracy even after four doublings of the context length.
Refer to caption: x2.png

Our work addresses these critical limitations by proposing a novel linear inference approach that extends the effective context length of sliding windows without increasing computational complexity or requiring additional training. We introduce a dynamic caching mechanism that organizes the key-value (KV) cache into cascading sub-caches, each designed to selectively retain tokens based on their historical importance. This approach allows for a more intelligent eviction policy, preserving tokens that are likely to play a crucial role in future predictions.

To demonstrate the effectiveness of our method, we present two compelling examples: First, we visualize the attention matrix on the PG19 test set, showing how our cascading KV cache retains context far beyond the reach of a static sliding window, while maintaining the same total cache size (Figure 1). Second, we provide an example on passkey retrieval up to 1M tokens given a total cache size of 65K (Figure 2), illustrating that our method consistently outperforms static window attention in terms of retrieval accuracy, maintaining superior performance over prior work even after four doublings of the context size. These results underscore the potential of our approach to transform the efficiency and accuracy of LLMs in both research and real-world applications.

Our contributions are as follows:

- •
We introduce a simple yet powerful training-free linear attention modification to sliding window attention on pretrained quadratic transformers that selectively retains important tokens, significantly extending the effective context length.
- •
Given the same total KV cache size, our method delivers substantial improvements of 12.13% average in long context benchmarks (LongBench), 0.4% in streaming perplexity (PG19), 4.48% in Book Summarization, and increases passkey retrieval accuracy by 24 percentage points (pp) at 1M token after four doublings of the context size and 18pp higher accuracy after 5 doublings of the context size.
- •
We provide a linear prefill strategy that avoids the restrictive quadratic prompt complexity of previous works. Our strategy reduces latency on 1M tokens by a factor of 6.8 compared to flash attention.
- •
We provide an efficient implementation of our KV cache, which achieves a more than two order of magnitude speedup over Streaming LLM.

Figure: Figure 3: Comparison of Streaming LLM (Xiao et al., ) and Cascading Cache (Ours). Top: Streaming LLM stores fixed sink tokens (red) along with a sliding window of $N$ recent tokens. Bottom: Our method segments the cache into smaller cascading sub-caches, where each successive sub-cache conditionally accepts a fraction of tokens based on the magnitude of past attention scores. This simple technique allows for important tokens to remain in the cache for a longer time instead of being naively evicted too early.
Refer to caption: x3.png

## 2 Related Work

Previous research has explored sparse attention patterns to reduce computational costs, such as in BigBird (Beltagy et al., ), where a sliding window with random sparse context is applied. While effective in reducing time complexity, this approach relies on uninformative random attention patterns that lack adaptability. Linear attention mechanisms, including locality-sensitive hashing (Kitaev et al., ), kernel-based approximations (Choromanski et al., ), attention matrix compression (Lee et al., ), and token compression techniques (Munkhdalai et al., ; Kim et al., ; Mohtashami & Jaggi, ), offer promising alternatives but often require extensive retraining, limiting their broad applicability to existing models.

Our work builds on insights from Streaming LLM (Xiao et al., ), which identified a key phenomenon in transformers: as features progress through layers, attention scores increasingly concentrate on the initial tokens, referred to as ”sink tokens.” This phenomenon presents a challenge in traditional fixed sliding windows, as evicting sink tokens from the cache disrupts the attention score distribution, forcing the model to redistribute probability mass onto less ideal tokens. Xiao et al. () demonstrated that this distribution shift leads to significant performance degradation unless the sink tokens are retained. We extend this insight by introducing an efficient, scalable method that increases the effective context window while maintaining a static cache size, improving throughput during the prompt stage, and enhancing performance on long-sequence tasks.

Subsequent approaches, such as SnapKV (Li et al., ) and H2O (Zhang et al., ), have introduced score-based KV cache eviction policies, but these focus primarily on quadratic prompts, with cache utilization limited to the decoding phase. These methods lack a linear prefill strategy, which results in the prefill attention operation remaining quadratic. Our method diverges from these previous approaches by offering a fast linear prefill strategy that is compatible with a cache that selectively retains important tokens through dynamic sparsity. This makes our approach not only more adaptive but also significantly more efficient, providing a practical path to extend transformer models’ context length with linear complexity and without incurring the high costs of retraining.

## 3 Method

Notation. We use the common notation of a boldface lowercase letter to denote a vector ${\bm{x}}$ and a boldface uppercase letter to denote a matrix ${\bm{X}}$. A superscript $(l)$ denotes that an object belongs to layer $l$, where $l\in[1,2,\dots,L]$. To simplify notation for attention, we omit the head dimension, output projections, and multi-layer perceptron (MLP) transformations, please see  Vaswani et al. () for an overview regarding those topics. We refer to a generic cache as ${\bm{C}}$, using subscripts ${\bm{C}}_{K}$ and ${\bm{C}}_{V}$ to refer to key and value caches, respectively.

Attention. Let $S\in{\mathbb{N}}$ represent a token sequence length, where each token at layer $l$ is represented by a vector ${\bm{x}}_{i}^{(l)}\in{\mathbb{R}}^{d}$. The collection of tokens in the sequence can be represented as a matrix ${\bm{X}}^{(l)}\in{\mathbb{R}}^{S\times d}$. With $\sigma$ being the softmax function, and different learnable query, key, and value matrices ${\bm{Q}}^{(l)},{\bm{K}}^{(l)},{\bm{V}}^{(l)}\in{\mathbb{R}}^{d\times d}$ the standard attention operation is as follows:

$$ ${\bm{X}}^{(l+1)}=\sigma\left(\frac{1}{\sqrt{d}}\left({\bm{X}}^{(l)}{\bm{Q}}^{( l)}\right)\left({\bm{X}}^{(l)}{\bm{K}}^{(l)}\right)^{\top}\right){\bm{X}}^{(l) }{\bm{V}}^{(l)}\in{\mathbb{R}}^{S\times d}$ (1) $$

Key-Value (KV) Caching. During LLM inference, a single token is generated at each time step. Combined with a causality condition such that token ${\bm{x}}_{i}$ cannot influence ${\bm{x}}_{j}$ iff $i>j$, it becomes more efficient to cache the $K^{(l)}$ and $V^{(l)}$ projected tokens in each layer $l$ rather than recomputing the full set of key, and value projections at each generation time step. Specifically, with $\cup$ representing a concatenation operation along the $S$ dimension, the calculation of the attention operation for the current token ${\bm{x}}_{j}$ during inference becomes,

$$ ${\bm{x}}_{j}^{(l+1)}=\sigma\left(\frac{1}{\sqrt{d}}\left({\bm{x}}_{j}^{(l)}{ \bm{Q}}^{(l)}\right)\left({\bm{C}}^{(l)}_{K}\cup\left({\bm{x}}_{j}^{(l)}{\bm{K }}^{(l)}\right)\right)^{\top}\right)\left({\bm{C}}^{(l)}_{V}\cup\left({\bm{x}} _{j}^{(l)}{\bm{V}}^{(l)}\right)\right)\in{\mathbb{R}}^{d}$ (2) $$

After the concatenation operation, ${\bm{x}}_{j}^{(l)}{\bm{K}}^{(l)}$ and ${\bm{x}}_{j}^{(l)}{\bm{V}}^{(l)}$ are considered to be added to the ${\bm{C}}_{K}$ and ${\bm{C}}_{V}$ caches for the next iteration in the sequence.

Sliding window attention, as used by Streaming LLM, treats the KV cache as a fixed size buffer which must evict tokens from the cache when the cache reaches its storage limit. Considering a cache which is fully populated (denoted by an overbar $\overline{{\bm{C}}}_{K}^{(t)}$ and omitting the layer index $l$), when a new token comes in, the cache must drop the oldest token in order to make space for the new token such that,

$$ $\overline{{\bm{C}}}_{K}^{(t+1)}=\overline{{\bm{C}}}^{(t)}_{K}\cup{\bm{x}}_{j}{ \bm{K}}=\left[\left\{{\bm{x}}_{i}{\bm{K}},{\bm{x}}_{i+1}{\bm{K}},\dots,{\bm{x} }_{j-1}{\bm{K}}\right\}\cup\left\{{\bm{x}}_{j}{\bm{K}}\right\}\right]\setminus \left\{{\bm{x}}_{i}{\bm{K}}\right\}.$ (3) $$

The only difference between standard sliding window attention and sink cache of Streaming LLM (Xiao et al., ) is that the sink cache perpetually retains the first $\alpha$ tokens in the sequence such that if a full cache starts from the first token ${\bm{x}}_{i}{\bm{K}}$, it is not token ${\bm{x}}_{i}{\bm{K}}$ which is evicted, but token ${\bm{x}}_{i+\alpha}{\bm{K}}$, with $\alpha\in{\mathbb{N}}$ being a fixed hyperparameter. Both sliding window attention and sink cache have the benefit of linear complexity, and a fixed overhead computation and memory cost for the window, as the size of the cache can only grow to a fixed, predetermined amount. However, this has the downside of constraining the information available in the cache, which may be needed for later predictions in the sequence. To illustrate, imagine streaming generation of an entire book with an LLM. If the cache can only contain tokens representing the average length of one chapter, then important information from previous chapters may be forgotten which could prove to be crucial for later steps of generation.

Figure: Algorithm 1 Strided Prefill

### 3.1 Linear Prefill

A crucial limitation of Streaming LLM (Xiao et al., ), H2O (Zhang et al., ), and SnapKV (Li et al., ) lies in the handling of the prompt/prefill stage of the LLM. Each of the aforementioned works only considers that either 1) tokens are processed one-by-one during the prompt or 2) the prompt is processed with full quadratic attention and then proceeds to apply the relevant caching strategy during the generation phase. Being limited in this way, even a linear model like Streaming LLM exhibits slower latency than a quadratic prompt utilizing flash attention in the non-asymptotic regime (see Figure 6(b)). To bridge this gap, our method utilizes an attention kernel which processes fixed-sized chunks (strides) of the prompt in a single operation before adding the keys and values to our cascading cache. Specifically, for a sequence length of $S$, quadratic attention can be seen as having a stride size of $S$, Streaming LLM can be seen as having a stride size of $1$, and our prefill method can be seen as having a stride size $K\in[1,S]$ that is somewhere in between. Algorithm 1 contains pseudocode for our strided prefill process, and Figure 5 contains an illustration. Our strided prefill allows for a more than two order of magnitude decrease in latency compares to single token processing and reduces latency by a factor of $6.8$ over quadratic flash attention when processing 1M tokens. This significant improvement delivers the benefits of linear attention on realistic, non-asymptotic sequence lengths.

Figure: Figure 4: Top: Each successive sub-cache window accepts a fraction of tokens evicted from the previous sub-cache. Bottom: At the boundaries between sub-caches, there are four possible cases where our method takes a different conditional action, creating a dynamic attention pattern. Circular buffers are not depicted for simplicity of visualization.
Refer to caption: x4.png

### 3.2 Cascading KV Cache

Our cache is a generalization of sliding window attention which allows for important historical tokens to be kept in the cache history for a longer period of time. To accomplish this, we view a fixed sized sliding window KV cache ${\bm{C}}$ with cardinality $|{\bm{C}}|$ as a collection of sub-caches ${\bm{C}}_{i}$ for $i\in[1,\dots,N]$, each with cardinality $|{\bm{C}}_{i}|\leq|{\bm{C}}|$ such that $|\bigcup_{i=1}^{N}{\bm{C}}_{i}|=|{\bm{C}}|$. Assuming a sub-cache is full, each full sub-cache $\overline{{\bm{C}}}_{i}$ evicts tokens when a new token is accepted into the sub-cache. However, each sub-cache accepts new tokens at a different rate, which in turn means that tokens may be discarded between the sub-caches and not solely at the end of the total cache window. An example of this process is depicted in Figure 4 (top), where the blue cache accepts all new incoming tokens. The green sub-cache, however, accepts only half of the tokens which are evicted from the blue sub-cache (every $2^{\text{nd}}$ iteration). The yellow sub-cache accepts half of the tokens evicted from the green sub-cache (every $4^{\text{th}}$ iteration) and so on.

Figure: Figure 5: Our strided prefill. We first compute attention for a chunk (stride) of new queries and new + cached keys, forming a rectangular slice of the attention matrix at each step.
Refer to caption: x5.png

As Xiao et al. () discovered, the initial attention sink tokens are vital for the stability of generation. Therefore, we too keep a fixed sub-cache of the initial tokens as attention sinks. With our method, a sink cache is a special case of a cascading cache, where the number of cascades (sub-caches) is set to 1, and it accepts all incoming tokens (*i.e*. using only the blue cache and red cache in Figure 4).

Token Selection. The process outlined in the preceding paragraph would result in a fixed heuristic pattern of tokens being dropped. However, such a heuristic pattern is not ideal, as it may be the case that the model naively holds onto tokens with limited value, while discarding important tokens. To remedy this, instead of discarding tokens naively, we dynamically select the most important tokens to retain by tracking the average attention score each token receives throughout time via an exponential moving average (EMA). We then selectively discard the token with the lower attention score EMA where possible. Given a hyperparameter $\gamma\in[0,1]$, and a vector of attention scores for all keys in the cache $\bm{s}_{k}^{(t)}$ at timestep $t$, we update the stored average attention scores $\bm{\mu}^{(t)}$ as, $\bm{\mu}^{(t+1)}=\gamma\bm{\mu}^{(t)}+(1-\gamma)\bm{s}_{k}^{(t)}$.

Consider the two sub-caches depicted in cases 2-4 of Figure 4. The blue sub-cache must accept all incoming tokens, while the green cache only accepts tokens every other iteration. Let the green cache be accepting tokens at the current timestep (case 2). When a new token comes in, it is added to the blue cache, which must then evict a token as the cache is full. The evicted token then goes to the green cache which accepts it unconditionally. The same process repeats on the next iteration, however, this time the green cache will not be accepting tokens (cases 3-4). At this step, when the blue cache accepts and subsequently evicts a token, we compare the attention score of the token evicted from the blue cache with the attention score of the newest token in the green cache. The token with the higher attention score is set (or remains) as the newest token in the green cache, while the token with the lower attention score is discarded. For a full pseudocode algorithm that covers cases 1-4 in Figure 4, see Algorithm 2 in the appendix.

Positional Encoding. We use the same positional encoding strategy as Streaming LLM, which reapplies positional encodings for each token relatively and consecutively, always starting with the 0 index. With $\text{PE}_{\text{index}}:{\mathbb{N}}\mapsto{\mathbb{N}}$ signifying a positional encoding index mutation function, and $i,j\in{\mathbb{N}}$ representing the original token positions in the input sequence, it follows that $i<j\iff\text{PE}_{\text{index}}(i)<\text{PE}_{\text{index}}(j)$. For example, if our cascading cache holds the token indices from the original sequence $[0,1,3,5,7,8]$, they would in turn receive positional encodings $[0,1,2,3,4,5]$.

Circular Buffers. Previous approaches have relied on tensor concatenation during cache add operations. However, this results in excessive copying operations, as each concatenation requires retrieving and storing all entries into a block of memory. We utilize more efficient circular buffers by tracking the location of the oldest token $\xi^{(t)}$ at timestep $t$ (where insertion should occur). We then increment $\xi$ after each insertion to the buffer such that $\xi^{(t+1)}=\left(\xi^{(t)}+1\right)\mod|{\bm{C}}_{i}|$. This way, at timestep $t$, we know that the oldest token in the buffer resides at $\xi^{(t)}$. When it is time to remove a token, we simply overwrite the content of the memory address of $\xi^{(t)}$ rather than performing a costly concatenation.

Cache Token Span. Given previously outlined cascading cache, we have lengthened the span of the tokens which currently reside in the cache. The process outlined above effectively extends sliding window context length by allowing older tokens to remain as keys and values for a longer time, with gaps between tokens. Assuming that each sub-cache has the same capacity (*i.e*. $|{\bm{C}}|=|\bigcup_{i}^{N}{\bm{C}}_{i}|=N|{\bm{C}}_{1}|$), and each will accept tokens with a different frequency function defined as $\frac{1}{f(i)}$, then the approximate context length will be a summation over the cache sizes and the inverse of the frequency functions. For example, if each successive sub-cache accepts half of the tokens evicted from the previous cache, the total span of the cache $\tilde{S}$ becomes,

$$ $\tilde{S}=\sum_{i=1}^{N}f(i)|{\bm{C}}_{i}|=|{\bm{C}}_{1}|\sum_{i=1}^{N}2^{i-1} =\frac{|{\bm{C}}|}{N}\sum_{i=1}^{N}2^{i-1}.$ (4) $$

We can then calculate the overall sparsity of the cache as $1-|C|/\tilde{S}$.

## 4 Experiments

Setup. We conduct experiments on streaming books (PG19) (Rae et al., ), Long Context Understanding (LongBench) (Bai et al., ), book summarization (Booksum) (Kryściński et al., ), and 1M token passkey retrieval. We evaluate our method with pretrained transformers from the Llama3.1 (Dubey et al., ) and Qwen2 (Bai et al., ) families of models. Please see Table 5 for model paths. In all our experiments, we keep the first 64 initial tokens as attention sinks. When considering the number of tokens in the cache, we always consider the sink tokens to be in addition to the cache size. Therefore a window size of $W$ has a total of $W+64$ tokens when accounting for the 64 sink tokens. We set the EMA parameter to $\gamma=0.9999$. We use the cascade token acceptance setting depicted in Figures 4 and 4, where each sub-cache accepts half of the tokens from the previous sub-cache. Unless otherwise indicated, our models use four cascading sub-caches. For experiments utilizing 8B model sizes, we use one NVIDIA A6000 (49GB) GPU, and for experiments utilizing 70B model sizes, we utilize 4 NVIDIA A100 GPUs. As Streaming LLM is a special case of our model, and 1 token per step is prohibitively slow, all results involving Streaming LLM results use our strided prefill strategy. Note that the strided prefill improves results over the original Streaming LLM with one token per step as shown in Figure 11.

Figure: (a) Cache Latency
Refer to caption: x6.png

### 4.1 Latency

We compare the latency of our cascading cache to the implementation from Xiao et al. () which uses tensor concatenation to add/evict tokens from the cache. Our implementation utilizes circular buffers and the Triton compiler (Tillet et al., ) to create an efficient CUDA kernel for the caching operation. To perform this experiment, we initialize a cache with 64 sink tokens and a total window size of $|{\bm{C}}|=16$K with 4 and 1 cascades, which are equivalent to our Cascading KV Cache and Streaming LLM, respectively. We also initialize the original Streaming LLM that uses concatenation. We then process a total of 16K tokens into the cache, and report the cumulative time spent on caching operations. Our method with one cascade (equivalent to Streaming LLM) takes just $0.01\%=1/10000$ of the total caching time of the original Streaming LLM, and our method with 4 cascades takes just $0.038\%=3.8/10000$ of the total caching time in Figure 6(a). In Figure 6(b), we show the overall attention latency, including caching, for a single attention layer processing 1M tokens. Our model uses a strided prefill of 4K with a total cache size of 16K. Our method is the only method which processes 1M tokens faster than flash attention, and takes only $14.8\%$ of the time of quadratic flash attention. In other words, flash attention is 6.8 times slower than ours for processing 1M tokens. We also study the effect of the size of our strided prefill on overall 1M token latency in Figure 6(b), finding that a stride of 4K delivers the lowest latency.

### 4.2 PG19

We measure perplexity on the PG19 (Rae et al., ) test set consisting of full-length books. Each book is streamed independently from start to finish without concatenation. We compare against a quadratic flash attention model, as well as BigBird and Streaming LLM, which are also training-free inference adaptations. We use three cache sizes of $16$K, $32$K, and $65$K with a strided prefill of 1K. Since our cache is equivalent to Streaming LLM sequence lengths that are less than the total cache size, we only run each experiment on the subset of books which exceed the given cache size. Flash Attention and Bigbird would exceed the GPU memory capacity, therefore, they are limited to processing books in chunk sizes equivalent to the total cache size. Results are displayed in Table 1. Our method delivers a consistently lower perplexity for all tested cache and model sizes. Additionally, we show examples of how our model behaves during the streaming process in Figure 7. After the cache size is exceeded, our eviction policy begins to differ from Streaming LLM, and our model tends to show lower perplexity due to the increased total token span in our cascading cache. Additional plots for Llama3.1 on all books exceeding 65K length can be seen in Figures 14 and 15.

**Table 1: PG19 Perplexity. Across all tested cache sizes, our cascading model maintains lower perplexity than baselines. Flash Attention and Bigbird operate by stepping through the entire sequence with a stride equivalent to the total cache size. They perform attention at each step until reaching the end of the sequence. The Qwen model is excluded from 65K cache size due to having only 32K positional embeddings.**
| Total<br>Cache Size | Num.<br>Books | Token<br>Count | Model | Methods / Perplexity ($\downarrow$) |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Flash Attn.<br>(strided) | Big Bird<br>(strided) | Streaming<br>LLM | Cascading KV<br>Cache (Ours) |  |  |  |  |  |  |  |
| 16384 | 91 | 9.78M | Qwen2${}_{\text{7B}}$ | 9.50 | (+0.36) | 10.65 | (+1.51) | 9.18 | (+0.04) | 9.14 |
| LLaMA3.1${}_{\text{7B}}$ | 8.08 | (+0.36) | 12.76 | (+5.04) | 7.78 | (+0.06) | 7.72 |  |  |  |
| 32768 | 77 | 9.42M | Qwen2${}_{\text{7B}}$ | 9.26 | (+0.26) | 10.38 | (+1.35) | 9.05 | (+0.02) | 9.03 |
| LLaMA3.1${}_{\text{8B}}$ | 7.86 | (+0.26) | 11.33 | (+3.73) | 7.65 | (+0.05) | 7.60 |  |  |  |
| 65536 | 62 | 8.25M | LLaMA3.1${}_{\text{8B}}$ | 7.73 | (+0.13) | OOM | (-) | 7.61 | (+0.01) | 7.60 |

Figure: Figure 7: PG19. For Llama with a total cache size of 65K, our cascade model stays equivalent to Streaming LLM until the cache size is exceeded and our eviction policy and token selection begin. Our intelligent eviction policy leads to better perplexity for the total stream. Additional figures can be seen in Figures 14 and 15.
Refer to caption: x10.png

**Table 2: Booksum book summarization. Among linear baseline models, our Cascading KV cache offers a consistent improvement. Averaged over all models and metrics, ours performs 4.48% better than linear baselines.**
| Method | Complexity | Llama 3.18B | Qwen 27B |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Rouge-1 | Rouge-2 | Rouge-L | Rouge-1 | Rouge-2 | Rouge-L |  |  |
| Vanilla Transformer | $\mathcal{O}(N^{2})$ | 35.37 | 7.29 | 21.24 | 31.70 | 5.70 | 17.8 |
| Snap KV | $\mathcal{O}(N^{2})$ | 36.80 | 8.11 | 21.63 | - | - | - |
| H2O | $\mathcal{O}(N^{2})$ | 35.62 | 7.42 | 21.16 | 31.10 | 5.47 | 17.58 |
| BigBird | $\mathcal{O}(N)$ | 33.36 | 6.55 | 18.59 | 20.83 | 2.31 | 12.62 |
| Streaming LLM | $\mathcal{O}(N)$ | 33.04 | 6.04 | 19.85 | 29.14 | 4.51 | 16.91 |
| Cascade (Ours) | $\mathcal{O}(N)$ | 34.47 | 6.63 | 20.52 | 30.34 | 5.02 | 17.54 |

Figure: (a) Ours 32K Total Cache Size
Refer to caption: x14.png

### 4.3 Passkey Retrieval

We perform passkey retrieval on both Streaming LLM and our Cascading KV Cache. For this experiment, we generate a random 5 digit passkey which is hidden in a random uniform point in the total sequence length. The rest of the text consists of random English words from the dictionary. We perform 20 trials for each insertion location range and sequence length for a total of 600 retrievals. We calculate accuracy for each digit, counting a correct digit prediction if it falls in the proper place in the output sequence. Therefore, a model which outputs random digits would receive an accuracy of 10%. We evaluate total cache sizes of 32K and 65K and sequence lengths with 8 cascades and a strided prefill of 4K. Context lengths start from 32K and double until we reach 1M tokens. Results are shown in Figure 8. For both 32K and 65K cache sizes, Streaming LLM begins to show near random accuracy after the first doubling, while our Cascading KV Cache is still better than random after 4 doublings of the context length.

### 4.4 LongBench

We evaluate our method against other linear scaling models on the same subset of tasks as (Xiao et al., ) in the LongBench long context understanding benchmark (Bai et al., ). We limit the total cache size of each model to be approximately 1/4 of the original prompt length $L$ of each input using the function $L^{\prime}=2^{\lfloor\log_{2}(L/4)\rfloor}$ and uses a strided prefill of 512. The results are displayed in Figure 9. Averaged over all datasets and tested models, our cascading cache improves performance over the next best model by 12.13%. For tabular results, please see Table 6.

Figure: (a) Llama 3.18B
Refer to caption: x18.png

### 4.5 Visualization

To visualize the effect of our method on the attention matrices, we reconstruct the full attention matrices of both Streaming LLM and our Cascading KV Cache using Llama3.1 8B on the first 8K tokens of the first book of the PG19 test set. We use a total cache size of 2048 and 4 cascades with a strided prefill of 256. The attention matrices are displayed in Figure 10. Naive sliding window attention Figure 10 (a,c) forms short static barrier where tokens are evicted regardless of their importance. Our method Figure 10 (b,d) maintains those tokens in the context history for a longer time where they may retain influence over future predictions, effectively increasing the available context.

We demonstrate the cascade boundaries of our proposed KV eviction policy in Figure 10(b), where the sparsity of each cascade gradually increases with the size of the cascade.
Additionally, our method preserves the attention patterns more thoroughly than Streaming LLM, such as the annotated preserved key value in Figure 10(d) which falls well outside of the range of the sliding window pattern. For more attention visualizations, please see Figure 17 in the appendix.

Figure: Figure 10: Attention matrix reconstruction for Sink Cache and our Cascading Cache. Both methods result in $\mathcal{O}(n)$ inference time complexity with the same total cache size ($|{\bm{C}}|=2048$).
Refer to caption: x23.png

**Table 3: The token selection process outlined in Section 3 is crucial for creating dynamic attention patterns.**
|  | $|C|=2048$ |
| --- | --- |
| KV Cache | LLaMA3.1${}_{\text{8B}}$ |
| Streaming LLM | 8.03 |
| Ours w/o token selection | 8.03 |
| Ours w/ token selection | $\mathbf{7.88}$ |

### 4.6 Ablation Study

To study the effect of different parts of our model, we provide three ablation studies including the effect of the strided prefill and token selection using the first book of PG19, and the effect of sparsity induced by the number of cascades. The effect of the strided prefill is shown in Figure 11. We find a decrease in perplexity with an increasing stride size. Intuitively, this comes from the fact that a larger stride provides a larger dense window at the leading edge of the attention matrix as shown in Figure 5. We study the token selection process outlined in Section 3 in Table 3 and find that without the token selection process, our model matches the performance of Streaming LLM, which highlights the importance of selecting higher scoring tokens. Lastly, we study the effect of the number of cascades, and thus overall sparsity, in Figure 12. For this experiment, we use a total cache size of 4K and consider context lengths from 4K which double until 262K. We average the passkey retrieval accuracy over all insertion locations. We find that accuracy steadily increases until the number of cascades exceeds 8 (more than 98% sparse).

Figure: Figure 11: Streaming LLM with our strided prefill achieved a progressively better perplexity and latency (Figure 6(c)) when increasing the stride due to a larger region of dense attention.
Refer to caption: x24.png

Interestingly, the token span in the cache remains a good predictor of accuracy for a moderate number of cascades. Given a token span, we may roughly calculate the expected accuracy in Figure 12(b) by considering the probability that the passkey falls within the span of the tokens. For example, with a token span of 1024, and a context size of 2048, we would expect an accuracy of approximately 50%. We find the token span to be a reliable predictor of accuracy until 4 cascades (73% sparsity).

Figure: (a) Actual Accuracy
Refer to caption: x25.png

## 5 Limitations & Future Work

As outlined in Equation 4 and visualized in Figure 12, our model’s overall sparsity increases as the number of cascades $N$ grows. This introduces a trade-off, as demonstrated in Figure 12(a), where performance improves up to a point, but diminishes once the number of cascades exceeds eight. Therefore, while our method enables substantial context length extrapolation, this process is not unbounded. Eventually, tokens must be discarded to maintain linear inference complexity, which inherently limits the scope of extrapolation. A promising direction for future work involves addressing the need to discard tokens while preserving linear complexity. One potential solution could involve developing methods for logarithmic complexity searches within the KV cache. By efficiently identifying the top-k relevant key-value pairs, such an approach could eliminate the need for token eviction and allow the model to maintain an overall complexity of $\mathcal{O}(N\log N)$. This would open new avenues for expanding transformer models’ context memory without compromising efficiency.

## 6 Conclusion

In this paper, we introduced a novel, training-free method for extending the context memory of streaming LLMs, offering significant improvements without increasing computational complexity. Our approach treats the fixed-size KV cache as a series of cascading sub-caches, allowing for dynamic token retention based on their historical importance. By selectively preserving high-impact tokens and evicting less critical ones, our method effectively extends the context window far beyond the limitations of traditional sliding windows. Our results demonstrate clear performance gains: a 12.13% average improvement on LongBench, a 4.48% boost in Book Summarization tasks, and superior passkey retrieval accuracy at 1M tokens, maintaining a significant edge even after four doublings of the context size. Additionally, our linear prefill strategy eliminates the quadratic complexity of previous approaches, achieving latency reduction by a factor of 6.8 compared to flash attention. These advancements highlight the potential of our method to significantly enhance the efficiency and accuracy of LLMs, making it a practical and impactful solution for both research and real-world applications that require long-context processing.

## 7 Reproducibility Statement

In order to aid in reproducibility of our experiments, we have provided our code which has been zipped into the supplementary file. We also provide exact pretrained model URL’s which are listed in Table 5. We provide an algorithm for our strided prefill method in Algorithm 1 and a full algorithm for our Cascading KV Cache in Algorithm 2. We have explained the parameters and computation budgets for all experiments in Section 4. As our method is deterministic and requires no stochastic training process, we have omitted error bars in our results.

## Appendix A Appendix

- •
Appendix B - Extra ablation study on head policy and head reduction function.
- •
Algorithm 2 - The algorithm of our cascading cache method
- •
Table 5 - Paths to exact pretrained models used in our experiments.
- •
Table 4 - MMLU experiment
- •
Table 6 - LongBench tabular data (displays the same data as Figure 9)
- •
Figures 14 and 15 - Extra PG19 plots for the subset of PG19 which exceeds 65K in length.
- •
Figure 16 - Quadratic Llama3.1 passkey results.
- •
Figure 17 - Additional attention matrix plots in higher resolution.

## Appendix B Head Policy and Head Reduction.

When making a decision for token selection, we may either apply the same homogeneous decision across all heads. Likewise, we may allow the heads to behave independently as illustrated in Figure 13(a). Additionally, as models may make use of Grouped Query Attention (GQA) (Ainslie et al., ), the number of attention heads may differ between queries and keys. Therefore, for both cases of homogeneous and independent heads, we need to select a head reduction function which will reduce the head dimension in the attention matrix to 1 (homogeneous heads) or K (key-value heads in GQA). We perform an ablation on the PG19 dataset to explore different options of head reduction functions and head policies in Figure 13(b). We find that in all cases, independent heads outperform homogeneous heads. Among the independent heads, we find that mean and max reductions resulted in similar performance, while a median reduction resulted in slightly worse performance.

Figure: (a) Head Policy Illustration
Refer to caption: x29.png

**Table 4: MMLU. We find that our Cascading KV Cache outperforms all linear models overall.**
| Model | Method | Humanities | Soc. Science | Other | STEM | Overall |
| --- | --- | --- | --- | --- | --- | --- |
| Llama 3.18B | Flash Attention (Full Context) | 61.59 | 76.47 | 73.09 | 56.01 | 67.00 |
| Vanilla (Truncated) | 61.23 | 76.08 | 73.06 | 55.79 | 66.80 |  |
| Streaming LLM | 61.57 | 76.37 | 73.19 | 55.92 | 66.98 |  |
| Cascade (Ours) | 61.45 | 76.63 | 73.25 | 56.23 | 67.11 |  |
| Qwen 27B | Flash Attention (Full Context) | 63.12 | 80.40 | 74.64 | 64.07 | 70.99 |
| Vanilla (Truncated) | 62.61 | 80.21 | 74.64 | 64.12 | 70.81 |  |
| Streaming LLM | 63.04 | 80.47 | 74.61 | 63.84 | 70.86 |  |
| Cascade (Ours) | 62.93 | 80.40 | 74.57 | 63.88 | 70.87 |  |

Figure: Algorithm 2 Cascading Sink Cache Algorithm (repeat for keys and values)

Figure: Figure 14: Additional plots for all of the books in the subset of PG19 books which exceeds 65K in length. The model used is Llama 3.1 8B. This chart is a complement to Figure 7
Refer to caption: x32.png

Figure: Figure 15: Additional plots for all of the books in the subset of PG19 books which exceeds 65K in length. The model used is Llama 3.1 8B. This chart is a complement to Figure 7
Refer to caption: x59.png

Figure: Figure 16: Passkey results for vanilla Llama3.1 as measured on an NVIDIA A6000 GPU with 49GM memory. According to the llama whitepaper (Dubey et al., ), it was finetuned for 131K positional embeddings until it achieved 100% accuracy on a passkey retrieval task. Therefore, given enough memory, it should achieve 100 for one more doubling of the context. Our model however, can extend high accuracy numbers past 131K with lower memory usage (see Figure 8).
Refer to caption: x86.png

**Table 5: Huggingface model paths used in our experiments.**
| Model | Huggingface Path | Experiment |
| --- | --- | --- |
| LLaMA3.1${}_{\text{8B}}$ (Dubey et al., ) | meta-llama/Meta-Llama-3-8B | PG19 |
| LLaMA3.1${}_{\text{8B}}$ Instruct (Dubey et al., ) | meta-llama/Meta-Llama-3-8B-Instruct | Booksum,Longbench,Ablation,Passkey |
| LLaMA3.1${}_{\text{70B}}$ Instruct (Dubey et al., ) | meta-llama/Meta-Llama-3-70B-Instruct | Longbench |
| Qwen2${}_{\text{7B}}$ (Bai et al., ) | Qwen/Qwen2-7B | PG19 |
| Qwen2${}_{\text{7B}}$ Instruct (Bai et al., ) | Qwen/Qwen2-7B-Instruct | Booksum,Longbench |
| Qwen2${}_{\text{72B}}$ Instruct (Bai et al., ) | Qwen/Qwen2-72B-Instruct | LongBench |

**Table 6: Tabular display of radar plot results from Figure 9. Higher scores are better ($\uparrow$). Total cache sizes are based on the length of the original prompt $L$.**
| Total Cache Size | Model | Cache | Narrative QA | HotPot QA | Qasper | Multi News | 2 Wiki MQA | Gov. Report | Mean |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| $2^{\lfloor\log_{2}(L/4)\rfloor}$ | LLaMA3.1${}_{\text{8B}}$ Instruct | Streaming LLM | 22.57 | 40.78 | 23.89 | 20.69 | 23.46 | 26.82 | 26.37 |
| Flash Attention (truncated) | 18.67 | 36.59 | 15.77 | 17.24 | 19.62 | 20.46 | 21.39 |  |  |
| Big Bird | 21.78 | 39.46 | 21.33 | 22.23 | 20.13 | 26.15 | 25.18 |  |  |
| Cascade (Ours) | 26.43 | 47.26 | 33.12 | 23.33 | 32.33 | 28.32 | 31.8 |  |  |
| $2^{\lfloor\log_{2}(L/4)\rfloor}$ | Qwen2${}_{\text{7B}}$ | Streaming LLM | 18.95 | 38.54 | 20.97 | 17.65 | 32.15 | 24.96 | 25.54 |
| Flash Attention (truncated) | 15.01 | 34.4 | 17.28 | 15.64 | 30.21 | 17.47 | 21.67 |  |  |
| Big Bird | 17.78 | 28.68 | 20.05 | 12.9 | 25.36 | 18.64 | 20.57 |  |  |
| Cascade (Ours) | 20.55 | 37.36 | 28.65 | 19.57 | 34.35 | 26.22 | 27.78 |  |  |
| $2^{\lfloor\log_{2}(L/4)\rfloor}$ | LLaMA3.1${}_{\text{70B}}$ | Streaming LLM | 25.72 | 42.39 | 24.5 | 20.93 | 31.78 | 27.67 | 28.83 |
| Flash Attention (truncated) | 24.62 | 39.77 | 19.93 | 17.26 | 24.23 | 20.51 | 24.39 |  |  |
| Big Bird | 27.02 | 52.0 | 28.63 | 23.0 | 31.51 | 27.53 | 31.61 |  |  |
| Cascade (Ours) | 30.3 | 52.61 | 36.97 | 23.04 | 46.82 | 29.3 | 36.51 |  |  |
| $2^{\lfloor\log_{2}(L/4)\rfloor}$ | Qwen2${}_{\text{72B}}$ | Streaming LLM | 20.8 | 50.08 | 20.68 | 17.97 | 43.83 | 27.83 | 30.2 |
| Flash Attention (truncated) | 17.69 | 43.25 | 14.96 | 15.68 | 40.08 | 21.35 | 25.5 |  |  |
| Big Bird | 23.55 | 48.41 | 25.89 | 18.54 | 44.58 | 30.35 | 31.89 |  |  |
| Cascade (Ours) | 25.0 | 53.78 | 29.48 | 20.17 | 48.22 | 29.4 | 34.34 |  |  |

Figure: (a) Layer 6
Refer to caption: x87.png

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 6
presentation: 2
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.
soundness: 2
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.
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.
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

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
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
soundness: 2
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
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.
weaknesses: The baseline may not be sufficient, most of the experiment are comparing with streamingLLM

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: I am wondering the perplexity results on more evaluation datasets.
soundness: 3
strengths: 1. Although the idea is simple, the experiment results are very good.
2. The idea is well implemented with careful code optimization.
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.
weaknesses: The experimental section would be better if the idea is evaluated on more data sets.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: Please refer to the weakness part for questions.
soundness: 3
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.
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.
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)

Meta-review: 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.

Decision: Accept (Poster)