CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration

We sincerely thank the reviewer for your valuable time and constructive feedback. In the following, we provide our responses to each question.

Weaknesses 1: The main concern I have with this paper is the issues of scalability. In the paper, the largest model being evaluated is 9B. It would be better if the authors could provide results on larger models (30B/70B).

We have tried our best to expand experiments on three larger models (30B&14B) in the following. Due to limited resources, we are unable to finish experiments on 70B models during this rebuttal period, but promise to include the results of 70B models in the final version.

The context length was set based on the maximum length supported by the model, with 128k for Qwen2.5-32B, Yi-34B-200K, and 16K for Phi-4-14B. The results in the above table demonstrate that CateKV scales effectively, delivering near-full-attention accuracy for 30B and 14B models, surpassing baseline methods such as SnapKV and PyramidKV.

Weaknesses 2: Some references [1-3] are missing in the related work for KV cache optimization.

Thank you very much for recommending these references, and we will carefully incorporate the discussions of these papers in §2.1 (KV Cache Eviction Algorithm), which is initially summarized as follows:

Keyformer proposes a score-based KV cache eviction algorithm that selectively retains only 'key' tokens with high attention weights, reducing cache size and memory bandwidth. Q-Hitter introduces a hybrid KV cache eviction criterion combining token importance (Heavy Hitters) and quantization-friendliness, enabling aggressive sparsification and low-bit quantization. ALISA designs a token-prioritization-based KV cache eviction algorithm using Sparse Window Attention (SWA) to dynamically reduce memory footprint, coupled with system-level optimizations for efficient caching-recomputation trade-offs. While these KV cache eviction methods demonstrate computational efficiency, they inherently incur information loss due to the reduction of attention across all heads or the quantization of data precision. In contrast, CateKV preserves the complete KV cache for adaptive heads, maintaining critical contextual information while still achieving efficiency through selective eviction in consistent heads.

[1] Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference, MLSys 2024.

[2] Q-Hitter: A Better Token Oracle for Efficient LLM Inference via Sparse-Quantized KV Cache, MLSys 2024.

[3] ALISA: Accelerating Large Language Model Inference via Sparsity-Aware KV Caching, ISCA 2024.

We sincerely thank the reviewer for your valuable time and constructive feedback. Below, we provide our responses point-by-point.

Experiments: Latency breakdown

We very appreciate the reviewer's constructive suggestion, and follow the advice to present a fine-grained latency analysis (here) for your reference.

Missing References

Thank you for your recommendation. We will carefully include MagicPIG and references suggested by other reviewers in the discussions of the submission.

W1: Assumption that ... may lead to inconsistencies.

Given a certain long-context input, we conjecture that the functionality of each head can be well characterized, even for the head where the focus may change. Informally, it resembles the functionality of human brain regions, which are categorized to play different roles. This pattern stability is well supported by empirical findings in DuoAttention, our cross-model consistency (here) and cross-task stability analysis (here), which we believe is inherently related to the working mechanism of LLMs. Regarding the implication of this phenomenon, we would like to give a plausible explanation: during inference, LLMs rely on both

• Heads that dynamically adapt to task-specific information, i.e., adaptive heads. For retrieval tasks with dynamic focus shifts, adaptive heads retaining more KV cache specifically handle such varying information.

• Heads that consistently attend to globally important tokens, i.e., consistent heads. Figure in this link demonstrates that the positions focused on by consistent heads exhibit strong similarity, indicating that some tokens are crucial for the model's understanding globally.

W2: Fixed Ratio of Adaptive vs. Consistent Heads.

First, the reason that we chose a fixed ratio here is that we focus on the task-agnostic scenario where previous works primarily considered. This ensures a fair comparison with the baselines under the same settings.

Second, CateKV can easily adapt to task-aware scenarios by setting task-adaptive head ratios. The table in this link exhibits different adaptive head ratios for tasks on RULER-128K (LLaMA-3.1-8B), where adjusting the ratio per task allows for lower ratios without sacrificing accuracy.

We will add a more detailed discussion in the revised submission.

W3: No Explicit Theoretical ... retention accuracy.

To address the concern, we present the following theoretical analysis of CateKV.

Lemma 1: Let $G$ denote our CV-based function hypothesis, $F$ denote the real-value class defined by the binary cross-entropy loss composite on $G$, and $N$ denote the sample number of the reference dataset. Then, with the probability $\delta$, we have the Rademacher complexity bound,

