Title: AttentionPredictor: Temporal Patterns Matter for KV Cache Compression
Venue: NeurIPS 2025
Accepted: yes
Keywords: Large Language Model

KV Cache Compression

Long Context
Abstract: With the development of large language models (LLMs), efficient inference through Key-Value (KV) cache compression has attracted considerable attention, especially for long-context generation.
To compress the KV cache, recent methods identify critical KV tokens through static modeling of attention scores. However, these methods often struggle to accurately determine critical tokens as they neglect the *temporal patterns* in attention scores, resulting in a noticeable degradation in LLM performance. 
To address this challenge, we propose **AttentionPredictor**, which is the **first learning-based method to directly predict attention patterns for KV cache compression and critical token identification**. 
Specifically, AttentionPredictor learns a lightweight, unified convolution model to dynamically capture spatiotemporal patterns and predict the next-token attention scores. An appealing feature of AttentionPredictor is that it accurately predicts the attention score and shares the unified prediction model, which consumes negligible memory, among all transformer layers. Moreover, we propose a cross-token critical cache prefetching framework that hides the token estimation time overhead to accelerate the decoding stage. By retaining most of the attention information, AttentionPredictor achieves **13$\times$** KV cache compression and **5.6$\times$** speedup in a cache offloading scenario with comparable LLM performance, significantly outperforming the state-of-the-arts. The code is available at https://github.com/MIRALab-USTC/LLM-AttentionPredictor.

AttentionPredictor: Temporal Patterns Matter for KV Cache Compression

Abstract

With the development of large language models (LLMs), efficient inference through Key-Value (KV) cache compression has attracted considerable attention, especially for long-context generation. To compress the KV cache, recent methods identify critical KV tokens through static modeling of attention scores. However, these methods often struggle to accurately determine critical tokens as they neglect the temporal patterns in attention scores, resulting in a noticeable degradation in LLM performance. To address this challenge, we propose AttentionPredictor , which is the first learning-based method to directly predict attention patterns for KV cache compression and critical token identification. Specifically, AttentionPredictor learns a lightweight, unified convolution model to dynamically capture spatiotemporal patterns and predict the next-token attention scores. An appealing feature of AttentionPredictor is that it accurately predicts the attention score and shares the unified prediction model, which consumes negligible memory, among all transformer layers. Moreover, we propose a cross-token critical cache prefetching framework that hides the token estimation time overhead to accelerate the decoding stage. By retaining most of the attention information, AttentionPredictor achieves 13 × KV cache compression and 5.6 × speedup in a cache offloading scenario with comparable LLM performance, significantly outperforming the stateof-the-arts. The code is available at https://github.com/MIRALab-USTC/LLMAttentionPredictor.

1 Introduction

Large language models (LLM) like OpenAI o1 [1] have shown impressive scalability and effectiveness in tackling complex tasks through long chain-of-thought (CoT) reasoning and multi-turn conversations [2-5]. However, these long-context tasks require LLMs to handle extremely lengthy contexts, presenting computational and memory challenges for LLMs [6]. Specifically, the key-value cache (KV cache), which holds the attention keys and values during generation to prevent re-computations, consumes huge GPU memory [7]. For example, for a model with 7 billion parameters, the parameters consume only 14 GB of memory whereas the KV cache requires around 72 GB with the 128K prompt length [8]. As the decoding latency and memory footprint scale with the KV cache, it is important to compress the KV cache as the prompt expands.

3

Spatiotemporal predictor

captures

dynamic

patterns

Temporal series

Spatial series

Ignore

direct

attention patterns

Statically

model

attention patterns

e.g. H2O

e.g. Quest

AttentionPredictor (Ours)

Sum of score

Query

t+1

t+1

0

t

t+1

✓

✓







✓

✓

✓

✓

Estimate

Predictor

Cache Eviction

Cache Retrieval

e.g. SeerAttention

Query

t+1

Training Based

Embedding

Layer

Embedding

Layer

Indirect

spatial-only

predictor

✓

 ✓

✓

Many attention sparsity-based methods have been proposed to compress the KV cache in the sequence dimension. Previous works have shown that a small portion of tokens dominate the attention distribution, substantially determining token generation precision [9-11]. Therefore, we can dramatically reduce the cache size by computing attention sparsely with only the critical tokens, while maintaining LLM performance. Cache eviction methods [11-13] use heuristic ranking with attention scores to identify critical keys and evict the less relevant ones. However, heuristic scoring methods can only model fixed patterns and struggle to identify critical tokens accurately, causing these methods to suffer from LLM performance degradation. Recently, SeerAttention [14] and Attention-Gate [15] use learnable modules to model dynamic patterns and retrieve critical tokens. However, they only represent keys or hidden states, rather than directly modeling the attention score distribution. As a result, they often exhibit limited accuracy after compressing KV cache (Table 8 for details). This motivates the need for a more effective solution that directly targets the prediction of attention scores for enhanced KV cache compression.

In our paper, we thus explore the importance of temporal patterns in attention scores for identifying critical KV cache. Previous research indicates that attention scores exhibit repetitive patterns intrinsic to LLMs [10, 16-18]. We further observe that attention patterns have temporal characteristics, such as re-access, sequential, and seasonal, as illustrated in Figure 2. Our theoretical analysis in Section 3.3 shows these patterns originate from intrinsic model properties such as query similarity and position encoding. These patterns are both stable and predictable across time, forming a two-dimensional (2D) temporal map with local continuity and translation-invariant characteristics.

KVcache compression with critical token identification can be formulated to attention score prediction which is a 2D time series prediction problem. Inspired by the theoretical understanding of attention mechanisms, we introduce AttentionPredictor , a time-series prediction approach to predict attention scores for critical token identification, as shown in Figure 1. Since temporal attention patterns are dynamic and difficult to capture with existing heuristic methods, we learn lightweight but accurate prediction models to address this challenge. Specifically, we design a convolutional module to capture 2D patterns in the input attention scores, allowing a single predictor to be shared across all transformer layers. These choices shrink our unified prediction model to roughly one millionth the size of the LLM (LLaMA-3.1-8B), resulting in negligible memory consumption. In contrast, SeerAttention's per-layer prediction model amounts to 101MB in total, whereas ours is only 21KB-just 0.02% of its size. Additionally, we employ distribution error calibration by periodically computing dense attention, and apply block-wise attention compression to further improve the efficiency of prediction. By accurately identifying critical tokens, AttentionPredictor retains most of the attention information after KV cache compression. To further accelerate LLM decoding, we apply AttentionPredictor to our proposed KV cache management framework, a cross-token KV cache prefetching system. In contrast to the existing cross-layer approach [19], our cross-token method can capitalize on longer transfer times, and adapt to more sophisticated and accurate methods for identifying critical tokens.

Our contributions: 1) Based on our observation of the temporal patterns in the attention score, we propose AttentionPredictor. To the best of our knowledge, it is the first learning-based method to directly predict attention patterns for KV cache compression and attention sparsity. 2) We propose the first cross-token prefetching framework, which effectively mitigates prediction delays and KV cache transfer latency in the LLM decoding stage. 3) Experiments demonstrate that our approach achieves 13 × KV cache compression with comparable LLM performance and 5.6 × speedup to full cache offloading.

2 Related works

2.1 Efficient LLM Inference

Several efforts have optimized pre-trained model inference efficiency from different perspectives. Speculative decoding methods [20-24] use smaller draft models to predict multiple future tokens for parallel validation by a target LLM. In contrast, our approach instead predicts the next token's attention scores to estimate KV cache compression. Inference systems [25-29] accelerate the prefilling stage with techniques like prefix cache. AttentionPredictor is compatible with them, for it is applied to the decoding stage. At the same time, our block-based sparsification can also be applied to paged attention [25] and other computational accelerations.

2.2 KV Cache Compression

Many methods compress the KV cache by leveraging the finding that attention scores are sparse, which allows for estimated sparse computation on high-score positions.

Cache eviction. These methods use heuristic approaches to identify key-value pairs of high importance and evict the less relevant ones. StreamingLLM [12] observes that the earliest tokens have high attention scores during inference, so it only retains the initial and the recent few tokens. H2O [11] accumulates all historical attention scores as an estimation, but suffers from the accumulation error of scores from the first few tokens being too frequent. SnapKV [13] accumulates attention scores within a recent window as an estimate and performs one-time filtering during the prefill stage. MInference [16] and FlexPrefill [30] inductively define several attention patterns and determine the compression strategy for each head to accelerate the prefilling stage. All these methods are heuristic-based and statistically model attention patterns. They struggle to capture the dynamic temporal patterns within attention scores accurately.

Cache retrieval. These methods, such as Quest [31], InfLLM [32], and PQCache [33], retrieve critical tokens by estimating attention weights over compressed key representations. However, such compression can degrade fidelity and introduce retrieval errors. In particular, Quest is sensitive to the page size, and the accuracy significantly drops with large page sizes and small budgets, as shown in Figure 8. Instead, our method employs a lightweight temporal predictor to estimate attention weights directly, yielding higher retrieval accuracy. Other methods involving KV cache quantization [18, 3436] and KV cache budget allocation [37, 8, 38, 39] are orthogonal to our token scoring approach and can be combined with our AttentionPredictor.

Training-based Approaches. Recently, a number of training-based cache compression approaches have been proposed. MoBA [40] extends the cache retrieval method by integrating sparse attention during training. Similarly, SeerAttention [14] and NSA [41] both train linear models to encode and retrieve key blocks more accurately. However, these approaches typically require training a separate model for each layer, which limits their scalability and generalization. In contrast, our framework employs a single plug-and-play module that can be applied uniformly to all layers and heads within the same LLM. Our model is only 0.02% the size of SeerAttention. Moreover, whereas NSA relies on the fine-tuning of LLMs to achieve optimal results, our method achieves comparable improvements without any additional fine-tuning, making it both more efficient and broadly applicable.

KV cache cross-layer prefetching. These methods hide part of the cache transfer time based on offloading the cache to the CPU. However, as the sequence length increases, the transfer time is too long to be hidden. InfiniGen [19] combines cache retrieval and prefetching by approximating the attention score for the next layer to load the critical cache. However, the estimation time increases significantly as the sequence grows, and the inference time for a single layer is insufficient to cover this. In contrast, our cross-token prefetching framework can hide longer estimation and transfer time within the per-token inference time.

3 Method

In this section, we introduce AttentionPredictor, the first learning-based method to directly predict attention patterns for KV cache compression and attention sparsity. , along with the cross-token prefetch framework for improved cache management. We begin with the problem formulation for attention prediction in Section 3.2, followed by the observation and theoretical analysis of temporal

attention patterns in Section 3.3. Then we introduce our novel AttentionPredictor in Section 3.4. Finally, Section 3.5 presents a cross-token prefetch framework that efficiently hides both evaluation and cache loading latencies.

3.1 Preliminary

In the language model decoding stage, we denote Q t ∈ R 1 × d , K ∈ R t × d as the query tensor and key tensor used for generate token t , respectively. Specifically, we denote K i ∈ R 1 × d , where i ∈ { 1 , 2 , . . . , t } , as the key tensor for token i , and K = K 1: t as the complete key tensor. The attention score at step t is calculated as: A t = Softmax ( 1 √ d Q t K ⊤ ) = [ a t, 1 , a t, 2 , . . . , a t,t ] ∈ R 1 × t .

The sparsity-based KV cache compression seeks to find a subset of keys with budget B that preserves the most important attention values. Specifically, the subset of key indices is S = { s 1 , s 2 , ..., s B } , where each s j is an index within the total t tokens. We define the attention recovery rate as:

which reflects the amount of information preserved after compression. A higher recovery rate R rec indicates less information loss caused by KV cache compression. Therefore, the goal of KV cache compression can be formulated as finding the index set S that maximizes R rec .

3.2 Problem Formulation

