PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Junxian Guo, Shang Yang, Haotian Tang, Yao Fu, Song Han

arXiv:2410.10819v1cs.CL

TL;DR

Long-context LLM inference is constrained by KV-cache memory and attention costs. DuoAttention assigns full KV caches to retrieval heads and constant-length caches to streaming heads, reducing memory and latency while preserving long-context capability, including 3.3 million-token decoding on one A100 with quantization.

  • Problem

    Long-context LLM inference requires costly attention computation and KV-cache memory that grow with sequence length.

  • Method

    DuoAttention applies full KV caches to retrieval heads and reduced constant-length caches to streaming heads identified by optimization on synthetic data.

  • Results

    DuoAttention reduces memory by up to 2.55× for MHA and 1.67× for GQA, while accelerating decoding and pre-filling with minimal accuracy loss.

  • Takeaways & Limitations

    With quantization, DuoAttention supports up to 3.30 million contextual tokens on a single A100 GPU, a 6.4× capacity increase over standard full-attention FP16 deployment.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Long-context LLM inference incurs substantial latency and KV-cache memory costs as context length grows. DuoAttention addresses this by assigning full attention only to retrieval heads while compressing streaming heads, enabling more efficient long-context deployment.

  • Motivation: Long-context inference has linearly increasing decoding latency, quadratically increasing pre-filling latency, and KV-cache memory that scales linearly with context length.These costs arise because full attention considers all preceding tokens and KV caching stores keys and values from them.
  • Key observation: Retrieval Heads are a fraction of attention heads requiring full attention, whereas Streaming Heads focus primarily on recent tokens and attention sinks.This head-level distinction motivates using different cache strategies within the same model.
  • Method: DuoAttention identifies non-compressible retrieval heads with a lightweight, optimization-based procedure using synthetic datasets.The method is designed to accelerate decoding and pre-filling while reducing memory footprints in long-context scenarios.
  • Compatibility: DuoAttention is compatible with Grouped-Query Attention and quantization, extending its applicability to existing inference optimizations.The introduction specifically highlights compatibility with GQA and quantization techniques.
  • Results: 3.3 million contextual tokens are supported by Llama-3-8B on a single A100 GPU, a 6.4× capacity increase over standard full-attention FP16 deployment.This result demonstrates million-level context handling with DuoAttention.

2 DUOATTENTION

DuoAttention distinguishes retrieval heads, which need full-token attention, from streaming heads, which can retain only attention sinks and recent tokens. It identifies these heads through gate optimization on synthetic passkey-retrieval data, then applies head-specific attention policies and chunked pre-filling to reduce memory and computation.

  • Synthetic Dataset for Identifying Retrieval Heads: The synthetic dataset embeds ten randomly generated passkey sequences in long text, providing supervision targeted at long-context retrieval and identifying compressible KV heads.Distillation loss focuses on the final passkey tokens, and L1 regularization encourages sparse gate values.
  • Head Types: Retrieval heads significantly alter model outputs when restricted to recent tokens and attention sinks, whereas streaming heads primarily attend to recent tokens and attention sinks.Pruning middle KV-cache tokens substantially harms retrieval-head passkey accuracy but has no significant impact for streaming heads.
  • Optimization-based Identification: DuoAttention assigns each KV head a gate α_i,j and optimizes the gates while keeping the LLM parameters fixed to assess the impact of streaming-cache compression.The forward pass mixes full and streaming attention using each gate as the mixing weight; GQA KV heads can correspond to multiple attention heads.
  • Binarizing Attention Implementations: At inference, DuoAttention binarizes optimized gate values using a sparsity-quantile threshold, applying full attention only to retrieval heads.Streaming heads therefore use the compressed attention policy rather than a full-token KV cache.
  • Chunked Pre-filling: Chunked pre-filling uses FlashAttention-2 and fixed-length prompt chunks to reduce peak memory by lowering intermediate activation size from sequence length to chunk size.The method pre-fills KV caches for both retrieval and streaming heads.

3 EXPERIMENTS