$\forall f\in F, P_{\text{head}}\left(\mathbb{E}[f]-\frac{1}{N}\sum_{n=1}^Nf_n \leq 2\mathcal{R}_N(F)+\sqrt{\frac{2\log{\frac{2}{\delta}}}{N}}\right)\geq 1-\delta$,

where $\mathcal{R}_N(F)$ is the conditional Rademacher average.

Let $P_1$, $P_2$ denote respectively the probabilities of correctly classified consistent heads and adaptive heads, and $\bar{P}\_1$, $\bar{P}\_2$ denote respectively the probabilities of misclassified consistent heads and adaptive heads, where we have $P\_1+P\_2=P\_{\text{head}}$ and $\bar{P}\_1+\bar{P}\_2=1-P\_{\text{head}}$. Then, we can decompose the probability in the above lemma with fine-grained analysis as follows.

Theorem 1: Let $\eta_1$, $\eta_2$ denote respectively the retention ratios of consistent heads and adaptive heads, and $\eta_1^*$, $\eta_2^*$ denote their optimal retention ratios correspondingly. Define the retion accuracy of different cases $r_{i,j}=\eta_i^*\mathbb{1}[\eta_j > \eta_i^*]+\eta_j(1-\mathbb{1}[\eta_j > \eta_i^*])$ by comparsing the retention budgets with the optimal budgets, and the hypothesis that the query attention score is the best description in order holds with the probability $\lambda$. Then, the token retention accuracy has the following inequalities,

$P\_{\text{token}} = \lambda(r_{1,1}P\_1 + r\_{2,2}P\_2 + r\_{2,1}\bar{P}\_1 + r\_{1,2}\bar{P}\_2)\geq \lambda \left(\min(r\_{2,1}, r\_{1,2}) + \left[\min(r\_{1,1}, r\_{2,2})-\min(r\_{2,1}, r\_{1,2})\right]P\_{\text{head}}\right)$

In this theorem, three factors, i.e., $\lambda$, budge control part and head identification accuracy $P_{\text{head}}$ make critical effect about the worst-tken retention accuracy in CateKV, which are actually also verified in the submission (like Figures 3, 4, 7 and Tables 2, 3). Due to space limitation, we will give more details and remarks on this theorem in the reviewer-author discussion phase.

We sincerely thank the reviewer for your valuable time and constructive feedback. In the following, we provide our responses point-by-point.

Claims: About the design of observation window

Here, we provide a more detailed explanation regarding the design of the observation window:

1. The observation window helps avoid quadratic memory and computational costs. In the pre-filling stage, the full attention matrix has size $n \times n$ (where $n$ is the context length). Analyzing the full attention matrix directly increases the computational complexity from $O(n)$ to $O(n^2)$, posing challenges in memory and computation, especially for long inputs. Using the full attention matrix for head classification led to an out-of-memory (OOM) error on a single A100 GPU for inputs larger than 8k tokens

2. The attention within the observation window is more representative. Due to the causal mask, the last query tokens capture global information similar to what is needed during decoding. Including intermediate query tokens may cause erroneous guidance, as these tokens have partial information. The table below shows that increasing the observation window size from 64 to 4k with $r = 0.4$ results in negative performance gains.

We will enrich our explanation and add the above analysis in the submission to well justify our choice. Thank you very much for pointing out the insufficiency of this claim.

Q1: Applicability of CateKV to traditional LLM Eval tasks and the role of long-context

We follow the research line of LLMs under long-context inference, which causes a significant bottleneck in memory consumption and computation latency. Regarding its practical value, some examples, like the RAG (retrieval-augmented generation) applications, long-document understanding, multi-turn interaction, etc., easily make the LLMs face long-context inference scenarios. In this case, it is critical to accelerate LLM inference to avoid the vanilla catastrophic response speed (it can be minute-level for A100 with 80G memory under 100k-long context).

Naturally, the existing long-context acceleration methods, including CateKV, are applicable to traditional short-context LLM evaluation datasets. We tested the performance of CateKV and baseline methods on MMLU and MMLU-Pro, with all methods retaining approximately 25% of the KV cache and utilizing a one-shot prompting approach. The results in the table below show that CateKV still maintains performance comparable to full attention. It is worth noting that in short-context scenarios, attention to recent tokens is more likely to retain complete contextual information, allowing simpler methods (such as StreamingLLM) to maintain much of the performance.

Essential Reference: [1] Model tells you what to discard: Adaptive kv cache compression for llms