Accurate next-step attention score modeling is critical to maintain high attention recovery rate with limited KV tokens and improve final LLM generation performance with KV cache compression. Therefore, KV Cache compression with token selection can be formulated as attention score prediction, which is a two-dimensional (2D) time series prediction problem. Specifically, we predict the attention in step t +1 as ˆ A t +1 and select the critical token positions with S = TopB ( ˆ A t +1 ) . Since stacking attention vectors over time yields a 2D spatiotemporal attention map A t = ⟨ A i ⟩ t i =1 , where the temporal axis corresponds to decoding steps and the spatial axis corresponds to token positions in context as Figure 1. We can thus frame attention score prediction as a time series prediction problem. The input to our predictor is a 2D spatiotemporal sequence A H = ⟨ A comp i ⟩ t i = t -H +1 , representing the H most recent block-wise compressed attention scores, where A comp i is generated via max-pooling on the original attention score A i (detailed in Sec. 3.4). We then train a lightweight model to predict the attention scores for step t +1 as ˆ A t +1 = F ( A H ) , where F ( · ) denotes the model function.

3.3 Attention Temporal Patterns

Observation. We observe that attention exhibits consistent and structured temporal behaviors, which we term attention patterns . At a local level, these patterns are remarkably stable. We have identified three fundamental patterns that reflect strong invariances under spatiotemporal shifts. However, these patterns reveal a complex dynamic globally, exhibiting long-term evolution. To fully understand this dynamic, we will first provide an analysis of the stable patterns.

Three Attention Patterns. Our work is motivated by the observation that attention exhibits three distinct temporal patterns-re-access, se-

quential, and seasonal, as shown in Figure 2. These patterns collectively suggest that attention evolves with both temporal continuity and structural regularity across tokens. We unify the three local spatiotemporal attention patterns to the following approximate invariance under joint temporal and spatial shifts:

where and δ t and δ i represent temporal and spatial offsets, respectively. (1) The re-access pattern corresponds to ( δ t = 1 , δ i = 0) , where the model repeatedly attends to the same tokens. (2) The sequential pattern corresponds to ( δ t = 1 , δ i = 1) , reflecting progressive token-wise attention. (3) The seasonal pattern corresponds to ( δ t > 1 , δ i = 0) , indicating periodic recurrence of attention to fixed positions. While our re-access and sequential patterns resemble shapes defined in MInference [16], our time-series analysis also reveals a novel seasonal pattern, and our dynamic predictor improves prediction accuracy by 5% over fixed-template approaches (see Appendix C.10).

To better understand why attention follows spatiotemporal patterns, we investigate the underlying theoretical mechanisms. We find out that the high continuity between queries and position embedding plays a central role in shaping attention behavior.

High Query Self-similarity. We observe that the query embedding sequence { q i } t i =1 exhibits strong temporal self-similarity, which is consistent with prior findings in [19]. Specifically, the one-step cosine autocorrelation is 0.87 as detailed in Appendix A.1, indicating a high degree of alignment between queries. As only the most recent KV cache is updated, high query self-similarity indicates high stability of the attention score over consecutive time steps , which is highly predictable.

For theoretical clarity, we define the raw attention scores as ˜ A t = q t k ⊤ 1: t omitting both the 1 / √ d factor and the softmax. This allows us to focus on the underlying structure of attention logits, as softmax is a monotonic function and does not alter the relative ranking of attention weights. For step t +1 , supposing the query vector evolves as q t +1 = q t +∆ q , the attention at step t +1 is:

Focusing on the first t entries of ˜ A t +1 , we obtain:

where ∆ ˜ A = ∆ qk ⊤ 1: t . Noting that ∥ ∆ q ∥ is small due to the high autocorrelation, and the key matrix k 1: t remains relatively stable, the term ∆ ˜ A is of small magnitude in norm:

Therefore, ˜ A t ≈ ˜ A t +1 holds pointwise to a good approximation when the query exhibits high self-similarity. This implies that the relative ranking and distribution of attention weights over past tokens change slowly across decoding steps. In other words, attention exhibits temporal smoothness and inter-step consistency, which form the basis of the observed re-access and sequential patterns.

Position Embedding. The attention patterns are closely related to position embedding, particularly Rotary Positional Embedding (RoPE), which applies position-dependent rotations to queries and keys. In RoPE, the query and key vectors are partitioned into d/ 2 groups along dimensions and apply rotational matrices with different frequencies θ to each group. The attention weights are calculated with RoPE encoding as [17, 42]:

where q ( m ) i , k ( m ) j ∈ R 2 denote the components of q i and k j in the m -th group, θ m is the frequency parameter, and ϕ ( m ) is the angle between q ( m ) i and k ( m ) j . For the query and key are stable across steps, the formulation reveals that the inner product depends primarily on the relative position j -i . When j -i = ∆ is fixed, the cosine terms remain stable across decoding steps, producing consistent attention along diagonals, showing as the sequential pattern . Furthermore, the periodicity of the cosine function explains the periodic slashes of sequential patterns observed in Figure 2.

Dynamic Patterns. These foundational patterns are not static and exhibit dynamic evolution even within a single generation sequence. As the decoding process progresses, patterns can shift, appear, or fade as illustrated in Figure 2. This dynamism is driven by the decay of query-key similarity from inherent query shifting, and is especially evident in long-output tasks like reasoning.

Unified

Lightweight

Convolutional

Model

Compressed

Attention

𝐴 𝑡

Update

Next Step

Attention

መ

𝐴𝑡+1

Top-K

Critical Block

Indices

𝑆

AttentionPredictor

Attention

Score

𝐴𝑡

Cross-token

Prefetch

LLM

Decoding

Token

t

Token

t+1

Attention

Predictor

✓

✓

Block

Predict

𝐴𝑡

Stored

Attention

Score

𝒜𝐻

Spatial

Dimension

Temporal

Dimension

Block Indices

𝑆

LLM

Layer 1

Layer 2

Layer 1

︙

Layer 1

Layer 2

︙

Layer 1

Layer 2

︙

Full Cache

K

V

Layer 2

︙

Predict Attention with Time Series Methods. Since the query similarity and position embedding are inherent features of LLMs and independent of input, the resulting attention pattern is also inherent to LLMs. This allows a single shared predictor that can generalize across heads and datasets. For the pattern also changes as the sequence grows, this dynamic requires the predictor not only to recover the stable structures but also to adapt to shifts. Building on theoretical insights into attention, we propose that the pattern can be treated as a two dimensional spatiotemporal process suited to time series methods. In the next section, we introduce an architecture that uses this view to capture and forecast the evolution of attention over time.

3.4 AttentionPredictor: A Spatiotemporal Predictor

Overall Process. As shown in Figure 3 and Algorithm 1, AttentionPredictor prepares an attention history sequence A H in the prefilling stage, and predicts attention during the decoding stage. First, the A t from the LLM is compressed to A comp t using block-wise attention compression. Next, A H is updated with A comp t . The next step attention ˆ A t +1 is then predicted with the pretrained model F . From ˆ A t +1 , the top-K positions are selected with a budget of B/b , since ˆ A t +1 is in compressed form. Finally, the indices are expanded with b to obtain the final critical token indices S .

Model Design. To capture spatiotemporal features, we introduce a convolutional neural network (CNN) composed of two 2D convolution layers followed by a 1D convolution layer. The 2D convolutions capture spatiotemporal features at multiple scales, while the 1D convolution focuses on the time dimension, extracting temporal patterns across time steps. By replacing the fully connected layer with a 1D convolution kernel, the model adapts to the increasing spatial dimension, without data segmentation or training multiple models. Compared to an auto-regressive LSTM [43], the CNN is more lightweight and offers faster predictions, main-

Algorithm 1 Identify Critical Tokens

Input : Attention scores A t , Attention history series A H , Block size b , KV budget B Output : Critical KV token indices set S

1: Pad A t to the nearest multiple of b with zero

3: A H ← Concat ( A H [ -H +1 :] , A comp t )

2: A comp t ← MaxPooling ( A t , b )

4: ˆ A t +1 ←F ( A H )

5: S comp ← Top-K ( ˆ A t +1 , B/b )

6: S ← ⋃ i ∈ S comp { i · b, i · b +1 , . . . , i · b +( b -1) } Return KV cache indices set S

taining a prediction time shorter than the single-token inference latency. Additionally, when compared to a MLP [44] on time-series dimension, the CNN is more effective at capturing spatial features, which improves prediction accuracy.

Training Data and Strategy. Our model is both data-efficient and generalizable. We train the model only on a small subset of attention data, specifically approximately 3% extracted from the dataset. The model performance on the entire dataset shows our model effectively captures the patterns (see Section 4.2). Additionally, the temporal patterns of attention are inherent in LLMs, a single model can generalize well across various datasets. For example, our model trained on LongBench also performs well on the GSM8K dataset, highlighting the generalization capability of AttentionPredictor.

Block-wise Attention Compression. To speed up prediction, we apply attention compression before computation. By taking advantage of the attention's locality, AttentionPredictor predict attention and identify critical tokens in blocks. Inspired by Quest [31], we use the maximum attention value in each block as its representative. Specifically, max-pooling is applied on A t with a kernel size equal to the block size b , as A comp t = Maxpooling ( A t , b ) , reducing prediction computation to roughly 1 /b .

Distribution Error Calibration. Due to the sparsity of attention computation, the distribution of attention history A H used for prediction may deviate from the distribution of dense attention. This deviation tends to accumulate over decoding, particularly as the output length increases. To mitigate this issue and enhance prediction accuracy, we introduce a distribution error calibration technique to correct these deviations. Specifically, we store the full attention score every M steps, effectively balancing accuracy with computational efficiency. Note that the full attention is only used to update the history. During LLM decoding, we still use the sparse KV cache, strictly adhere to the budget. Additionally, we predict and update the position S every several steps for efficiency.

3.5 KV Cache Cross-token Prefetching

To address the increased memory cost of longer contexts, current LLM systems offload the KV cache to the CPU, but I/O transfer latency becomes a significant bottleneck in inference. KV cache prefetching offers a solution by asynchronously loading the sparse cache in advance. We introduce the cross-token KV cache prefetching framework, which differs from the cross-layer method [19] by leveraging longer transfer time budgets and enhancing data integration. They are helpful for improving prediction accuracy with larger models and improving transfering efficiency. Specifically, our implementation involves a prefetching process for each layer. As illustrated in Figure 3, during the decoding phase, AttentionPredictor forecasts the critical token indices S for the next step. The framework then prefetches the sparse KV cache with S for the next step from the full cache to the GPU. At the same time, LLM inference of other layers is also ongoing. Subsequently, the prefetched sparse KV cache is used to calculate the sparse attention of the next token. The timeline of cross-token prefetching can be seen in Figure 10 in Appendix B.1.

4 Experimental Results

4.1 Settings

Tasks. 1) We use the LongBench [45], InfiniteBench [46] and RULER QA [47] dataset for long context evaluation. LongBench is a widely used benchmark for long-context LLM inference. It consists of 16 datasets covering 6 tasks, including single- and multi-document QA (SQA and MQA), summarization, few-shot learning, synthetic tasks, and code completion. The metrics for each task are in Appendix B.3. InfiniteBench has up to extreme 124K average context length. 2) We employ the AIME [48] dataset, a tough mathematics dataset, for evaluating the performance of our method with the long-reasoning output scenario. 3) We use GSM8K math dataset to evaluate the long n-shot CoT task. 4) We use MMLU [49] and GPQA [50] dataset to evaluate the multi-choice task. 5) We use the Needle In A Haystack [51] experiment to evaluate the in-context retrieval capabilities.

Baselines and LLMs. We select four sparse attention approaches as our baselines. These include four cache eviction methods, StreamingLLM [12] and its advanced method Cascading [52], accumulative approaches H2O [11] and SnapKV [13], and two cache retrieval method Quest [31] and learningbased SeerAttention [14]. We use the official implementations of all baseline experiments. A detailed description of each baseline is provided in Appendix B.2. We choose two widely used long-context models for our evaluation: LongChat-v1.5-7b-32k [53] with 32K context length and LLaMA-3.1-8B-Instruct [54] with 128K context length.

