MagicPIG: LSH Sampling for Efficient LLM Generation

## Abstract

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 ∼ 3.9 × 1.9\sim 3.9\times 1.9 ∼ 3.9 × 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 https://github.com/Infini-AI-Lab/MagicPIG .

## 1 Introduction

Figure: Figure 1: While $\mathrm{TopK}$ attention performs well on retrieval tasks (niah) where the useful information reduces to a few words, it degrades severely in aggregated tasks like word extraction (cwe, fwe). x-axis: proportion of attention keys used for $\mathrm{TopK}$ attention.
Refer to caption: x2.png

Large language models (LLMs) with long context windows, such as GPT [], Llama [], and Gemini [], have gained significant attention for their ability to enhance applications like chatbots [], search engines [], and video analysis []. However, serving long-context LLMs is highly challenging due to the unique bottleneck in auto-regressive generation—the key-value (KV) cache, which stores intermediate attention keys and values to avoid re-computation [, ].
Specifically, the KV cache grows linearly with both the batch size and sequence length, occupying substantial GPU memory and increasing decoding time.
Moreover, the KV cache makes LLM generation extremely memory-bound, leading to underutilization of GPU computational power.
For instance, an NVIDIA A100-40GB GPU can only handle a single request for Llama with a 128k context length, with nearly half of the decoding time spent accessing the KV cache, and poor GPU utilization  [].

Leveraging the common insight that attention is naturally sparse, dynamic sparse or $\mathrm{TopK}$-based approximation has been extensively studied [, , , ], but three major challenges prevent a wide adoption in LLM serving systems.
(1) Quality Degradation. They usually propose various strategies to approximate a subset of KV cache that yields the highest attention scores. However, $\mathrm{TopK}$ attention itself is a biased attention approximation and lacks theoretical guarantees. Figure 1 shows that even exact $\mathrm{TopK}$ attention results significantly degrade the accuracy of certain downstream tasks.
(2) High Overhead. There is a large overhead to identify $\mathrm{TopK}$ attention, which becomes the bottleneck rather than the attention computation. For example, as studied in Liu et al. [], naively applying a search algorithm like IVF [] requires access over $30\%$ key states to obtain the exact $\mathrm{TopK}$, showing an unsatisfying trade-off between search accuracy and cost.
(3) No Memory Saving. Although saving KV cache loading time, they cannot reduce the total memory occupied by the KV cache, which limits the maximum context and batch sizes when VRAM is scarce.

An ideal sparse attention approximation approach should (1) preserve full accuracy for a diverse set of downstream tasks with guarantees, (2) involve low-cost overhead for KV cache selection, and (3) save GPU memory. The following observations, together with the performance drop shown in Figure 1 suggest that to achieve such demanding requirements, we need to go beyond $\mathrm{TopK}$ attention:

- •
Attention is not always sparse. Contradictory to previous belief [, , , ], we observe that attention is not always sparse, especially for tasks that leverage the full context. As shown in Figure 2a, in some layers, attention distribution can be very long-tailed, i.e., the $\mathrm{Top}20\%$ attention can only cover $70\%$ of the total attention scores.
- •
Seemingly high sparsity is usually a consequence of an attention sink. Most of the attention scores concentrate on initial tokens (attention sink phenomenon) [], making the distribution look sparser. However, as shown in Figure 2b, attention scores are distributed more uniformly among tokens except for the sink. According to the geometrical interpretation of sink, keys, and queries shown in  Figure 2c, the attention sink, which we found surprisingly almost static regardless of the input token, is just for imposing sparsity on the attention distribution.
- •
It is hard to find $\mathrm{TopK}$ attention. Figure 2c also shows why searching for the Top-K keys is intrinsically costly. The keys and queries usually lie within two narrow cones with nearly opposite orientations, except for the attention sink. This significant mismatch between query and data distributions causes nearest-neighbor search methods to perform poorly.

Figure: (a) Long tailed phenomena
Refer to caption: x3.png

These limitations of $\mathrm{TopK}$ attention require rethinking the sparse attention approximation.
Rather than only using the keys and values with the highest scores, leveraging information on the distribution can make the estimation more accurate.
We approach this as a bias correction problem in sampling.
Unbiased and efficient sampling has been long studied in biology [], sociology [] as well as machine learning [, , ], with theoretical guarantees.

Figure 3 shows that sampling values according to their corresponding attention score (we call this oracle sampling) achieves a much lower (up to $4\times$) estimation error than the naive $\mathrm{TopK}$ selection.
Deploying sampling estimation in attention is promising, but three challenges remain.
First, how a reduction of the attention error can make a difference in downstream performance is unclear [, ].
Second, modeling the attention score distribution is necessary for efficient sampling, but inferring the distribution parameters requires expensive computations.
Third, fully leveraging the resources of modern hardware, GPU and CPU, with a theoretically efficient algorithm is non-trivial.

Figure: Figure 3: $\mathrm{TopK}$ v.s. Sampling, 16k total context
Refer to caption: x6.png

This paper proposes Magic samPlIng for Generation (MagicPIG), which leverages Locality sensitive hashing (LSH) sampling for efficient LLM generation. LSH is employed for sampling to approximate the attention score distribution and estimate attention output. By computing hash functions on GPU and conducting sampling on CPU, MagicPIG can allow massive hash tables and hash functions compared to prior work [, ], which are of vital importance for accurate estimation []. Following the practice of Aminabadi et al. [], He and Zhai [], we offload the KV cache computation, which is memory bound, to CPU to allow a larger batch or longer context. Specifically,

- •
In Section 3, we analyze the failures of $\mathrm{TopK}$ attention. Moreover, we study sampling-based attention estimation assuming an oracle for the key distribution (Oracle Sampling Estimation) and empirically demonstrate that it is consistently more effective both for distribution estimation and downstream tasks.
- •
In Sections 4.1, 4.2 and 4.3, we present a sampling algorithm to approximate oracle sampling for attention estimation based on locality sensitive hashing and the intuition and motivation from statistic perspectives. To our best knowledge, MagicPIG is the first to leverage LSH sampling in self-attention in decoder-only LLM generation.
- •
In Section 4.4, we present our system design to efficiently offload attention computation on the CPU, breaking the memory limit of the GPU for serving larger batches or longer contexts. We also overcome the new challenges of computation and memory size raised by our sampling algorithm to support a larger scale of hashing tables beyond prior work [, ].

