DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads

We sincerely thank the reviewer for their thoughtful feedback and constructive suggestions. Below, we address the key points raised:

1. Baseline Comparisons

We acknowledge the reviewer's concern about baselines. In response, we conducted additional experiments with recent KV cache compression methods, such as SnapKV and PyramidKV, and pre-filling acceleration techniques like MInference. These comparisons are detailed in Appendix Sections A.6 and A.7.

a. Comparison with KV Cache Compression Methods (SnapKV and PyramidKV)

We evaluated DuoAttention against SnapKV and PyramidKV on the Needle-in-a-Haystack (NIAH) benchmark using Llama-2-7B-32K and Llama-3-8B-1048K models.

• Methodological Differences: SnapKV and PyramidKV prune KV tokens based on attention scores computed over a small observation window (e.g., the last 8–64 tokens). While effective for tasks where queries appear at the end of the context, this assumption limits their applicability in scenarios such as multi-turn dialogues or tasks with earlier query positioning. DuoAttention, by contrast, does not rely on positional assumptions, making it more versatile.
• Performance Results: As detailed in Figures 17 and 18:
  • Without simulation of the last tokens as generated tokens, DuoAttention matches or surpasses these baselines across all benchmarks.
  • With simulation (e.g., the last 50 tokens treated as generated inputs to mimic multi-turn dialogues), SnapKV and PyramidKV experience severe accuracy drops due to their reliance on final token positions for pruning. DuoAttention remains robust, maintaining high accuracy even under these conditions.

These results highlight DuoAttention’s adaptability and reliability across diverse scenarios, making it a more general solution for KV cache compression.

b. Comparison with Pre-filling Acceleration Methods (MInference)

While MInference accelerates pre-filling using sparsity patterns (e.g., block-sparse or vertical-slash), it does not optimize KV cache storage or decoding latency. Thus, it is orthogonal to DuoAttention rather than directly comparable.

• Integration with DuoAttention: As shown in Appendix Section A.7, MInference can complement DuoAttention by further accelerating the pre-filling stage. By applying MInference kernels to retrieval heads identified by DuoAttention, we achieve:
  • Comparable accuracy to pre-filling all heads with MInference alone.
  • Significant reductions in KV cache size (nearly halved) and decoding overhead.

This demonstrates the compatibility and efficiency of combining DuoAttention with MInference for enhanced optimization.

2. Inaccurate Statements in the Manuscript

We appreciate the reviewer pointing out inaccuracies, and we have revised the manuscript accordingly:

1. GQA Compatibility: We revised our statement in the introduction regarding baseline methods and GQA. While DuoAttention explicitly considers GQA in its design, we acknowledge that baseline methods such as StreamingLLM can also be adapted for GQA. However, we note that in some baselines implementations (e.g., PyramidKV and SnapKV), GQA compatibility is achieved by effectively converting GQA models into MHA models (e.g., by repeating KV caches before pruning and storing), which actually inflates reported compression ratios. DuoAttention achieves true GQA compatibility without such workarounds, providing robust compression for both MHA and GQA models.
2. StreamingLLM Pre-filling: We acknowledge that StreamingLLM can accelerate pre-filling. However, its accuracy is limited, comparable to pre-filling only the final segment of the context. In contrast, DuoAttention achieves both pre-filling and decoding acceleration while delivering significant memory reduction and maintaining highly competitive accuracy. We have revised the related work section to present this distinction with greater clarity and rigor.

3. Comparison with SnapKV Follow-Up Works

We appreciate the reviewer’s insights and agree that methods like SnapKV and its follow-ups address KV cache compression using dynamic sparsity and sample-specific budget allocation. However:

• Positional Assumptions: These methods rely on the hypothesis that critical queries are short and located at the end of the context. This is not always valid in real-world scenarios, such as multi-turn dialogues or earlier query positioning, which reduces their generalizability.
• DuoAttention’s Robustness: DuoAttention avoids such assumptions and operates as a more general and robust framework, adaptable to diverse tasks and sequences.

We also envision future work combining these dynamic sparsity techniques with DuoAttention, such as applying dynamic compression for retrieval heads when query positions can be anticipated. This hybrid approach could further enhance the performance of LLMs.

Thank you for your prompt follow-up question! We sincerely appreciate your thoughtful observations and the opportunity to clarify further.

