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
Primary area: generative models
Keywords: locality sensitive hashing, randomized algorithms, llm inference, kv cache

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 up to $5\times$ across various GPU hardware and achieve 54ms decoding latency on a single RTX 4090 for Llama-3.1-8B-Instruct model with a context of 96k tokens.

Decision: accept (spotlight)
Reviewer scores: 8, 8, 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1730411439149

thread_type:
official_review

<thread_object>
creator: reviewer_qu5o

<content>
summary:
This paper introduced a novel method dubbed "MagicPIG" to reduce the computation cost of self-attention in long context. Specifically, MagicPIG utilizes Locality-sensitive hashing to approximate the attention score distribution and estimate the attention output. While not decreasing the overall cache size required to store Keys and Values, MagicPIG sampled only a fraction of Keys and Values to calculate the attention scores, reducing the overall computation cost.

strengths:
This paper is exceptionally well-written and clearly presented, making it tremendously helpful for understanding complex topics. Concepts, definitions, and proofs are structured logically, with clear and concise writing. The proposed approach is intuitive and relatively straightforward, with much of the intuition supported by prior explanations. Additionally, the empirical results are strong.

weaknesses:
I have a few questions:
1. what is the intuition for selecting (K, L) for the hash table size? 
2. For Table 1/2/3, why is latency not included in the comparison? 
3. Does the author believe further improvement can be made by combining this approach with PEFT?

questions:
See weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730625747183

thread_type:
official_review

<thread_object>
creator: reviewer_q8fb

<content>
summary:
The authors introduce MAGICPIG, a heterogeneous system leveraging LSH (Locality-Sensitive Hashing) sampling to estimate a complete attention distribution, overcoming limitations of traditional Top-K sparse attention methods, which can underperform in certain downstream tasks.

strengths:
1.	MAGICPIG addresses shortcomings of traditional Top-K attention methods in LLMs, which often assume sparsity and suffer in some downstream applications. Using LSH-based sampling, MAGICPIG more accurately estimates the attention distribution, mitigating the bias found in Top-K approximations. The approach is backed by theoretical guarantees and empirical evidence, underscoring its effectiveness in sparse attention acceleration.

2.	MAGICPIG overcomes GPU VRAM constraints by offloading parts of the computation, including hash table operations, to the CPU. This approach is pivotal for scaling LLMs with LSH-based sampling in resource-constrained, practical environments.

weaknesses:
1.	While the authors discuss CPU-GPU collaboration, they provide limited data on the effects of PCIe bandwidth and CPU-GPU data transfer overhead. This omission may hinder understanding MAGICPIG’s real-world performance across different hardware configurations.

2.	The paper lacks a detailed analysis of the overhead associated with hash tables. As noted by the authors, hash tables could introduce significant memory and computational costs. Therefore, a more thorough evaluation of these overheads would better illustrate the trade-offs of the proposed method.

questions:
1.	Is the size of the hash table related to model size and sequence length? How does the size of the hash table affect the performance?

2.	What is the time overhead of constructing hash tables, and which factors influence this overhead?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730660924321

thread_type:
official_review

<thread_object>
creator: reviewer_epkx

<content>
summary:
This paper introduce a novel approach that leverages LSH sampling to approximate the oracle sampling. Empirical evaluation shows improvement over the baseline.

strengths:
S1. The problem and the solution is well-motivated in the paper.
S2. The CPU-GPU co-design enables the storage of large LSH index table.
S3. The empirical results outperforms baselines.

weaknesses:
Weak Points
----
W1. In the design of the proposed system, the authors claim that putting the retrieval stage in the CPU side would allow large hash tables. I wonder if moving the full system into GPU would reduce the latency when the GPU memory is sufficiently large to fit the hash table.

W2. The author discussed a few KV Cache reduction methods in Section 2. However, only quest is considered as the baselilne in the experiments. I would suggest the author to add a reasonable justification or add more baselines.

W3. Another direction of accelerating the inference is to quantize the model. How does the proposed method work on quantized LLM is not discussed.

W3. No code is provided. It might be hard for readers to reproduce the results.

Presentation
----
P1. In the abstract, without any notes, the author claims "achieve 110ms decoding latency on a single RTX 4090" while not actually running the code on RTX 4090. I believe this is a false claim without mentioning the simulation.
P2. Although it might be obvious for readers with retrieval and word extraction background, the acronym niah, cwe, and fwe and not explained before usage. 
P3. The numbers in Figure 6 might be a bit outdated. In addition, the connection between CPU and GPU could be faster SXM.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730698435768