Thank you for the recommendation and we will incorporate the discussion of FastGen [1] in Section 2.2. For the comparison, we need to clarify that in our practice, FastGen easily incurs OOM when going beyond 8K. This is mainly because of the usage of the complete attention matrix during its pre-filling stage. We will add more explanations in the updated version.

Q2: How is the method generalize without reference dataset ?

While CateKV uses a reference dataset to achieve static identification, the method proposed in Section 3.2 actually can perform without a reference dataset, yielding the dynamic identification. In the following table, we show that the dynamic identification method exhibits similar performance as the reference-based static identification method, both comparable to or surpassing full attention, highlighting the generalizability of CateKV.

However, the following drawbacks arise without a reference dataset:

• At the pre-filling stage, the extra sample-wise computation needed for head identification increases computational overhead.
• The need for a complete observation matrix in the pre-filling stage to assist classification prevents integration with sparsification-based pre-filling acceleration methods.

We will include the above analysis in the submission to improve the clarity.

We appreciate the supportive comments and the comprehensive evaluation of our method. In the following, we provide our responses point-by-point.

Weakness 1: About setting the best threshold

We select the threshold based on a fixed ratio of consistent-to-adaptive heads, as while CV scores vary significantly across inputs, the proportion of heads exhibiting consistent patterns remains remarkably stable within the model architecture.

In task-agnostic scenarios, determining the optimal threshold or adaptive head ratio for each sample is challenging. One potential approach is to incorporate attention weight recall as an additional criterion in the head classification process. If the attention weight recall exceeds a threshold (denoted as $t$), these heads can also adopt the consistent pattern.

As shown in the table below (LLaMA-3-8B,$r=0.4$,$t=0.98$), this recall-threshold method enables input-adaptive ratio adjustment while maintaining performance (Avg 84.32% vs 84.61% CateKV). However, the marginal gains suggest this approach may not be universally optimal. We will further explore dynamic threshold adjustment methods in future work.

For task-aware scenarios, please refer to our response to Reviewer 3ByN'Weakness 2, where we discuss dynamic ratio selection in task-aware contexts.

Weakness 2: About the meaning of Topk curve.

Figure 4 compares different methods' ability to approximate the ground-truth top-k attention weights during decoding. The top-k curve represents the upper bound of achievable recall, as it directly uses the actual highest attention weights (which are unknown in advance during real inference).

This does not imply top-k is universally "better," but rather that it serves as the theoretical optimum. The key advantage of CateKV/Quest is that they avoid computing full attention, making them practical for real-world deployment.

Weakness 3 and Question: Comparison Between CateKV and DuoAttention

1. Differences between CateKV and DuoAttention: From a conceptual perspective, consistent heads emphasize the overall similarity of attention across different sequences, while Streaming Heads prioritize whether initial and recent tokens cover most of the attention. Thus, in terms of token selection, streaming heads only focus on the initial and recent tokens for attention, whereas consistent heads have a broader search space. To further illustrate this, we randomly selected 100 samples from WikiText and performed inference using LLama3-8B. Our analysis shows that approximately 72% of the tokens selected by consistent heads do not belong to the initial or recent tokens, demonstrating greater flexibility compared to streaming heads.

2. Analysis of the rapid drop: From the perspective of attention recall, the number of streaming heads within a model is inherently limited. As shown in Figure 4, a larger proportion of heads recall less than 0.8 of the attention weights under the streaming pattern. An overreliance on the streaming pattern results in significant information loss, which explains the rapid drop in the performance of DuoAttention when the ratio is below 0.4. In contrast, even with an aggressive increase in the proportion of consistent heads, the consistent pattern effectively preserves critical tokens, including those that capture global information from prior attention. These retained tokens enhance the robustness of CateKV, making it more effective than DuoAttention.

Suggestion: Using Soft Attention Weight Matrix for Method Design

Utilizing a soft attention weight matrix is an intuitive approach. However, we believe that the values in the soft attention weight matrix reflect the importance of individual tokens, whereas the binarized observation matrix used in CateKV more effectively captures the relative importance of tokens. Since sparse attention methods aim to identify the top-k important positions, sequence similarity should be evaluated by the alignment of their top-k attention indices, which reflects a relative relationship. Additionally, the values of attention weights fluctuate with the query, and attention sinks in some heads complicate the use of a soft attention matrix for classification.

We also conducted an experiment where we removed the binarization step from the identification process (a soft version). When $r = 0.5$, on the RULER-128K niah-multikey3 task, the performance dropped by 14.58 points (from 63.54 to 48.96) compared to full attention on LLama-3.1-8B.