DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads
Large Language Models; Long Context; Efficiency;
Yes
No ethics review needed.
One of the concurrent works, RazorAttention, also utilizes a very similar recipe. The authors compared it in Figure 13(1), but I find the reading to be very different from RazorAttention's own reporting. Any insight of why?
1. The proposed method is extremely performant. It is rare to see a prefilled compressible token dropping method capable of doing NIAH. 

2. The efficiency optimization is solid, and it is supported by a thorough efficiency evaluation.

3. The optimization-based head selection approach offers a different avenue for attention head pattern matching.

4. The paper is nicely written.
DuoAttention leverages finding around retrieval heads. By allocating more KV budget to such heads and compressing the rests, the method claims to have excellent long context performance.
1. The method mainly involves retrieval head + token dropping (in this case, StreamingLLM), which is slim in novelty. I don't think it is much of an issue given the impressive performance boost achieved while remaining practical, but this is worth mentioning.

2. The featured compared methods — SteamingLLM, H2O, and TOVA — are very dated. I'd like to see how it stands off with more modern KV cache compassion/sparse inference methods, e.g., SnapKV and MInference.

3. The NIAH setting is unclear. What does the needle look like? What background is used? The author should better clarify this, as different settings can lead to very different results. 

4. Following #3, it looks like LongBench is the only real long context test evaluated. I'd like to see more coverage on other long context datasets, e.g., $\infty$Bench and RULER.

I believe the proposed method marks a breakthrough in eviction-based methods. Its execution is solid and the presentation is clear. I am willing to boost my score should the coverage/clarity of performance evaluation be improved.
Yes
No ethics review needed.
Please address the questions in the weaknesses.
1.	It saves the GPU memory and enhances the inference efficiency for both MHA and GQA.
This paper introduces a lightweight, optimization-based method DuoAttention that only applies a full KV cache to retrieval heads while using a constant-length KV cache for streaming heads. The proposed method achieves good performance in both MHA and GQA.
1.	There are some related works identifying the different patterns of retrieval and non-retrieval heads, such as RazorAttention, Retrieval Head Mechanistically Explains Long-Context Factuality, MInference, etc. It would be better to clarify the novelty of the proposed method.
2.	It can only be applied to multiple KV cache scenarios. For only single KV cache model structure, such as YOCO, it cannot be applied.
3.	The baselines are not comprehensive. It would be beneficial if the authors could compare their results with more advanced baselines.
4.	It seems the proposed method cannot achieve the best performance in all the benchmarks in Figure 7. It would be beneficial to give detailed analysis on the worse cases.
Yes
No ethics review needed.
* How effective is this method for reasoning-based tasks? Do the retrieval heads remain the same for question-answering type tasks as well as reasoning type tasks?
* Rather than merely illustrating the distinction between retrieval and streaming heads through specific examples, it would be more compelling to present statistically significant evidence demonstrating how the identified retrieval heads generalize across various sequences and tasks.
* Is there any way to extend this method to multi-GPU setting?
* To what extent can the memory savings lead to increased throughput? (Just a suggestion)
* The paper is well-structured and the claims are easy to follow.
* This work addresses a significant research problem in LLM inference - KV cache optimization. 
* The authors perform comprehensive ablation studies to motivate each aspect of their proposed approach. The advantage of their proposed novel retrieval head selection method over the existing methods is well-supported by empirical results.
This work presents methods to reduce memory footprint and computational demands of Large Language Model's (LLM) attention layers to enhance inference speed and to accommodate longer context lengths. The authors build upon the observation that only a few attention heads within an LLM, also known as retrieval heads, have considerable impact on long-context performance. They propose a novel, optimization-based approach to identify these less critical heads, unlike previous methods that relied on heuristic techniques. This optimization process utilizes synthetic data designed specifically to evaluate long-context capabilities. Their results indicate superior model performance compared to current state-of-the-art retrieval head allocation techniques and certain KV sparsification methods, while maintaining comparable KV cache budget.
* While tensor parallelism has proven effective in reducing latency and memory usage per GPU, this paper omits any discussion on how its proposed methods could be adapted to such model-parallel settings. For example, in an 8-way tensor-parallel configuration, the retrieval heads might become the performance bottleneck, potentially negating the gains from this work.

* Some state-of-the-art KV sparsification methods are missing in the comparative analysis. For instance, SnapKV, PyramidKV, AdaKV are some notable omissions.
   1. 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, 2024d. *Arxiv /abs/2404.14469*
   2. Zefan Cai, Yichi Zhang, Bofei Gao, Yuliang Liu, Tianyu Liu, Keming Lu, Wayne Xiong, Yue Dong, Baobao Chang, Junjie Hu, and Wen Xiao. Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling, 2024. *Arxiv /abs/2406.02069*
   3. 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, 2024. *Arxiv /abs/2407.11550*