1. SnapKV and PyramidKV's Limitations in Chunked Pre-filling and Multi-round Dialogue: The previously conducted NIAH experiments effectively demonstrate the inherent limitations of SnapKV and PyramidKV in chunked pre-filling and multi-round dialogue scenarios. These settings simulate cases where the context is compressed without access to the question, followed by the question being input subsequently. The results clearly show that SnapKV and PyramidKV fail to retain critical information in such scenarios due to their reliance on heuristic assumptions about query placement. In contrast, DuoAttention remains robust, demonstrating its capability to handle real-world tasks without such constraints.
2. New Experiment with Query Positioned at the Middle: Following your insightful suggestion, we conducted an additional experiment to evaluate query positioning more comprehensively. The setup includes:
   • An instruction placed at the beginning of the context: "This is a very long book with a question embedded. Please answer the embedded question at the end of the book."
   • The question inserted immediately before the needle in the middle of the haystack (e.g., "Q: What is the best thing to do in San Francisco?\n\nA: The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.").
   • A partial answer prompt at the end of the context: "Answer: The best", used to elicit the model’s response.

For this evaluation, we tested SnapKV, PyramidKV, and DuoAttention on NIAH without simulating the last tokens. The entire context—including the instruction, question, haystack, and partial answer—was provided before applying KV cache compression.

Results and Analysis: The findings from this experiment are detailed in the updated manuscript Appendix Section 12 and visualized in Figure 23. The key results are as follows:

• DuoAttention: Achieved accuracy (score = 0.901) comparable to full attention (score = 0.906), demonstrating its robustness in scenarios with variable query positioning.
• SnapKV and PyramidKV: Experienced significant accuracy degradation (SnapKV: 0.694, PyramidKV: 0.772), highlighting their dependency on the assumption that queries are positioned at specific locations (e.g., the end of the context).

These results further validate DuoAttention’s resilience to query position variability, which sets it apart from SnapKV and PyramidKV. By not relying on assumptions about query placement, DuoAttention proves to be a more robust and generalizable solution for KV cache compression.

We sincerely appreciate your detailed follow-up concerns and suggestions. Below, we address each point thoroughly and have made plans for additional experiments to further support our arguments.

1. Readability of Revisions

To improve the clarity of changes, we have marked all rebuttal-revised text in red in the updated manuscript. This should make it easier to identify and evaluate the modifications.

2. Query Positioning and 50-Token Simulation

We appreciate the opportunity to clarify this point. While SnapKV and PyramidKV do not rely on global attention weights, their heuristic assumes that queries are positioned at the end of the context. The purpose of simulating the last 50 tokens is to evaluate the robustness of these methods in scenarios where this assumption does not hold, such as multi-turn dialogue or non-linear input structures.

For example, in the Needle-in-a-Haystack (NIAH) benchmark, the query (e.g., "What is the best thing to do in San Francisco?") is placed at the end of the input. If the last 50 tokens are treated as simulated decoding (e.g., input in a new dialogue round), the remaining tokens (input[:-50]) used for pre-filling exclude the query. SnapKV and PyramidKV, relying only on a window of prior tokens (e.g., input[-58:-50]), mistakenly prune important tokens while retaining irrelevant ones (the haystack). This simulation reveals the methods’ inability to handle real-world scenarios where queries are not placed predictably.

As demonstrated in PyramidKV, even when compressing 32K to 128 with Mistral-7B-Instruct, both SnapKV and PyramidKV exhibit minimal performance degradation. However, this is only attainable based on knowing the location of the query and using the query as the observation tokens to prune the KV cache. DuoAttention doesn't rely on this heuristic. As demonstrated in our updated NIAH results, both SnapKV and PyramidKV fail when the observation tokens are not the query tokens, even at a high retention ratio like 50%, which is the same as DuoAttention.

3. SnapKV’s Claim on Query Position Insensitivity

SnapKV’s statement about query positioning insensitivity is based on hit rates for critical tokens but does not provide corresponding end-to-end accuracy results. Our NIAH experiments demonstrate that when not using the queries as the observation tokens, i.e., inputing the queries after the KV cache is pruned, these methods fail to identify and retain critical tokens, leading to degraded performance.

We acknowledge the importance of additional benchmarks to strengthen this argument and will include quantitative results within two days to provide further evidence.

4. Compatibility with Chunked Pre-Filling

SnapKV and PyramidKV depend on tokens at the end of the context for pruning, which becomes problematic in chunked pre-filling scenarios. After each chunk’s pre-filling, the observation window is limited to tokens at the end of the chunk, which are unlikely to include the critical query. This limits their effectiveness when pruning KV cache during chunked pre-filling.

Our NIAH results serve as an example of chunked pre-filling: the first chunk includes input[:-50] (without the query), and the second chunk contains the query (input[-50:]). In this scenario, both SnapKV and PyramidKV fail, highlighting their incompatibility with chunked pre-filling.

5. Clarification on DuoAttention’s LongBench Results

We respectfully clarify that our DuoAttention LongBench experiments differ from SnapKV’s setup. While we did not use the chunked pre-filling strategy in these experiments, we did simulate the last 50 tokens as generated tokens to ensure fair comparison with baselines like H2O and TOVA. DuoAttention demonstrated robust performance in this setup, maintaining comparable accuracy even when the query is input after KV cache pruning.

