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

• URL: http://arxiv.org/abs/2410.10819v1
• Date: Mon, 14 Oct 2024 17:59:58 GMT
• Title: DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads
• Authors: Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, Song Han,
• Abstract summary: Caching all Key and Value states across all attention heads consumes substantial memory. 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. Our method significantly reduces long-context inference memory by up to 2.55x for MHA and 1.67x for GQA models.
• Score: 22.462489968597
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.55x for MHA and 1.67x for GQA models while speeding up decoding by up to 2.18x and 1.50x and accelerating pre-filling by up to 1.73x and 1.63x 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.3 million context length on a single A100 GPU. Code is provided in https://github.com/mit-han-lab/duo-attention.

Related papers

• ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference [25.638980944695728]
  ShadowKV is an efficient long-context large language models (LLMs) inference system. It stores the low-rank key cache and offloads the value cache to reduce the memory footprint for larger batch sizes and longer sequences. It can support up to 6$times$ larger batch sizes and boost throughput by up to 3.04$times$ on an A100 GPU.
  arXiv  Detail & Related papers  (2024-10-28T19:08:12Z)
• Not All Heads Matter: A Head-Level KV Cache Compression Method with Integrated Retrieval and Reasoning [19.942402563256962]
  Key-Value (KV) caching is a common technique to enhance the computational efficiency of Large Language Models (LLMs) We propose HeadKV, a head-level KV cache compression method, and Head KV-R2, which leverages a novel contextual reasoning ability estimation for compression. Our method retains just 1.5% of the KV cache while achieving 97% of the performance of the full KV cache on the contextual question answering benchmark.
  arXiv  Detail & Related papers  (2024-10-25T02:22:00Z)
• A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts [38.867323730365406]
  LongGen finetunes a pretrained LLM into an efficient architecture during length extension. LongGen achieves 1.55x training speedup and reduces wall-clock time by 36%, compared to a full-attention baseline. During inference, LongGen reduces KV cache memory by 62%, achieving 1.67x prefilling speedup and 1.41x decoding speedup.
  arXiv  Detail & Related papers  (2024-10-02T12:35:53Z)
• ThinK: Thinner Key Cache by Query-Driven Pruning [63.13363917871414]
  Large Language Models (LLMs) have revolutionized the field of natural language processing, achieving unprecedented performance across a variety of applications. This paper focuses on the long-context scenario, addressing the inefficiencies in KV cache memory consumption during inference. We propose ThinK, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels.
  arXiv  Detail & Related papers  (2024-07-30T17:59:08Z)
• vTensor: Flexible Virtual Tensor Management for Efficient LLM Serving [53.972175896814505]
  Large Language Models (LLMs) are widely used across various domains, processing millions of daily requests. Large Language Models (LLMs) are widely used across various domains, processing millions of daily requests.
  arXiv  Detail & Related papers  (2024-07-22T14:37:58Z)
• DeFT: Decoding with Flash Tree-attention for Efficient Tree-structured LLM Inference [22.684773338989007]
  Large language models (LLMs) are increasingly employed for complex tasks that process multiple generation calls in a tree structure with shared prefixes of tokens. Existing inference systems for tree-based applications are inefficient due to improper partitioning of queries and KV cache during attention calculation. We propose DeFT, a hardware-efficient attention algorithm with prefix-aware and load-balanced KV cache partitions.
  arXiv  Detail & Related papers  (2024-03-30T04:34:54Z)
• CHAI: Clustered Head Attention for Efficient LLM Inference [11.056313961175702]
  Clustered Head Attention (CHAI) is able to reduce the memory requirements for storing K,V cache by up to 21.4% and inference time latency by up to 1.73x without any fine-tuning required. We show that CHAI is able to reduce the memory requirements for storing K,V cache by up to 21.4% and inference time latency by up to 1.73x without any fine-tuning required.
  arXiv  Detail & Related papers  (2024-03-12T20:10:04Z)
• Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM Inference [78.65321721142624]
  We focus on a memory bottleneck imposed by the key-value ( KV) cache. Existing KV cache methods approach this problem by pruning or evicting large swaths of relatively less important KV pairs. We propose LESS, a simple integration of a constant sized cache with eviction-based cache methods.
  arXiv  Detail & Related papers  (2024-02-14T18:54:56Z)
• KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache [67.9776980972508]
  We develop a tuning-free 2bit KV cache quantization algorithm named KIVI. KIVI can enable Llama, Falcon, and Mistral models to maintain almost the same quality while using $mathbf2.6times$ less peak memory.
  arXiv  Detail & Related papers  (2024-02-05T06:06:47Z)
• Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs [82.08922896531618]
  We introduce adaptive KV cache compression, a plug-and-play method that reduces the memory footprint of generative inference for Large Language Models (LLMs) We conduct targeted profiling to discern the intrinsic structure of attention modules. Based on the recognized structure, we then construct the KV cache in an adaptive manner: evicting long-range contexts on attention heads emphasizing local contexts, discarding non-special tokens on attention heads centered on special tokens, and only employing the standard KV cache for attention heads that broadly attend to all tokens.
  arXiv  Detail & Related papers  (2023-10-03T05:17:08Z)
• Efficient Streaming Language Models with Attention Sinks [72.20260088848987]
  StreamingLLM is an efficient framework that enables Large Language Models to generalize to infinite sequence lengths without any fine-tuning. We show that StreamingLLM can enable Llama-2, MPT, Falcon, and Pythia to perform stable and efficient language modeling with up to 4 million tokens and more.
  arXiv  Detail & Related papers  (2023-09-29T17:59:56Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.