thread_type:
official_review

<thread_object>
creator: reviewer_uyfg

<content>
summary:
This paper proposes MAGICPIG for efficient attention score sampling to resolve the large KV-cache for LLM inference. The observation is that exact top-k attention sampling may not perform well. The proposal is to conduct sampling using locality sensitive hashing (LSH) and use importance sampling to obtain unbiased estimations. Empirical results show that MAGICPIG achieves high accuracy with low computation cost.

strengths:
1.	The observation that top-k attention does not perform well is interesting.

2.	The idea of using LSH to conduct importance sampling is novel.

3.	MAGICPIG partitions the computation reasonably between GPU and CPU, i.e., the hashing (which involves matrix operations) is conducted on the GPU while attention computation is conducted on CPU.

weaknesses:
1.	Lacks an intuitive explanation why LSH-based importance sampling works better than exact top-k attention. For the theorical view, I get it that importance sampling provides unbiased estimation while exact top-k attention does not. However, both importance sampling and top-k selects some attention scores to compute. Is it because (i) importance sampling select some scores that top-k will not select or (ii) once sampled, importance sampling assigns higher weights to scores with low sampling probabilities? It will be good if an ablation study can be conducted. For instance, if the case is (i), will it work if combine top-k sampling and sampling some random tokens (or some tokens at regular intervals of the sequence, for a good representation of the sequence)?

2.	The parameter configurations for LSH can be discussed, which involves the number of hash table (H), the number of hash functions for a hash table (L), the number of collisions for a token to be considered as a candidate for attention computation (T). Currently, T is fixed at 2. I understand that to sample a fixed number of attention scores, when H is increased, L should be reduced. We can also increase both H and L, but reduce T. Please provide some insights on how these parameters should be set.     

3.	What are the current execution statistics of the system? When the CPU is computing the sampled attention scores, is the GPU idle? GPU or CPU has a longer running time? If we use a pipeline (e.g., by switching between two mini-batches) to overlap GPU and CPU computation, which one will be the straggler?

questions:
See the weakness part
</content>
</thread_object>

</thread>
<thread>
creation date:
1730715357326

thread_type:
official_review

<thread_object>
creator: reviewer_kbnm

<content>
summary:
This paper studies the optimization of long-context LLM inference. Unlike most existing approaches that mainly adopt TopK selection for attention calculation, this paper presents a novel method based on importance sampling, where SimHash is used for estimation. Experiments on a set of benchmarks demonstrate the effectiveness of the proposed method and its superiority over a state-of-the-art TopK selection approach.

strengths:
S1. This paper studies from a new perspective of estimating attention scores, in contrast to TopK selection that has been widely targeted in existing works. The proposed approach showcases its potential in dealing with the non-sparse case of attention. 

S2. Experiments are highly promising, outperforming Quest in both accuracy and efficiency. 

S3. System design is discussed, with the jobs of CPU and GPU clearly depicted in the figure.

weaknesses:
W1. The context length seems to be short in the evaluation.

W2. Some parameter evaluations are missing in the experiments. 

W3. Only Llama series models are evaluated.

questions:
Q1. In Figure 3, did you mean even the exact TopK selection yields a higher relative error than oracle sampling? For oracle sampling, I suppose you estimate the weight of each value vector in the attention. As such, a better result can be obtained, especially for the case when attention is not sparse, where TopK selection treats all non-TopK values as zero-weights. 

Q2. For LSH, why SimHash was chosen? The method proposed by Andoni et al. (Practical and optimal LSH for angular distance, NeurIPS 2015) is a better approach than SimHash and has been used in Reformer. 

Q3. How does the budget B relate to K and L? For each hash probe out of L, there could be multiple k_i's having a hash collision with q. I suppose the number of retrieved k_i's in L hash probes should be reflected to the budget. 

Q4. Following Q3, hash collision could be a problem when the context goes long. In this case, K and L can be adjusted to strike a balance, but it is unclear how they are affected by the context length (the context used in the paper seems to be short, see Q5). 

Q5. What is the maximum context length used in the experiments? It seems to be 96K. It is encouraged to see what if the context goes longer, e.g., 1M, which has been evaluated in some TopK-based approaches such as InfLLM.