Implementation details. We set the history step H to 64, the block size b to 16, and the calibration step M to 5. Performance analysis of these hyperparameters is discussed in Section 4.6. Follow Quest [31], we did not apply our method or any other algorithms to the first two layers of the LLM. Following the settings of H2O and StreamingLLM, We allocated the budget equally to the prefix and local tokens, assigning 64 tokens each. The remaining KV budget is allocated to intermediate tokens, determined by the prediction model. We conducted experiments on NVIDIA A800 (80GB) GPUs.

4.2 Main Results

Results on LongBench. We evaluate our method on various long-context tasks in the LongBench benchmark, with the KV cache budgets ranging from 1024 to 4096. We restrict H2O's attention to the past 64 steps to ensure fairness and denote this variant as H2O+. As shown in Table 1, AttentionPredictor surpasses the performance of all SOTA KV cache eviction and retrieval methods across various KV budgets and LLMs. Notably, the average performance loss compared to the full cache is less than 0.5% across all cache budgets. This demonstrates the ability of our method to effectively model attention patterns and precisely predict the locations of critical tokens. Furthermore, in an extremely sparse budget setting of 8% (1K/13K), our approach achieves full cache performance, while all other methods suffer at least a 1.25% reduction. Additional results of SeerAttention, long-context benchmark InfiniteBench and RULER QA are in Appendix C.2, C.3, and C.4.

1024

2048

KV Cache Budget

0

20

40

60

80

100

Accuracy (%)

30.0

46.7

16.7

43.3

53.2

Full Cache (80%)

H2O

Quest

AttentionPredictor

76.7

4K

8K

16K

32K

Context Length

0

50

100

150

200

250

300

Decode Latency (ms)

59

58

49

48

47

83

64

49 49

49

135

67

50

50

50

279

85

51

50

50

FlashAttention2

H2O-1/4 Budget

Ours-1/4 Budget

Ours-1/8 Budget

Ours-1/16 Budget

5.6x

Results on CoT reasoning. Reasoning significantly enhances the problem-solving capabilities of LLMs. For instance, QwQ-32B achieves an accuracy of 76.7% on the challenging AIME benchmark with the aid of detailed reasoning. However, generating complex answers leads to output length of up to 31.5K tokens. This sharply contrasts with typical long-context scenarios, which involve long inputs but short outputs. Such a shift introduces distinct challenges for KV cache compression during decoding. Under this demanding setting, we evaluate AttentionPredictor against two decoding-phase acceleration baselines, H2O and Quest. As shown in Figure 4, AttentionPredictor consistently outperforms both baselines. It achieves comparable accuracy of 76.7% with a KV budget ratio of 0.13 (2K/15K average), whereas Quest drops sharply to 43.3%. Moreover, our approach attains this accuracy with an average output of only 15K tokens, offering significantly higher efficiency compared to H2O (21K) and Quest (20K). Additional results of GSM8K, multi-choice benchmarks MMLUand GPQA are in Appendix C.1, C.5, and C.6.

4.3 Needle In A HayStack

We conduct the "Fact Retrieval Across Context Lengths" (Needle In A Haystack) experiment. Our approach achieved a perfect 100% retrieval accuracy with a 64K-token context and a 1/64 cache compression ratio, surpassing all existing methods. This outcome underscores the effectiveness of our method in preserving critical information necessary for accurate retrieval.

1000

3000

5000

7000

9000

11000

13000

15000

17000

19000

21000

23000

25000

27000

29000

31000

33000

35000

37000

39000

41000

43000

45000

47000

49000

51000

53000

55000

57000

59000

61000

63000

Token Limit

0.0

11.0

22.0

33.0

44.0

56.0

67.0

78.0

89.0

100.0

H2O, Acc= 60.0%

0.0

0.2

0.4

0.6

0.8

1.0

1000

3000

5000

7000

9000

11000

13000

15000

17000

19000

21000

23000

25000

27000

29000

31000

33000

35000

37000

39000

41000

43000

45000

47000

49000

51000

53000

55000

57000

59000

61000

63000

Token Limit

0.0

11.0

22.0

33.0

44.0

56.0

67.0

78.0

89.0

100.0

SnapKV, Acc= 84.5%

0.0

0.2

0.4

0.6

0.8

1.0

1000

3000

5000

7000

9000

11000

13000

15000

17000

19000

21000

23000

25000

27000

29000

31000

33000

35000

37000

39000

41000

43000

45000

47000

49000

51000

53000

55000

57000

59000

61000

63000

0.0

11.0

22.0

33.0

44.0

56.0

67.0

78.0

89.0

100.0

Quest, Acc= 98.7%

0.0

0.2

0.4

0.6

0.8

1.0

1000

3000

5000

7000

9000

11000

13000

15000

17000

19000

21000

23000

25000

27000

29000

31000

33000

35000

37000

39000

41000

43000

45000

47000

49000

51000

53000

55000

57000

59000

61000

63000

0.0

11.0

22.0

33.0

44.0

56.0

67.0

78.0

89.0

100.0

AttentionPredictor, Acc= 100.0%

0.0

0.2

0.4

0.6

0.8

1.0

Token Limit

Token Limit

4.4 Efficiency

Decode latency. We evaluate decoding latency in a cache offloading setup [55], essential for long contexts where the KV cache exceeds GPU memory. In this scenario, the standard approach (FlashAttention2) creates a bottleneck by prefetching the entire KV cache from CPU to GPU, which our method optimizes. Further details on the setup are in Appendix B.4. As shown in Figure 5, our method consistently outperforms both FlashAttention2 and H2O across all context lengths. Notably, at 32K tokens, we achieve a 5.6 × acceleration compared to FlashAttention2, demonstrating the dramatic acceleration enabled by loading only a small fraction of the cache. Moreover, we attain a 1.6 × speedup over the calculation-efficient H2O, showing the effectiveness of our cross-token prefetch system.

Overload breakdown. Table 2 details the per-layer overhead of our prefetching system. A key result is that this total overhead remains remarkably stable, increasing only from 47.3ms to 50.0ms as context length grows from 4K to 32K. This stability is a core advantage of our asynchronous design. The growing latency from Predict Attention and Transfer Cache is effectively hidden by being absorbed into the Wait for Main LLM period, during which the main model is busy processing the current token's inference. Due to parallel execution across layers, this per-layer overhead is equivalent to the final per-token latency.

4.5 Prediction Accuracy

Compare with Baselines. We evaluate the prediction accuracy to evaluate our predictor, and the metric is R pred rec /R target rec . On the three representative tasks-QA, summary, and mathematical reasoning-with different KV cache budgets, AttentionPredictor consistently achieves a higher average recovery rate compared to H2O and Quest, as shown in Table 3. This demonstrates that our method accurately identifies the positions of critical tokens, thereby minimizing information loss. Notably, with the KV budget of 512, AttentionPredictor shows a significant advantage over Quest, achieving a 7% higher average recovery rate, emphasizing our robustness under the extremely high (20 × ) compression ratio.

Generalization across datasets. We evaluate our model's generalization by training exclusively on a single dataset and then evaluating prediction accuracy on all remaining, unseen tasks within LongBench. As shown in Figure 7, the model consistently achieves over 95% accuracy across every task, demonstrating strong cross-task transferability and generalization.

4.6 Ablation Study

The effects of hyperparameter block size on AttentionPredictor performance are studied. Other hyperparameters, predictor structure and training efficiency are shown in App. C.8, C.7, and D.1.

Block size b . Figure 8 shows that AttentionPredictor maintains superior average performance with block sizes 8-64. While larger block sizes degrade performance from coarser token positioning, AttentionPredictor declines more mildly than Quest's block-wise KV cache retrieval. Quest's performance drop may stem from losing total attention score information when its minmax based per-block attention score upper bound estimation is used for retrieval, causing inaccurate block identification. In contrast, AttentionPredictor accurately captures attention

8

16

32

64

30

35

40

45

Accuracy (%)

Full Cache

AttentionPredictor

Quest

patterns, improving critical block identification and mitigating block-wise retrieval drawbacks.

5 Conclusion

We present AttentionPredictor, the first learning-based method to directly predict attention patterns for KV compression and attention sparsity method. With a lightweight convolution model, AttentionPredictor captures spatiotemporal patterns of attention score and predicts the next-token attention score accurately. Additionally, we propose the first cross-token prefetching framework, which effectively mitigates prediction and transfer delays in the LLM decoding stage. Experiments on the long-context datasets demonstrate that AttentionPredictor achieves 13 × KV cache compression and 5.6 × speedup with comparable accuracy of LLM.

6 Acknowledgments

This work was supported in part by National Key R&D Program of China under contract 2022ZD0119801, National Nature Science Foundations of China grants U23A20388, 62021001, U19B2026, and U19B2044. This work was supported in part by Huawei as well. We would like to thank all the anonymous reviewers for their insightful comments. This research was also supported by the advanced computing resources provided by the Supercomputing Center of the USTC.

MQA

SQA

Summary

Few-shot

Synthetic

Code

Test Dataset

MQA

SQA

Summary

Few-shot

Synthetic

Code

Train Dataset

97.5

97.7

97.2

97.1

98.1

96.7

97.3

97.6

97.1

96.5

97.1

95.5

97.4

97.6

97.2

96.7

97.1

95.6

96.9

97.2

96.7

96.5

96.8

95.1

96.7

96.9

96.2

95.7

97.2

94.7

97.2

97.4

97.0

96.5

97.2

95.5

References

OpenAI. Learning to reason with LLMs. https://openai.com/index/ learning-to-reason-with-llms/ , 2024. [Accessed 30-01-2025].

Runquan Gui, Zhihai Wang, Jie Wang, Chi Ma, Huiling Zhen, Mingxuan Yuan, Jianye HAO, Defu Lian, Enhong Chen, and Feng Wu. Hypertree planning: Enhancing LLM reasoning via hierarchical thinking. In Forty-second International Conference on Machine Learning , 2025. URL https://openreview.net/forum?id=45he3Ri6JP .

Shervin Minaee, Tomas Mikolov, Narjes Nikzad, Meysam Chenaghlu, Richard Socher, Xavier Amatriain, and Jianfeng Gao. Large language models: A survey. arXiv preprint arXiv:2402.06196 , 2024.

Xu Huang, Weiwen Liu, Xiaolong Chen, Xingmei Wang, Hao Wang, Defu Lian, Yasheng Wang, Ruiming Tang, and Enhong Chen. Understanding the planning of llm agents: A survey. arXiv preprint arXiv:2402.02716 , 2024.

Haoyang Liu, Jie Wang, Yuyang Cai, Xiongwei Han, Yufei Kuang, and Jianye HAO. Optitree: Hierarchical thoughts generation with tree search for LLM optimization modeling. In The Thirty-ninth Annual Conference on Neural Information Processing Systems , 2025. URL https: //openreview.net/forum?id=Ej20yjWMCj .

Zixuan Zhou, Xuefei Ning, Ke Hong, Tianyu Fu, Jiaming Xu, Shiyao Li, Yuming Lou, Luning Wang, Zhihang Yuan, Xiuhong Li, et al. A survey on efficient inference for large language models. arXiv preprint arXiv:2404.14294 , 2024.

Haoyang Li, Yiming Li, Anxin Tian, Tianhao Tang, Zhanchao Xu, Xuejia Chen, Nicole Hu, Wei Dong, Qing Li, and Lei Chen. A survey on large language model acceleration based on kv cache management. arXiv preprint arXiv:2412.19442 , 2024.

Dongjie Yang, Xiaodong Han, Yan Gao, Yao Hu, Shilin Zhang, and Hai Zhao. Pyramidinfer: Pyramid KV cache compression for high-throughput LLM inference. In Findings of the Association for Computational Linguistics, ACL 2024, Bangkok, Thailand and virtual meeting, August 11-16, 2024 , pages 3258-3270, 2024.

Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, et al. Deja vu: Contextual sparsity for efficient llms at inference time. In International Conference on Machine Learning , pages 22137-22176. PMLR, 2023.

Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive KV cache compression for llms. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024 , 2024.

Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, and Beidi Chen. H2O: heavy-hitter oracle for efficient generative inference of large language models. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurIPS 2023, New Orleans, LA, USA, December 10 16, 2023 , 2023.

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024 . OpenReview.net, 2024. URL https://openreview.net/forum?id=NG7sS51zVF .

Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation. arXiv preprint arXiv:2404.14469 , 2024.

Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Peiyuan Zhou, Jiaxing Qi, Junjie Lai, Hayden Kwok-Hay So, Ting Cao, Fan Yang, et al. Seerattention: Learning intrinsic sparse attention in your llms. arXiv preprint arXiv:2410.13276 , 2024.