Yes
No ethics review needed.
1) Is the proportion of Retrieval Heads the same for different models? Why choose 25% Retrieval Heads for MHA and 50% for GQA?
1) The paper is well written and well motivated.
2) DuoAttention is a plug-and-play solution compatible with FlashAttention.
3) DuoAttention can accelerate inference during both the prefill and decoding stages.
To tackle KV Cache problem, the paper introduces a new approach called DuoAttention, which classifies attention heads in LLMs into two types: Retrieval Heads and Streaming Heads. Retrieval Heads are essential for processing long contexts and require full KV cache storage, while Streaming Heads focus on recent tokens and can operate with reduced KV caches. DuoAttention uses a lightweight optimization process to identify non-compressible Retrieval Heads, allowing for efficient memory usage and faster processing. This method integrates easily with existing optimization techniques and significantly reduces the memory footprint and decoding time. When combined with quantization, DuoAttention enables models like Llama-3-8B to handle up to 3.33 million tokens on a single GPU, achieving a 6.4× increase in capacity compared to standard deployments.
1) Lack the experiment with the accuracy results after combining DuoAttention and KV quantization.
2) Lack the comparisons with new KV Cache compression technologies such as Minference.
3) Using different compression rates for different heads may lead to uneven computation across cards during parallel inference, potentially affecting performance.
Yes
No ethics review needed.
1. the offline method to find the retrieval heads is more important for practical application
2. the compression ratio of KV cache should be improved
1. the experiments is enough and confident
2. the writing and clarity are clear
3. the concept of retrieval heads is significant for static KV cache compression
The DuoAttention uses the retrieval heads and local streaming heads, realizing the KV cache compression and acceleration:
1) the retrieval heads preserve the major information of the long-context, which can not be compressed;
2) the non-retrieval heads are compressed by locality, which is accelerated with StreamingLLM;
3) the retrieval heads are detected with importance learnable weights;
4) Finally, the prefill is accelerated by chunked prefilling, and the decoding is accelerated with compressed KV cache.
1. the detection of retrieval heads is based on the fine-tuning of importance weights, the usability should be enhanced
2. the compression ratio is limited, if the retrieval heads preserve the full context information
Yes
No ethics review needed.
1. The core of this paper focuses on prioritizing the retention of the KV cache in important attention heads while attempting to discard less important cache in other heads, based on offline detection results. Some follow-up works on SnapKV seems to align closely with this approach[3][4]. For example, [4] employs a similar strategy by identifying important attention heads through online analysis of "altering model outputs," subsequently allocating more budget to these key heads and reducing the budget for others. What do you think is the relationship between these budget allocation strategies and the detection of retrieval heads?

2. In Equation 2, why is the index i set to $L$? Shouldn't it be $N$ instead?

3. How do you control the cache budget to a specifical ratio like 50% in the experiments? It seems challenging to precisely manage the cache budget within this approach.

[1] Li, Y., Huang, Y., Yang, B., Venkitesh, B., Locatelli, A., Ye, H., ... & Chen, D. (2024). Snapkv: Llm knows what you are looking for before generation. arXiv preprint arXiv:2404.14469.

[2] Zhang, Y., Gao, B., Liu, T., Lu, K., Xiong, W., Dong, Y., ... & Xiao, W. (2024). PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling. arXiv preprint arXiv:2406.02069.

[3] Rehg, I. (2024). KV-Compress: Paged KV-Cache Compression with Variable Compression Rates per Attention Head. arXiv preprint arXiv:2410.00161.

[4] Feng, Y., Lv, J., Cao, Y., Xie, X., & Zhou, S.K. (2024). Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference. ArXiv, abs/2407.11550.

While this paper is engaging and provides extensive evaluations, several limitations hold me back from giving it a higher score.  If the authors can address these concerns during rebuttal phase, I believe it would greatly enhance the paper’s quality, and I’d be glad to reconsider my score.
1. The paper is well-written and easy to follow, with a novel methodology.
2. The work is solid, having been evaluated on multiple benchmarks. It also combines the orthogonal technique of quantization, further enhancing its results.
3. The evaluations demonstrate promising results.
Overall, this is an interesting and solid work. It categorizes all attention heads of LLMs into two types—Retrieval Heads and Streaming Heads—based on an offline analysis. By pruning the KV cache of Streaming Heads, this study achieves more efficient inference.
1. **Choice of Baseline**: DuoAttention is essentially a KV cache compression algorithm, and all the performance gains come from reducing the amount of KV cache that was previously stored, whether in the prefilling stage or the decoding stage. While the authors compare their method to H2O, TOVA, and StreamingLLM, these approaches are no longer state-of-the-art . For instance, SnapKV[1], which was published in April and is now widely regarded as a more powerful KV cache compression method, should have been considered. In the commonly used LongBench evaluation, SnapKV outperforms H2O with only one-quarter of the cache budget[1]. Additionally, in needle-in-a-haystack tasks, SnapKV benefits from its pooling operations and significantly improves accuracy compared to H2O. When compressing an 8K sequence to 128 tokens, SnapKV incurs minimal performance loss, outperforming H2O by a large margin, as shown in [2]. In the experimental section, the authors highlight a limitation of H2O: “Since the original designs of H2O and TOVA do not support long contexts, we modify their algorithms by replacing the pre-filling stage with FlashAttention and simulating decoding for the last 50 tokens of the input, following Tang et al. (2024b).” This limitation arises because H2O requires the accumulation of global attention weights, which is incompatible with FlashAttention. However, SnapKV and its successors only observe attention scores within a small observation window for compression, making them highly compatible with FlashAttention. This allows recalculating only a small portion of attention weights  for compression when combined with chunked prefill. Therefore, SnapKV would be a more appropriate baseline for this paper.

