RazorAttention: Efficient KV Cache Compression Through Retrieval Heads

Q3: The retrieval pattern among different queries.

As clarified in Q1, the selection of retrieval heads is consistent across queries, indicating that these patterns are primarily model-driven rather than query-specific.

Q4: The effectiveness of the compensation token.

Our ablation studies (Figure 7 in the main paper) show that compensation tokens significantly improve accuracy, even with identical KV cache allocations. This evidence highlights the critical role of compensation tokens in mitigating information loss due to KV cache reduction. Exploring the combination of compensation tokens with other baseline algorithms presents an exciting direction for future work, and we aim to explore this in subsequent studies.

Q5: The KV cache budget for different algorithms.

To ensure fairness, all baseline methods were evaluated under the same total compression ratio, as detailed in Table 2 of the main paper. While the absolute KV cache budget may vary proportionally with input length due to RazorAttention’s dynamic compression strategy, the total compression ratio remains consistent across all methods for comparison purposes.

We sincerely appreciate your constructive feedback of our paper, and we will make certain updates w.r.t. your concerns in our revised version. Below are our responses to your concerns.

Q1: Whether the retrieval pattern holds in general query settings.

We appreciate this observation. Retrieval patterns remain highly consistent across diverse input queries, confirming they are predominantly model-based rather than query-specific. As detailed in Section 2 of our revised supplementary materials, we demonstrate the stability of retrieval head selection using varied “Needle in a Haystack” queries, where the identified heads remain unchanged across inputs.

Further, we validate these findings on complex datasets such as LongBench and InfiniBench/Ruler (see results below), which encompass diverse and challenging task types. RazorAttention consistently retains the original performance. This evidence supports the robustness of the observed retrieval patterns.

Q2: Comparison with SnapKV.

Thank you for this suggestion. SnapKV is indeed a prominent method for KV cache compression. Below, we provide a comparative analysis of RazorAttention and SnapKV in both query-aware and query-agnostic settings. While SnapKV excels in query-aware scenarios, its performance degrades significantly in query-agnostic tasks. RazorAttention’s training-free design and headwise sparse patterns ensure superior robustness across diverse queries.

Additionally, we include complete results for LongBench (Llama3.1-8B-instruct) to demonstrate the practical advantages of RazorAttention over SnapKV, especially under query-agnostic settings. This discussion will be incorporated into the revised paper.

We sincerely appreciate your constructive feedback of our paper, and we will make certain updates w.r.t. your concerns in our revised version. Below are our responses to your concerns.

Q1: Comparison with [1].

Although [1] does provide a similar finding as ours, directly retaining the KV cache in induction heads introduces a significant accuracy drop for tasks such as Needle In A Haystack. We introduce two essential strategies that go beyond the use of induction heads to achieve lossless performance:

1. Expanded Definition of Retrieval Heads: The concept of induction heads was initially introduced in [1], where these heads were identified as attention heads that follow a token-retrieval pattern (attending to previously seen tokens). In [2], the authors extended this observation to long-context scenarios, showing that induction heads are more critical than other heads under such conditions. However, these works considered only induction heads as retrieval heads. Our analysis reveals that echo heads are equally crucial for retaining model performance. For instance, when tested on the Needle in a Haystack dataset, using the definition of retrieval heads from [1] leads to an approximate 10% accuracy drop. By redefining retrieval heads to include both induction and copy heads, we retain the full performance of the model (see Figure 5 in our paper).
2. Compensation Token Strategy: Directly discarding remote tokens in non-retrieval heads results in severe performance degradation, with accuracy drops exceeding 30%. To address this, we designed the compensation token strategy, which condenses the dropped information from these heads into a compact form. This strategy ensures that essential information is preserved while enabling efficient KV cache compression.

Q2: Comparison with [2].

An obvious advantage of the training-free algorithm is that it can be used as a plug-and-play component for advanced LLM serving system, such as TRT-LLM and Triton, and training-free lossless compression is usually a much more challenging task compared to the training-based algorithms. Moreover, we believe the compensation token we introduced and our head selection method (inductive + echo heads) provide valuable insights into the underlying functionality of LLMs in long-context scenarios.

Q3: How to quickly determine the suitable proportion of Induction Heads and Echo Heads for different models.

We find that for most open-source models, about 15%~20% induction heads+1% echo heads shall be enough without training. However, if there still exists certain performance loss, our recipe is to increase about 3% of the induction heads at one step. We will provide a discussion section about this in our revised version.

Q4: General theoretical analysis explaining why retrieval heads are widely present in various models.

We believe this is a quite important and challenging topic, and we will leave this for future study.

We sincerely appreciate your constructive feedback of our paper, and we will make certain updates w.r.t. your concerns in our revised version. Below are our responses to your concerns.