Zihao Zeng, Bokai Lin, Tianqi Hou, Hao Zhang, and Zhijie Deng. In-context kv-cache eviction for llms via attention-gate. arXiv preprint arXiv:2410.12876 , 2024.

Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin-Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention. arXiv preprint arXiv:2407.02490 , 2024.

Federico Barbero, Alex Vitvitskyi, Christos Perivolaropoulos, Razvan Pascanu, and Petar Veliˇ ckovi´ c. Round and round we go! what makes rotary positional encodings useful? In The Thirteenth International Conference on Learning Representations , 2025.

Xing Li, Zeyu Xing, Yiming Li, Linping Qu, Hui-Ling Zhen, Wulong Liu, Yiwu Yao, Sinno Jialin Pan, and Mingxuan Yuan. Kvtuner: Sensitivity-aware layer-wise mixed precision kv cache quantization for efficient and nearly lossless llm inference, 2025. URL https://arxiv.org/abs/2502.04420 .

Wonbeom Lee, Jungi Lee, Junghwan Seo, and Jaewoong Sim. { InfiniGen } : Efficient generative inference of large language models with dynamic { KV } cache management. In 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24) , pages 155-172, 2024.

Zhihai Wang, Jie Wang, Jilai Pan, Xilin Xia, Huiling Zhen, Mingxuan Yuan, Jianye HAO, and Feng Wu. Accelerating large language model reasoning via speculative search. In Forty-second International Conference on Machine Learning , 2025. URL https://openreview.net/ forum?id=oq0t5BXilT .

Yuhui Li, Fangyun Wei, Chao Zhang, and Hongyang Zhang. EAGLE: speculative sampling requires rethinking feature uncertainty. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024 , 2024.

Hanshi Sun, Zhuoming Chen, Xinyu Yang, Yuandong Tian, and Beidi Chen. Triforce: Lossless acceleration of long sequence generation with hierarchical speculative decoding. arXiv preprint arXiv:2404.11912 , 2024.

Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437 , 2024.

Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng, Jason D. Lee, Deming Chen, and Tri Dao. Medusa: Simple LLM inference acceleration framework with multiple decoding heads. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024 , 2024.

Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Symposium on Operating Systems Principles , pages 611-626, 2023.

Mingcong Song, Xinru Tang, Fengfan Hou, Jing Li, Wei Wei, Yipeng Ma, Runqiu Xiao, Hongjie Si, Dingcheng Jiang, Shouyi Yin, et al. Tackling the dynamicity in a production llm serving system with sota optimizations via hybrid prefill/decode/verify scheduling on efficient meta-kernels. arXiv preprint arXiv:2412.18106 , 2024.

LMDeploy Contributors. Lmdeploy: A toolkit for compressing, deploying, and serving llm. https://github.com/InternLM/lmdeploy , 2023.

Lianmin Zheng, Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E Gonzalez, et al. Sglang: Efficient execution of structured language model programs. arXiv preprint arXiv:2312.07104 , 2024.

Zihao Ye, Lequn Chen, Ruihang Lai, Wuwei Lin, Yineng Zhang, Stephanie Wang, Tianqi Chen, Baris Kasikci, Vinod Grover, Arvind Krishnamurthy, et al. Flashinfer: Efficient and customizable attention engine for llm inference serving. arXiv preprint arXiv:2501.01005 , 2025.

Xunhao Lai, Jianqiao Lu, Yao Luo, Yiyuan Ma, and Xun Zhou. Flexprefill: A contextaware sparse attention mechanism for efficient long-sequence inference. In The Thirteenth International Conference on Learning Representations , 2025.

Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. QUEST: query-aware sparsity for efficient long-context LLM inference. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024 , 2024. URL https://openreview.net/forum?id=KzACYw0MTV .

Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, and Maosong Sun. Infllm: Training-free long-context extrapolation for llms with an efficient context memory. In The Thirty-eighth Annual Conference on Neural Information Processing Systems , 2024.

Hailin Zhang, Xiaodong Ji, Yilin Chen, Fangcheng Fu, Xupeng Miao, Xiaonan Nie, Weipeng Chen, and Bin Cui. Pqcache: Product quantization-based kvcache for long context llm inference. arXiv preprint arXiv:2407.12820 , 2024.

Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. KIVI: A tuning-free asymmetric 2bit quantization for KV cache. In Fortyfirst International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024 , 2024.

Coleman Hooper, Sehoon Kim, Hiva Mohammadzadeh, Michael W Mahoney, Yakun Sophia Shao, Kurt Keutzer, and Amir Gholami. Kvquant: Towards 10 million context length llm inference with kv cache quantization. arXiv preprint arXiv:2401.18079 , 2024.

Zhenyu Zhang, Shiwei Liu, Runjin Chen, Bhavya Kailkhura, Beidi Chen, and Atlas Wang. Q-hitter: A better token oracle for efficient llm inference via sparse-quantized kv cache. Proceedings of Machine Learning and Systems , 6:381-394, 2024.

Zijie Geng, Jie Wang, Ziqi Liu, Feng Ju, Yiming Li, Xing Li, Mingxuan Yuan, Jianye Hao, Defu Lian, Enhong Chen, and Feng Wu. Accurate KV cache eviction via anchor direction projection for efficient LLM inference. In The Thirty-ninth Annual Conference on Neural Information Processing Systems , 2025.

Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, et al. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling. arXiv preprint arXiv:2406.02069 , 2024.

Yuan Feng, Junlin Lv, Yukun Cao, Xike Xie, and S Kevin Zhou. Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference. arXiv preprint arXiv:2407.11550 , 2024.

Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, et al. Moba: Mixture of block attention for long-context llms. arXiv preprint arXiv:2502.13189 , 2025.

Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. arXiv preprint arXiv:2502.11089 , 2025.

André Jonasson. Rotary outliers and rotary offset features in large language models. arXiv preprint arXiv:2503.01832 , 2025.

Alex Graves and Alex Graves. Long short-term memory. Supervised sequence labelling with recurrent neural networks , pages 37-45, 2012.

David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams. Learning representations by back-propagating errors. nature , 323(6088):533-536, 1986.

Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. LongBench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 3119-3137, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10.18653/v1/2024.acl-long.172. URL https://aclanthology.org/2024.acl-long.172 .

Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, et al. ∞ bench: Extending long context evaluation beyond 100k tokens. arXiv preprint arXiv:2402.13718 , 2024.

Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? In First Conference on Language Modeling , 2024.

MAA. American invitational mathematics examination-AIME. https://maa.org/ math-competitions/american-invitational-mathematics-examination-aime , feb 2024. Accessed February 2024 from American Invitational Mathematics ExaminationAIME 2024.

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In International Conference on Learning Representations , 2021.

David Rein, Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang, Julien Dirani, Julian Michael, and Samuel R Bowman. Gpqa: A graduate-level google-proof q&a benchmark. In First Conference on Language Modeling , 2024.

G. Kamradt. Llmtest needleinahaystack,GitHub repository, 2023. https://github.com/ gkamradt/LLMTest_NeedleInAHaystack , Last accessed on 2025-5-15.

Jeffrey Willette, Heejun Lee, Youngwan Lee, Myeongjae Jeon, and Sung Ju Hwang. Trainingfree exponential context extension via cascading kv cache. In The Thirteenth International Conference on Learning Representations , 2025.

Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Jason Flinn, Margo I. Seltzer, Peter Druschel, Antoine Kaufmann, and Jonathan Mace, editors, Proceedings of the 29th Symposium on Operating Systems Principles, SOSP 2023, Koblenz, Germany, October 23-26, 2023 , pages 611-626. ACM, 2023. doi: 10.1145/3600006.3613165. URL https://doi.org/10.1145/3600006. 3613165 .

Meta AI. Introducing Llama 3.1. https://ai.meta.com/blog/meta-llama-3-1/ , 2024. [Accessed 09-01-2025].

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-theart natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 38-45, Online, October 2020. Association for Computational Linguistics. URL https://www.aclweb.org/anthology/ 2020.emnlp-demos.6 .

Tri Dao. Flashattention-2: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations , 2024.

Zhongchao Yi, Zhengyang Zhou, Qihe Huang, Yanjiang Chen, Liheng Yu, Xu Wang, and Yang Wang. Get rid of isolation: A continuous multi-task spatio-temporal learning framework. Advances in Neural Information Processing Systems , 37:136701-136726, 2024.

Zhengyang Zhou, Qihe Huang, Binwu Wang, Jianpeng Hou, Kuo Yang, Yuxuan Liang, Yu Zheng, and Yang Wang. Coms2t: A complementary spatiotemporal learning system for dataadaptive model evolution. IEEE Transactions on Pattern Analysis and Machine Intelligence , 2025.

Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye, Yang Zhou, Jianyu Zhang, Niklas Nolte, Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, and Beidi Chen. MagicPIG: LSH sampling for efficient LLM generation. In Adaptive Foundation Models: Evolving AI for Personalized and Efficient Learning , 2024. URL https://openreview.net/forum?id= SFsvqfNUsh .

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2110.14168 , 2021.

A High Similarity of Queries and Keys

A.1 Qeury Similarity

Our model predicts the attention at step t +1 from a recent history. Since the temporal dimension corresponds to the query sequence (one query per step), the similarity between adjacent queries ( q t and q t +1 ) is the most critical factor for demonstrating step-by-step temporal continuity. Therefore, we evaluate the similarity of consecutive queries. Specifically, we extract query data from LLM using LongBench and compute the cosine autocorrelation of adjacent queries with a lag of 1:

As shown in Figure 9, the heatmap illustrates the query similarity across each layer and head of the LLM. The average similarity is 86% for Longchat and 87% for LLaMA-3.1, indicating strong continuity. Furthermore, the similarity in LLaMA is more consistent across layers, while Longchat shows higher similarity in the shallower layers.

0

5

10

15

20

25

30

Mean

0

5

10

15

20

25

30

Mean

Layer Index

87 79 79 100 95

77 81 100 76

91 91 100 88

97 87 100 85

82 86 100 80

97 92 75 100100 84 96 80 88 99 93 89

87 80 98 99 98 99 95 98 96 99 82 97 99 99 96 85 92 96 74 92 98 99 97 86 91 88 88 92 92 99 95 99 93

95 98 94 95 91 96 98 95 94 89 95 96 94 95 90 97 95 96 92 96 93 94 93 94 95 97 95 95 93 97 96 95 95

87 96 92 93 91 92 78 93 88 92 88 94 91 89 92 92 93 91 91 88 94 90 94 91 87 89 78 91 85 91 94 73 90

93 88 95 92 89 93 92 91 93 90 93 93 83 92 94 91 84 92 93 92 90 92 94 92 88 93 96 94 91 93 90 91 91

95 92 92 92 95 93 90 84 91 93 93 91 91 90 93 88 94 89 92 92 90 84 92 89 92 85 91 90 92 83 91 92 91

91 89 89 92 91 89 93 87 90 74 90 79 82 90 82 87 71 86 89 92 87 86 91 91 89 90 90 89 88 92 67 86 87

85 86 87 91 79 89 74 90 89 88 75 91 78 77 84 85 87 90 92 90 88 89 89 83 89 95 89 85 77 86 84 79 86

80 88 89 89 81 81 84 87 87 91 87 92 89 87 84 88 89 88 89 85 88 84 79 90 85 75 66 87 89 90 90 75 85

88 87 87 86 85 88 87 88 71 86 84 84 87 84 88 92 81 83 88 88 85 89 88 83 88 84 88 93 88 87 84 84 86

87 82 84 90 82 89 86 86 89 88 81 85 83 82 76 88 87 91 80 88 85 87 88 88 82 89 87 87 88 79 86 86 86