Our NIAH results further show that SnapKV and PyramidKV degrade significantly under similar conditions, failing to retain the necessary tokens when the query is not included in the initial pre-filling. We will further support this with results on LongBench in two days.

Dear Reviewer JVPb,

Thank you for your continued engagement and valuable feedback. We are pleased to share the results of additional experiments using another model, Llama-7B-32K-Instruct, on LongBench with a 25% KV cache budget. The detailed results are as follows:

On average, DuoAttention achieves the highest average score (34.49) across datasets, outperforming SnapKV (32.98), PyramidKV (34.35), and AdaKV (33.43).

Addressing Your Concerns:

1. Results Across Multiple Models: We now present results for two different models, Llama-3-8B-Instruct-Gradient-1048k and Llama-7B-32K-Instruct, on LongBench. DuoAttention consistently outperforms baseline methods, further supporting its effectiveness.
2. Broader Comparisons and Robustness: These results confirm DuoAttention’s superior accuracy, GQA compatibility, and resilience to real-world complexities, including varying query positions and continuous pre-filling scenarios.

Moving Forward:

We will incorporate these additional results into the main paper in future revisions, along with detailed experimental settings and analyses. Conducting these extensive comparisons required considerable time and effort, but they validate DuoAttention’s superiority comprehensively.

We kindly ask you to reconsider your rating in light of these extensive experiments, which directly address your concerns and further demonstrate DuoAttention’s merits.

Thank you again for your engagement and constructive suggestions. We sincerely appreciate your contributions to strengthening our work.

Best regards,

Authors

We sincerely thank the reviewer for their thoughtful feedback and constructive suggestions. Below, we address the key points raised:

1. Detection of Retrieval Heads

We appreciate the reviewer’s interest in the retrieval head detection process. This detection is performed offline prior to deployment and is designed to be lightweight and efficient. Specifically:

• Optimization Process: The detection process involves freezing the entire model weights and fine-tuning only the gated values assigned to each attention head, which range from 256 (8 heads x 32 layers) to 1,024 (32 heads x 32 layers) values in our experiments.
• Resource Requirements: All experiments were conducted on a single 8×A100 GPU node, and the retrieval head identification process takes less than 4 hours.
• Task Independence: This process is model-specific and does not need to be repeated for every new task, similar to the calibration stages of quantization methods like SmoothQuant, AWQ, and GPTQ.

These characteristics make DuoAttention highly usable in practical applications. The lightweight nature of the detection process ensures minimal overhead, and its independence from task-specific requirements further enhances its applicability.

2. Compression Ratio of KV Cache

DuoAttention achieves significant KV cache compression, with 75% compression for MHA models and 50% compression for GQA models, representing a substantial improvement over existing methods. During the rebuttal stage, we conducted comparisons with recent KV cache compression techniques, including SnapKV and PyramidKV, as detailed in Appendix Section A.6.

• Comparison with Recent Methods: While SnapKV and PyramidKV reportedly achieve higher compression ratios, they rely on a critical assumption that query-relevant information resides at the end of the context. This assumption often fails in real-world scenarios, such as multi-turn dialogues or tasks with queries positioned earlier in the context.

  • Methodological Differences: SnapKV and PyramidKV prune KV tokens based on attention scores computed over a small observation window (e.g., the last 8–64 tokens). While effective for scenarios with end-positioned queries, this approach lacks generalizability to tasks with diverse query positions. In contrast, DuoAttention operates without any assumptions about token placement, making it more versatile and robust.
  • Performance Results: As demonstrated in Figures 17 and 18 of the manuscript:
    • Without simulation of the last tokens as generated tokens, DuoAttention matches or surpasses these baselines across all benchmarks.
    • With simulation (e.g., the last 50 tokens treated as generated inputs to mimic multi-turn dialogues), SnapKV and PyramidKV experience severe accuracy drops due to their reliance on final token positions for pruning. DuoAttention remains robust, maintaining high accuracy even under these conditions.

  These findings highlight DuoAttention’s robustness and adaptability, outperforming these baselines in scenarios with diverse token distributions.

• Generalization and Robustness: DuoAttention’s independence from token position assumptions ensures its applicability to a broader range of real-world scenarios. Its design supports various tasks and sequences without compromising accuracy or performance.

Considering these points, we believe that DuoAttention’s compression ratios are not only competitive but also highly practical for real-world applications. Its universal applicability across diverse tasks further reinforces its robustness and value.

We hope this response clarifies your concerns and highlights DuoAttention's practical advantages. Please let us know if you have further questions or suggestions; we are happy to provide additional information.

We sincerely thank the reviewer for their thoughtful feedback and constructive suggestions. Below, we address the key points raised:

1. Experiment with Accuracy Results After Combining DuoAttention and KV Quantization

We acknowledge the reviewer’s concern about the lack of experimental results combining DuoAttention with KV quantization. In response, we have conducted additional experiments and included the results in Appendix Section A.9 of the updated manuscript. The experiments evaluate the performance of combining DuoAttention (50% sparsity) with INT4 KV Pre-rope quantization, as proposed in KIVI [1].

• Baseline (INT4 KV Pre-rope Quantization): The original model achieves an overall score of 86.7%, demonstrating a slight performance drop compared with full attention.
• DuoAttention + INT4 KV Quantization: When DuoAttention (50% sparsity) is combined with INT4 KV quantization, the overall score is 85.1%, reflecting a marginal difference of only 1.6% compared to the baseline.

These results demonstrate that DuoAttention introduces negligible accuracy loss when combined with KV quantization while retaining its computational benefits. This validates the compatibility and efficiency of the combined approach.

[1] KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache. Liu, et al.

2. Comparison with MInference

We would like to clarify that MInference is not a KV cache compression method. It accelerates pre-filling for long-context LLMs using sparse attention patterns but does not reduce KV storage requirements or decoding latency. Thus, a direct comparison with DuoAttention is not appropriate. However, we explored the compatibility of MInference with DuoAttention by combining the two as orthogonal optimizations, as detailed in the updated Appendix Section A.7.

• Results: As detailed in Appendix Section A.7, combining MInference with DuoAttention enhances pre-filling efficiency while maintaining accuracy. Specifically:
  • Applying MInference alone introduces some accuracy degradation compared to full attention or pure DuoAttention.
  • By combining MInference with DuoAttention, we utilize MInference kernels to pre-fill retrieval heads identified by DuoAttention. This combined approach maintains comparable accuracy while significantly reducing both KV cache size (nearly halved) and decoding overhead.

These findings highlight that MInference and DuoAttention can complement each other as orthogonal optimizations, further extending the utility of both methods.

3. Addressing Load Imbalance in Multi-GPU Settings

We appreciate the reviewer’s observation regarding potential load imbalance when using different compression rates for different heads in multi-GPU settings. The imbalance primarily arises when the number of retrieval heads per layer is not divisible by the number of GPUs (e.g., 3 retrieval heads with 4 GPUs).

• Proposed Solution: Pre-training the model with a hybrid configuration specifically designed for such scenarios can mitigate this issue. For example, configuring a model with 8 retrieval heads and 32 streaming heads per layer ensures better alignment with common tensor-parallel setups, minimizing imbalance.
• Compatibility with Parallelism Techniques: DuoAttention is compatible with sequence parallelism (e.g., Ring Attention) and pipeline parallelism, both of which can further enhance scalability and performance in multi-GPU environments.

Although we have not conducted specific pre-training studies for this purpose due to resource constraints, we consider this a promising avenue for future research.

4. Proportion of Retrieval Heads

As shown in Figure 4 of our paper, the optimal proportion of retrieval heads varies by model architecture， mainly related to the attention architecture used:

• MHA Models: These models can be more aggressively compressed, requiring only 25% retrieval heads.
• GQA Models: Due to KV cache sharing across multiple queries (e.g., in Llama-3-8B, 4 query heads share one KV head), pruning a single KV head affects multiple queries. To maintain performance, GQA models require a higher proportion of retrieval heads, at 50%.

These optimal ratios were determined through extensive experiments across multiple benchmarks. The 25% retrieval head ratio for MHA and 50% for GQA represent safe and effective configurations, achieving both high accuracy and significant memory reduction, along with pre-filling and decoding acceleration.

We hope our responses have addressed your concerns and clarified the contributions of our work. Please let us know if there are further questions or suggestions, as we are happy to provide additional information.

We sincerely thank the reviewer for their thoughtful feedback and constructive suggestions. Below, we address the key points raised:

1. Tensor Parallelism and Multi-GPU Scalability

We appreciate the reviewer's concern regarding the scalability of DuoAttention in tensor-parallel configurations. In multi-GPU settings, retrieval heads can become a bottleneck, especially in configurations like 8-way tensor parallelism. To address this issue, aligning the number of retrieval heads with the number of tensor-parallel GPUs is critical to avoid load imbalances. For example, imbalances may occur when the number of retrieval heads per layer (e.g., 3 retrieval heads) is not evenly divisible by the number of GPUs (e.g., 4 GPUs).

To mitigate this, a feasible approach involves pre-training the model using a hybrid strategy tailored for such deployment scenarios. For instance:

• Configuring a model with 8 retrieval heads and 32 streaming heads per layer ensures better alignment with common tensor-parallel setups.
• This configuration reduces potential imbalances and facilitates efficient GPU utilization.