Q1: The efficiency evaluation.

We evaluate RazorAttention’s decoding latency and throughput on GLM-9B-1M model on 8 Ascend 910B NPUs under different input lengths for both prefill and decoding

Below is the maximum throughput improvement of using RazorAttention under different input lengths:

Q2: Overhead Analysis.

After computing the KV Cache at each layer, we directly calculate the average of the tokens that need to be compressed. Only the averaging operation is introduced, and the capacity of the compressed KV Cache is significantly reduced. When transferring the data back from the compute unit to the storage unit, the data transfer overhead is also reduced. In summary, this operation does not introduce significant performance overhead, so our overall performance improves in long-sequence scenarios for the prefill stage as well.

Q3: Comparison with other algorithms.

Below we evaluated the accuracy of Snapkv at the same compression rate using LongBench on the Llama3.1-8B-instruct model.

However, when the future queries are unknown for SnapKV, its performance drops significantly. Below, we present the results of SnapKV in both query-aware and query-agnostic settings for Needle In A Haystack: (see detailed Figures in Section.1 of our supplementary updated):

Only RazorAttention is still feasible in this case. This resilience is due to its training-free design and headwise sparse pattern, which ensure minimal information loss across various query types.

Our algorithm can be perfectly combined with current token-dropping algorithms, such as H2O, SnapKV, MInference. Since we only constrain the cache budget in the short heads while in long heads the KV cache can be further compressed using these algorithms.

Q4: Results of RazorAttention under different compression ratios.

The reason why we only consider one specific ratio is that we found the ratio of the necessary retrieval heads is about 15%. Using fewer retrieval heads would introduce a certain accuracy drop (see Table.4 in our paper) while using more heads gives no more accuracy gain. Therefore we only report the comparison result under this compression ratio for all baselines. We will clarify this in our revised version.

Thanks for the answers to my questions; I've increased the rating. I have some follow-ups here:

• For the efficiency tests, please also use GPUs like A100, H100, etc.
• I would like to see performance when combining with other methods like token-dropping, quantization, etc. To see how far can we go in terms of KV cache compression.
• For figure 4 in the supplement, what is the model? And can you further elaborate on how to interpret the result and what heads are retrieval heads?

We sincerely appreciate your constructive feedback of our paper, and we will make certain updates w.r.t. your concerns in our revised version. Below are our responses to your concerns.

Q1: Novelty of RazorAttention.

Our work is the first to apply induction heads for KV cache compression, and we introduce two essential strategies that go beyond the use of induction heads to achieve lossless performance:

1. Expanded Definition of Retrieval Heads: The concept of induction heads was initially introduced in [1], where these heads were identified as attention heads that follow a token-retrieval pattern (attending to previously seen tokens). In [2], the authors extended this observation to long-context scenarios, showing that induction heads are more critical than other heads under such conditions. However, these works considered only induction heads as retrieval heads. Our analysis reveals that echo heads are equally crucial for retaining model performance. For instance, when tested on the Needle in a Haystack dataset, using the definition of retrieval heads from [1] leads to an approximate 10% accuracy drop. By redefining retrieval heads to include both induction and copy heads, we retain the full performance of the model (see Figure 5 in our paper).
2. Compensation Token Strategy: Directly discarding remote tokens in non-retrieval heads results in severe performance degradation, with accuracy drops exceeding 30%. To address this, we designed the compensation token strategy, which condenses the dropped information from these heads into a compact form. This strategy ensures that essential information is preserved while enabling efficient KV cache compression.

Q2: More comparison.

According to your suggestion, here we also include the comparison of RazorAttention with Snapkv at the same compression rate, under both query-aware and query-agnostic settings. While SnapKV excels in query-aware scenarios, its performance degrades significantly in query-agnostic tasks. RazorAttention’s training-free design and headwise sparse patterns ensure superior robustness across diverse queries.

Additionally, we include complete results for LongBench (Llama3.1-8B-instruct) to demonstrate the practical advantages of RazorAttention over SnapKV, especially under query-agnostic settings. This discussion will be incorporated into the revised paper.

We sincerely appreciate your constructive feedback of our paper, and we will make certain updates w.r.t. your concerns in our revised version. Below are our responses to your concerns.

Q1: Clarify the experimental setup for H2O.

For H2O, we ensured a fair comparison by setting the KV cache budget to match that of RazorAttention for each sentence. We appreciate your suggestion and will clarify this in our revised version.

Q2: Efficiency evaluation results.

We evaluate RazorAttention’s decoding latency and throughput on GLM-9B-1M model on 8 Ascend 910B NPUs under different input lengths for both prefill and decoding