91 87 86 87 85 86 83 84 87 86 85 88 88 91 84 76 87 83 76 82 88 85 83 90 89 88 85 89 86 88 83 84 86

88 88 72 88 83 89 86 86 89 85 85 86 88 87 86 86 77 88 82 86 90 86 85 87 87 83 74 91 84 87 85 90 86

87 89 84 84 86 85 85 86 87 82 87 73 87 88 86 88 88 88 89 84 85 88 88 76 85 84 89 88 89 85 84 86 86

85 80 87 84 85 86 85 82 90 84 85 87 84 83 89 83 85 87 84 86 87 89 89 87 86 85 86 85 88 84 90 84 86

87 85 91 88 88 83 92 89 88 88 86 91 73 89 85 88 92 88 87 88 88 86 87 90 87 83 85 86 88 85 88 87 87

85 81 88 89 88 88 89 88 89 88 88 91 91 88 82 90 89 86 86 76 86 82 84 90 80 90 87 86 87 86 84 87 87

87 85 88 86 89 85 84 86 85 89 87 82 82 84 87 91 84 88 85 91 86 85 80 86 88 87 87 88 89 84 89 83 86

89 83 89 88 89 82 91 89 86 83 82 89 90 85 85 87 88 87 84 87 82 90 85 88 86 84 87 90 89 89 82 84 87

88 86 88 87 83 87 88 86 87 83 79 83 79 86 83 84 91 84 86 84 80 86 82 84 88 87 88 86 84 80 83 87 85

79 85 83 86 89 88 81 89 84 81 82 87 88 87 88 87 81 85 89 83 88 86 89 84 82 84 82 86 81 84 83 90 85

86 85 85 85 85 84 82 83 84 88 86 84 86 86 86 85 82 83 84 81 83 83 90 85 86 83 83 83 81 85 76 87 84

82 84 80 83 88 82 86 87 83 84 87 89 83 82 85 86 86 85 81 84 84 87 80 87 85 83 91 84 82 85 82 80 84

85 86 82 85 87 83 79 83 84 80 87 83 84 81 82 84 86 83 89 85 82 82 81 81 83 80 85 80 85 77 85 85 83

81 86 83 81 86 85 85 85 82 82 87 82 84 79 86 83 85 83 88 82 80 84 80 88 85 79 86 83 86 75 82 79 83

85 81 82 79 84 85 78 77 84 83 81 81 82 80 82 83 81 83 81 80 83 85 83 83 80 81 83 82 84 79 80 76 82

80 81 80 85 86 80 84 80 82 85 87 80 80 87 82 85 85 83 84 78 84 84 87 83 85 83 84 81 79 82 80 81 83

81 85 86 82 84 84 87 84 79 85 85 81 88 81 84 86 86 83 82 83 81 80 79 79 82 78 82 86 82 85 81 84 83

82 83 81 82 84 83 86 81 82 85 85 84 81 78 82 84 83 81 78 82 81 83 81 85 80 87 83 85 82 80 78 81 82

84 82 81 82 86 79 84 79 85 84 84 82 83 81 82 74 82 81 82 81 81 81 80 83 81 81 78 85 84 87 88 81 82

83 82 86 79 84 86 83 86 80 77 87 82 84 85 82 80 85 86 83 83 85 82 83 84 80 82 85 82 81 81 81 84 83

87 83 83 87 85 80 83 84 83 82 81 90 80 88 81 82 84 81 78 80 85 84 78 83 82 85 86 85 78 81 78 87 83

86 85 86 87 87 86 86 87 86 86 86 87 85 86 85 86 86 86 85 86 86 86 86 86 86 86 85 87 85 85 85 85 86

Query Similarity (lag=1) of Longchat

0

5

10

15

20

25

30

Mean

0

5

10

15

20

25

30

Mean

Layer Index

95 88 93 92 92 87 85 87 86 86 88 89 87 86 84 93 87 86 85 91 91 93 87 91 79 93 93 91 87 81 89 96 89

94 92 97 92 91 91 91 95 97 96 98 90 94 91 94 92 92 94 88 94 91 95 94 78 92 93 93 95 93 96 90 95 93

91 87 87 91 95 89 88 93 89 91 94 92 78 82 86 77 92 89 91 94 69 69 74 83 87 75 78 80 89 89 92 93 86

91 92 92 89 84 93 92 91 95 90 89 88 87 93 91 89 89 90 87 90 93 90 92 90 92 92 90 83 91 91 90 91 90

86 91 92 91 89 91 90 90 92 93 90 93 88 94 91 91 65 92 90 78 94 89 96 92 91 92 89 94 92 84 90 90 90

87 85 90 89 90 80 87 89 77 77 77 78 84 89 88 89 92 88 91 90 89 83 83 88 88 90 98 91 90 90 89 87 87

88 87 84 83 89 89 87 88 89 85 91 89 90 91 91 93 87 89 84 91 89 84 83 88 79 89 79 85 84 78 89 77 86

91 89 89 78 90 82 85 83 90 90 87 89 88 89 87 88 88 90 91 90 92 92 83 80 84 76 88 86 89 89 89 89 87

78 73 83 81 87 88 88 88 76 83 85 78 82 84 80 86 83 82 85 82 89 92 85 90 84 86 86 85 83 82 86 81 84

85 84 83 85 88 86 89 89 88 87 84 90 79 77 82 84 87 87 87 82 79 86 86 85 86 81 83 77 83 82 90 77 84

77 73 72 80 87 82 79 78 86 88 87 88 73 75 72 86 86 85 82 81 83 84 85 84 88 85 90 83 78 83 84 79 82

86 80 83 87 80 78 78 81 82 84 87 90 71 82 81 79 88 89 84 86 85 83 83 82 86 87 83 85 78 88 82 88 83

93 91 89 89 87 88 85 90 90 88 81 89 91 90 82 92 89 88 88 87 89 88 87 88 89 89 89 89 86 88 86 90 88

85 82 84 82 83 85 82 86 81 85 85 84 87 88 85 86 85 83 82 83 85 85 89 87 90 89 84 81 88 91 87 90 85

86 86 86 87 84 81 87 87 90 92 88 91 86 81 84 87 77 74 81 62 81 90 84 88 87 88 92 89 86 82 85 84 85

88 82 83 83 91 87 92 91 89 90 90 85 89 86 88 89 85 83 87 93 89 86 87 87 88 89 88 91 71 82 75 89 87

83 81 82 83 89 88 90 88 86 89 90 89 89 89 88 89 91 90 91 89 70 85 85 71 88 84 86 85 87 92 90 95 87

90 91 89 89 89 89 87 90 77 86 87 91 90 92 89 92 88 91 87 89 91 86 87 87 82 83 82 84 87 83 91 88 88

91 89 90 88 78 89 90 89 85 86 93 89 92 90 90 90 84 89 91 87 81 87 86 84 88 90 90 92 84 88 84 87 88

83 87 79 84 84 87 84 91 90 83 90 89 80 80 85 79 89 92 91 91 86 89 89 86 90 89 87 90 91 88 90 89 87

84 83 85 78 88 84 87 86 91 87 90 88 87 84 83 84 86 86 87 86 83 87 88 80 79 81 83 87 92 89 86 90 86

88 83 79 83 91 90 92 89 84 87 87 83 89 90 88 90 90 89 89 91 89 89 85 86 90 85 86 85 91 92 92 89 88

89 91 90 88 87 88 90 88 86 86 87 88 82 87 84 85 88 90 87 88 87 88 87 91 85 85 85 85 86 84 85 92 87

88 85 86 81 88 89 90 86 89 90 88 88 92 90 83 86 87 88 85 90 87 85 87 86 87 85 90 87 89 89 87 88 87

87 88 88 88 81 80 83 84 84 83 84 84 89 88 87 87 90 91 90 87 91 88 92 84 86 89 81 83 89 89 89 87 87

85 90 86 83 92 90 91 92 86 87 87 91 92 90 91 91 88 87 84 89 86 86 86 85 89 88 88 89 86 84 86 85 88

76 84 79 79 84 83 85 84 84 86 87 86 90 84 90 80 84 85 84 82 81 87 88 87 89 88 89 83 87 87 84 91 85

86 87 85 84 91 83 86 82 85 87 84 85 86 82 90 86 82 85 86 86 87 85 91 87 85 88 88 88 78 87 84 83 86

83 85 84 86 86 83 84 86 92 89 83 83 88 92 89 84 82 84 82 83 82 84 86 82 87 85 85 88 86 83 82 87 85

87 86 85 87 87 88 87 88 86 76 83 81 87 87 87 86 86 87 84 85 82 78 82 83 84 84 82 79 87 87 86 85 85

82 80 81 83 84 86 81 84 86 79 84 83 85 84 85 84 85 91 78 93 85 85 87 86 80 81 77 81 87 88 82 82 84

88 87 89 88 88 90 87 84 91 90 90 77 89 85 82 90 85 85 91 82 87 91 90 87 90 91 87 83 87 86 83 94 87

87 86 86 85 87 86 87 87 87 87 87 87 86 87 86 87 86 87 87 87 86 87 87 85 87 87 86 86 86 87 87 88 87

Query Similarity (lag=1) of LLaMA-3.1

Head Index

Head Index

To provide a more detailed analysis of the similarity of queries at a far distance, we measure the average query cosine similarity across various distances (lags). The results are presented in Table 4. The query similarity gradually decreases as the distance increases, which is expected. However, the similarity remains relatively high even at longer distances. This finding suggests that while the influence of the most recent tokens is strongest, the attention patterns exhibit a strong local stability that decays slowly over time. This supports our model's effectiveness in capturing these evolving yet persistent patterns using a recent history window.

A.2 Key Similarity

In addition to queries, we also investigate the similarity of keys over time. When analyzing the attention scores between pairs such as A t,i and A t +1 ,i , both attend to the same position i . Consequently, the keys involved in the two computations are exactly the same ( k i ), and therefore no issue of key self-similarity arises in this case.

In contrast, when analyzing patterns like the "Periodic Slashes" pattern, the comparison involves pairs such as A t,i and A t +1 ,i +1 , where the attention scores are influenced by different keys ( k i and

k i +1 ). In this case, the self-similarity of keys becomes relevant. To quantify this, we compute the cosine similarity between keys at different positions with various lags:

where τ denotes the distance between keys. The results are also presented in Table 4, and we can find that they indeed exhibit strong self-similarity. These consistently high similarity values indicate that keys vary slowly over time. Our findings are also consistent with prior works such as KIVI [34] and KVQuant [35], which reported that keys often contain fixed, massive channels, suggesting limited variation across positions.

B Implement Details

B.1 Implement Details of AttentionPredictor

Model Architecture Details. The architecture of the CNN-based AttentionPredictor model is summarized in Table 5. This design ensures efficiency and adaptability to varying sequence lengths, while maintaining sufficient capacity to capture temporal patterns.

Training Details. We capture the attention scores during full-cache inference as raw data. We package the attention scores of H steps, together with the ( H +1) -th step's score, as input-output pairs. The training data is block-compressed to align with the usage during LLM model inference. Specifically, the input history attention is A H ∈ R H × t b padding with zero for short attentions and the output attention score is C ∈ R 1 × t b . Attention scores with insufficient lengths are padded with zeros at the end to ensure that all data within a package has the same length. For the attention at step t +1 , we only predict the first t positions because the last position is derived from the newly generated key, which does not need to be compressed or transferred. Notably, only the attention scores from the decoding phase are used as output data. As a result, the training dataset includes the final H steps of the pre-filling attention and all attention scores from the decoding stage. To train the model, we calculate the discrepancy between the predicted and real attention scores using mean squared error (MSE) loss. The training epoch is set to 30, and we select the model with the best attention prediction accuracy as shown in Equation 1.

Cross-token Prefetching. We leverage GPU parallel streams and CPU multi-threading to parallelize the token retrieval and cache transfer across different layers. Consistent with the main experiment, we maintain uncompressed caches for the first two layers to ensure higher inference accuracy. The timeline of prefetching framework is in Figure 10.

B.2 Implement Details of Baselines