In Section 5, we show the empirical evaluation results of the performance of MagicPIG, demonstrating the accuracy and efficiency. While maintaining high accuracy for diverse tasks, MagicPIG can improve serving throughput by $1.9\sim 3.9\times$ (A100, L40, RTX 4090) and can achieve 110ms decoding latency on a single RTX 4090 for Llama-3.1-8B-Instruct  [] with 96K context. More importantly, we show that MagicPIG already outperforms $\mathrm{TopK}$ attention in the two aggregation tasks in Figure 1, suggesting that sampling indeed goes beyond $\mathrm{TopK}$ attention.

## 2 Background

In this section, we formulate the targeted attention estimation problem and related works.

### 2.1 Problem formulation

In LLM decoding phase, self-attention part calculates a weighted average of previous values by

$$ $o=\mathrm{Softmax}(\frac{qK^{T}}{\sqrt{d}})V=wV\quad q\in\mathbb{R}^{1\times d }\quad K,V\in\mathbb{R}^{n\times d}\quad w\in\mathbb{R}^{1\times n}$ (1) $$

where $d$ is the head dimension and $n$ is the context size. $K=[k_{1},k_{2},...,k_{n}],V=[v_{1},v_{2},...,v_{n}],k_{i},v_{i}\in\mathbb{R}^{
1\times d}$ is KV cache. Normalized attention weight $w=\mathrm{Softmax}(\frac{qK^{T}}{\sqrt{d}})\in\mathbb{R}^{1\times n}$ is also called attention (score) distribution. Our target is to find sampling matrix $\Pi\in\mathbb{R}^{n\times m}$ and diagonal matrix $D\in\mathbb{R}^{m\times m}$ which minimize

$$ $\delta=||wV-w\Pi D\Pi^{T}V||$ (2) $$

where $m\ll n$ is computation budget. For $\mathrm{TopK}$ attention, suppose $w_{r_{1}}>...>w_{r_{m}}>...>w_{r_{n}}$, then