Experiments show that DuoAttention preserves long- and short-context performance while improving KV-cache efficiency, decoding, and pre-filling. Its optimization-based retrieval-head identification outperforms profiling and language-modeling alternatives, with sink-plus-recent attention important for deployment.

  • Evaluation Setup: DuoAttention is evaluated on NIAH and LongBench for long-context tasks, and MMLU, MBPP, and MT-Bench for short-context knowledge, coding, and helpfulness.The evaluation includes Llama-2-7B-32K-Instruct, Llama-3-8B-Instruct-Gradient-1048k, Llama-3-[8,70]B-Instruct, and Mistral-7B-v0.2-Instruct, with comparisons to H2O, TOVA, FastGen, and StreamingLLM.
  • Long-Context Accuracy: DuoAttention handles sequences up to 1048K tokens on NIAH, while baseline methods fail across sequence depths because they discard KV states containing needed information.It preserves full KV caches in retrieval heads and discards them only in streaming heads.
  • Long-Context Accuracy: DuoAttention provides a superior KV-budget–accuracy trade-off on most of 14 LongBench tasks and achieves performance comparable to full attention.The comparison uses the same KV cache budget for DuoAttention, H2O, TOVA, and StreamingLLM.
  • Decoding Efficiency: 2.55× and 1.67× memory reductions, alongside 2.18× and 1.50× latency reductions, are achieved for MHA and GQA decoding, respectively.Memory and latency decrease linearly as the retrieval-head ratio is reduced, approaching the inverse ratio at longer contexts.
  • Pre-filling Efficiency: 1.73× and 1.63× latency reductions are achieved for MHA and GQA pre-filling, respectively, with savings increasing as the pre-filling chunk size decreases.The reduction comes from lower time and memory complexity for streaming heads at smaller chunk sizes.
  • Retrieval-Head Identification: Optimization-based identification significantly outperforms attention profiling and traditional language modeling, while combining sink and recent attention is necessary for effective retrieval-head identification.Deployment performance plateaus at 16 sink tokens and 64 recent tokens, with further increases yielding marginal improvements.

4 RELATED WORK

Long-context efficiency methods span architectural changes, approximate attention, KV-cache quantization, and system optimizations. Existing approaches reduce memory or computation in specific ways, but their limitations leave KV-cache reduction and decoding efficiency as complementary targets for DuoAttention.

  • Long-context efficiency methods fall into four categories: model architectures, approximate attention, KV-cache quantization, and system-level optimizations.
  • Model Architecture: MQA and GQA reduce KV-cache size by sharing KV heads across query heads, but require architecture-specific pre-training and do not reduce computational costs.Linear attention reduces memory usage but tends to underperform on long-context tasks.
  • Approximate Attention: Sparse Transformer, LongFormer, and BigBird reduce attention complexity through local, block, or combined local-global patterns, but often require custom GPU kernels or retraining.H2O and TOVA simplify attention by discarding tokens based on query patterns.
  • KV Cache Quantization: 8-bit and 4-bit KV-cache quantization reduce cache size but do not address attention-kernel computational overhead, making them complementary to DuoAttention.
  • System Optimizations: vLLM, FlashAttention, FlashDecoding, and RingAttention improve computational performance, but do not reduce KV-cache size and therefore complement DuoAttention.These methods optimize batch processing, GPU memory hierarchies, decoding speed, or sequence-level parallelism.
  • Recent Works: Wu et al. introduce retrieval heads to explain long-context capabilities but do not compress non-retrieval-head KV caches, while MInference accelerates pre-filling without optimizing KV-cache storage or decoding latency.

5 CONCLUSION

DuoAttention optimizes long-context LLM inference by assigning full KV caches only to Retrieval Heads while using Streaming Heads to reduce memory and latency. It reports substantial memory and decoding-speed gains for both MHA and GQA models.

  • DuoAttention distinguishes Retrieval Heads from Streaming Heads and applies a full KV cache only to Retrieval Heads.This framework targets memory and computational resources in long-context LLM applications.
  • 2.55× memory reduction is achieved for MHA models, compared with 1.67× for GQA models.
  • 2.18× decoding-speed improvement is achieved for MHA models, compared with 1.50× for GQA models.

A APPENDIX · A.1 EXPERIMENTAL DETAILS