Our experiments are based on LLMs using 16-bit floating points, and we utilize FlashAttention2 [56] during the prefilling stage to reduce GPU memory consumption. Since FlashAttention2 does not output intermediate attention scores, for methods that rely on historical attention, we additionally compute partial prefill attention to obtain the scores. Our batch size is fixed at 1. For all baselines, we use the official code implementations to ensure the effectiveness of the methods. Additionally, for all methods, we skip compress the first two layers of the LLM.

Token

𝒕

Token

𝒕 + 𝟏

Layer 1

Layer 2

…

Layer 1

Layer 2

…

…

Token

𝒕

…

Token

𝒕 + 𝟏

Fetch KV

Cross-token

prefetch KV

Timeline

Asynchronous tasks

hide latency

①

②

③

④

①

②

③

④

①

②

③

④

①

②

③

④

Comp

Query

Comp

Query

Comp

Query

Comp

Top-k

Comp

Top-k

Comp

Top-k

Fetch

KV

Fetch

KV

Fetch

KV

Attn

FFN

Attn

FFN

Attn

FFN

Attn

FFN

Attn

FFN

Attn

FFN

Attn

FFN

Comp

Query

Comp

Query

Comp

Query

Comp

Query

Attention

Predictor

Attention

Predictor

Prefetch

KV

Prefetch

KV

StreamingLLM [12] finds that the sink token is very important, so it retains the initial and the most recent few tokens. For StreamingLLM, we evenly distribute the budget between sink tokens and local tokens.

H2O [11] accumulates all historical attention scores as an estimation. We evenly distribute the budget between the heavy budget and the recent budget. Since computing the full historical data can cause OOM with longer inputs, we use the last 64 steps of historical information to calculate the heavy hitters, and denote this as H2O+. Note that accumulating all prefill attention scores causes H2O to focus on initial tokens, so this is an improvement to the method's effectiveness.

SnapKV [13] accumulates attention scores within a recent window as an estimate and performs one-time filtering during the prefill stage. We set the time window to a size of 64 to align with our method. Since SnapKV only performs cache compression once during the prefill stage, the number of KV tokens it uses continues to grow during decoding, resulting in a relatively lower compression ratio compared to other methods.

Quest [31] uses the current query and paged key to approximate the attention score. We use a chunk size of 16, corresponding to the block size in our method, which is also the parameter value used in the original paper.

B.3 Details of LongBench Dataset

LongBench [45] is a carefully crafted benchmark suite designed to evaluate the capabilities of language models in processing extended documents and complex information sequences. It was developed for multi-task assessment of long-context inputs. The details of the metrics, number of words, language, and data used in LongBench are presented in Table 6.

B.4 Details of Efficiency Experiment Setup

Our efficiency evaluation is conducted under the cache offloading scenario, which is essential for long-context generation when the KV cache size exceeds the available GPU memory. Unlike GPUbased eviction/compression, offloading keeps the full cache on CPU, preserving model accuracy. This setting is also adopted in recent works [32, 19, 33].

We conduct experiments under this offloading setup on FlashAttention2, H2O, and our proposed AttentionPredictor:

FlashAttention2 [56]: We utilize the standard KV cache offloading implementation provided by the Hugging Face library. In this setup, at each decoding layer, the entire KV cache for the upcoming layers is prefetched from CPU to GPU to compute the full attention.

H2O [11]: Also uses cross-layer prefetching but transfers sparse KV cache, making the comparison a more direct evaluation of the prefetching framework's efficiency.

AttentionPredictor: Runs on an offloading setup but prefetches across tokens, improving efficiency.

C Additional Experimental Results on Accuracy

C.1 Results on Long N-shot CoT GSM8K.

We evaluate the reasoning task under a long-input scenario, which is different with long output scenario in Section 4.2. As shown in Table 7, we evaluate our method on the mathematical reasoning dataset GSM8K with LLaMA-3.1.

To simulate CoT tasks within long-context, we increased the number of few-shot examples. Specifically, we randomly selected a fixed number of questions and standard CoT answer pairs as prompts, along with the questions to be tested. We chose 25, 47, and 97 few-shot examples, resulting in input lengths of approximately 4K, 8K, and 16K tokens respectively. The few-shot data were sourced from the GSM8K training set. Since the test set does not overlap with the training set, the answers remain undisclosed to the LLM.

Unlike long-context tasks LongBench, CoT mathematical reasoning tasks present distinct challenges. For example, the retrieval-based SOTA method, Quest, excels on long-context benchmarks but performs poorly on CoT tasks, showing a 16.91% accuracy drop at a sequence length of 16K. In contrast, AttentionPredictor achieves a significantly smaller accuracy loss of just 0.91%. Moreover, across all sequence lengths, our method outperforms baselines in most cases.

Single-DocumentQA

Multi-DocumentQA

Summary

Few-shot Learning

Synthetic

Code

Budget

Method

NrtvQA

Qasper

MF-en

HotpotQA

2WikiMQA

Musique

GovReport

QMSum

MultiNews

TREC

TriviaQA

SAMSum

PCount

PRe

Lcc

RB-P

Average

↑

Longchat-v1.5-7b-32k

Full

Full cache

20.70

29.41

43.09

33.05

24.14

14.66

30.84

22.79

26.61

66.50

83.99

40.90

0.00

30.50

52.94

56.78

36.06

StreamingLLM

15.03

17.71

27.14

27.00

21.52

9.09

21.51

19.39

21.97

55.00

81.07

37.07

0.50

5.00

46.40

53.05

28.65

Cascade

16.56

19.02

24.16

25.57

22.13

7.88

11.76

19.50

15.77

51.50

82.44

36.73

1.00

7.50

6.07

12.45

22.50

H2O+

19.01

24.00

33.08

30.60

22.47

12.23

21.01

21.59

22.07

50.50

83.42

39.62

0.00

27.50

51.35

54.63

32.07

SnapKV

18.73

25.21

38.18

33.26

23.79

13.27

21.84

21.15

22.58

59.00

84.20

38.41

0.00

25.50

54.39

57.43

33.56

Quest

16.80

31.43

38.01

27.95

23.73

10.03

27.91

22.58

25.68

60.50

81.27

39.95

0.00

41.50

51.95

50.22

34.34

512

AttentionPredictor

17.69

28.49

41.63

34.13

23.67

14.18

28.40

21.82

25.52

64.00

84.66

40.56

1.00

25.50

54.42

55.89

35.10

StreamingLLM

16.30

20.06

32.30

28.60

21.37

9.56

25.52

20.36

23.51

60.50

82.15

39.29

1.50

6.00

52.74

53.71

30.84

Cascade

19.30

21.21

27.47

30.24

22.29

10.78

14.60

20.63

19.19

57.00

83.43

39.10

2.50

9.50

6.69

12.83

24.80

H2O+

18.92

24.65

37.38

32.40

22.28

12.46

22.72

21.62

24.22

59.00

84.39

39.90

0.00

28.00

53.74

55.60

33.58

SnapKV

19.46

26.97

38.96

34.01

23.22

14.33

23.17

21.46

24.50

64.00

84.33

39.93

0.00

28.00

53.96

57.72

34.63

Quest

18.51

31.09

41.32

29.92

22.95

11.35

30.16

22.67

26.52

66.00

83.35

40.86

0.00

35.00

50.77

54.16

35.29

1024

AttentionPredictor

19.36

29.50

41.67

33.64

24.30

14.91

29.92

22.30

26.09

66.00

84.85

40.58

0.00

28.00

52.22

56.54

35.62

StreamingLLM

15.52

22.66

36.20

29.03

22.15

10.13

27.59

21.04

26.20

64.00

81.93

38.62

2.00

9.50

49.78

55.46

31.99

Cascade

19.49

24.80

30.63

33.48

22.68

12.37

16.82

20.82

20.76

63.00

83.51

40.32

3.00

11.00

7.28

13.38

26.46

H2O+

19.10

26.62

36.82

33.23

22.66

12.83

25.13

21.93

21.09

62.50

84.65

40.60

0.00

29.00

53.33

55.99

34.09

SnapKV

20.02

28.69

40.61

34.13

24.02

14.67

26.05

21.82

25.59

65.00

84.71

40.89

0.00

28.50

54.93

58.44

35.50

Quest

19.12

32.08

43.72

31.91

24.97

14.83

31.56

22.65

26.36

66.50

85.09

41.69

0.00

31.50

50.53

55.73

36.14

2048

AttentionPredictor

20.07

30.06

42.60

34.26

24.20

14.64

30.53

22.07

26.61

67.00

84.94

41.86

0.00

30.50

52.45

56.59

36.15

StreamingLLM

19.45

27.90

41.53

28.54

20.50

10.99

26.86

20.73

26.54

58.00

75.71

34.82

2.50

13.00

54.31

56.41

32.36

Cascade

20.47

27.98

35.38

32.45

23.90

12.15

17.87

21.33

21.46

64.50

84.99

41.55

0.50

16.00

7.48

13.09

27.57

H2O+

20.53

28.28

41.41

32.78

23.87

14.52

27.81

21.76

21.71

64.50

84.09

40.90

0.00

29.50

52.71

56.61

35.06

SnapKV

20.30

28.90

42.36

33.49

24.46

15.11

28.37

22.54

26.54

66.00

84.19

41.20

0.00

29.00

54.70

58.69

35.99

Quest

19.67

29.26

43.90

32.72

23.61

12.61

31.49

22.95

26.27

67.00

84.38

41.18

0.00

31.00

51.62

55.61

35.83

4096

AttentionPredictor

20.12

29.33

42.88

34.12

24.27

15.10

30.74

22.70

26.59

67.00

84.14

41.42

0.00

29.50

53.00

56.69

36.10

LLaMA-3.1-8B-Instruct

Full

Full cache

29.28

45.36

56.12

56.67

48.90

32.43

33.77

25.20

27.08

74.00

91.48

40.85

5.07

99.50

56.86

48.15

48.17

StreamingLLM

22.49

24.93

38.64

48.48

41.41

29.43

24.08

20.24

23.38

58.50

81.60

38.22

8.00

98.50

55.32

45.49

41.17

Cascade

25.49

30.09

39.83

52.05

44.75

22.92

24.16

21.49

24.57

48.00

87.49

40.61

8.50

94.50

45.58

36.33

40.40

H2O+

25.93

29.59

45.53

50.59

45.75

29.24

24.50

23.28

23.56

48.50

88.90

41.39

8.16

99.50

56.06

47.88

43.02

SnapKV

26.51

39.81

52.05

54.60

46.31

26.92

24.27

23.22

24.16

65.50

90.93

40.31

4.39

99.00

55.66

46.64

45.02

SeerAttention

20.46

43.07

52.82

48.86

41.54

26.32

29.86

22.88

26.37

61.50

88.10

42.51

4.50

54.00

40.49

33.32

39.79

Quest

25.82

44.01

53.99

58.08

43.96

28.84

32.90

24.50

26.80

70.00

88.55

39.42

11.16

99.50

51.65

42.03

46.33

512

AttentionPredictor

27.53

43.92

56.58

57.78

45.30

27.00

30.21

24.92

26.54

67.50

92.28

40.45

7.50

99.50

59.11

46.81

47.06

StreamingLLM

24.69

27.84

42.65

48.20

37.82

26.50

26.80

21.55

22.40

64.00

89.42

40.65

8.00

94.50

56.18

46.33

42.35

Cascade

27.14

37.17

42.51

53.06

37.75

25.73

26.74

22.25

25.61

58.50

89.70

40.73

6.50

96.00

46.88

38.27

42.16

H2O+

26.24

36.30

48.47

53.97

45.78

27.08

26.11

23.59

25.05

53.00

88.54

40.88

9.05

99.50

55.20

49.22

44.25

SnapKV

27.63

41.86

52.37

55.83

46.39

27.69

26.39

24.06

25.49

70.00

90.66

40.76

8.12

99.50

55.55

49.58

46.37

SeerAttention

23.64

46.13

54.41

54.61

38.78

28.13

31.47

24.28

26.69

65.50

89.21

41.79

6.00

82.00