Q6. Despite evaluating the importance of centering, Figure 8(a) can be seen also as an evaluation of the impact of L. However, I didn't find the evaluation of K. I wonder how K = 8-10 was determined in the experiment. 

Q7. On LongBench and RULER, the performance is even higher when a smaller set of (K, L) is used, e.g. (8, 75) and (9, 120), in comparison to (10, 150). Why?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732231461251

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all reviewers

comment:
We thank all the reviewers [**R1 (KBnM), R2 (Uyfg), R3 (ePKx), R4 (Q8fb), R5 (qU5o)**] for their thoughtful and highly supportive feedback! We were glad that the reviewers found the work **novel and promising [R1, R2, R3, R5]**, with **empirically solid results [R1, R2, R3, R4, R5]** and **theoretical guarantees [R4]** felt the founding is **interesting [R2]**, and believed our approach is **pivotal [R4]** for scaling LLMs in resource-constrained settings.

We have updated the paper to incorporate constructive suggestions, as shown in the revision. We summarize the major changes:
- **Evaluations on longer contexts and various models [R1]**:  We extended the maximum evaluated context lengths to 256K and evaluated two additional models, i.e., MegaBeam-Mistral-7B-512K and Llama3-Prolong-512K, showing 3.6% to 4.3% absolute accuracy improvement on average, compared to SOTA (i.e., Quest), with significantly smaller computation cost. We added the results in $\text{\textcolor{blue}{Appendix D.1 (Pg. 18, Table 4)}}$.
- **More baselines [R3]**: We added another dynamic KV cache sparsity algorithm, Loki, to our comparison, showing that our methods can significantly outperform Loki by over 25% on average. The results are added in $\text{\textcolor{blue}{Section 5.1 (Pg. 9, Table 3)}}$.
- **Analysis of LSH overhead [R4]**:  We analyzed the memory and computation overhead introduced by LSH sampling. The analysis is presented in $\text{\textcolor{blue}{Appendix E.2 (Pg. 19, Table 6)}}$.
- **LSH hyper-parameters (K, L) configuration [R1, R2, R4, R5]**: We analyzed the questions of how (K, L) is related to computation cost, accuracy, and how to select (K, L) with detailed ablation study presented in $\text{\textcolor{blue}{Appendix E (Pg. 18 - 20)}}$. We also add a note discussing LSH hyper-parameters as a reference to reviewers. 
- **Impact of model sizes [R4]**:  We evaluated our approach with Llama-3.1-70B-Instruct, showing that the same set of LSH hyper-parameters (K, L) work well when scaling to larger models.  We added the results in $\text{\textcolor{blue}{Appendix D.2 (Pg. 18, Table 5)}}$.  We also present the size and overhead of hash tables for different sizes of models in $\text{\textcolor{blue}{Appendix E.2 (Pg. 19, Table 6)}}$.
</content>
</thread_object>

</thread>
<thread>
creation date:
1740601102357

thread_type:
public_comment

<thread_object>
creator: ~anastasiia_filippova2

<content>
comment:
Dear Authors,

Thank you very much for your work! It gave a really solid understanding of the problems behind TopK eviction strategy. 

I have a question regarding Algorithm 1: MagicPIG Decoding (page 8). If I understand correctly, you use two types of key-value pairs to compute the final attention: $K_s, V_s$ and $K_T, V_T$. You refer to $K_T, V_T$ as the static KV cache. However, I could not find an explicit reference to static cache in the paper. If I understand correctly, you are referring to the attention sink tokens as well as local tokens. Could you confirm whether this interpretation is correct? If so, could you kindly help me find the ratio of static vs. sampled keys in the paper?

Additionally, you mentioned:
> Sink tokens (the first several tokens) and local tokens are more likely to be sampled according to their high similarity to the query. To further reduce CPU workload, MagicPIG stores these tokens on GPU and does not apply LSH sampling to them.

Correct me if I am wrong: you store sink keys as well as local keys on the GPU and **do not apply hashing to them**.

I could be wrong, but given the above, does this not imply that the estimator you apply is no longer unbiased?

I look forward to your response! This is a really interesting contribution, and I want to be sure that I fully understand the details.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies the optimization of long-context LLM inference and instead of using TopK selection for attention calculation, presents a method based on LSH and importance sampling. Experiments are encouraging.