Below is the maximum throughput improvement of using RazorAttention under different input length:

Q3: Comparison with previous work [1].

We acknowledge the similarities with [1] and highlight the following key distinctions:

1. Broader Definition of Retrieval Heads: In [1], the authors define retrieval heads as only the induction heads. In contrast, our analysis shows that beyond the inductive heads, about 1% of echo heads play an equally essential role in retaining model performance (refer to Figure 5 in our paper). Using the narrower definition in [1] results in an approximate 10% accuracy drop in the Needle In A Haystack benchmark, emphasizing the importance of our algorithm design.
2. Compensation Token Design: Directly applying a truncation on the non-retrieval heads' KV cache introduces a significant accuracy drop (>30%). Our compensation token design mitigates this issue effectively and ensures minimal information loss. In summary, while [1] provides foundational insights into retrieval heads, our work extends these findings to achieve lossless KV cache compression. We will add a detailed discussion of these distinctions in the revised version.

Q4: Elaborate on the “Retrieve and Process” hypothesis.

The “Retrieve and Process” hypothesis originates from the observed behavior of attention heads in LLMs. Most attention heads focus primarily on recent tokens, while only a few retrieval heads attend to remote tokens to extract relevant information. This mirrors the human cognitive process: recalling pertinent information before generating a response. We will clarify this point in our revised version.

Q5: What does "Protect" mean.

In line 234, “protect” refers to preserving the KV cache in selected heads while discarding remote tokens in others.

We sincerely appreciate your constructive feedback of our paper, and we will make certain updates w.r.t. your concerns in our revised version. Below are our responses to your concerns.

Q1: More experimental results.

As you suggest, we have included the results of infinite bench and ruler below with Llama3.1-8B-instruct under the same setting of our paper. Due to time constraints, only a subset of the experiments have been included. We will present additional results upon completion of the remaining experiments.

Meanwhile, here we include the full results of RazorAttention on Longbench with Llama3.1-8B-instruct .

Q2: The exact compression ratio.

We compress 85% of the non-retrieval heads, with the compressed size set to max(4000, N/5), where $N$ represents the length of the input sequence. The compression ratio on each subset of LoneBench is correlated with the dataset's sequence length, and we ensured a fair comparison for H2O by setting the KV cache budget to match that of RazorAttention for each sentence.

We sincerely appreciate your constructive feedback of our paper, and we will make certain updates w.r.t. your concerns in our revised version. Below are our responses to your concerns.

Q1: Comparison with SnapKV.

We agree that SnapKV is a competitive approach for KV cache compression, especially when user queries are known before generation. Below, we present the results of SnapKV in both query-aware and query-agnostic settings: • Query-aware: SnapKV demonstrates impressive compression ratios and accuracy under known query scenarios. • Query-agnostic: In scenarios where the user query is not pre-defined, SnapKV’s performance deteriorates significantly, as shown by the following results of Needle in a Haystack (see Section.1 in our supplementary updated):

This resilience is due to its training-free design and headwise sparse pattern, which ensure minimal information loss across various query types. We believe this discussion is crucial and will incorporate it into the revised version to emphasize the practical advantages of RazorAttention over query-dependent methods like SnapKV.

Q2: Detailed setting of experiments and more long-context benchmarks.

We acknowledge the variation in Needle in a Haystack accuracy (approximately 2%) under different prompts. We will clarify the specific prompt used for RazorAttention and all baselines in our revised version. Meanwhile, Here we include the full results of RazorAttention on Longbench with Model Llama3.1-8B-instruct under the same setting with our paper. We also report the accuracy of Snapkv at the same compression rate.

As you suggest, we have included the results of infinite-bench and ruler below. Due to time constraints, only a subset of the experiments have been included. We will present additional results upon completion of the remaining experiments.

Q3: The efficiency evaluation.

We evaluate RazorAttention’s decoding latency and throughput on GLM-9B-1M model on 8 Ascend 910B NPUs under different input length for both prefill and decoding:

Below is the maximum throughput improvement of using RazorAttention under different input length:

Q4: Comparison with [1] and [2].

We respectfully emphasize that RazorAttention is the first method to utilize a headwise sparse pattern for KV cache compression. Below, we summarize the key differences between RazorAttention and the recent works [1] and [2]:

1. Comparison with [2]: In [2], the authors rely on query-based importance metric to select the “important” tokens (Section 3.3 in their paper). While this approach can perform well when the query is known, it introduces the risk of information loss for future queries, just like SnapKV.
2. Comparison with [1]: Unlike [1], which employs continual training and knowledge distillation to recover model performance after compression, RazorAttention achieves a nearly lossless compression without any additional training. This training-free design is a more challenging and practical approach. We believe the compensation token we introduced and our head selection method (inductive + echo heads) provide valuable insights into the underlying functionality of LLMs in long-context scenarios.
3. Regarding Figure 13(1) in [2]: In [2] authors did not provide specific details about this experiment. We think the discrepancy might be due to their exclusion of echo heads, which we have found to play a critical role in our results. We are confident in the reproducibility of our results.

