title: DuoAttention: Efficient Long-Context LLM Inference with Retrieval and
    Streaming Heads
summary: 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.
paper_authors._id: 670f0ee0125e4ced02d514e4
paper_authors.name: Guangxuan Xiao
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-10-18T08:25:16.976Z
paper_authors._id: 6362fefe19cf373a5fc5b39e
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/6362fefe19cf373a5fc5b39e/v4uJ5bzjpZJOxqGUHYPM2.jpeg
paper_authors.fullname: Guangxuan Xiao
paper_authors.type: user
paper_authors.user: Guangxuan-Xiao
paper_authors._id: 670f0ee0125e4ced02d514e5
paper_authors.name: Jiaming Tang
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-10-21T07:48:42.023Z
paper_authors._id: 63f30d28f4e30ffd2bda9aff
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/1676872959244-noauth.jpeg
paper_authors.fullname: Jiaming Tang
paper_authors.type: user
paper_authors.user: Sakits
paper_authors._id: 670f0ee0125e4ced02d514e6
paper_authors.name: Jingwei Zuo
paper_authors._id: 670f0ee0125e4ced02d514e7
paper_authors.name: Junxian Guo
paper_authors._id: 670f0ee0125e4ced02d514e8
paper_authors.name: Shang Yang
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2025-02-21T10:00:33.554Z
paper_authors._id: 641d8bacd526196afc12766d
paper_authors.avatarUrl: /avatars/73f7b2d86a7bf27940bec2b1f199d71b.svg
paper_authors.fullname: Shang Yang
paper_authors.type: user
paper_authors.user: Shangy
paper_authors._id: 670f0ee0125e4ced02d514e9
paper_authors.name: Haotian Tang
paper_authors._id: 670f0ee0125e4ced02d514ea
paper_authors.name: Yao Fu
paper_authors._id: 670f0ee0125e4ced02d514eb
paper_authors.name: Song Han
paper_title: DuoAttention: Efficient Long-Context LLM Inference with Retrieval and
    Streaming Heads
paper_summary: 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.
paper_ai_summary: DuoAttention reduces memory and latency for long-context LLMs by using a full KV cache only for critical retrieval heads while employing a lightweight cache for streaming heads, improving efficiency without sacrificing performance.
paper_ai_keywords: long-context large language models LLMs Key and Value states KV cache pruning attention heads Retrieval Heads Streaming Heads DuoAttention full KV cache lightweight KV cache optimization-based algorithm synthetic data long-context inference memory reduction decoding speed pre-filling speed accuracy loss quantization Llama-3-8B A100 GPU