The experiments use distributed training and sequence-parallel infrastructure for long sequences, together with a block-sparse approximation for streaming attention. Model-specific maximum sequence lengths and training hyperparameters are reported in the appendix.

  • A.1 EXPERIMENTAL DETAILS: FSDP2 in PyTorch is used for model training.
  • A.1 EXPERIMENTAL DETAILS: DeepSpeed Ulysses sequence parallelism supports long-sequence training.The cited implementation is Jacobs et al. (2023).
  • A.1 EXPERIMENTAL DETAILS: Training uses an efficient block-sparse approximation of Λ-like attention for streaming attention.
  • A.1 EXPERIMENTAL DETAILS: The streaming-attention approximation follows the implementation of Guo et al. (2024).
  • A.1 EXPERIMENTAL DETAILS: Maximum sequence lengths vary across models and are detailed in Table 2.Table 2 is titled “Training Hyperparameters.”
  • A.1 EXPERIMENTAL DETAILS: Figure 14 illustrates the block-sparse approximation of Λ-like attention.

A.2 FULL LONGBENCH RESULTS · A.3 IMPLEMENTATION OF H2O AND TOVA ON LONG-CONTEXT BENCHMARKS

The original H2O and TOVA designs cannot support FlashAttention pre-filling because they require attention scores that FlashAttention does not materialize, causing OOM in long-context evaluations. In full LongBench results with Llama-3-8B-Instruct-1048K, DuoAttention performs best on most datasets at a 50% KV cache budget.

  • A.3 IMPLEMENTATION OF H2O AND TOVA ON LONG-CONTEXT BENCHMARKS: H2O and TOVA rely on attention scores to evict tokens.Their original designs use attention scores during token eviction.
  • A.3 IMPLEMENTATION OF H2O AND TOVA ON LONG-CONTEXT BENCHMARKS: FlashAttention does not materialize attention scores during pre-filling.This incompatibility prevents the original H2O and TOVA algorithms from operating in FlashAttention pre-filling.
  • A.3 IMPLEMENTATION OF H2O AND TOVA ON LONG-CONTEXT BENCHMARKS: H2O and TOVA cannot be evaluated on long-context benchmarks because they cause OOM during context pre-filling.The passage specifically names needle-in-the-haystack and LongBench as affected evaluations.
  • A.2 FULL LONGBENCH RESULTS: DuoAttention achieves the best performance on most Full LongBench datasets.The comparison uses Llama-3-8B-Instruct-1048K.
  • A.2 FULL LONGBENCH RESULTS: DuoAttention reaches this result with a 50% KV cache budget.The table caption identifies this as the cache budget associated with its best performance on most datasets.
  • A.2 FULL LONGBENCH RESULTS: The Full LongBench comparison uses Llama-3-8B-Instruct-1048K.This model designation is specified in the Table 3 caption.

A.4 NIAH RESULTS ON MISTRAL MODELS

The paper reports NIAH results for the Mistral-7B-Instruct-v0.2 and Mistral-7B-Instruct-v0.3 models.

  • Figure 15 presents NIAH results for Mistral-7B-Instruct-v0.2.
  • Figure 16 presents NIAH results for Mistral-7B-Instruct-v0.3.

A.5 IMPLEMENTATION OF FASTGEN ON LONG-CONTEXT BENCHMARKS

This section evaluates DuoAttention against FastGen on LongBench using reproduced FastGen results and reports that DuoAttention consistently outperforms FastGen on feasible datasets. On Llama-2-7B-Instruct-32K, DuoAttention achieves the best performance with a 25% KV cache budget on most datasets.

  • FastGen is reproduced using a community codebase referenced by its official repository because no official implementation is available.
  • FastGen sparsity is controlled through the recovery ratio T, while average KV cache usage across test cases measures overall sparsity.
  • DuoAttention achieves the best performance with a 25% KV cache budget on most Llama-2-7B-Instruct-32K LongBench datasets.
  • DuoAttention consistently outperforms FastGen on LongBench datasets for both evaluated models.Because FastGen has quadratic memory consumption, results are limited to datasets feasible on 8x A100-80G GPUs.
Loading 2410.10819v1…