45.97

35.49

43.38

Quest

27.46

43.58

54.95

57.22

44.53

29.86

33.64

25.37

26.68

72.00

88.61

40.05

9.55

99.50

53.92

43.82

46.92

1024

AttentionPredictor

28.30

46.66

55.81

56.87

50.67

30.10

32.34

25.08

26.66

74.00

92.11

42.15

7.17

100.00

59.12

50.26

48.58

StreamingLLM

25.79

32.35

45.03

49.84

36.71

23.29

29.31

22.43

26.83

69.00

86.87

38.53

7.25

83.00

43.32

46.66

41.64

Cascade

28.49

40.74

46.05

51.86

45.54

26.58

28.98

22.42

26.49

63.50

90.55

41.85

10.50

94.00

47.03

39.13

43.98

H2O+

28.07

40.35

53.22

54.81

44.75

31.77

28.20

23.50

26.37

64.50

90.71

41.32

6.71

100.00

57.97

49.19

46.34

SnapKV

28.00

45.34

54.98

57.97

45.89

30.50

28.92

24.29

26.32

71.50

89.36

40.71

6.64

100.00

58.79

49.90

47.44

SeerAttention

27.56

46.05

55.70

55.83

47.07

29.09

32.02

24.25

26.82

68.50

88.24

40.14

6.00

93.50

46.94

37.74

45.34

Quest

29.25

46.05

55.90

57.91

45.51

30.82

34.05

24.89

26.93

74.00

92.41

41.77

8.50

100.00

54.29

45.72

48.00

2048

AttentionPredictor

29.52

45.72

54.20

56.50

46.93

33.45

33.28

24.68

26.88

72.00

91.79

42.20

10.08

99.50

57.47

49.09

48.33

StreamingLLM

25.88

43.45

55.55

49.90

37.78

28.79

28.13

22.18

26.70

62.00

78.53

36.32

8.29

74.00

58.63

45.91

42.63

Cascade

28.43

46.21

50.77

55.02

47.68

30.29

31.47

23.75

26.93

66.50

90.03

41.70

8.67

94.50

47.14

40.02

45.57

H2O+

28.37

45.33

55.72

56.27

47.74

31.63

30.91

24.15

26.89

65.50

90.96

41.97

7.08

99.50

57.52

49.70

47.45

SnapKV

29.17

44.03

55.40

56.14

49.01

30.86

31.36

24.73

26.85

72.00

90.31

40.58

8.68

99.50

57.75

49.03

47.84

SeerAttention

30.43

46.06

56.00

55.67

45.73

31.08

32.60

24.49

26.82

68.50

89.44

42.94

10.62

99.50

46.13

39.46

46.59

Quest

29.30

46.43

54.66

55.88

47.01

30.65

34.54

25.23

26.31

71.00

89.76

43.92

6.42

100.00

55.25

45.06

47.59

4096

AttentionPredictor

28.67

47.07

55.87

57.46

48.35

27.91

33.12

24.76

26.62

73.50

90.30

43.40

6.75

100.00

58.03

48.85

48.17

C.2 Additional Results on Longbench Benchmark

Due to space constraints, we presented the test results after aggregating task types in Section 4.2. Here, we present the individual results for all tasks. The Table 8 shows that our method surpasses the majority of the SOTAs on most tasks, and the average performance under all budgets and LLMs exceeds all compared methods, demonstrating the effectiveness of our approach.

SeerAttention utilizes a learnable module to aid cache compression. We use the officially released SeerAttention model for our tests. Since the official pre-trained model of SeerAttention for LongChat is unavailable, we evaluate its performance on LLaMA-3.1-8B-Instruct.

Table 8 indicates that our method outperforms SeerAttention, which also employs a learnable module. In particular, with a cache budget of 512, our method achieves an average accuracy 7.27% higher than SeerAttention. This could be attributed to SeerAttention encoding keys and queries without directly modeling attention scores, thus failing to accurately retrieve important token blocks under extreme resource constraints. Moreover, by employing a single unified model for all LLM layers, our

approach results in a total model size that is merely 0.02% of SeerAttention's, further attesting to the effectiveness of our method.

C.3 Results on InfiniteBench

To evaluate the effectiveness of our proposed method, we conduct experiments on selected tasks from the InfiniBench [46] benchmark, an extremely long-context benchmark with an average context length 214K. We use LLaMA-3-8B as the base model and focus on the Math Find and Choice tasks, which require numerical reasoning and multi-option selection, respectively. We compare our approach against several state-of-the-art KV compression baselines under an 8K token KV cache budget.

As shown in Table 9, AttentionPredictor achieves the highest average performance (44.9), outperforming Quest (43.4), SnapKV (42.5), and InfLLM (33.7). Notably, on the Math Find task, which is particularly sensitive to the preservation of precise attention information, our method attains 34.3-significantly higher than all baselines. These results suggest that AttentionPredictor is more effective at identifying and preserving the most relevant attention information under constrained memory, enabling better retention of long-range dependencies.

C.4 Results on RULER QA Benchmark

We further evaluate our method on the RULER QA benchmark [47], which is designed to assess long-context understanding under varying context lengths and constrained KV cache budgets. The results are summarized in Table 10. The results demonstrate that with a highly constrained budget of only 1K tokens, our AttentionPredictor consistently and significantly outperforms the SnapKV baseline across all tested context lengths.

C.5 Results on MMLU Benchmark

To further assess the effectiveness of our approach, we conduct experiments on the MMLU benchmark [49]. Since the default evaluation mode in the official lm-eval library is a single-step prediction task, it does not fully reflect the benefits of our method, which is designed to optimize multi-step decoding. Therefore, we adopt the cot_fewshot configuration, which applies an n-shot Chain-of-Thought (CoT) prompting strategy and involves a multi-step reasoning process. This setup provides a more suitable testbed for evaluating our decoding-stage optimizations.

As shown in Table 11, our proposed AttentionPredictor achieves performance very close to that of the full-cache model while consistently outperforming the SnapKV baseline under different KV cache

budgets. These results demonstrate that our method remains effective even on relatively short input tasks, provided that the decoding process requires multi-step reasoning.

C.6 Results on GPQA Benchmark

To further validate our approach on expert-level tasks, we also tested it on the GPQA benchmark [50]. GPQA is a challenging dataset of graduate-level multiple-choice questions curated by domain experts across biology, physics, chemistry, and philosophy. It is similar in format to MMLU but significantly more difficult, requiring deep expert-level reasoning.

As shown in Table 12, our method again proves effective, maintaining high accuracy in a scenario requiring deep reasoning.

C.7 Prediction Model

Our prediction task can be framed as modeling spatio-temporal information inherent in the token sequences, a problem for which various architectures have been explored [43, 44, 57, 58]. We evaluate the performance of various prediction model implementations. MLP is applied to the sequence dimension to handle variations in token length. For LSTM, attention is divided into 16 width blocks and predictions are independent for each block, referred to as LSTM-block in Table 13. Similarly, CNN-block utilizes a block-wise prediction approach, employing a fully connected layer for fixed block lengths. Finally, the CNN model, as the primary setting in our experiments, predicts all attention scores simultaneously. As reported in Table 13,

CNN outperforms other models, achieving the highest attention prediction accuracy by effectively capturing attention patterns. MLP failed to account for neighboring token interactions, while LSTMblock and CNN-block were restricted to block-level information without global context. Notably, CNN also required the fewest parameters and was the most memory-efficient during inference.

C.8 Ablation Study Results of Hyperparameters

We evaluate the impact of hyperparameters on our method. Specifically, we train the predictor using the default parameters of the main experiment, i.e., b = 16 , H = 64 . Weemploy full attention without sparse computation during training, which is equivalent to performing distribution error calibration at each step, i.e., M = 1 . Subsequently, we test the performance under different hyperparameters using the trained model. Other parameters during testing remain consistent with the main experiment, with a KV budget of 1K. We evaluate six representative tasks in Longbench. The results are shown in Table 14.

Calibration step M . Table 14 shows the average performance of AttentionPredictor with calibration steps ranging from 1 to 20. Performance improves with shorter calibration intervals, indicating that the calibration scheme reduces cumulative errors caused by differences between sparse and original attention distributions. However, higher calibration frequencies increase computational costs, creating a trade-off between accuracy and efficiency.

History step H . As depicted in Table 14, we evaluate the average accuracy of AttentionPredictor over a range of history steps. Overall, the performance gap compared to the full cache remains consistently small. Notably, the middle value of H maximizes performance by providing the necessary information for pattern recognition, while mitigating redundancy caused by the decaying self-correlation of attention scores.

C.9 Necessity of the Learnable Attention Predictor

A pertinent question arises regarding the necessity of a learnable attention predictor: given the potential for high similarity between attention distributions of adjacent tokens or layers, could one simply utilize the attention map from a preceding token or layer as a direct estimate? To address this, we conducted a comparative experiment evaluating our proposed AttentionPredictor against two such heuristic baselines: (1) Previous-Layer Attention , which uses the attention map from the corresponding position in the immediately preceding layer, and (2) Previous-Token Attention , which employs the attention map of the immediately preceding token within the same layer.

The results of this evaluation, performed on the LongBench benchmark using Llama3.1-8B-Instruct, are presented in Table 15. Employing previous-layer attention yields an average score of 42.51%, which is a substantial 5.29% lower than our AttentionPredictor's average of 47.80%. This performance decrement underscores the limitation of direct cross-layer attention transference, primarily because attention patterns often exhibit considerable variation across different layers, making the previous layer's attention a less reliable proxy for identifying critical tokens in the current layer.

The previous-token attention strategy achieves a higher average score of 46.82%. This is 4.31% better than the previous-layer approach and aligns with the observation that attention mechanisms frequently display sequential patterns along the token dimension within the same layer. However, despite this improvement, relying solely on the single immediately preceding token is still suboptimal.

Our learnable AttentionPredictor, with an average score of 47.80%, outperforms the previous-token attention baseline by 0.98%. This superiority stems from our model's capacity to learn and leverage more complex, dynamic attention patterns from several previous tokens, rather than being confined to the heuristic of the single last token, which may not adequately capture evolving contextual dependencies.

Therefore, these comparative results affirm that while attention similarities exist, simplistic heuristic estimations are insufficient . The empirical evidence highlights the necessity and distinct advantage of our proposed learnable AttentionPredictor for achieving more accurate anticipation of attention distributions.

C.10 Comparison between MInference and AttentionPredictor

To further evaluate the effectiveness of our proposed AttentionPredictor, we compare its performance against MInference [16]. MInference operates by categorizing attention heads into three predefined sparse patterns (A-shape, Vertical-Slash, and Block-Sparse) and uses a kernel-aware search to assign the optimal pattern to each head offline. During inference, it dynamically builds sparse indices based on the assigned pattern and the specific input.

We consider two variants of MInference for a comprehensive comparison: (1) MInference (prefill) : The standard version where the pattern-fitting algorithm is run only during the pre-filling stage. (2) MInference (decode) : A stronger, more dynamic variant where the computationally expensive pattern-fitting algorithm is re-run at every decoding step to adapt to new contexts.

As shown in Table 16, AttentionPredictor consistently outperforms both MInference variants across all tasks, achieving an average accuracy of 94.25%. This demonstrates that by learning from recent history instead of relying on fixed pattern templates, AttentionPredictor more accurately captures the dynamic nature of attention during decoding.

C.11 Top-K Selection Strategy Analysis