Although we have not yet conducted specific pre-training studies for such scenarios due to resource constraints, this remains a promising avenue for future research to further enhance DuoAttention’s scalability and performance.

Additionally, DuoAttention’s design is inherently compatible with other parallelism techniques such as sequence parallelism (e.g., Ring Attention) and pipeline parallelism. These methods can complement DuoAttention to further improve scalability and performance in multi-GPU environments.

2. Baseline Comparisons

We appreciate the reviewer's suggestion to include comparisons with additional state-of-the-art KV sparsification methods, such as SnapKV, PyramidKV, and AdaKV. In response, we have conducted experiments with SnapKV and PyramidKV, detailed in Appendix Sections A.6.

• Methodological Differences: SnapKV and PyramidKV rely on attention scores computed over a small observation window (e.g., the last 8–64 tokens) to prune KV tokens. While effective for scenarios with queries located at the end of the context, this assumption does not generalize to tasks such as multi-turn dialogues or cases where queries are positioned earlier in the sequence. DuoAttention, in contrast, operates without assumptions about token positions, making it more general and robust.
• Performance Results: As shown in Figures 17 and 18, DuoAttention achieves comparable or better performance than SnapKV and PyramidKV under equivalent KV cache budget constraints. For example:
  • Without simulation of the last tokens as generated tokens, DuoAttention matches or surpasses these baselines across all benchmarks.
  • With simulation (e.g., the last 50 tokens treated as generated inputs to mimic multi-turn dialogues), SnapKV and PyramidKV experience severe accuracy drops due to their reliance on final token positions for pruning. DuoAttention remains robust, maintaining high accuracy even under these conditions.

These results underscore DuoAttention’s versatility and reliability compared to the baselines, making it suitable for a broader range of real-world scenarios.

3. Effectiveness for Reasoning-Based Tasks

DuoAttention is highly effective for reasoning-based tasks, as demonstrated by its performance on the HotpotQA benchmark in LongBench. For instance:

• With the Llama-2-7B-32K model, DuoAttention achieves comparable results to full attention at a 25% retrieval head ratio.
• With the Llama-3-8B-1048K model, DuoAttention achieves comparable results at a 37.5% retrieval head ratio.

These results indicate that the retrieval heads identified by DuoAttention’s optimization-based approach are well-suited for reasoning tasks as well as question-answering tasks. Additional results across various reasoning-based datasets further validate this effectiveness.

4. Generalization of Retrieval Heads

We appreciate the reviewer’s request for statistical evidence of retrieval head generalization across tasks and sequences. To address this, we conducted extensive experiments on diverse benchmarks, including LongBench, RULER, and NIAH, as well as shorter-context tasks like MMLU, MBPP, and MT-Bench.

The results show that the retrieval heads identified by DuoAttention consistently outperform existing methods across different tasks and sequences. This consistency highlights DuoAttention’s ability to generalize effectively across a wide range of contexts and applications.

We sincerely thank the reviewer for their constructive feedback and suggestions. Below, we address the key points raised:

1. Novelty

We respectfully disagree with the assessment that the method lacks novelty. DuoAttention represents one of the first methods to leverage retrieval heads specifically for KV cache optimization. DuoAttention employs an optimization-based strategy to identify retrieval heads, leading to superior accuracy and computational efficiency in long-context tasks.

While concurrent work RazorAttention partitions heads into retrieval and non-retrieval categories, it relies on profiling techniques that our experiments show are less precise than DuoAttention’s optimization-based approach. Additionally, RazorAttention does not address pre-filling optimization, limiting its scope.

Other related works, such as:

• MInference, accelerate pre-filling through sparsity patterns but do not address KV cache compression or decoding efficiency.
• Retrieval Head Mechanistically Explains Long-Context Factuality focuses on analyzing the functionality of retrieval heads rather than improving memory efficiency or computational cost.

By contrast, DuoAttention delivers a comprehensive solution for KV cache management, achieving higher compression rates and improved performance during both pre-filling and decoding stages. These contributions are clarified in the related work section of the revised manuscript.

2. Applicability to Single KV Cache Models

We acknowledge the reviewer’s concern about the applicability of DuoAttention to single KV cache models, such as YOCO and MQA. However, we argue that these models are not widely adopted in practice. The majority of large language models (LLMs) in use today, including GQA models (e.g., Llama-3, Mistral, Qwen) and **MHA models **, feature multiple KV caches. These architectures are the focus of our work.

DuoAttention has demonstrated significant reductions in memory footprint and computational demands for GQA and MHA models while preserving accuracy. This underscores the practical significance of our method for widely used architectures.

3. Baseline Comparisons

We acknowledge the importance of comparing against state-of-the-art methods and have expanded our experiments to include SnapKV and PyramidKV, two recently introduced KV cache compression techniques, as well as pre-filling acceleration methods like MInference. We updated our manuscripts's Appendix Section A.6 and A.7 to include these experiments. Below are the key findings:

a. Comparison with KV Cache Compression Methods (SnapKV and PyramidKV):

We evaluated DuoAttention against SnapKV and PyramidKV on the Needle-in-a-Haystack (NIAH) benchmark using the Llama-2-7B-32K and Llama-3-8B-1048K models. The implementations of these methods were sourced from their official repositories.

• Methodological Differences: SnapKV and PyramidKV rely on computing attention scores for a small window of tokens (e.g., the last 8–64) and pruning KV cache tokens based on these scores. While effective in scenarios where queries appear at the end of the context, this assumption does not generalize well to real-world tasks, such as multi-turn dialogues or earlier query positioning.
• Performance Results: As shown in Figures 17 and 18 of the updated paper, DuoAttention performs comparably or better than these methods under equivalent KV cache budget constraints. For instance:
  • In the absence of simulated last tokens in prompts as generated tokens, DuoAttention matches or surpasses SnapKV and PyramidKV, and all methods achieve good results.
  • When simulating the last 50 tokens as generated tokens (similar to a new round of dialogue to ask the model), SnapKV and PyramidKV exhibit significant accuracy drops due to their reliance on the final tokens for pruning guidance. In contrast, DuoAttention remains robust under these stress tests.

These results highlight DuoAttention's generality and flexibility, making it a more reliable solution across diverse contexts.

3. Needle-in-a-Haystack (NIAH) Setting

We recognize the need for greater clarity regarding the NIAH setting. In response, we have updated Appendix Section A.10 to provide comprehensive details, including:

• Needle Description: The "needle" consists of a unique keyword or phrase inserted within a long sequence of irrelevant "haystack" tokens.
• Background Description: The "haystack" is composed of unrelated, randomly sampled tokens designed to mimic real-world long-context scenarios. Variations in the needle's location and the surrounding haystack are also described.

This additional context ensures that the experimental setup is transparent and reproducible.

4. RULER results

We agree with the reviewer’s suggestion to expand beyond LongBench and have included new results on the RULER benchmark. These results, detailed in Appendix Section A.8, further validate DuoAttention’s efficacy across diverse long-context datasets. Specifically:

• DuoAttention achieves comparable accuracy with full attention on RULER benchmarks across all context lengths.

Results on RULER

RULER is a synthetic dataset designed to rigorously evaluate long-context language models with configurable sequence lengths and task complexities. It includes 13 tasks spanning 4 categories, assessing long-context capabilities beyond simple in-context recall.

The table above presents the average accuracy of full attention and DuoAttention (50% sparsity) across different context lengths using the Llama-3-8B-Instruct-Gradient-1048k model for sequences up to 128K. DuoAttention demonstrates accuracy scores comparable to full attention across all context lengths, achieving an average performance increase of 0.05%.

These findings confirm DuoAttention’s effectiveness in maintaining strong accuracy on a rigorous benchmark, even under challenging long-context evaluation settings. We hope our response has clarified your concerns and can improve your rating of our paper. Please let us know if there are further concerns, as we are happy to respond.

5. Comparison with RazorAttention

We appreciate the reviewer’s observation regarding the differences in performance between DuoAttention and RazorAttention. While we lack access to the RazorAttention codebase for a direct comparison, we hypothesize the following:

• Dataset Variations: Differences in the datasets used for attention profiling may account for discrepancies in trade-off curves.
• Performance Comparison: Despite these differences, DuoAttention demonstrates superior performance in long-context tasks with lower KV cache budgets. For instance, DuoAttention achieves all-green NIAH results with 25% retrieval heads on Llama-2 (MHA), whereas RazorAttention requires 30% retrieval heads and exhibits accuracy drops (as per Figure 1 in the RazorAttention paper).

These findings underscore the strength of DuoAttention’s optimization-based retrieval head identification, which provides a more effective allocation of computational resources.

We thank the reviewer again for their valuable insights. The additional experiments and clarifications incorporated into the revised manuscript directly address the concerns raised and highlight the robustness and generality of DuoAttention. We hope our responses has clarified your concerns and can improve your rating of our paper. Please let us know if there are further concerns, as we are happy to respond!

We sincerely thank the reviewer for their thoughtful feedback and constructive suggestions. Below, we address the key points raised:

1. Novelty

We respectfully disagree with the assessment that the method is "slim in novelty." DuoAttention represents a pioneering approach by leveraging retrieval heads specifically for KV cache optimization, moving beyond traditional heuristics such as attention score distribution profiling. By employing an optimization-based strategy, DuoAttention identifies retrieval heads with greater precision, resulting in substantial performance gains in long-context tasks. While a concurrent work, RazorAttention, also explores retrieval head utilization, DuoAttention's novel optimization framework consistently delivers superior results in both accuracy and computational efficiency, as evidenced by the comparisons detailed in Section 3.5 of our paper and the reported NIAH results in the RazorAttention paper. These advancements underscore DuoAttention's distinctiveness and innovative contribution.