Q5: Full results of LongBench.

Below we present the full results of Longbench with Llama3.1-8B-instruct:

[1]. DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads

[2]. Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning

We sincerely appreciate your valuable feedback. In response, we will incorporate the efficiency results tested on A100 and the performance results on SharedContextBench (SCB) in the updated version of our manuscript. We have conducted the following experiments to address the concerns about the performance of RazorAttention and SnapKV under different compression ratios, and the combination of RazorAttention with other methods.

1. Performance of RazorAttention and SanpKV under different compression ratios:

Below we present the results for these two methods under 15%, 20%, and 40% KV cache budget for Needle In A Haystack. We use Llama3.1-70b and masked out the queries for both methods (to simulate the situation where the user queries are unknown).

2. Combination of RazorAttention and Quantization:

Below we further compress the KV cache under RazorAttention into MX-FP4 on Llama3.1-8B and report its results on LongBench.

3. How does DuoAttention choose the gate for each head:

In DuoAttention, authors use continual training together with distillation before compression. This training process determines a score for each head, and they set the heads with a larger score as full-cache attention heads afterward in compression.

I appreciate the authors for providing a comprehensive rebuttal. With a proper comparison to SnapKV and broader evaluation coverage, many of my concerns have been alleviated. However, there are still a few areas the authors should better address:

• Efficiency Evaluation. The efficiency metrics, gathered on non-standard NPUs, offer very limited insights to the community. It is essential that the authors provide GPU-based efficiency reports, especially for an efficiency-focused paper. My assumption is that the authors currently lack a complete kernel implementation for the proposed method. This is typically a disqualifying factor, as it leaves the end-to-end efficiency of the proposed approach unverified. That said, given that DuoAttention achieves notable efficiency gains under thorough evaluation and has been open-sourced, there is little doubt RazorAttention could achieve similar results, given their nearly identical inference patterns. The authors should consider borrowing the kernel implementation of DuoAttention, crediting it appropriately, and providing comparable efficiency results.

• Query Dependency vs. SnapKV. According to the newly added results, there is only a slight performance difference between SnapKV and the proposed method. The authors correctly observe that SnapKV's dropping strategy heavily depends on query awareness (e.g., whether it is within the window). However, the current evaluation does not adequately support this claim. To my knowledge, SharedContextBench (SCB) is the only dataset that evaluates multi-round scenarios in a long-context setting. I recommend that the authors include results from SCB in the final version. SCB has already demonstrated the performance degradation of SnapKV, and given the (seemingly) task-agnostic mechanism of RazorAttention’s retrieval/induction heads, I am confident it would outperform SnapKV.

Normally, I would request these results before adjusting my score. However, in this case, I am inclined to increase my score now due to 1) the limited time available (partially due to my own late response), 2) the confidence that supportive results can be achieved based on findings from concurrent work, 3) SCB’s ongoing review at ICLR, which may delay access to finalized supporting materials, and 4) most importantly, RazorAttention is the first work (and predates others by a considerable margin) to leverage induction/retrieval head findings for KV cache compression. That said, I ask the authors to confirm that both of these improvements will be included in their camera-ready version or, if SCB’s availability blocks progress, in an arXiv version released soon after.

---

Finally, I’d like to share my thoughts on several standing issues raised by other reviewers:

• Comparison with DuoAttention. The two methods indeed share many similarities. However, RazorAttention predates Duo/Not All Heads Matter by a significant margin. While the methods are concurrent and likely developed independently, this alone should shield RazorAttention from criticism in this regard.

• Limited Compression Ratio. The evaluation already demonstrates that RazorAttention remains effective at a relatively high compression ratio (~3x). While finer-grained compression is a fair request, I see it more as a bonus rather than a core requirement as near-lossless performance at ~3x is already pretty usable. That said, the authors should consider expanding Table 4 to include results with >15% protected heads for completeness.

• Combining with Other KV Cache Compression Methods. KV cache compression spans multiple dimensions, and methods addressing each dimension are still evolving. While combining approaches is possible, such combinations are likely out of scope for this paper and, most importantly, probably too heavy of an ask for a paper with seven reviewers already demanding numerous additional results.

---

One additional question: does DuoAttention "employ continual training and knowledge distillation to recover model performance after compression"? I was under the impression that Duo does distillation to update its gates, identifying heads worthy of a full cache, but not after the compression. Could the authors please double-check this?