WhyTop-K Outperforms Sampling? While prior works such as MagicPIG [59] suggest that Top-K may introduce bias in certain tasks, our experiments reveal a critical insight: prediction accuracy matters more than the design of the selection policy. As shown in Table 17, we compare the performance of our AttentionPredictor and MagicPIG on the LongBench benchmark using Llama3.1-8B-Instruct (MagicPIG's results are taken from the original paper). The results demonstrate that when combined with our high-precision AttentionPredictor, even a simple Top-K strategy achieves superior performance compared to MagicPIG's more sophisticated sampling approach. This indicates that high-quality attention predictions can enable basic selection policies to outperform more complex alternatives. Furthermore, it is important to note that our AttentionPredictor is decoupled from the selection policy and can be flexibly integrated with various methods , including both Top-K and importance sampling.

D Additional Experimental Results on Efficiency

D.1 Training Efficiency

We evaluate the impact of the training data ratio on prediction accuracy on the hotpotqa dataset of Longbench. The predictor's performance improved only slightly, from 96.0% to 97.5%, as the proportion of training samples increased from 3% to 70%. This efficiency can be attributed to two factors. First, the predictor only needs to capture the relative patterns rather than precisely predict attention scores, which is easier. Second, since attention patterns are inherent characteristics of the LLM, the predictor exhibits strong generalization capabilities. These findings suggest that a relatively small amount of training data is adequate for the predictor to effectively learn the underlying attention dynamics.

D.2 Computational Overhead of AttentionPredictor.

We have quantified the inference overhead of our predictor model. As shown in Table 19, both the latency and memory footprint are minimal, ensuring that the prediction process does not become a bottleneck.

D.3 Memory Efficiency of Cross-Token Prefetching

We have evaluated the memory footprint of our proposed cross-token prefetching framework during decoding. As shown in Table 20, the framework consistently consumes less GPU memory compared to the standard cross-layer approach. This advantage arises because the memory savings from loading only a small, sparse KV cache outweighs the potential overhead from the cross-token prefetching mechanism.

E Limitations

Our current investigation is primarily centered on a specific range of LLM architectures and sizes, and the generalizability of observed attention patterns to vastly different model types warrants further study. The necessary conditions for a specific attention pattern can be explored in the future work.

F Boarder Impacts

Academic Impact: Our study of LLM attention score patterns and their predictability may offer insights into attention dynamics, potentially informing future research on these characteristics. By exploring attention score patterns, this work could also be relevant to advancements in KV cache compression, possibly contributing to the development of more memory-efficient model designs.

Societal Impact: Through the potential implications for KV cache compression via attention predictability, this research could contribute to more efficient LLM inference. This might lead to reduced latency in AI applications and lower energy consumption. Such improvements could potentially enhance the accessibility of LLMs and support more sustainable AI practices.

G License

We include the following licenses for the code, models, and benchmarks we used in this paper.

Benchmarks: LongBench [45]: License, AIME 2024 [48]: License, GSM8K [60]: License, Needle In A Haystack [51]: License, InfiniteBench [46]: License. Models: LLaMA-3.1-8B-Instruct [54]: License, LongChat [53]: License. Code: H2O [11]: License, Minference [16]:License, SeerAttention [14]: License.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: This paper analyzes the temporal patterns in the attention score and proposes AttentionPredictor to predict attention scores and dynamically identify critical tokens. Moreover, a cross-token prefeching system is proposed to accelerate the decoding stage. Experiments on long-context benchmarks shows AttentionPredictor achieves 13x cache compression with comparable performance.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: The limitations of our AttentionPredictor are detailed in Appendix E. Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: The theoretical results are given in Section 3.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The implementation of our approach are detailed in Section 3 and Appendix B.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide the code and dataset in the supplemental material.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: The details about experiment setting are in Section 4.1 and Appendix B.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: The results are in Section C.8 and Section 4.6.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: The details about hardware specification are in Section 4.1.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have reviewed and conformed with the Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: The details of broader impacts are included in Section F.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our model and data are focused on KV Cache compression, with no risk of misuse.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: Licenses are in Appendix G.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: The details of the model are detailed in this paper.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: The core method development in this research does not involve LLMs as any important, original, or non-standard components.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a design named AttentionPredictor, a learning-based method for Key-Value (KV) cache compression in large language models (LLMs), enabling efficient long-context inference. It utilizes a lightweight, unified convolutional model to capture spatiotemporal patterns in attention scores, predicting next-step attention scores, and integrates a cross-token prefetching framework to mitigate prediction and cache loading latencies.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\nThe solution introduces a learning-based approach for predicting direct attention scores and a cross-token prefetching framework, which differs from heuristic-based methods and per-layer models. Its lightweight unified prediction model enhances scalability and efficiency.\n\nIt also reduces memory and computational demands for long-context tasks, with implications for multi-turn dialogues, complex reasoning, and resource-constrained settings.\n\nThe paper articulates the theoretical basis, model design, and implementation details.\n\nWeaknesses:\n\nWhile standard deviations are provided, the paper does not specify the method for error computation or the distributional assumptions. Experiments are limited to LLaMA-3.1-8B and QwQ-32B, lacking validation on other architectures.\n\nWhile societal impacts, such as reduced latency and energy consumption, are noted, no deployment data for edge devices is provided.\nHyperparameter choices are empirically determined without theoretical justification or sensitivity analysis. The communication overhead analysis for cross-token prefetching lacks specific latency data.\n\nQUESTIONS:\nSee weaknesses above. Also other minors:\n\nSections C.5 and 4.6 provide standard deviations but do not specify the error computation method (e.g., closed-form, bootstrap) or distributional assumptions (e.g., normality). Please provide these details.\n\nCalibration step M and history step H (Table 9) are empirically determined (lines 631-647). Please provide a theoretical justification or sensitivity analysis for their impact on performance.\n\nThe cross-token prefetching framework claims to hide latency (Figure 5), but lacks a detailed breakdown of overhead components (e.g., prediction time, data transfer, synchronization). Can you quantify the contribution of each component to latency and explore potential optimizations?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-28"}

comments: SUMMARY:
This work finds Temporal Patterns in the attention scores of RoPE transformers. Based on this Temporal Patterns, this work use CNNs to predict critical tokens of time t+1 based on Attention scores patterns before t. This work computes the full attention score every M steps to avoid error accumulation. The author uses LongBench and simple Needle-in-a-HayStack tasks to evaluate the proposed method.

STRENGTHS AND WEAKNESSES:
## Strengths: 
- The discovered Temporal Patterns in attention scores of RoPE transformers is interesting. 
- This work also provides a reaonable explaination of the sequential patterns from RoPE viewpoint. 
- Related works are well discussed. 

## Weakness:
- The writing remains a large improvement space. E.g. line 146 mentioned the notion $A^{comp}$, but doesn't tell what is $A^{comp}$. 
- Some explaination are still questionable. E.g. line 171 High query self-similarity.  We know that softmax attention is a normalized score. This author only show the high similarity (0.87) of consecutive queries. To support the explaination in line 171, one need to also show the similarity of queries at far distance. Moreover, ther author doesn't explain the reason of Seasonal Pattern. 
- the calibration step M is 5, a small number. AKA, every 5 steps the proposed method needs to 1) compute attention scores on the whole sequence; 2) fetch KV cache of the whole sequence. The former requires high computation cost, while the later requires high memory cost. 
- The evaluation are conducted on narrow tasks, a simple needle-in-a-haystack task and a longbench task.

QUESTIONS:
- What is the reason of seasonal pattern? 
- What is the performance of the proposed method on other tasks? e.g. RULER qa tasks, GSM8K, MMLU?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-06-30

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses the high memory and latency costs of maintaining the full Key–Value (KV) cache during long-context decoding in large language models. It proposes AttentionPredictor, the first learning-based method for predicting future attention scores via a lightweight, shared convolutional network. The method captures patterns in past attention history to identify critical tokens for KV cache retention. When combined with cross-token prefetching, which overlaps prediction and I/O, the approach maintains accuracy close to that of the dense model under high compression, resulting in significant generation speed-ups.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n* The paper proposes an efficient, learning-based method to predict attention scores, which performs well on downstream tasks and in CoT settings.\n* The authors discover a periodic seasonal pattern in attention score distributions.\n* Cross-token prefetching is used effectively to hide the overhead introduced by the AttentionPredictor.\n\nWeaknesses:\n\nMotivation and Contributions:\n* In Section 3.1, you state that \"attention recovery rate reflects the amount of information preserved after compression.\" Could you please define information formally here? The goal of KV compression is arguably to minimize the error in model output caused by compression, which may or may not correlate with attention recovery rate. MagicPig (Chen et al., 2024), in Appendix F, discuss why Top-K is a biased estimator of attention output.\n* You claim to discover temporal patterns in attention scores, but MInference, which you cite, has already identified the sequential and re-access patterns. In Lines 91–93, you write: \"They struggle to capture the dynamic temporal patterns within attention scores accurately.\" Could you provide evidence to support this claim?\n* Some theoretical explanations are unclear. In Lines 190–196, could you please further elaborate further how the formulation of RoPE leads to consistent attention along diagonals and how the periodicity of the cosine function explains the periodic slashes of the sequential patterns?\n\nEfficiency Claims and Measurements:\n* The abstract claims a 5.6× speed-up and states that the method \"significantly outperforms state-of-the-arts,\" suggesting a general breakthrough. However, this speed-up is achieved in a specific cache-offloading setup, which is not clearly defined or explained in Section 4.3.\n* The memory impact of cross-token prefetching is not discussed. My understanding is that prefetching the KV cache for the next token prediction increases GPU VRAM usage. If so, this would be unfair compared to baselines that use cross-layer prefetching.\n* There is no compute breakdown provided to understand the overhead caused by the AttentionPredictor.\n* The provided code is incomplete and only includes the implementation of the proposed method.\n\nPerformance Claims:\n* In Line 289, the authors claim that their method surpasses the performance of other methods. However, the paper lacks statistical analysis, which is required to establish the significance of the results.\n* The same issue applies to the CoT reasoning results. This is especially important given that (unless I'm mistaken) the benchmark uses a small number of samples.\n\nConference Checklist:\n* Q4–Q5: The authors do not release the code required to reproduce baselines, which limits understanding and reproducibility.\n* Q7: Standard deviations and statistical significance is not discussed.\n* Q8: Experimental setup details in Section 4.1 are minimal. The only information provided is: \"We conducted experiments on NVIDIA A800 (80GB) GPUs.\"\n\nQUESTIONS:\nClarifications:\n* Line 35: Could you clarify what is meant by static patterns? My interpretation is that these are patterns constant across inputs, but methods like H2O are dynamic, selecting different critical tokens per input. If so, I don't understand this point.\n* Lines 36–38: What exactly are dynamic patterns? What does it mean to say \"they only represent keys or hidden states\"? Based on this, I don’t fully understand why this motivates modelling the distribution of attention scores.\n\nRequests for Additional Discussion:\n* Could you more clearly describe the efficiency experiment setup, and include a comparison where offloading is not used?\n* Could you explicitly discuss how token prefetching affects memory requirements?\n\nLIMITATIONS:\nNo, the authors didn't properly addressed the limitations of their work which I discuss in Strengths and Weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents AttentionPredictor, an innovative approach for KV cache compression in large language models (LLMs) by predicting next-token attention scores during the decoding phase. The method employs a lightweight convolutional model to identify and retain only critical tokens, significantly reducing memory usage while maintaining model performance. Additionally, the authors propose a cross-token prefetching framework to hide the computational overhead of attention score prediction. The approach achieves impressive results, including 13$\\times$ KV cache compression and 5.6$\\times$ speedup while maintaining comparable performance to full cache inference.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1) The idea of learning-based attention predictor is novel. It also surprises me that the learned predictor generalizes well on unseen datasets.\n2) The unified convolutional model is remarkably lightweight (only 21KB) compared to alternatives like SeerAttention (101MB), making it practical for deployment.\n3) Comprehensive evaluations across multiple benchmarks demonstrate consistent performance improvements over state-of-the-art methods.\n4) The paper is well-motivated and clearly written.\n\nWeaknesses:\n1) It is a little confusing that since we get the indices of critical tokens from the attention predictor (see Algorithm 1), does it mean we need to keep all KV caches? If so, it is inappropriate to claim it is a method for KV cache compression.\n2) While the training process is described in Appendix B.1, more details about the CNN architecture (number of layers, kernel sizes, etc.) would be helpful.\n3) The computational overhead of the prediction model itself could be better quantified, especially since it runs at every decoding step.\n\nQUESTIONS:\nSee weaknesses.\n\nLIMITATIONS:\nSee weaknesses.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 4, "ORIGINALITY": 3, "DATE": "2025-07-01"}