RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

• URL: http://arxiv.org/abs/2409.10516v2
• Date: Wed, 18 Sep 2024 13:11:13 GMT
• Title: RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval
• Authors: Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chengruidong Zhang, Bailu Ding, Kai Zhang, Chen Chen, Fan Yang, Yuqing Yang, Lili Qiu,
• Abstract summary: RetrievalAttention is a training-free approach to both accelerate attention computation and reduce GPU memory consumption. Our evaluation shows that RetrievalAttention only needs to access 1--3% of data while maintaining high model accuracy.
• Score: 24.472784635757016
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Transformer-based Large Language Models (LLMs) have become increasingly important. However, due to the quadratic time complexity of attention computation, scaling LLMs to longer contexts incurs extremely slow inference latency and high GPU memory consumption for caching key-value (KV) vectors. This paper proposes RetrievalAttention, a training-free approach to both accelerate attention computation and reduce GPU memory consumption. By leveraging the dynamic sparsity of attention mechanism, RetrievalAttention proposes to use approximate nearest neighbor search (ANNS) indexes for KV vectors in CPU memory and retrieves the most relevant ones with vector search during generation. Unfortunately, we observe that the off-the-shelf ANNS indexes are often ineffective for such retrieval tasks due to the out-of-distribution (OOD) between query vectors and key vectors in attention mechanism. RetrievalAttention addresses the OOD challenge by designing an attention-aware vector search algorithm that can adapt to the distribution of query vectors. Our evaluation shows that RetrievalAttention only needs to access 1--3% of data while maintaining high model accuracy. This leads to significant reduction in the inference cost of long-context LLMs with much lower GPU memory footprint. In particular, RetrievalAttention only needs a single NVIDIA RTX4090 (24GB) for serving 128K tokens in LLMs with 8B parameters, which is capable of generating one token in 0.188 seconds.

Related papers

• MagicPIG: LSH Sampling for Efficient LLM Generation [41.75038064509643]
  We show that TopK attention itself suffers from quality degradation in certain downstream tasks because attention is not always as sparse as expected. We propose MagicPIG, a heterogeneous system based on Locality Sensitive Hashing (LSH) MagicPIG significantly reduces the workload of attention while preserving high accuracy for diverse tasks.
  arXiv  Detail & Related papers  (2024-10-21T16:44:51Z)
• 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)
• Training-Free Exponential Context Extension via Cascading KV Cache [49.608367376911694]
  We introduce a novel mechanism that leverages cascading sub-cache buffers to selectively retain the most relevant tokens. Our method reduces prefill stage latency by a factor of 6.8 when compared to flash attention on 1M tokens.
  arXiv  Detail & Related papers  (2024-06-24T03:59:17Z)
• A Training-free Sub-quadratic Cost Transformer Model Serving Framework With Hierarchically Pruned Attention [43.211427581302715]
  We propose Hierarchically Pruned Attention (HiP) to increase context length in large language models. HiP reduces the time complexity of the attention mechanism to $O(T log T)$ and the space complexity to $O(T)$, where $T$ is the sequence length. We show that HiP significantly reduces both prefill and decoding latencies, as well as memory usage, while maintaining high-quality generation with minimal degradation.
  arXiv  Detail & Related papers  (2024-06-14T08:32:45Z)
• CORM: Cache Optimization with Recent Message for Large Language Model Inference [57.109354287786154]
  We introduce an innovative method for optimizing the KV cache, which considerably minimizes its memory footprint. CORM, a KV cache eviction policy, dynamically retains essential key-value pairs for inference without the need for model fine-tuning. Our validation shows that CORM reduces the inference memory usage of KV cache by up to 70% with negligible performance degradation across six tasks in LongBench.
  arXiv  Detail & Related papers  (2024-04-24T16:11:54Z)
• Towards Model-Size Agnostic, Compute-Free, Memorization-based Inference of Deep Learning [5.41530201129053]
  This paper proposes a novel memorization-based inference (MBI) that is compute free and only requires lookups. Specifically, our work capitalizes on the inference mechanism of the recurrent attention model (RAM) By leveraging the low-dimensionality of glimpse, our inference procedure stores key value pairs comprising of glimpse location, patch vector, etc. in a table. The computations are obviated during inference by utilizing the table to read out key-value pairs and performing compute-free inference by memorization.
  arXiv  Detail & Related papers  (2023-07-14T21:01:59Z)
• Similarity search in the blink of an eye with compressed indices [3.39271933237479]
  Graph-based indices are currently the best performing techniques for billion-scale similarity search. We present new techniques and systems for creating faster and smaller graph-based indices.
  arXiv  Detail & Related papers  (2023-04-07T23:10:39Z)
• ETAD: A Unified Framework for Efficient Temporal Action Detection [70.21104995731085]
  Untrimmed video understanding such as temporal action detection (TAD) often suffers from the pain of huge demand for computing resources. We build a unified framework for efficient end-to-end temporal action detection (ETAD) ETAD achieves state-of-the-art performance on both THUMOS-14 and ActivityNet-1.3.
  arXiv  Detail & Related papers  (2022-05-14T21:16:21Z)
• SreaMRAK a Streaming Multi-Resolution Adaptive Kernel Algorithm [60.61943386819384]
  Existing implementations of KRR require that all the data is stored in the main memory. We propose StreaMRAK - a streaming version of KRR. We present a showcase study on two synthetic problems and the prediction of the trajectory of a double pendulum.
  arXiv  Detail & Related papers  (2021-08-23T21:03:09Z)
• IRLI: Iterative Re-partitioning for Learning to Index [104.72641345738425]
  Methods have to trade between obtaining high accuracy while maintaining load balance and scalability in distributed settings. We propose a novel approach called IRLI, which iteratively partitions the items by learning the relevant buckets directly from the query-item relevance data. We mathematically show that IRLI retrieves the correct item with high probability under very natural assumptions and provides superior load balancing.
  arXiv  Detail & Related papers  (2021-03-17T23:13:25Z)

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.