MagicPIG: LSH Sampling for Efficient LLM Generation

• URL: http://arxiv.org/abs/2410.16179v2
• Date: Mon, 28 Oct 2024 14:44:22 GMT
• Title: MagicPIG: LSH Sampling for Efficient LLM Generation
• Authors: Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye, Yang Zhou, Jianyu Zhang, Niklas Nolte, Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, Beidi Chen,
• Abstract summary: 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.
• Score: 41.75038064509643
• License:
• Abstract: Large language models (LLMs) with long context windows have gained significant attention. However, the KV cache, stored to avoid re-computation, becomes a bottleneck. Various dynamic sparse or TopK-based attention approximation methods have been proposed to leverage the common insight that attention is sparse. In this paper, we first show that TopK attention itself suffers from quality degradation in certain downstream tasks because attention is not always as sparse as expected. Rather than selecting the keys and values with the highest attention scores, sampling with theoretical guarantees can provide a better estimation for attention output. To make the sampling-based approximation practical in LLM generation, we propose MagicPIG, a heterogeneous system based on Locality Sensitive Hashing (LSH). MagicPIG significantly reduces the workload of attention computation while preserving high accuracy for diverse tasks. MagicPIG stores the LSH hash tables and runs the attention computation on the CPU, which allows it to serve longer contexts and larger batch sizes with high approximation accuracy. MagicPIG can improve decoding throughput by $1.9\sim3.9\times$ across various GPU hardware and achieve 110ms decoding latency on a single RTX 4090 for Llama-3.1-8B-Instruct model with a context of 96k tokens. The code is available at \url{https://github.com/Infini-AI-Lab/MagicPIG}.

Related papers

• Squeezed Attention: Accelerating Long Context Length LLM Inference [64.11145320159126]
  We propose Squeezed Attention as a mechanism to accelerate LLM applications where a large portion of the input prompt is fixed. We use K-means clustering offline to group the keys for the fixed context based on semantic similarity and represent each cluster with a single centroid value. We then compute exact attention using only these important keys from the fixed context, thereby reducing bandwidth and computational costs.
  arXiv  Detail & Related papers  (2024-11-14T18:54:19Z)
• Anchor Attention, Small Cache: Code Generation with Large Language Models [15.94784908771546]
  Current practices in NLP often use sparse attention which may, unfortunately, lead to substantial inaccuracies, or hallucinations, in code generation tasks. We propose a novel approach, AnchorCoder, which features token-wise anchor attention designed to extract and compress contextual information. It can consistently achieve a significant (at least 70%) reduction in KV cache requirements, while preserving the majority of model's performance.
  arXiv  Detail & Related papers  (2024-11-11T02:47:05Z)
• RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval [24.472784635757016]
  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.
  arXiv  Detail & Related papers  (2024-09-16T17:59:52Z)
• 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)
• MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention [36.49445805074941]
  MInference (Milliontokens Inference) is a sparse calculation method designed to accelerate pre-filling of long-sequence processing. We demonstrate that MInference effectively reduces inference latency by up to 10x for pre-filling on an A100, while maintaining accuracy.
  arXiv  Detail & Related papers  (2024-07-02T17:59:56Z)
• 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)
• Enabling High-Sparsity Foundational Llama Models with Efficient Pretraining and Deployment [56.44025052765861]
  Large language models (LLMs) have revolutionized Natural Language Processing (NLP), but their size creates computational bottlenecks. We introduce a novel approach to create accurate, sparse foundational versions of performant LLMs. We show a total speedup on CPUs for sparse-quantized LLaMA models of up to 8.6x.
  arXiv  Detail & Related papers  (2024-05-06T16:03:32Z)
• Sketch and shift: a robust decoder for compressive clustering [17.627195350266796]
  Compressive learning is an emerging approach to drastically reduce the memory footprint of large-scale learning. We propose an alternative decoder offering substantial improvements over CL-OMPR. The proposed algorithm can extract clustering information from a sketch of the MNIST dataset that is 10 times smaller than previously.
  arXiv  Detail & Related papers  (2023-12-15T16:53:55Z)
• 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)
• SMYRF: Efficient Attention using Asymmetric Clustering [103.47647577048782]
  We propose a novel type of balanced clustering algorithm to approximate attention. SMYRF can be used as a drop-in replacement for dense attention layers without any retraining. We show that SMYRF can be used interchangeably with dense attention before and after training.
  arXiv  Detail & Related papers  (2020-10-11T18:49:17Z)

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.