2. Baselines

We acknowledge the importance of comparing against state-of-the-art methods and have expanded our experiments to include SnapKV and PyramidKV, two recently introduced KV cache compression techniques, as well as pre-filling acceleration methods like MInference. We updated our manuscripts's Appendix Section A.6 and A.7 to include these experiments. Below are the key findings:

a. Comparison with KV Cache Compression Methods (SnapKV and PyramidKV):

We evaluated DuoAttention against SnapKV and PyramidKV on the Needle-in-a-Haystack (NIAH) benchmark using the Llama-2-7B-32K and Llama-3-8B-1048K models. The implementations of these methods were sourced from their official repositories.

• Methodological Differences: SnapKV and PyramidKV rely on computing attention scores for a small window of tokens (e.g., the last 8–64) and pruning KV cache tokens based on these scores. While effective in scenarios where queries appear at the end of the context, this assumption does not generalize well to real-world tasks, such as multi-turn dialogues or earlier query positioning.
• Performance Results: As shown in Figures 17 and 18 of the updated paper, DuoAttention performs comparably or better than these methods under equivalent KV cache budget constraints. For instance:
  • In the absence of simulated last tokens in prompts as generated tokens, DuoAttention matches or surpasses SnapKV and PyramidKV, and all methods achieve good results.
  • When simulating the last 50 tokens as generated tokens (similar to a new round of dialogue to ask the model), SnapKV and PyramidKV exhibit significant accuracy drops due to their reliance on the final tokens for pruning guidance. In contrast, DuoAttention remains robust under these stress tests.

These results highlight DuoAttention's generality and flexibility, making it a more reliable solution across diverse contexts.

b. Comparison with Pre-filling Acceleration Methods (MInference):

MInference employs sparsity patterns to accelerate pre-filling but does not address KV cache compression or decoding speed optimization. Thus, MInference is on a different optimization track and is not directly comparable to DuoAttention. However, we explored the compatibility of MInference with DuoAttention by combining the two as orthogonal optimizations, as detailed in the updated Appendix Section A.7.

• Results: Combining MInference with DuoAttention enhances pre-filling efficiency without compromising accuracy, while simultaneously achieving significant reductions in KV cache size and decoding overhead. These complementary strengths demonstrate the synergy between the two methods.

• Performance Analysis: As illustrated in Figures 19 and 20, applying MInference alone on the NIAH benchmark results in some accuracy degradation compared to full attention or pure DuoAttention (refer to Figure 19).

  By integrating MInference into DuoAttention, we utilize MInference kernels to pre-fill all retrieval heads identified by DuoAttention. This combined approach maintains accuracy levels comparable to pre-filling all heads with MInference while nearly halving the KV cache size and significantly reducing decoding overhead. These findings underscore the compatibility and efficiency of the DuoAttention-MInference integration, further extending their utility in real-world applications.