$$ $\displaystyle\Pi_{i,j}=\left\{\begin{aligned} 1&,&\text{if }i=r_{j},\\ 0&,&\text{otherwise}.\end{aligned}\right.\quad D_{ii}=\frac{1}{\sum_{i=1}^{m}w _{r_{i}}}$ (3) $$

### 2.2 Related works

Efficient Attention.
Attention approximation has been long studied. Reformer [], KDEformer [] and ScatterBrain []
tackle the problem via locality sensitive hashing. These methods work in training and encoder models like BigGAN []. Theoretically, the error bounds and minimal workload required are continuously improved [, ] but have not proven to be practical for wall-clock acceleration in LLM decoding. Besides, flash-attention [, , ], flash-decoding [, ] and SlimAttention [] losslessly accelerate scaled product attention operator by maximizing the utilization of hardware, which is orthogonal to our approach.

Locality sensitive hashing.
Locality sensitive hashing (LSH) [, ] is a family of hashing functions which assigns the same hash codes for similar inputs with higher probability than others [, ]. LSH uses two hyper-parameters, $(K,L)$. $L$ hash tables are independently built. Each hash table has its own function $H$ which projects a high-dimension vector to an integer by concatenating $K$ random independent hash functions. In the sampling process, all vectors that share hash codes in at least one hash table with a query will be collected.
SimHash [] is the LSH family based on cosine similarity. For a vector $x\in\mathbb{R}^{d}$, SimHash generates a random hyperplane $w$ and returns $\mathrm{Sign}(w^{T}x)$. Vectors share the same sign if and only if the random projection is not in between them. For a random projection, all angles are equally likely, thus the probability that
two vectors $x$, $y$ share the same sign is $p=1-\frac{\theta}{\pi}$, where $\theta=\arccos{\frac{xy^{T}}{||x||\cdot||y||}}$. If we have $L$ hash tables each with $K$ random hash functions, the probability of $y$ to be retrieved by query $x$ is $1-(1-p^{K})^{L}$.

KV Cache reduction.
To get rid of memory bound introduced by KV cache thus enabling a larger batch size or serving a longer prompt, many methods are proposed to reduce the volume of KV cache. For example, H2O [], SnapKV [] and Keyformer [] calculate heuristics during the prefilling phase to decide which tokens to preserve for decoding phase. Quest [] and Loki [] do not evict KV cache but apply dynamic sparsity to reduce KV Cache loading at inference time.
Besides the reduction along the dimension of sequence length, methods like KIVI [] and QServe [] reduce the size of KV Cache by quantization.

## 3 Rethinking attention sparsity

In this section, we examine $\mathrm{TopK}$ attention, which is the theoretical upper bound of prior search-based algorithms, including both static methods [, ] and dynamic methods [, , ].
We show that $\mathrm{TopK}$ is *sub-optimal* and present another attention approximation based on sampling and estimation with an oracle that improves the accuracy and/or the computation cost.

### 3.1 Achilles’ heel of TopK attention

Figure: Figure 4: $\mathrm{TopK}$ estimation error for a KV-cache of 16k tokens.
Refer to caption: x7.png

As it is defined, $\mathrm{TopK}$ attention only computes the weighted average on elements with the highest attention scores. To quantify its performance, the computation budget of $\mathrm{TopK}$ attention is defined as the number of selected tokens, i.e., the $\mathrm{K}$ of $\mathrm{TopK}$.
Searching-based sparse attention algorithms, like [, , ], are approximations for $\mathrm{TopK}$ attention by replacing the true $\mathrm{TopK}$ keys with the ones found by approximate searching algorithms.

However, we find significant performance degradation in downstream tasks caused by $\mathrm{TopK}$ attention as shown in Figure 1.
Although $\mathrm{TopK}$ attention preserves accuracy for retrieval tasks that only require a minimal subset of the context (needle-in-a-haystack single/multikey []), it severely degrades for aggregation tasks that leverage the full context (common word extraction and frequent word extraction []).
Intuitively, the information is distributed more broadly for aggregation tasks, which results in less peak attention score distribution.

$\mathrm{TopK}$ attention is biased and inaccurate, especially when the distribution of attention scores is long-tailed and the computation budget or density (i.e., $K$) is limited.
Unfortunately, long-tailed phenomena do occur in LLMs across all layers (prior works [, , ] usually skip the first two layers to maintain accuracy) as presented in Figure 2a.
$\mathrm{Top}20\%$ tokens can only cover $70\sim 80\%$ attention scores, leaving a large proportion of keys and values not considered, which is translated into a non-negligible ($15\sim 20\%$) estimation error in Figure 4.

To better understand the attention distribution, we study the geometry of $q,k$ and make the following three observations. (1) Key states of the initial token (also known as attention sink, denoted by $k_{sink}$) remain almost the same for arbitrary input. In Figure 5a, we randomly draw $32$ samples from the vocabulary and measure the mutual cosine similarity of key states. Surprisingly, we find that the orientations of the key states of different input tokens are almost identical with a similarity $>0.99$.
(2) The orientation of the center of key states (i.e. $k_{avg}=\frac{1}{n}\sum_{i=1}^{n}k_{i}$) remains stable for different input sentences. In Figure 5b, we measure the mutual cosine similarity of $k_{avg}$ of $50$ different input sentences. Although variance exists, the similarity of $k_{avg}$ is over $0.9$. (3) The orientations of $k_{avg}$ and $k_{sink}$ are almost opposite. In Figure 5c, we find that for each head, $k_{sink}$ and $k_{avg}$ has a cosine similarity between $-0.9\sim-0.8$.

These observations shape the geometry as shown in Figure 2c. The attention sink, which is static regardless of input, produces high sparsity in the attention distribution, whereas other parts are more uniformly distributed. Simply applying $\mathrm{TopK}$ will place even more weight on the sink token, thus losing contextual information. In addition, the misalignment of $q$ and $k$ also causes difficulty in search [].

Figure: (a) $k_{sink}$
Refer to caption: x8.png

### 3.2 Estimate attention with sampling

Figure: (a) Estimation error Layer11
Refer to caption: x11.png

Existing TopK attention mechanisms ignore tokens in the KV cache with low attention scores, which introduces a bias since the ignored tokens comprise a large proportion of attention scores (Figure 2a).
As a result, TopK attention achieves suboptimal performance for long-context tasks, such as information aggregation (Figure 1).
Increasing the computation budget for TopK attention does help reduce the estimation error (Figure 4) since it will involve more elements in computing. However, the following question is posed:

> Can we improve the estimation quality with low computational budgets?

Inspired by mark and recapture [, , , ], we show in the following that attention output can be estimated with sampling.
Using notations from Section 2.1 we can re-write attention output $o$ as the expectation of $v_{i},1\leq i\leq n$ from distribution $w$, i.e. $o=\mathbb{E}_{i\sim w}(v_{i})$, which can be estimated by the following method.

###### Definition 3.1 (Oracle Sampling Estimation) .

Given a sampling budget $\mathcal{B}$ and normalized attention score $w$, $\mathcal{B}$ elements are sampled independently from $w$ (i.e. $i_{1},i_{2},...,i_{\mathcal{B}}\overset{\mathrm{iid}}{\sim}w$). Then the attention output is estimated as

$$ $\bar{o}=\frac{1}{\mathcal{B}}\sum_{j=1}^{\mathcal{B}}v_{i_{j}}$ (4) $$

This is not the lowest variance estimator but has better downstream performance (see Appendix B).
We call it “oracle” because it assumes that the exact attention vector $w$ is known, which is not true for sparse attention approximations.

###### Theorem 3.2 .

Oracle sampling estimation is unbiased, and the trace of covariance monotonically decreases with $\mathcal{B}$.

This theorem (proved in Appendix A) theoretically guarantees a low estimation error of oracle sampling.
We also present an empirical comparison between oracle sampling estimation and $\mathrm{TopK}$ attention in Figures 6a and 6b.
In summary, oracle sampling estimation can reduce relative error by up to $4\times$.

Note that the sampling budget $\mathcal{B}$ is not the actual computation cost for oracle sampling estimation:
duplicate $X_{i}$ need to be computed/loaded only once, so $\bar{o}$ can be computed by

$$ $\bar{o}=\sum_{i\in\mathcal{S}}\frac{f_{i}}{\mathcal{B}}v_{i}\quad{S}=\mathrm{ Unique}(\{i_{1\leq i\leq\mathcal{B}}\})$ (5) $$

where $f_{i}$ is the number of duplicates of $X_{i}$.
Intuitively, if $w$ has an peaked distribution (e.g. $w_{i}>99\%$), then almost all samples in $\{i_{1},...,i_{\mathcal{B}}\}$ are identical to $i$. The actual computation cost of oracle sampling estimation is $|{S}|$, the number of *unique* samples, which we bound in the following:

###### Theorem 3.3 .

The expected computation budget ($\mathbb{E}({|{S}|}))$ has an upper bound of $1+\mathcal{B}\epsilon$, where $\epsilon=1-\max_{i}w_{i}$.

This theorem (proved in Appendix A) shows that the computation cost of oracle sampling is usually far less than the sampling budget.
In Figure 6c, we present the downstream accuracy comparison between oracle sampling estimation and $\mathrm{TopK}$ attention.
The former preserves high accuracy for both tasks, even with a very small computation cost ($0.002\%$ out of 16k context, which is approximately $32$).

###### Definition 3.1 (Oracle Sampling Estimation) .

###### Theorem 3.2 .

###### Theorem 3.3 .

## 4 MagicPIG

Section 3.2 demonstrates the potential of sampling-based estimation. In Sections 4.1 and 4.2, we present how we arrive at Locality sensitive hashing to unleash this potential from a statistical perspective.
In Section 4.3, we show the practical algorithm.
Finally, in Section 4.4, we demonstrate our system co-design for accurate and efficient LLM decoding through GPU-CPU collaboration.

Note that most of the derivations in this section might be classical and can even be found in textbooks, but our goal is to leverage them to motivate MagicPIG design and precisely demonstrate the power of a rigorously sound algorithm with system co-design in deep generative models.

### 4.1 Self-normalized importance sampling for attention estimation

Oracle sampling estimation cannot go beyond $2\times$ wall clock speed up because obtaining distribution $w$ requires full computation of all $qk_{i}^{T}$, thereby only saving the $wV$ computation.

Fortunately, importance sampling [, , ] allows us to estimate unknown distribution $w$ by sampling from a proposed distribution $u$.
In our problem setting, the normalization factor of $w$, i.e. $Z=\sum_{i=1}^{n}\exp{\frac{qk_{i}^{T}}{\sqrt{d}}}$ is also unknown because computing it requires evaluating all $qk_{i}^{T}$.
However, we do have access to unnormalized weights $\widetilde{w_{i}}=e^{\frac{qk_{i}^{T}}{\sqrt{d}}}$ for sampled indices $i$.
Hence, by employing a variant of importance sampling, self-normalized importance sampling [], we sample indices $i_{1},i_{2},...,i_{\mathcal{B}}$ from a proposed distribution $u$ and the resulting estimator is

$$ $X^{\mathrm{IS}}=\frac{1}{\widetilde{Z}}\sum_{j=1}^{\mathcal{B}}\frac{ \widetilde{w_{i_{j}}}}{u_{i_{j}}}v_{i_{j}}\quad\text{where}\quad\widetilde{Z}= \sum_{j=1}^{\mathcal{B}}\frac{\widetilde{w_{i_{j}}}}{u_{i_{j}}}$ (6) $$

which has a very nice property for accurately estimating attention output that $\mathbb{P}[\lim_{k\rightarrow\infty}X^{\mathrm{IS}}=o]=1$. Its variance(^1^11We assume head dimension $d=1$ here for simplicity. Higher dimensions have similar formulations and analyses by replacing variance with the trace of covariance.) is related to the distribution $u$, and can be approximated by

$$ $\widetilde{\mathrm{Var}}(X^{\mathrm{IS}})=\frac{1}{\mathcal{B}}\mathbb{E}_{i \sim u}[\frac{w_{i}^{2}}{u_{i}^{2}}(v_{i}-o)^{2}]=\frac{1}{\mathcal{B}Z^{2}} \mathbb{E}_{i\sim u}[\frac{\widetilde{w_{i}}^{2}}{u_{i}^{2}}(v_{i}-o)^{2}]$ (7) $$

To minimize the variance, $u$ should satisfy $u\propto\widetilde{w_{i}}|v_{i}-o|$ [].
The variance will be high if $u_{i}$ and $\widetilde{w_{i}}|v_{i}-o|$ assign a high probability mass to different regions of the sample space or have different modes.
Therefore, the challenge is computing a distribution $u$ aligned with $\widetilde{w_{i}}|v_{i}-o|$ without accessing too many $\widetilde{w_{i}}$. Besides, Equation 6 requires that sampling probability $u$ can be computed and $u_{i}>0$, which is not satisfied by many deterministic approximations like $\mathrm{TopK}$.

### 4.2 Variance reduction with LSH

We decompose $\widetilde{w_{i}}|v_{i}-o|=\exp({\frac{qk_{i}^{T}}{\sqrt{d}}+\log|v_{i}-o|}$).
We observe empirically (Figure 10 in the appendix) that $\log|v_{i}-o|$ does not fluctuate significantly compared to $\frac{qk_{i}^{T}}{\sqrt{d}}$. Hence, we simplify the requirement of $u$ to share the same peaks with $qk_{i}^{T}$. By the following transformation,

$$ $\displaystyle r=\max_{1\leq i\leq n}|k_{i}|\quad\bar{q}=[q,0]\quad\bar{k_{i}}= [k_{i},\sqrt{r^{2}-|k_{i}|^{2}}]$ (8) $$

we further transfer the inner product $qk_{i}^{T}$ to cosine similarity between $\bar{q}$ and $\bar{k_{i}}$ (which is a common practice in Maximum Inner Product Search []).

Inspired by prior work [, ], we leverage Locality sensitive hashing-based sampling for this estimation problem.
Specifically, leveraging a hash function $h$ in the LSH family that preserves cosine similarity such as SimHash [], we can sample from probability distribution $u_{i}=\mathbb{P}[{h}(q)=h(k_{i})]$ which is monotonic to $\cos{\frac{qk_{i}^{T}}{|q|\cdot|k_{i}|}}$.

### 4.3 Algorithm implementation

To make this estimation practical, MagicPIG is implemented by the following specific design.

Estimator approximation. Self-normalized important sampling Equation 6 requires $i_{1},i_{2},...,i_{k}$ $\mathrm{iid}$ sampled, but the probabilities provided by hashing are not normalized.
Hence we adapt our estimator: After obtaining $S$ with probability $u$, MagicPIG computes

$$ $X=\frac{\sum_{i=1}^{n}\frac{\widetilde{w_{i}}}{u_{i}}v_{i}\mathbf{1}_{i\in S}} {\sum_{i=1}^{n}\frac{\widetilde{w_{i}}}{u_{i}}\mathbf{1}_{i\in S}}=\frac{\sum_ {i\in S}\frac{\widetilde{w_{i}}}{u_{i}}v_{i}}{\sum_{i\in S}\frac{\widetilde{w_ {i}}}{u_{i}}}$ (9) $$

Hash function selection. MagicPIG leverages SimHash [], that draws with $K\times L$ random vectors.
For each of the $L$ hash tables, the ${q}$ and ${k_{i}}$s vectors are projected on $K$ directions, and only the sign of the projection is kept, which yields a $K$-bit hash value.
Key ${k_{i}}$ is sampled only if there exist at least two hash tables where ${k_{i}}$ shares the hash value with ${q}$.
The corresponding probability is

$$ $u_{i}=\mathbb{P}[\text{${k_{i}}$ is sampled}]=1-(1-p^{K})^{L}-Lp^{K}(1-p^{K})^ {L-1}\quad\text{where}\quad p=1-\frac{1}{\pi}\arccos{\frac{{q}{k_{i}}^{T}}{|{q }|\cdot|{k_{i}}|}}$ (10) $$

Figure: Algorithm 1 MagicPIG Decoding

Data pre-processing. Before building hash tables, MagicPIG centers the $k_{i}$ vectors.
As shown in Figure 2c, keys are almost always concentrated on one side of the queries, except the initial token.
In this case, random projections cannot effectively distinguish keys, resulting in uniform sampled probabilities. $\mathrm{Softmax}$ is translation invariant.
Centering ($\bar{k_{i}}=k_{i}-\frac{1}{n}\sum_{i=1}^{n}k_{i}$) distributed the keys better and remains computationally equivalent.

Combining Equations 9 and 10 gives a closed form of the MagicPIG attention estimation. Assuming sample set $S$ is obtained with LSH,

$$ $\displaystyle\bar{o}$ $\displaystyle=\sum_{i\in S}\frac{\exp{(\frac{qk_{i}^{T}}{\sqrt{d}}-\log u_{i}) }}{\sum_{i\in S}\exp{(\frac{qk_{i}^{T}}{\sqrt{d}}-\log u_{i})}}v_{i}$ $\displaystyle u_{i}$ $\displaystyle=1-(1-p_{i}^{K})^{L}-Lp_{i}^{K}(1-p_{i}^{K})^{L-1}$ $\displaystyle p_{i}$ $\displaystyle=1-\frac{1}{\pi}\arccos{\frac{{q}{k_{i}}^{T}}{|{q}|\cdot|{k_{i}}|}}$ (11) $$

### 4.4 System co-design

Figure: (a) Memory Hierarchy
Refer to caption: x14.png

The memory size of KV cache remains a bottleneck for long-context LLM decoding, especially when GPU VRAM is limited. DRAM on the CPU side offers sufficient memory capacity with $100-200$GB/s bandwidth, which is usually $10-20\%$ of GPU VRAM bandwidth (see Figure 7a).
Ideally, this gap can be mitigated by $5-10\times$ sparsity.
To make CPU DRAM an aggregated memory for GPU, the workload must be partitioned.
In our experiments, $K=9$ or 10, and $L$ is a few hundred.

Our system design extends prior work [, ] by splitting LLM decoding into three parts.
(1) Parameter computations, i.e., all linear projectors including MLP and $W_{Q},W_{K},W_{V},andW_{O}$ in the self-attention module run on GPU.
(2) Attention computation, which involves $o=\mathrm{Softmax}(\frac{qK^{T}}{\sqrt{d}})V$, runs on CPU.
(3) Random projections. At generation time, for each $q$, $K\times L$ random projections are conducted to obtain the hash codes.
Since all heads can share the same random projectors, the memory overhead is limited (25 MB in our implementation), so this step is compute-bound.
Therefore, the projection is placed on GPU.
(4) Retrieval. The hash codes of $q$, need to be looked up in $L$ hash tables, which is negligible computationally. However, the pre-built hash tables for $k_{i}$s can occupy considerable memory, making it a better fit for the CPU.
With the above partition, we are able to support hash tables with $K$ and $L$ beyond the scale of prior work [, , ] without worrying about computation for hash codes as well as the storage of hash tables.

On-device cache. 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. We leverage the recursive attention technique [] to merge the attention output from CPU and GPU.

Our algorithm applies to a single attention head, see Algorithm 1.
The details of Encode, Query, as well as the hash table construction, are described in prior work [, ].

## 5 Evaluation

In this section, we aim to demonstrate that MagicPIG can speed up LLM decoding while preserving high accuracy. We first present MagicPIG’s accuracy in downstream tasks, followed by our end-to-end system results showing wall-clock performance.

- •
In Section 5.1, we demonstrate MagicPIG preserves high accuracy (less than $2\%$ degradation) across moderate to long context tasks with computation cost $2\%\sim 5\%$ of full attention.
- •
In Figure 8, we demonstrate the system performance of MagicPIG, which achieves up to $3.9\times$ throughput improvement and 110ms decoding latency on a single RTX 4090 for Llama-3.1-8B-Instruct with 96K context.
- •
In Section 5.3, we verify the effectiveness of centering, which is of vital importance for the success of sampling. Also, we demonstrate that MagicPIG already outperforms $\mathrm{TopK}$ attention in the two aggregation tasks in Figure 1, indicating that sampling indeed goes beyond $\mathrm{TopK}$ attention.

### 5.1 MagicPIG Preserves Accuracy

We demonstrate that MagicPIG can preserve accuracy in diverse tasks with less than $5\%$ computation.

Setup. Our experiments are based on Llama [, , ] models. Three types of tasks are included, which are 3 mid-context comprehensive tasks from lm-eval-harness [] (GSM8K-CoT [], MMLU-Flan-Cot-Fewshot [] and COQA []), and 6 long context tasks from [] (QASPER [], LCC, Repobench-P [], TriviaQA [], PRE and TREC [, ]) and 13 synthetic tasks from RULER [] (with 50 examples per task).

Baselines. Besides full attention, Quest [] and its variants are used as baselines.
In its default setting, Quest uses a “page size” of 16, i.e. 1/16 of the full attention cost.
To compare the methods fairly in the low computation budget regime, we also evaluate Quest with page size 32 and 64 and make sure at least one page is selected in every test example.
The initial 4 tokens and local 64 (for LongBench [] and RULER []) or 24 (for lm-eval-harness []) tokens as well as layer-$\{0,16\}$ are statically preserved.
We do not use the theoretical transformations in Equation 8 in our implementations, as we do not find them to contribute to accuracy improvements.

Cost. The cost for the attention approximation consists of two parts:
$\text{Cost}_{1}$ is the sampling/search cost to obtain $S$ in Equation 11,
$\text{Cost}_{2}$ is the attention computation cost, see Equation 11.
We report the ratio of the number of FLOPs compared to the full attention computation.
For MagicPIG, $\text{Cost}_{1}\simeq 0$ and $\text{Cost}_{2}$ is empirically measured for different LSH hyper-parameters. For Quest with page size $K$, $\text{Cost}_{1}=\frac{1}{K}$ and $\text{Cost}_{2}$ is controlled manually.

Analysis. From Tables 1, 2 and 3,
(1) MagicPIG preserves high accuracy (degradation less than $2\%$) for all kinds of tasks, with a computation cost of $2\%\sim 5\%$.
(2) Compared with Quest, which also shows reasonable performance on long context tasks, MagicPIG also demonstrates good performance on tasks with moderate context sizes in lm-eval-harness [], indicating a more robust performance in general serving. (3) With LSH sampling, which introduces an order of magnitude lower sampling/searching cost ($\text{Cost}_{1}$), MagicPIG can achieve equivalent or better accuracy with only half of the computation cost.

**Table 1: Comprehensive tasks on lm-eval-harness []. MagicPIG significantly outperforms other methods with lower computation. The config (K, L) is a hyper-parameter of LSH for MagicPIG or page size and ratio of selected pages for Quest []. $\text{Cost}_{1}$, $\text{Cost}_{2}$ represents the cost for searching/sampling and sparse attention computation.**
| Methods | Config | GSM | COQA | MMLU | Avg. | $\text{Cost}_{1}$ | $\text{Cost}_{2}$ | $\text{Cost}_{total}$. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-2-7b-chat | Full | 22.4 | 75.8 | 49.2 | 49.1 | 0.00 | 1.00 | 1.00 |
| MagicPIG | (10,220) | 17.3 | 76.4 | 48.6 | 47.4 | 0.00 | 0.04 | 0.04 |
| MagicPIG | (8,90) | 18.7 | 75.0 | 47.9 | 47.2 | 0.00 | 0.08 | 0.08 |
| Quest | (16,0.05) | 13.0 | 69.4 | 41.4 | 41.3 | 0.06 | 0.05 | 0.11 |
| Quest | (32,0.1) | 15.7 | 70.2 | 44.0 | 43.3 | 0.03 | 0.10 | 0.13 |
| Llama-3.1-8B-Instruct | Full | 77.6 | 78.5 | 65.2 | 73.7 | 0.00 | 1.00 | 1.00 |
| MagicPIG | (10,220) | 72.7 | 78.1 | 62.7 | 71.2 | 0.00 | 0.03 | 0.03 |
| MagicPIG | (8,90) | 71.0 | 78.0 | 61.3 | 70.1 | 0.00 | 0.07 | 0.07 |
| Quest | (16,0.05) | 57.9 | 64.6 | 42.5 | 55.0 | 0.06 | 0.05 | 0.11 |
| Quest | (32,0.1) | 64.5 | 65.0 | 48.0 | 59.2 | 0.03 | 0.10 | 0.13 |

**Table 2: Long context tasks on LongBench []. MagicPIG preserves high accuracy with low computation. Config and cost are defined as in Table 1. Code models are only evaluated by Repobench-P and LCC.**
| Methods | Config | QaS | RbP | LCC | PrE | TrC | TrQ | Avg. | $\text{Cost}_{1}$ | $\text{Cost}_{2}$ | $\text{Cost}_{total}$. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-3.1-8B-Instruct | Full | 44.9 | 52.1 | 66.8 | 100.0 | 71.3 | 91.8 | 71.2 | 0.00 | 1.00 | 1.00 |
| MagicPIG | (10,150) | 43.2 | 50.2 | 64.4 | 100.0 | 71.3 | 92.2 | 70.3 | 0.00 | 0.02 | 0.02 |
| MagicPIG | (8,75) | 43.5 | 50.4 | 67.0 | 100.0 | 71.7 | 91.7 | 70.7 | 0.00 | 0.05 | 0.05 |
| Quest | (16,0.05) | 45.7 | 49.7 | 64.9 | 100.0 | 71.7 | 91.5 | 70.6 | 0.06 | 0.05 | 0.11 |
| Quest | (32,0.1) | 44.4 | 50.5 | 65.1 | 100.0 | 71.3 | 91.6 | 70.5 | 0.03 | 0.10 | 0.13 |
| Code-Llama-13b-16K | Full |  | 58.5 | 74.7 |  |  |  | 66.6 | 0.00 | 1.00 | 1.00 |
| MagicPIG | (10,150) |  | 56.9 | 74.0 |  |  |  | 65.5 | 0.00 | 0.03 | 0.03 |
| Quest | (16,0.05) |  | 56.4 | 74.4 |  |  |  | 65.4 | 0.06 | 0.05 | 0.11 |

**Table 3: Synthesized tasks on RULER []. MagicPIG preserves high accuracy with low computation. Config and cost are defined as in Table 1.**
| Methods | Config | 16K | 32K | 64K | 96K | Avg. | $\text{Cost}_{1}$ | $\text{Cost}_{2}$ | $\text{Cost}_{total}$. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-3.1-8B-Instruct | Full | 94.2 | 91.5 | 86.1 | 83.0 | 88.7 | 0.00 | 1.00 | 1.00 |
| MagicPIG | (10,150) | 91.8 | 88.9 | 84.8 | 80.0 | 86.4 | 0.00 | 0.02 | 0.02 |
| MagicPIG | (9,120) | 93.4 | 90.6 | 84.7 | 81.5 | 87.6 | 0.00 | 0.04 | 0.04 |
| MagicPIG | (8,75) | 92.9 | 90.2 | 84.9 | 81.7 | 87.4 | 0.00 | 0.05 | 0.05 |
| Quest | (16,0.04) | 86.3 | 85.4 | 81.9 | 74.9 | 82.1 | 0.06 | 0.04 | 0.10 |
| Quest | (32,0.06) | 84.3 | 84.0 | 80.1 | 74.4 | 80.7 | 0.03 | 0.06 | 0.09 |
| Quest | (64,0.08) | 85.2 | 84.3 | 77.0 | 74.2 | 80.2 | 0.02 | 0.08 | 0.10 |
| MegaBeam-Mistral-7B-512K | Full | 91.7 | 88.1 | 83.5 | 83.7 | 86.8 | 0.00 | 1.00 | 1.00 |
| MagicPIG | (10,150) | 89.8 | 86.5 | 81.7 | 80.7 | 84.7 | 0.00 | 0.02 | 0.02 |
| MagicPIG | (9,120) | 90.7 | 88.5 | 82.9 | 82.4 | 86.1 | 0.00 | 0.04 | 0.04 |
| MagicPIG | (8,75) | 90.6 | 86.4 | 82.8 | 81.6 | 85.4 | 0.00 | 0.05 | 0.05 |
| Quest | (16,0.04) | 83.3 | 83.2 | 79.3 | 78.6 | 81.1 | 0.06 | 0.04 | 0.10 |
| Quest | (32,0.06) | 81.5 | 80.8 | 76.7 | 74.4 | 78.4 | 0.03 | 0.06 | 0.09 |
| Quest | (64,0.08) | 79.6 | 77.5 | 73.8 | 73.7 | 76.1 | 0.02 | 0.08 | 0.10 |

### 5.2 MagicPIG Shows Impressive Efficiency across Various Hardware Settings

We show MagicPIG can bring up to $3.9\times$ wall clock speed up and reduce GPU memory consumption on different models and hardware settings (A100, L40, RTX4090).

Figure: (a) A100 with 34B model
Refer to caption: x16.png

Setup. We evaluate our system performance on 3 serving settings. (1) 80GB GPU (A100) and 34B model (CodeLlama-34B) [] with 16K contexts; (2) 48GB GPU (L40) and 13B model (CodeLlama-13B) [] with 16K contexts; (3) 24GB GPU(^2^22We simulate 24GB GPU by setting memory limit with L40. As the bandwidth of L40 (864GB/s) is less than RTX 4090 (1TB/s), the real speed of our system should be slightly faster than the simulation.) (e.g. RTX 4090) and 8B model (Llama-3.1-8B) [] with 96K contexts.

Baselines. Our baselines for (1) and (2) are full attention on GPU, and for (3) is full attention on CPU with theoretical estimated bandwidth. Our system’s GPU part is implemented in native Pytorch [] and the CPU part in FBGEMM [] in bfloat16 precision. Our CPU is Intel(R) Xeon(R) Platinum 8480+ for A100 and Intel(R) Xeon(R) Gold 6338 CPU @ 2.00GHz for L40. In the last setting, the CPU bandwidth is estimated at 100GB/s, above the empirical bandwidth we measure when running a group query attention of size 4.

Analysis. In Figures 8a, 8c and 8b, we demonstrate (1) MagicPIG significantly improves decoding throughput for all three scenarios (A100: $1.9\times$, L40: $3.9\times$, RTX 4090: $2.9\times$) and can achieve a latency of 110ms for single request generation with 96K context for RTX 4090. (2) With KV cache offloading, MagicPIG can fit much larger batches than GPU full attention baselines ($15\sim 18\times$).

### 5.3 Ablation Study

Figure: (a) Importance of centering
Refer to caption: x19.png

In this section, we empirically validate our two previous observations.

Centering is important for good performance. In Section 4.3, we use a translation to center the keys before applying LSH sampling. Empirical results show this to be important for downstream tasks as shown in Figure 9a. Without centering, the accuracy drops to almost zero in retrieval (NIAH) and degrades to $65\%$ in FWE. We find almost no keys (less than $0.1\%$) can be sampled by the query without centering, as their orientation is almost opposite, as shown in Figure 2c.

Sampling goes beyond $\mathrm{TopK}$. In Figures 9b and 9c, We compare the performance of MagicPIG and $\mathrm{TopK}$ attention in two aggregated tasks (CWE, FWE) where $\mathrm{TopK}$ attention experiences significant performance degradation (Figure 1). MagicPIG can even beat exact $\mathrm{TopK}$ attention in these two tasks by a margin up to $3\%$ and $8\%$ respectively, demonstrating that sampling improves the ceiling of $\mathrm{TopK}$, which is impossible for a search-only algorithm.

## 6 Conclusion

In this work, we first present the limitation of $\mathrm{TopK}$ attention approximation for addressing the computational and memory challenges of long-context LLM generation.
Then we show oracle sampling can go beyond $\mathrm{TopK}$ and introduce MagicPIG, a novel approach that leverages LSH sampling to approximate the oracle sampling.
MagicPIG significantly reduces the workload of attention computation while preserving high accuracy across diverse tasks.
MagicPIG relies on LSH sampling and a system co-design that offloads hash tables and reduced attention computation to the CPU.
Our experimental results demonstrate that MagicPIG substantially improves throughput and latency across multiple hardware configurations, outperforming traditional $\mathrm{TopK}$ attention mechanisms.
The theoretical soundness, robustness, and scalability of MagicPIG open up new opportunities in both attention approximation methods and algorithm-hardware co-design.

## Appendix A Proofs for theorems

### A.1 Proof for Theorem 3.2

###### Proof.

$$ $\displaystyle\mathbb{E}(\bar{o})=\frac{1}{\mathcal{B}}\sum_{j=1}^{\mathcal{B}} \mathbb{E}[v_{i_{j}}]=\frac{1}{\mathcal{B}}\sum_{i=1}^{{n}}w_{i}v_{i}=o$ (12) $$

Assume $\Sigma_{1}$ is the covariance matrix of $\bar{o}$, $\Sigma_{2}$ is the covariance matrix of $v_{i}$

$$ $\displaystyle\text{Tr}(\Sigma_{1})=\frac{1}{\mathcal{B}}\text{Tr}(\Sigma_{2})= \frac{1}{\mathcal{B}}(\mathbb{E}[||v_{i}||^{2}]-||\mathbb{E}[v_{i}]||^{2})= \frac{1}{\mathcal{B}}(\mathbb{E}[||v_{i}||^{2}]-||o||^{2})$ (13) $$

$\mathbb{E}[||v_{X}||^{2}]-||o||^{2}$ is a constant, so the trace of covariance matrix monotonically decreases with $\mathcal{B}$.
∎

###### Proof.

### A.2 Proof for Theorem 3.3

###### Proof.

$$ $\displaystyle\mathbb{E}[|S|]=\mathbb{E}\Big{[}\sum_{i=1}^{n}\textbf{1}_{i\in S }\Big{]}=\sum_{i=1}^{n}\mathbb{E}[\textbf{1}_{i\in S}]=\sum_{i=1}^{n}(1-(1-w_{ i})^{\mathcal{B}})=n-\sum_{i=1}^{n}(1-w_{i})^{\mathcal{B}}$ (14) $$

Without loss of generality, let $a_{i}=1-w_{i}$ and $a_{1}=\min_{1\leq i\leq n}a_{i}={\epsilon}$, then

$$ $\displaystyle\mathbb{E}[|S|]$ $\displaystyle=n-\sum_{i=1}^{n}a_{i}^{\mathcal{B}}=n-a_{1}^{\mathcal{B}}-\sum_{ i=2}^{n}a_{i}^{\mathcal{B}}$ (15) $\displaystyle=n-{\epsilon}^{\mathcal{B}}-\sum_{i=2}^{n}a_{i}^{\mathcal{B}}$ (16) $$

$f(x)=x^{\mathcal{B}}$ is convex function with $\mathcal{B}\geq 1$ and $x\geq 0$. Then with Jensen’s inequality, we have

$$ $\displaystyle\sum_{i=2}^{n}a_{i}^{\mathcal{B}}$ $\displaystyle\geq(n-1)\Big{(}\frac{\sum_{i=2}^{n}a_{i}}{n-1}\Big{)}^{\mathcal{ B}}=(n-1)\Big{(}\frac{(\sum_{i=1}^{n}a_{i})-a_{1}}{n-1}\Big{)}^{\mathcal{B}}$ (17) $\displaystyle=(n-1)(\frac{n-1-{\epsilon}}{n-1})^{\mathcal{B}}=(n-1)(1-\frac{{ \epsilon}}{n-1})^{\mathcal{B}}$ (18) $$

Let $g(x)=(1-x)^{\mathcal{B}}+\mathcal{B}x-1$. We can prove $g(x)\geq 0$ for any $x\in(0,1),\mathcal{B}\geq 1$.
Then we have

$$ $\displaystyle\sum_{i=2}^{n}a_{i}^{\mathcal{B}}\geq(n-1)(1-\frac{{\epsilon} \mathcal{B}}{n-1})=n-1-{\epsilon}\mathcal{B}$ (19) $$

Then we finally have

$$ $\displaystyle\mathbb{E}[|S|]=n-{\epsilon}^{\mathcal{B}}-\sum_{i=2}^{n}a_{i}^{ \mathcal{B}}\leq 1+{\epsilon}\mathcal{B}$ (20) $$

∎

###### Proof.

## Appendix B Oracle sampling

The optimal sampling probability to guarantee estimation is unbiased in terms of lowest variance is not directly using attention score distribution $w_{i}$, but $u_{i}^{\prime}\propto w_{i}||v_{i}||$. However, this sampling probability is not optimal in terms of downstream accuracy and efficiency. We attribute this to two reasons. First, we observe the value norm of the sink token is significantly smaller than others (Figure 11), given its lower probability of being sampled, which may influence the functionality of attention. Second, due to the same reason, $u_{i}^{\prime}\propto w_{i}||v_{i}||$ is flatter than $w_{i}$, resulting larger computation cost (as analyzed by Theorem 3.3).

## Appendix C Supplementary analysis

Figure: Figure 10: The range of fluctuation of $\log|v_{i}-o|$ and $\frac{qk_{i}^{T}}{\sqrt{d}}$ in a single decoding step. Compared to $\frac{qk_{i}^{T}}{\sqrt{d}}$, $\log|v_{i}-o|$ is stable, hence we do not consider $\log|v_{i}-o|$ in our proposed sampling probability.
Refer to caption: x22.png

Figure: Figure 11: The $y$-axis is the norm of values states $\|v_{i}\|$ for token $i$ (on the x-axis). We observe that the value norm $\|v_{0}\|$ of the attention sink is significantly smaller than others.
Refer to caption: x23.png

Figure 10 shows that compared to $\frac{qk_{i}^{T}}{\sqrt{d}}$, $\log|v_{i}-o|$ is stable in a decoding step. Figure 11 shows that the norm of the value states of attention sink is smaller than others.

## Appendix D Additional Experiments

### D.1 Scaling-up

We provide accuracy results for larger models (meta-llama/Meta-Llama-3.1-70B-Instruct []).

**Table 4: Synthesized tasks from RULER []. MagicPIG preserves high accuracy with low computation for 70B level models. 4 layers {0,16,32,48} are preserved. Config and cost are defined as in Table 1.**
| Methods | Config | 16K | 32K | 64K | Avg. | $\text{Cost}_{1}$ | $\text{Cost}_{2}$ | $\text{Cost}_{total}$. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Llama-3.1-70B-Instruct | Full | 96.4 | 94.6 | 88.4 | 93.1 | 0.00 | 1.00 | 1.00 |
| MagicPIG | (10,150) | 95.5 | 93.4 | 88.2 | 92.4 | 0.00 | 0.02 | 0.02 |

### D.2 Reasoning

In mathematical reasoning tasks infini_igsm [, ], MagicPIG consistently outperforms Quest [] across all complexity (in terms of operators). We also find $\mathrm{TopK}$ attention suffers from significant performance degradation while Oracle Sampling can maintain high accuracy.

**Table 5: Tasks from infini_igsm [, ]. MagicPIG preserves high accuracy for reasoning tasks. Config and cost for MagicPIG and Quest are defined as in Table 1. Config denotes the ratio of selected tokens for $\mathrm{TopK}$ and sampled tokens for oracle sampling. For oracle sampling, massive duplication exists in sampled tokens, so $\text{Cost}_{2}$ is significantly lower than the ratio of sampled tokens Theorem 3.3.**
| Task | Methods | Config | 2-Ops | 4-Ops | 5-Ops | $\text{Cost}_{1}$ | $\text{Cost}_{2}$ | $\text{Cost}_{total}$. |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 4K close [] | Llama-3.1-8B-Instruct | Full | 87.4 | 71.4 | 26.8 | 0.00 | 1.00 | 1.00 |
| MagicPIG | (10,300) | 83.1 | 67.2 | 20.7 | 0.00 | 0.06 | 0.06 |  |
| MagicPIG | (10,220) | 79.8 | 58.9 | 17.9 | 0.00 | 0.04 | 0.04 |  |
| MagicPIG | (10,150) | 68.3 | 43.5 | 11.7 | 0.00 | 0.02 | 0.02 |  |
| $\mathrm{TopK}$ | 0.06 | 78.6 | 62.9 | 20.8 | 0.50 | 0.06 | 0.56 |  |
| $\mathrm{TopK}$ | 0.04 | 76.2 | 59.0 | 19.2 | 0.50 | 0.04 | 0.54 |  |
| $\mathrm{TopK}$ | 0.02 | 71.5 | 44.0 | 11.3 | 0.50 | 0.02 | 0.52 |  |
| Oracle Sampling | 0.3 | 88.1 | 72.4 | 27.6 | 0.50 | 0.02 | 0.52 |  |
| Oracle Sampling | 0.1 | 88.5 | 69.2 | 26.2 | 0.50 | 0.01 | 0.51 |  |
| Oracle Sampling | 0.02 | 83.1 | 57.9 | 11.9 | 0.50 | 0.005 | 0.505 |  |
| Quest | (16,0.06) | 55.8 | 23.2 | 5.2 | 0.06 | 0.06 | 0.12 |  |
| 8K close [] | Llama-3.1-8B-Instruct | Full | 80.2 | 68.8 | 26.0 | 0.00 | 1.00 | 1.00 |
| MagicPIG | (10,300) | 78.6 | 61.5 | 25.2 | 0.00 | 0.06 | 0.06 |  |
| MagicPIG | (10,220) | 72.2 | 60.7 | 20.4 | 0.00 | 0.04 | 0.04 |  |
| MagicPIG | (10,150) | 67.1 | 44.0 | 11.9 | 0.00 | 0.02 | 0.02 |  |
| $\mathrm{TopK}$ | 0.06 | 70.2 | 61.1 | 22.3 | 0.50 | 0.06 | 0.56 |  |
| $\mathrm{TopK}$ | 0.04 | 66.9 | 55.2 | 20.6 | 0.50 | 0.04 | 0.54 |  |
| $\mathrm{TopK}$ | 0.02 | 64.7 | 47.2 | 15.9 | 0.50 | 0.02 | 0.52 |  |
| Oracle Sampling | 0.3 | 80.0 | 67.3 | 26.2 | 0.50 | 0.02 | 0.52 |  |
| Oracle Sampling | 0.1 | 76.6 | 64.1 | 25.4 | 0.50 | 0.01 | 0.51 |  |
| Oracle Sampling | 0.02 | 79.0 | 60.3 | 20.4 | 0.50 | 0.005 | 0.505 |  |
| Quest | (16,0.06) | 54.8 | 30.0 | 11.1 | 0.06 | 0.06 | 0.12 |  |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
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?
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
questions: See the weakness part
soundness: 3
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.
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.
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?

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 2
questions: See weaknesses.
soundness: 3
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.
summary: This paper introduce a novel approach that leverages LSH sampling to approximate the oracle sampling. Empirical evaluation shows improvement over the baseline.
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.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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?
soundness: 3
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.
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.
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.

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 4
questions: See weaknesses
soundness: 3
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.
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.
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?

Meta-review: 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.

Decision: Accept (Spotlight)