2. **Inappropriate statements:**

a.  _“Moreover, DuoAttention is **fully compatible** with important optimization techniques like GQA and quantization.” “Approximative attention methods, such as H2O (Zhang et al., 2023b), StreamingLLM (Xiao et al., 2023b), TOVA (Oren et al., 2024), and FastGen (Ge et al., 2024), often compromise accuracy in long-context applications and are **incompatible** with essential KV cache optimization techniques like GQA"_ 

I believe the compatibility of DuoAttention with GQA is not significantly different from other methods. DuoAttention forcedly classifies a group of attention heads under GQA into the same category through offline analysis, thus upporting GQA. However, other methods could achieve the same effect by accumulating weights for a group of kv cache during cache eviction. Therefore, GQA compatibility seems more like an implementation detail rather than a unique feature of the algorithm. In fact, StreamingLLM does not encounter any GQA compatibility issues. The lack of GQA compatibility in earlier works, in my opinion, stems from the fact that GQA was not widely adopted when those methods were initially proposed. Later methods, for the sake of comparison, maintained this approach without further integrating GQA, which could have been easily addressed at the code implementation level. A recent study[3] applying cache eviction in paged attention has demonstrated that this compatibility can be easily achieved in practice.

b._“Also, these methods (H2O, TOVA, StreamingLLM) cannot reduce the prefilling cost of long-context LLMs.”_

 Previous evaluations of cache compression often compress the KV cache after the prefilling stage to ensure comparability. Since the KV cache for each layer can be compressed immediately after the completion of that layer's prefilling computation, this already substantially reduces the peak memory usage during the prefilling process. If further reduction in computational cost during the prefilling stage or a more significant decrease in peak memory usage for very long input texts is desired, combining these methods with chunked prefilling can provide additional acceleration—a straightforward solution. This is particularly applicable to StreamingLLM, which is same with the same Streaming Head in this paper, and thus faces no obstacles in applying such methods. If one argues that the H2O method combined with chunked prefilling may require additional accumulation of global attention weights, the additional baseline SnapKV can effectively address this issue.
We significantly reduce both pre-filling and decoding memory and latency for long-context LLMs without sacrificing their long-context abilities.
@inproceedings{
anonymous2024duoattention,
title={DuoAttention: Efficient Long-Context {LLM} Inference with Retrieval and Streaming Heads},
author={Anonymous},
booktitle={Submitted to The Thirteenth International Conference on Learning Representations},
year={2024},
url={https://openreview.net/forum?id=cFu7ze7xUm},
note={under review}
}
Deploying long-context large language models (LLMs) is essential but poses significant computational and memory challenges.
Caching all Key and Value (KV) states across all attention heads consumes substantial memory.
Existing KV cache pruning methods either damage the long-context capabilities of LLMs or offer only limited efficiency improvements.
In this paper, we identify that only a fraction of attention heads, a.k.a, Retrieval Heads, are critical for processing long contexts and require full attention across all tokens.
In contrast, all other heads, which primarily focus on recent tokens and attention sinks—referred to as Streaming Heads—do not require full attention.
Based on this insight, we introduce DuoAttention, a framework that only applies a full KV cache to retrieval heads while using a light-weight, constant-length KV cache for streaming heads, which reduces both LLM's decoding and pre-filling memory and latency without compromising its long-context abilities.
DuoAttention uses a lightweight, optimization-based algorithm with synthetic data to identify retrieval heads accurately.
Our method significantly reduces long-context inference memory by up to 2.55$\times$ for MHA and 1.67$\times$ for GQA models while speeding up decoding by up to 2.18$\times$ and 1.50$\times$ and accelerating pre-filling by up to 1.73$\times$ and 1.63$\times$ for MHA and GQA models, respectively, with minimal accuracy loss compared to full attention.
Notably, combined with quantization, DuoAttention enables Llama-3-8B decoding with 3.33 million context length measured on a single A100 GPU. Code and dataset will be released upon publication.
I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.
I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Large Language Models; Long Context; Efficiency;
I certify that there is no acknowledgement section in this submission for double blind review.
/pdf/93ded0fa15cb617e2e11700cc0e4771e35e68451.pdf
foundation or frontier models, including LLMs
I understand the reciprocal reviewing requirement as described on https://iclr.cc/Conferences/2025/CallForPapers. If none of the authors are registered as a reviewer, it may result in a desk rejection at the discretion of the program chairs. To request an exception, please complete this form at https://forms.gle/Huojr6VjkFxiQsUp6.
I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.
DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads
ICLR 2025 Conference Submission
ICLR.cc/2025/Conference/Submission