I thank the authors for the response; much of my concerns (#2, #3, and partially #4) are resolved. The authors' needle setting seems proper, and achieving good results on this (and by extension, RULER) is a significant advancement per benchmark literature like [1], [2], or method papers with thorough experimental execution like MInference. It seems the general community lacks an understanding of how challenging a proper needle test is for many established KV cache compression methods, so it might be worth further highlighting this in the presentation.

That being said, I still have a few standing concerns/questions:

1. The $\infty$Bench request has not been fulfilled. While I believe the featured experiments already showcase the performant nature of DuoAttention, this should be added. This is especially important if the authors are featuring some long context-capable models.

2. On the note of #1, I honestly do not understand the decision to feature models like Llama-3-8B-Instruct-Gradient-1048k and Llama-2-7B-32K-Instruct — these are third-party finetuned models with limited adaptation in the community. They were reasonable choices before Llama 3.1 as there wasn't a long context-capable Llama, but with Llama 3.1 being available and widely tested, it should clearly be the base model to avoid reaching very third-party recipe-specific conclusions (especially when the authors are not doing much, if any, >128k evaluations). I recommend featuring this model instead or in addition, should your work reaches camera-ready

3. It is indeed true that MInference is a prefill-only compression method, but I believe the comparison is still reasonably warranted as DuoAttention is also advertised for its prefill compression capability. While I appreciate the additional results in A.7, I am still not quite sure how Figure 20 is achieved even after careful reading. The authors stated they "utilize MInference kernels to prefill all retrieval heads" in contrast to "prefilling all heads with MInference." First, what budget does Figure 19 have, and is it aligned with Figure 6? More specifically, are the authors only applying MInference sparse patterns to the retrieval heads? What drives that decision, and what would happen if you let MInference determine the patterns of non-retrieval heads? Would there be any improvement compared to making them all StreamingLLM heads? Also, MInference seems to exhibit some degradation on short-context tasks — can this be alleviated through some combination design?

These are not core questions to the proposal of DuoAttention, and I generally hesitate to ask extensive questions about combining two methods due to the large scope such potential combination shall entail. However, many close followers of the field would appreciate these findings, so I leave it to the authors to decide whether to explore these studies. Some clarification on A.7 should be added regardless.

4. Regarding RazorAttention, did you consider the echo heads proposed by the authors?

5. What does "retrieval head ratio" mean? Is it the ratio of retrieval heads being preserved or the total compression ratio? This distinction might lead to different conclusions regarding your 25% vs. 30% comment on RazorAttention.

Lastly, I respectfully present a hard disagreement regarding the novelty of the methods. The core ingredient of the proposed method is leveraging the existence of induction/retrieval heads, where the importance of such heads has been well-studied in prior work. DuoAttention is not the first to explore assigning full cache to retrieval heads (though other works may be considered concurrent), and its gate-based optimization approach is heavily rooted in gating network studies from the MoE and pruning realms. For example, [3] — a widely regarded classic on MoE — also employs trainable gate values to determine the activation of different experts. Similarly, DuoAttention employs trainable gate values to determine head categorization, where the resemblance is strong. Various pruning methods also assign trainable importance scores to guide component dropping.

As I detailed in my initial review, I don't think this matters much given the "impressive performance boost achieved while remaining practical," but I encourage the authors to tone down the novelty claims and provide a discussion of such prior gating-related work.

I am bumping my score to 6 as I believe this work is well above the acceptance threshold, though further improvements are pending in the authors' response.

(I also recommend the authors replace Figure 7 with a table reporting 6 LongBench subtasks — small diagrams with varying y-axes are difficult to read and cross-reference.)

[1] KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark... [2] A Controlled Study on Long Context Extension and Generalization in LLMs
[3] Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer

4. RazorAttention and Echo Heads

We have not explored echo heads, as DuoAttention achieves strong performance with its current framework. That said, integrating echo heads may be a promising optimization for future work.

The retrieval head ratio refers to the proportion of retrieval heads retained in the KV cache, rather than the total compression ratio. For example:

• A model with 128 streaming heads and 128 retrieval heads achieves an approximate 50% KV cache budget when the input sequence is sufficiently long.

To provide further clarity, consider a scenario with a 64K token input, 64 sink tokens, and 256 recent tokens for streaming heads. The KV cache budget ratio is calculated as follows:

• $\text{Budget Ratio} = \frac{(64K\times 0.5) + (320 \times 0.5)}{64K} = 0.5025 \approx 0.5$

This definition supports accurate and consistent budget allocation in our experiments.

Regarding the 25% vs. 30% comparison with RazorAttention, we believe the discrepancy stems from methodological differences rather than budget misalignment. The Figure 1 of RazorAttention evaluates up to a 64K context length. For our Llama-2-7B-32K experiments using a 25% retrieval head ratio, the total KV budget under a 32K context length is:

$\text{Budget Ratio} = \frac{(32K\times 0.25) + (320 \times 0.25)}{32K} = 0.2525$

This budget is still clearly smaller than the 30% budget used in RazorAttention’s experiments. Despite this, our 25% ratio results show no degradation, whereas RazorAttention’s 30% ratio results exhibit performance drops. This underscores DuoAttention's stronger efficiency and robustness under smaller budget constraints.

We will clarify these points further in the future revision of our manuscript to ensure transparency and address any potential ambiguities.

5. Novelty and Gating-Related Methods

We appreciate your feedback on DuoAttention’s novelty. While its optimization-based approach marks a significant advance, we acknowledge its connection to prior gating-related works, such as sparsely gated mixture-of-experts (MoE). We will revise the manuscript to better contextualize DuoAttention within this lineage and adjust novelty claims accordingly.

6. Presentation Suggestions

We agree that Figure 7 can be improved. We will include a table summarizing key LongBench results for better clarity and usability.

7. Query-Dependency of SnapKV-Like Methods

We concur that query-location sensitivity is a design limitation of SnapKV-like methods. DuoAttention’s query-position robustness offers a clear advantage, which we will further validate with datasets like SharedContextBench in future work.

8. Compression Granularity

DuoAttention delivers practical gains by achieving near-full KV performance at 50% budget, with breaking points illustrated in Figure 7 for various tasks. We will explore these limits further will guide future optimizations.

Final Remarks

Thank you again for your valuable insights and constructive feedback. We are committed to addressing your suggestions in our revisions and future work.

Best regards,

Authors