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
Venue: ICLR 2025 Spotlight (2025)
Keywords: locality sensitive hashing
randomized algorithms
llm inference
kv cache
Primary area: generative models
TL;DR: We leverage Locality Sensitive Hashing (LSH) sampling to accelerate LLM decoding.

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.

Published as a conference paper at ICLR 2025
MAGICPIG: LSH SAMPLING FOR EFFICIENT LLM GENERATION
Zhuoming Chen', Ranajoy Sadhukhan', Zihao Ye², Yang Zhou', Jianyu Zhang",, Niklas Nolte Yuandong Tian', Matthijs Douze', Leon Bottou', Zhihao Jia', Beidi Chen
CCamegie Mellon University UUniversity of Washington 3 New York University <Meta Al {zhuominc; nc, rsadhukh, yangzhuo6, zhihao beidic @andrew cmu, edu zhyeêcs washington. edu {j [jianyuu Jjanyuzhang, nolte, yuandong, matthi Leonel@eaa.eecmceemeooo
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- attention approximation methods have been proposed 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 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, 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 L.S 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 5x across various GPÜ hardware and achieve 54ms decoding latency on a single RTX 4090 for Llama-3.. 8-88--nstruct model with context of 96k tokens.
always sparse
INTRODUCTION
Large language models (LLMs) with long context windows, such as GPT (Achiam et al., 2023). Llama (Dubey al. 2024), and Gemini (Team et al. 2023), have gained signif- icant attention for their ability enhance applications like chatbots (Chiang et al.) 2024), search engines (Wang et al. 2024), and video analysis (Cheng al. 2024). However, serv- ing long-context LLMs is highly challenging due to the unique bottleneck in auto-regressive generation- key-value (KV)
90
80
niah single
niah multikey
60
50
cache, which stores intermediate attention keys and values to avoid re- co-mputtaiion (Pope et al., 2022 Zhang et al. 2023b).
Density
Figure 1: While TopK attention per- forms well on information retrieval tasks (niah) where the useful information re- duces few words, it degrades severely harder aggregated tasks like word ex- traction (cwe, fwe). x-axis: proportion of attention keys used for Topk attention.
Specifically, the KV cache grows linearly with both the batch size and sequence length, occupying substantial GPU mem- ory and increasing decoding time. Moreover, the KV cache makes L.LM generation extremely memory-bound, leading to underutilization of GPU computational power. For instance, an NVIDIA A A100-40GB GPU can only handle a single request for Llama with 128k context length, with nearly half of the decoding time spent accessing the K cache, and poor GPU utilization (He & Zhai 2024).
Leveraging the common insight that attention is naturally sparse, dynamic sparse or TopK-based approximation has been extensively studied (Tang et al. [0224] Singhania et al. 20224] Zhang et al. 2024 et al. (022a)), but three major challenges prevent wide adoption in LLM serving systems. (1) Quality Degradation. They usually propose various strategies to approximate a subset KV cache that yields the highest attention scores. However, TopK attention itself biased attention

---

Published as a conference paper at ICLR 2025
CDF
CDF
1.0
1.0
0.8
0.8
Ravg
0.6
0.6
0.4
Layer14 cwe
0.4
Attn w/ sink
Altn sink
0.2
Layer24 niah
20%
0.0
20%
0.0
5000
10000
15000
20
60 100
Ritsk
(a) Long tailed phenomena
(b) Attention sink reshapes sparsity
(c) Geometry of attention
Figure 2: Left: Examples of long tailed distribution in L.M.M The x-axis is the fraction (or number of tokens) used the TopK, a.k.a. the sampling budger. Mid: Sink tokens make attention score look sparser. Right: The geometry of attention. The key states of attention sink Kn is almost opposite to other tokens and its orientation is surprisingly invariant of input tokens. Query states lie close to ko, thus forming attention sink and Figure 2b) k usually lie in narrow cone that is far away from q. In certain heads, this geometry will result long tailed distribution of attention score as well the difficulty search for the TopK keys.
approximation and lacks theoretical guarantees. F that even exact TopK attention results significantly degrade the accuracy of certain downstream tasks. (2) High Overhead. There is a large overhead to identify TopK attention, which becomes the bottleneck rather than the attention computation. For example, as studied in Liu et al. (2024a), naively applying search algorithms like IVF (ouze et al. 2024) requires to access over 30% key states to obtain the exact Topk, showing an unsatisfying trade-ofi 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 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] ]
suggest that to achieve such demanding requirements, we need to go beyond TopK attention:
Attention is not always sparse. Contradictory to previous belief (Zhang et al. 2023b 2024 Tang
et al. 2024] Liu et al. (2024)), we observe that attention is not always sparse, especially for tasks which leverage the full context. As shown in Figure 2a in some layers, attention distribution can be very long-tailed, i.e., the Top20% attention can only cover 70% of the total attention scores.
Seemingly high sparsity usually a consequence of an attention sink. Most of the attention scores concentrate on initial tokens (attention sink phenomenon) (Xiao al. [2023), making the distribution look sparser. However, as shown 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 the attention distribution.
It is hard to find Topk attention. Figure also also shows why searching for the Top-K keys is intrin- sically 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.
These limitations of TopK attention requires rethinking the sparse attention approximation. Rather than only using the keys and values with highest scores, leveraging information on the distribution can make the estimation more accurate. We ap- proach this as as hias correction problem in sampling. Unbiased and efficient sampling has been long studied in biology (Lukacs) 2009), sociology (Chen et al. (2018) as well as machine learn- ing (Backurs et. al. 3019/Chen at.l. 2019 Zandieh et al. 2023 with theoretical guarantees.
Topk Duude Sampring
0.4
Figure 3 shows that sampling values according to their corre- sponding attention score (we call this oracle sampling) achieves much lower (up to 4x) estimation error than the naive TopK selection. Deploying sampling estimation in attention is promising, but three challenges remain. First, how reduction of the attention error can make difference in downstream performance is
Sampling Budget
Figure 3: TopK v.s. Sampling, 16k total context
2

---

Published as a conference paper at ICLR 2025
unclear (Backurs et al. 2019 2018). Second, modeling the attention scores 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 theoretically efficient algorithm is non-trivial.
In this paper, we propose Magic samPling for Generation (MAGICPIG), which leverages Locality sensitive hashing (LSH) sampling for efficient LLM generation. LSH employed for sampling to approximate the attention scores 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 (Kitaev et al. 2020: Chen et al. (222)) which are of vital importance for accurate estimation (Backurs et 2018). Following the practice of Aminabaai et al. (2022) He & Zhai 2024), 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 TopK attention. Moreover, we study sampling based attention estimation assuming an oracle for the key distribution (Oracle Sampling Estimation) in detail and empirically demonstrate that it is consistently more effective both for in distribution estimation and on downstream tasks.
In Sections 4. 4.34 we present a sampling algorithm to approximate oracle sampling for attention estimation based locality sensitive hashing and the intuition and motivation from statistic perspec- tives. To our best knowledge, MAGICPIG is the first to leverage LSH sampling in self attention in decoder-only LLM generation.
Section 4.4 we present our system design to efficiently offlood attention computation on CPU, breaking the memory limit of GPU for serving larger batch 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 (Chen et al. 2021 Kitaev et al. (0220).
In S5ction55 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.5 ~ 5x (A100, L20, RTX 4090) and can achieve 54ms decoding latency on single RTX 4090 for Llama-3. 1-88--BBnn Instrust (Dubey et al., (024) with 96K context. More importantly, we show that MAGICPIG already outperforms TopK attention in the two aggregation tasks in Figure[] suggesting that sampling indeed goes beyond TopK attention.
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 an weighted of previous values by
gkT
o= Softmax
V=
GERIxd KVVER"xe
W EER IEERIx 1xn
where d is the head dimension and n is the context size. K= ke..,, = U, Un] } k,,,% E
g is KV cache. Normalized attention weight w =ofomaaax qk* R 1xn is also called attention
vd
(score) distribution. Our target is to find sampling matrix IT E Rx and diagonal matrix D E R mm which minimize
(2) where m < is computation budget. For TopK attention, suppose Wr, >..>> Wrm > Wr,, then
8 wIllI
1.
if otherwise.
11
=== m W Wr:
Dú
0,
2.2 RELATED WORKS
Efficient Attention. Attention approximation has been long studied. Reformer (Kitaev et al. 2020). KDEformer (Zandieh al. 2023). HyperAttention (Han et al. 2023) and ScatterBrain (Chen et al.
2021) tackle the problem via locality sensitive hashing. These methods work in training and encoding (e.g., BigGAN (Brock et al. 2019)) Theoretically, the error bounds and minimal workload required
are continuously improved Brand et al. 2023 Alman & Song 2023) but not proven to be practical for
wall-clock acceleration in LLM decoding. Flash-attention (Dao al. 2022b: Dao 2023
Dao et al.
3

---

Published as a conference paper at ICLR 2025
2022a), fash-deeodingg Ye et al. 2024 Hong al. 2024) and SlimAttention et 2024 accelerate attention operator by maximizing hardware utilization, which is orthogonal to our approach.
Locality sensitive hashing. Locality sensitive hashing (LSH) (Backurs et al. 2019; 2018) is a family of hashing functions which assigns the same hash codes for similar imputs with higher probability than others (Chen et al. 2020b; Jafari et al., 2021). LSH uses two syper-parameters. (K, L). L hash tables are independently built. Each hash table has its own function / which projects a high dimension vector to an integer by concatenating K random independent hash functions. In the sampling process, all vectors which share hash codes in at least one hash tables with query will be collected. SimHash (Charikar 2002) is the LSH family based on cosine similarity. For a vector E RR SimHash generates a random hyperplane w and returns Sign(w" x). Vectors share the same sign if and only the random projection is not in-between them. For random projection, all angles are equally likely, xy" thus the probability that two vectors 2, y share the same sign for is p= 1- -, where 0 arccos |||| 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-1ppk)
KV Cache reduction. To get rid of memory bound introduced by KV cache thus enabling a larger batch size or serving longer prompt, many methods are proposed to reduce the volume of KV cache. For example, H₂O (Zhang et al. 2023b), SnapKV (Li et al.. 2024) and Keyformer (Adnan et al. 2024) calculates heuristics during prefilling phase to decide w which tokens to preserve for decoding phase. Quest (Tang et al., 2024) and Loki (Singhania et al.) 2024) do not evict KV cache but apply dynamic sparsity to reduce KV Cache loading at inference time. StreamingLLMM (Xiao al. 2023) and InfLLMM Xiao et al. 2024) reduce attention computation and support extremely long context by context window extrapolation. Methods like KIVI (Liu al. 2024b) and QServe (Lin et al.
2024
reduce the size of KV Cache by quantization.
3 RETHINKING ATTENTION SPARSITY
In this section, we examine TopK attention, which the theoretical upper bound of prior search- based algorithms including both static methods (Zhang al. [023bb Li et al., (2024) and dynamic methods (Tang et al. [2024] Singhania et al. [024 Mao et al. 2024). We show that Topk is sub-optimal and present another attention approximation based sampling and estimation with oracle, that improves the accuracy and/or the computation cost.
3.1 ACHILLES" HEEL OF TOPK ATTENTION
Relative Emor
As it is defined, TopK attention only computes the weighted average on elements with highest attention scores. To quantify its performance, the computation budget of TopK attention, is defined as the number of selected tokens, i.e. the K of TopK. Searching-based sparse attention algorithms, like (Tang al.) 2024; Singhania et al.) 2024, Wu et al., (004)) are approxi- mations for TopK attention by replacing the true Topk keys with the ones found by approximate searching algorithms.
Layar10-Head2
0.5
Layaris-Hedd22
0.4
Layerl 3-
Heydr12-Head20
0.3
0.2
0.1
0.0
2000
6000
12000
Budget
However, we do find significant performance degradation in downstream tasks caused by TopK attention as shown cache of 16k tokens.
Figure 4: Topk estimation crror for a KV-
in Figure I Although TopK attention preserves accuracy for retrieval tasks that only require minimal subset of the context -naall--i---hhyttack niah, single/multikey (Hsieh et al. (024)). it severely degrades for aggregation tasks that leverage the full context (common word extraction, ewe, and
frequent word extraction, fwe (Hsich et al. 2024 ). Intuitively, the information is distributed more broadly for aggregation tasks, which results in less peaky attention score distribution.
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. Long tail phenomena do occur in LLMs across all layers as presented in Figure 2a Top20% tokens can only cover 70~ 80% attention scores, leaving large proportion of keys and values not considered, which is translated into non-negligible (15 (15~20% 20%) estimation error in Figure More observations are presented in Appendix
3.2 ESTIMATE ATTENTION WITH SAMPLING
Existing TopK attention mechanisms ignore tokens in the KV cache with low attention scores, which introduces a bias, since the ignored tokens sum up to a large proportion of attention scores (Figure As a result, TopK attention achieves suboptimal performance for long-context tasks, such as information
4

---

Published as a conference paper at ICLR 2025
100
Topk
Tork
044
Orace Sanping
044
Crscce Samping
0.3
Accuracy
02
022
70
rah mul key Tapk.
60
Rah PuR key Orn Samplin
0.0
00
50
2000
6020
1200
2000
5000
12000
Computation Budget Ratio
(a) Estimation error Layer11
(b) Estimation error Layer19
(c) Downstream Comparison
Oracle sampling estimation can significantly reduce numerical error compared to TopK attention. The evaluated context size is 16k. The r-axis is sampling budget for oracle sampling and computation budget for TopK attention. Notice that the estimation error of TopK attention will cross oracle sampling after certain large budget (12k in figures). The reason is that oracle sampling will repetitively sample the same subset of tokens with high probability while Topp will not. Theorem furrtterr explains this. Right: Downstream comparison for oracle sampling estimation and Topk attention. The r-axis for both methods is computation budget ratio, i.e. the fraction of selected/sampled tokens.
Figure 5: Left and Middle:
aggregation (Figure Tp. Increasing the computation budget for TopK attention does help reduce the estimation error (Figure since 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 (Lukacs 2009 Owen 2013 Lohr 2021 Chen et al. 2018), we show in
the following that attention output can be estimated with sampling. Using notations from Section 2. we can re- write attention output as the expectation of o,, 1,,1 < Sn from distribution w, i.e. o= E (t)).
which can be estimated by the following method.
Definition 3.1 (Oracle Sampling Estimation). Given a sampling budget B and normalized attention
iid score w, B elements are sampled independently from w (i.e. i1,22,..., ~ w). Then the attention output estimated
B B
(4)
This is not the lowest variance estimator but has a better downstream performance (see Appendix C 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 is monotonically decreasing with 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 TopK attention in Figures and 5b In summary, oracle sampling estimation can reduce relative error by up to 4x.
Note that the sampling budget B is not the actual computation cost for oracle sampling estimation: duplicate X, need to be computed/loded only once, so ō can be computed by
Sa B
Ō=
S=Unique((issvss
(5)
FES
where f/ is the number of duplicates of X4. Intuitively, if w has an peaked distribution (e.g. w; > 99%), then almost all samples in (i..)...g) } 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 budger (E((S))) has an upper bound of + Be, where 1 - max,w;,
This theroem (proved in Appendix A) shows that the computation cost of oracle sampling is usually far less than the sampling budget. In Figure present the downstream accuracy comparison between oracle sampling estimation and TopK attention. The former preserves high accuracy for both tasks, even with very small computation cost (0.002% out of 16k context, which is approximately 32).
5

---

Published as a conference paper at ICLR 2025
4 MAGICPIG
Section 3.2 demonstrates the potential of sampling-basse basedg estimation. In Sections 4.1 and we present how arrives at Locality sensitive hashing to unleash this potential from 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 collaborrt.....
Note that most of the derivations 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 2x wall clock speed up because obtaining distribution w requires full computation of all qk? and thereby only saving the wV computation.
Fortunately, importance sampling (Kloek & Van Dijkj 1978. Owen 2013 Lohr 2021 allows us to perform estimation for unknown distribution by sampling from proposed distribution u. In
our problem setting, the normalization factor of w. i.e. Z = >i= computing requires evaluating all qh."- However, we do have access to unnormalized weights
exp va
ak! is also unknown because
Wi - for sampled indices i. Hence, by employing a variant of importance sampling, self- normalized importance sampling (Owen) 2013), we sample indices i .1, from a proposed
distribution u and the resulting estimator is
Wis Uiy
Wiy
X
where
7
1=1
which has very nice property for accurately estimating attention output that xo Its variance is related to the distribution u. and can be approximated by
= =
=1.
Var(x IS
W
= BZ2
UU 0
(7)
B
To minimize the variance, u should satisfy x wi 02 (Hesterberg 2003 The variance will be high if u, and (0/00 assign high probability mass different regions of the sample space have different modes. Therefore, the challenge is compute distribution aligned with wp/r; 02 o without accessing too many wi. Besides, Equation (6) requires that sampling probability u can be computed and u, > 0, which is not satisfied by many deterministic approximations like TopK.
4.2 VARIANCE REDUCTION WITH LSH
We decompose We| 0 ol = exp( gh. + +lgg es ol). We observe emprically (Figure i0 the appendix)
that log U ol does not fluctuate significantly compared to of share the same peaks with qk:" By the following transformation,
g
Hence, we simplify the requirement
Vd
r max [Rxa] [= k= [k - [k,
(8)
we further transfer inner product gk? to cosine similarity between y and ky (which is common practice in Maximum Inner Product Search (Shrivastava & Li 2014)).
Chen al. 2020a). leverage Locality sensitive hashing- massing-basedd sampling for this estimation problem. Specifically lly, leveraging hash function h the LSH family that preserves cosine similarity such as SimHash (Sadowski] 2007), we can sample from
Inspired by prior work (Spring& & Shrivastava 201
probability distribution u4 =P (P(h(g) = ((:):)) which is monotonic to cos [g]1k
gk
4.3 ALGORITHM IMPLEMENTATION
To make this estimation practical, MAGICPIG is implemented by the following specific design.
Estimator approximation. Self normanormalized important sampling Equation (6) requires é ,,,2.....44 iid sampled but the probabilities provided by hashing are not normalized. Hence we adapt our estimator:
We assume head dimension d=1 1 here for simplicity. Higher dimension has similar formulations and analysis by replacing variance with trace of covariance.

---

Published as a conference paper at ICLR 2025
Model 1
Random Projectors
GPU E
VRAM 1000GB/s
GPU
Parameters
PCIE 315GB/s
Hash 4
CPU
DRAM 512GB
100GB/ss
Tables
KV Cache
CPL
(a) Memory Hierarchy
(b) Workload Partition
Figure 6: Left: Memory hierarchy of hardware. GPU VRAM has high bandwidth but is limited. CPU DRAM is sufficient but is relatively slow. The limited bandwidth of PCIE forbids farge-scald data transferring. Right: Workload partition MAGICPIG. Linear projections and hash function computation (by random projection) are done on GPU, while sampling with hash tables and attention are done CPU. The execution order D
After obtaining S with probability u, MAGICPIG computes
X -E
Eies my
(9)
LIES
Eies
Hash function selection. MAGICPIG leverages SimHash (Sadowski 2007 that draws with K L random vectors. For each of the L hash tables, the and k;s vectors are projected on K directions and only the sign of the projection kept, which yields K-bit hash value. Key k, is sampled only there exist at least two hash tables where k; shares the hash value with q. The corresponding probability is
1
qk;
u, = P[k, sampled] (1- (1- Lp (1 11-pK)L-
where p=
arccos
(10)
π
[ |g
Data pre-processing. Before building hash tables, MAGICPIG centers the k; vectors. As shown in Figure 2c] keys are almost always concentrated on one side of the queries, except the initial token. Random projections cannot effectively distinguish keys in this case, result- ing in uniform sampled probabilities. Luckily, Softmax is translation invariant. Centering (k, = k4 1 - kj) distributed the keys bet- ter and remains computationally equivalent.
Algorithm 1: MAGICPIG Decoding
Input: K,V Raxd q E Rlxd, random
projectors W E (kdx(KXXx) hash tables HT,
static KV cache KT, VT E Rexd
Compute hash code for new query
Gende = Encode (q, W)
Query hash tables to sample S in Equation 9
= Query HTT .qood), k== Vs V[S]
Compute inner product for q and sampled K
Combining Equations 9 and (10), gives a closed form of the MAGICPIG attention es-
Wg w-=0p qK wT wT=qKK
Compute collision probability for each hash function
timation. Assuming sample set S is obtained with LSH,
p 1 - arccos( w/ /[laal (Ilg|l-11/5 ((s))
Q4? exp( logu,) Vd
Compute sampling probability
u= 4= ( (1-p**)
iS EiEs exp( logu,;)
9h:""
U/3
Compute attention output estimation
v7
Ō Softmax
[ws wr! wwwwwwlooogggggg logg((u.,)))) Vr.
4. 1 (1 - pi
K
Lpp (1 (1-p)) -pp 1
Return Vd
1
qk:" T
arecos
(11)
[ql Ikk
Our algorithm applies to single attention head, Algorithm The details of Encode, Query as well as the hash table construction are described in prior work (Sadowski 2007 Chen et al. 2020b).
4.4 SYSTEM CO-DESIGN
The memory size of KV cache remains a bottleneck for long-contexs LLM decoding, especially when GPU VRAM is limited. DRAM on the CPU side offers sufficient memory capacity with 100- 200GB/s bandwidth, which usually 10 20% of GPU VRAM bandwidth (see Figure Ideally, this gap can be mitigated by 5-100 sparsity. To make CPU DRAM an aggregated memory for GPU, the workload must be partitioned. In our experiments = 9 or 10 and is few hundreds.
1000-200
Our system design extends prior work (He & Zhai 2024 Aminabadi et al., 2022) by splitting LLM decoding into four parts. (1) Parameter computations, ie. all linear projectors including MIP and
7

---

Published as a conference paper at ICLR 2025
Wq. Wx. Wv, Wo in the self-attention module runs GPU. (2) Attention computation, which involves o- Softmaxx gk" )v, ruuss runs on CPU. (3) Random projections, Al generation time, for each KxL random projections conducted to obtain the hash codes. Since all heads can share the same random projectors the memory overhead is limited (400 KB in our implementation), so this step is compute-bound. Therefore, the projection is placed on GPU. (4) Retrieval. The hash codes of need to be looked up in L hash tables, which is negligible computationally. However, the pre-built hash tables for k;s can occupy considerable memory, making better fit for CPU. With the above partition, we are able to support hash tables with K and L beyond the scale of prior work (Kitaev et al. [2020] Chen et al. 2021 Zandieh et al., 2023) without worrying about computation for hash codes well as the storage of hash tables.
vd
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 s accuracy in downstream tasks, followed by our end-to-end system results showing wall-clock performance.
we demonstrate MAGICPIG preserves high accuracy (less than 2% degradation) across moderate to long context tasks with computation cost 2% ~ 5% of full attention.
In Section 5.1
In Section we demonstrate the efficiency of MAGICPIG, which achieves up to 5x throughput improvement and 54ms decoding latency on RTX 4090 for Llama-3. 1-11--8B-Inntee 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 TopK attention in the two indicating that sampling indeed goes beyond TopK attention.
aggregation tasks in Figure I
5.1 MAGICPIG PRESERVES ACCURACY
We demonstrate MAGICPIG can preserve accuracy in diverse tasks with less than 5% computation.
Setup. Our experimentts are based on Llama Met@ 2024 Dubey et al. 2024 Touvron et al. 2023
models. Three types of tasks are included, which are mid-context comprehensive tasks from Im-eval-
harness (Gao et al.) (202T) (GSM8K-CoT -CTT (Cobbe et al. 2021). MMLU-Flan-COO. Caraan---oo---hhhh (Hendrycks)
et al., 2020) and COQA (Reddy et al.) 2019)). and 6 long context tasks from (Bai et al. 2023)
(QASPER (Dasigi et al. 2021). LCC, Repobench- (Liu al. 2023), TriviaQA (Joshi et al. 2017). PRE and TREC PREC(Li& & Roth (00i))202220000 Hovy etal. 200T)) and 13 synthetic tasks from RULER (Hsieh et al.
2024) (with 50 examples per task).
Baselines. Besides full attention, Quest (Tang et al. 2024) and its variants are used as baselines. In its default setting, Quest uses a "page size" of 16, i.c. 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 least one page is selected in every test example. The initial 4 tokens and local 64 (for LongBench (Bai et al. 2023) and RULER (Hsieh et al., 2024)) or 24 (for in-eval-harned (Gao
et al.
2021 tokens as well as layer-(0,16) (0,16) are statically preserved. We do not use the transformations in Equation
K in our implementations, we do not find them to contribute accuracy improvements.
Cost. The cost for the attention approximation consists of two parts: Cost, is the sampling/search cost to obtain S in Equation [I]. Cost2 is the attention computation cost, see Equation (II). We report the ratio of number of FL.OPs compared of the full attention computation. For M MAGICTGG Cost, 2 and Costa is empirically measured for different LSH prrper-parametes... For Quest with page size K, Cost, A and Cost, is controlled manually.
Analysis. From Tables I to 3, (1) MAGICPIG preserves high accuracy (degradation less than 2%) for all kinds of tasks, with a computation cost of 2% ~ 5%. (2) Compared with Quest, which also shows reasonable performance on long context tasks, MAGIC also demonstrates good performance on tasks with moderate context sizes in Im-eval-harness (Gao et al. 2021), indicating more robust performance in general serving. (3) With LSH sampling which introduces an order of magnitude lower sampling/scarching cost (Cost,), MAGICPIG can achieve equivalent or better accuracy with only half of the computation cost.
5.2 MAGICPIG SHOWS IMPRESSIVE EFFICIENCY ACROSS VARIOUS HARDWARE SETTINGS
We show MAGICPIG can bring up to 5x wall clock speed up and reduce GPU memory consumption on different models and hardware settings (A 100, 1.20, RTX4090).
8

---

Published as a conference paper at ICLR 2025
Table I: Comprchensive tasks on cal-eval-hareess et 3. MAGICPIG significandly outperforms other methods with lower computation. The config is hyper parameter of LSH for MAGICPIG or page size and ratio of selected pages for Quest (Tang et al. 2024 Costi, Costa represents cost for parching/sampling and sparse attention computation respectively.
Methods
Config
GSM
COQA
MMLU
Avg.
Cost:
Cost2
Cosleotal-
Llama-2-77 2-aa------hht
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.----3--------- 1-8-IInsttuctt
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 Bai ct al. 2023). MAGICPIG preserves high accuracy with low
computation. Config and cost are defined as in Table [ Code models are only evaluated by Repobench- and LCC.
Methods
Config
QaS
RbP
LCC
PrE
TrC
TrQ
Avg.
Cost
Costs
Costrotal-
Llama-3.1-8 1-8B-Innstucc
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 1 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 Hsiich et al. 2024).
MAGICPIG preserves high accuracy with low
computation. Config and cost of MAGIC PIG and Quest are defined as in Table We evaluate Loki with rank 32 and select 3% takens for attention computation (Singhania et al. 2024).
Methods
Config
16K
32K
64K
96K
Avg.
Cost
Costz
Cosleotal-
A.ma--311-B -B-Insttuct
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
Loki
(32,0.03)
80.0
63.6
61.9
34.7
60.1
0.12
0.03
0.15
Setup. We evaluate our system performance on serving settings. (1) 80GB GPU (A (00)) and 34B model (CodeLlama-34)
(Rozière et al. (022)) with 16K contexts; (2) 48GB GPU (L20) and 13B
model (CodeLlama- lama--3B) (Rozière et al. 2024) with 16K contexts; (3) 24GB G (e.g. RTX 4090) and 8B model (Llama-3.1 1.BBBBBBBBBBBB--- (Dubey et al. 2024) with 96K contexts.
Baselines. Our baselines for (1) and (2) are full attention GPU, and for (3) full attention on CPU with theoretical estimated bandwidth. Our system's 's GPU part is implemented in native Pytorch (Pazkke) al. 2019) and the CPU part FBGEMM (Khudia al. 202I) in b precision. Our CPU Intel Platinum 8480+ for A100 and Intel 8563C for L20. In the last setting, the CPU bandwidth is estimated 150GB/s, above the empirical bandwidth we measure when running a group query attention of size
Analysis. In Figures 7a to 7c we demonstrate (1) MAGICPIG significantly improves decoding throughput for all three scenarios (A100: 1.5x, L20: 5.0x, RTX 4090: 3.3 and can achieve latency of 54ms 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 (over 12x). The ablation study of decoding throughput with different LSH hyper-parameters is presented in Table
We simulate 24GB GPU by setting memory limit with 1.20. As the bandwidth of 1.20 (864GB//) is less than RTX 4090 (1TB/s), the real speed of system should be slightly faster than the simulation.

---

Published as a conference paper at ICLR 2025
Durs Raceline
Sastiie
20
30
1akenss
10
40
20
22
15
22
24
23
32
35
40
0
10 12
(a) A100 with 34B model
(b) L20 with 13B model
(c) RTX 4090 with 8B model
Figure 7: evaluate MAGICPIG on three serving scenarios. Left: A A100 serves 34B model with 16K context.
MAGICPIG achieves 1.5x throughput improvement. Mid: 1.20 serves 13B model with 16K context. MAGICPIG achieves 5.0x throughput improvement. Right: Simulated RTX 4090 serves 8B model with 96K context. MAGICPIG achieves a latency of 54ms in a single request serving and can improve the throughput of baseline by up to x. The dashed lines denote the highest throughput of baselines. With KV cache offloading, MAGICPIG fit much lareer hatch size. which contributes to the throuehout imnrovement.
100
90
88
80
Accurscy
centaring nith Af contering rinn contaring fee Kl cencering fee
Accuracy
844
86
84
80
78
76
0.01
Budget
Budget
(a) Importance of centering
(b) CWE
(c) FWE
Figure 8: Left: Accuracy comparison for with and without centering. Here we fix K and vary L for the two settings. Mid and Right: In the two aggregated tasks, sampling based MAGICPIG can even beat the exact Topk attention. The experiments done RULER (Hsich et al. 2024 with 16K context size.
5.3 ABLATION STUDY
In this section, we empirically validate our two previous observations.
Centering important for good performance. In Section 4.3) we use a translation center the keys before applying LSH sampling. Empirical results show this to be important for downstream tasks as shown in Figure Sal Without centering, the accuracy drops to almost zero retrieval (NIAH) and degrades 65% in FWE. We find almost none keys (less than 0.1%) can be sampled by query without centering, as their orientation is almost opposite as shown in Figure 2c
Sampling goes beyond TopK. In Figures and Sc We compare the performance of MAGICPIG and TopK attention in two aggregated tasks (CWE, FWE) where TopK attention experiences significant performance degradation (Figure D. MAGICPIG can even beat exact TopK attention in these two tasks by a margin up to 3% and 8% respectively, demonstrating that sampling improves the ceiling of TopK, which is impossible for search-only algorithm.
CONCLUSION
In this work, first present the limitation of TopK attention approximation for addressing the computational and memory challenges of long-context LLM generation. Then we show oracle sampling can go beyond TopK and introduce MAGICPIG, 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 system co-design that offloads hash tables and reduced attention computation to the CPU. Our experimental results demonstrate that MAGICPIG achieves substantial improvements in throughput and latency across multiple hardware configurations, outperforming traditional TopK attention mechanisms. The theoretical soundness, robustness, and scalability of MAGICPIG open up new opportunities in both attention approximation methods and algorithm-hardwee co-design.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank the anonymous reviewers for their helpful discussions and feedback on the paper. This work was partially supported by the National Science Foundation under grant numbers CNS-2147909 CNS-221 1882, and 12S-2233351, along with gift awards from Li Auto, Amazon, Cisco, Google, Intel, Meta, Moffet AI, Oracle, Qualcomm, and Samsung.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Iige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08771 2023.
Muhammad Adnan, Akhil Arunkumar, Gaurav Jain, Prashant Nair, llya Soloveychik, and Purushotham Kamath. Keyformer: Kv cache reduction through key tokens selection for efficient generative inference. Proceedings of Machine Learning and Systems, 6: 6:114-127,
AltMeta. Ilama 3 model card. 2024. URL
lama/ lama3/
b blob obooam iee/ii/MMEELCCC
md
Josh Alman and Zhao Song. Fast attention requires bounded entries. In Oh, Nau- mann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 63117-63135, Curran Associates, Inc., 2023. URL https::/pproceedinges wwwp:iiiiiinnn.uuueueeee.ee..eee nwwrips wttpcp//pppppeeeeeeeeeeeeeeppece Files/paperr2023/ httpss//ppooeeednssseeee.
[122:778881111555001313322222231 Paper wonferenneeeecee Poaprpeeeee.ppaaeeeeeooreeeeeee.
Reza Yazdani Aminabadi, Samyam Rajbhandari, Minjia Zhang, Ammar Ahmad Awan, Cheng Li, Du Elton Zheng, Jeff Rasley, Shaden Smith, Olatunji Ruwase, et Deepspeed inference: Enabling efficient inference of transformer models at unprecedented scale. arXiv preprint arXiv:2207.00032, 2022.
Alexandr Andoni and Ilya Razenshteyn. Optimal data-dependent hashing for approximate near neighbors. In Proceedings the forty-seventh annual ACM symposium Theorr computing, pp. 793-801, 2015.
Alexandr Andoni, Piotr Indyk, Thijs Laarhoven, llya Razenshteyn, and Ludwig Schmidt. Practical and optimal LSH for angular distance. Proceedings of the 28th International Conference on Neural Information Processing Systems- Volume 1, pp. 1225-1233, 2015.
Arturs Backurs, Moses Charikar, Piotr Indyk, and Paris Siminelakis. Efficient density evaluation for smooth kernels. In 2018 IEEE 59th Annual Symposium on Foundations of Computer Science WTT.////..20...0000..
(FOCS), pp. 615-626, 2018. doi:
Arturs Backurs, Piotr Indyk, and Tal Wagner. Space and time efficient kernel density estimation in high dimensions. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché-Bue Buc, F. Fox, and R. Garnett (eds.), Advances in Neural Dnemattttni mewreiiiiiiiin...uepp cc/paper fiwes/ htlles/pappe/22019 Inc., 2019. [Lee2/22eee110622230002233233223
PrciiiiionnPPoPeeee Processing Systems, volume 32. Curran Associates,
Yushi Bai, Xin Lv, Jiajie Zhang. Hongchang Lyu. Jiankai Tang. Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, Yuxiao Dong, Jie Tang, and Juanzi Li. Longbench: bilingual, multitask benchmark for long context understanding. arXiv preprint 233v 14508, 2023.
Jan van den Brand, Zhao Song, and Tianyi Zhou Algorithm and hardness for dynamic attention maintenance in large language models. arXiv preprint 2324000000.00
Andrew Brock, Jeff Donahue, and Karen Simonyan. Large scale gan training for high fidelity natural http://aaxvvv... org/ abs/ 1809 11096
image synthesis. 2019.
Moses S. Charikar. Similarity estimation techniques from rounding algorithms. In Proceedings of the Thiry-Fourth Fourth Annual ACM Symposium on Theory Computing, STOC '02, pp. 380-388, New York, NY, USA. 2002. Association for Computing Machinery. ISBN 158 134959. doi: 10.1 001114//59997.5099909
URLtttte pictps://caen 1tt.//......../00000000 509965
11

---

Published as a conference paper at ICLR 2025
Beidi Chen, Anshumali Shrivastava, and Rebecca C Steorts. Unique entity estimation with application to the syrian conflict. The Annals of Applied Statistics, ((22):033--167, 1067,2012 2018.
Beidi Chen, Yingchen Xu, and Anshumali Shrivastava. Fast and accurate stochastic gradient estimation. Advances in Neural Information Processing Systems, 32, 2019.
Beidi Chen, Tharun Medini, James Farwell, sameh gobriel, Charlie Tai, and Anshumali Shrivastava. Slide: In defense of smart algorithms over hardware acceleration for large- deep learning sys- tems. In I. Dhillon, D. Papailiopooloss and V. Sze (eds.), Proceedings of Machine Learning and Sys- fems, volume pp. -0000 2020a. URL UtLhhttss::/grrcceedingg ml sys org org/papee paper.
fi files/paper/// 2020/ fi / waw.eeaa3444000 /tt./////00000000000000000003330 PPaper pdf
Beidi Chen, Tharun Medini, James Farwell, Charlie Tai, Anshumali Shrivastava, et al. SLIDE: In defense of smart algorithms over hardware acceleration for large-scale deep learning systems. Proceedings of Machine Learning and Systems, 2:291 1-36,, 2020b.
Beidi Chen, Tri Dao, Eric Winsor, Zhao Song, Atri Rudra, and Christopher Ré. Scatterbrain: Unifying sparse and low-rank attention. Advances in Neural Information Processing Systems, 34:17413- 17426, 2021.
Zhuoming Chen, Avner May, Ruslan Svirschevski, Yuhsun Huang, Max Ryabinin, Zhihao Jia, and Beidi Chen. Sequoia: Scalable, robust, and hardware-awe aware speculative decoding. arXiv preprint 1rXiv:2402.123774 2024.
Zesen Cheng. Sicong Leng, Hang Zhang, Yifei Xin, Xin Li, Guanzheng Chen, Yongxin Zhu, Wenqi Zhang, Ziyang Luo, Deli Zhao, al. Videollama 2: Advancing spatial-temports modeling and audio understanding in video-Ilms. arXiv preprini arXiv:2406.07476, 2024.
Wei-Lin Chiang, Lianmin Zheng. Ying Sheng, Anastasios Nikolas Angelopoulos, Tianle Li, Dacheng Li, Hao Zhang, Banghua Zhu, Michael Jordan, Joseph E Gonzalez, al. Chatbot arena: An open platform for evaluating Ilms by human preference. arXiv preprint arXiv:2403.041 04:::222..4 32, 2024.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021.
Tri Dao. F Faster attention with better parallelism and work partitiioning. CoRR, abs/2307.08661, 2023. doi: 10.48550ARXXV2200..6666 48550/arXiv wttp:xxvvvv.v.7..0
URL htttps://doi.orq///
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35: 16344-16359, 2022a.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. In Sanmi Koyejo. S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information P Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28. December 9, 2022, 2022b.
Pradeep Dasigi, Kyle Lo, Iz Beltagy, Arman Cohan, Noah A Smith, and Matt Gardner. A dataset of informaaiin--seeking questions and answers anchored in research papers. arXiv preprint arxiv:2105.030/1 2002.
Matthijs Douze, Alexandr Guzhva, Chengqi Deng, Jeff Johnson, Gergely Szilvasy, Pierre-Emmanuel Mazaré, Maria Lomeli, Lucas Hosseini, and Hervé Jégou. The faiss library. arXiv preprint arXiv:2401.08288, 2024.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The laaa herd of models. arXiv preprint arXiv:2407.22178. 2024.
Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell. Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for shot language model evaluation, September 2021. URL RRl.httss:://dd http://d..oor////// 10.. /enooo 5371
Niklas
12

---

Published as a conference paper at ICLR 2025
Tianyu Gao, Alexander Wettig. Howard Yen, and Danqi Chen. How to train long-contess language models (effectively). arXiv preprint arXiv:2410.02660, 2024.
Insu Han, Rajesh Jayaram, Amin Karbasi, Vahab Mirrokni, David Woodruff, and Amir Zandieh. Hyperattention: Long-context attention in near-linear time. arXiv preprint arXiv:2310.0588% 2023,
Jiaao He and Jidong Zhai. Fastdecode: High-throughput gpu-efficient Ilm serving using heterogeneous pipelines. arXiv preprint arXiv:2403. 1421111220 2024.
Pujiang He, Shan Zhou, Wenhuan Huang, Changqing Li, Duyi Wang. Bin Guo, Chen Meng, Sheng Gui, Weifei Yu, and Yi Xie. Inference perforaacce optimization for large language models on cpus, 2024. BRLLtttps://axxttt.
org/ absaaass//22/ 2407.07304
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300, 2020.
Timothy Hesterberg. Advances in importance sampling. 01 2003.
Ke Hong, Guohao Dai, Jiaming Xu, Qiuli Mao, Xiuhong Li, Jun Liu, Kangdi Chen, Yuhan Dong, and Yu Wang. Flashdecodingt+: Faster large language model inference on gpus, 2024. URL
01282
Eduard Hovy, Laurie Gerber, Ulf Hermjakob, Chin-Yer Lin, and Deepak Ravichandran. Toward semantics- answer pinpointing. In Proceedings the First International Conference on Human Language Technology Research, 2001. URL https://wwww a www.wcwww..b...... org/ htwp://owww..e.oooo/acoooooooooo
H01-
Cheng-Ping Hsich, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? arXiv preprint arXiv:240..06633 2024.
Omid Jafari, Preeti Maurya, Parth Nagarkar, Khandker Mushfiqul Islam, and Chidambaram Cru- shev. A survey locality sensitive hashing algorithms and their applications. arXiv preprint arXiv:2102.08942, 2021.
Mandar Joshi, Eunsol Choi, Daniel S. Weld, and Luke Zeulemoyer. Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension, 2017. URL httpt:////// /ttp://aaivvvvv/vv org/ abs/1705 03551
Daya Khudia, Jianyu Huang, Protonu Basu, Summer Deng, Haixin Liu, Jongsoo Park, and Mikhail Smelyanskiy. Fbgemm: Enabling high-performance low-precision deep learning inference. arXiv preprint arXiv:2101. 05615, 2021.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001 044511, 2020.
Teun Kloek and Herman K Van Dijk. Bayesian estimates of equation system parameters: an application of integration by monte carlo. Econometrica: Journal of the Econometric Society, pp. 1-19, 1978.
Xin Li and Dan Roth. Learning question classifiers. In COLING 2002: The 19th International Confer- ence on Computational Linguistics, 2002. URL BttLhttps: //wp::::/wwww... ac awweb. org/ hwtggloyhhhhgy//////o
[02-1150
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation. arXiv preprint arXiv:2404. 14469, 2024.
Yujun Lin, Haotian Tang, Shang Yang, Zhekai Zhang, Guangxuan Xiao, Chuang Gan, and Song Han. Qserve: W4a8kv4 quantization and system co-design for efficient Imm serving. arXiv preprins 04X55:2405.04532, 2024.
Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang. Zhenhua Han, Qianxi Zhang, Qi Chen, Chengruidong Zhang, Bailu Ding, Kai Zhang, et al. Retrievalattentinn: Accelerating long-context 1lm inference via vector retrieval. arXiv preprint a409 105 16, 2024a.
13

---

Published as a conference paper at ICLR 2025
Tianyang Liu, Canwen Xu, and Julian McAuley. Repobench: Benchmarking repository-level code auto-completion systems, 2023. /
03091
Zirui Liu, Jiayi Yuan, Hongye Jin, Shaochen Zhong, Zhaozhuo Xu, Vladimir Braverman, Beidi Chen, and Xia Hu. Kivi: A tuning-free asymmetric 2bit quantization for kv cache. arXiv preprint arXiv:2402.02750, 2024b.
Sharon L Lohr. Sampling: design and analysis. Chapman and Hall/CRC, 2021.
Paul Lukacs. Closed population capture- models. Program MARK: a gentle introduction, 2009.
Qin Lv, William Josephson, Zhe Wang. Moses Charikar, and Kai Li. Intelligent probing for locality sensitive hashing: multi-probe ish and beyond. Proc. VLDB Endow., 10(12):001-20024, August 2017. ISSN 2150-8097. doi: 10.1477//13376553333333
URL https::/aoi.org///
14778/3137768 3137836
Yuzhen Mao, Martin Ester, and Ke Li. Iceformer: Accelerated inference with long-sequence trans- formers on cpus. arXiv preprint wrxiv:2405.028-22 2024.
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Rac Ying Yee Wong, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerat- ing generative Ilm serving with speculative inference and token tree verification. arXiv preprint arXiv:2305.09781, 2023.
Toan Nguyen Mau and Yasushi Inoguchi. Locality-sensitive hashing for information retrieval system on multiple gpgpu devices. Applied Sciences, 10(7), 2020. ISSN 2076-3417. doi: 10.3390/ app10072539. BtLhhtps:: /www. mdwii com/2076 com.2//..cccm/2666666 //tp:///////////////////////////
3417/10/7/2539
Art B. Owen. Monte Carlo theory, methods and examples, httpp////////ee......uuuu... domaiins mc/ 2013.
Zaifeng Pan, Feng Zhang, Hourun Li, Chenyang Zhang. Xiaoyong Du, and Dong Deng. G-slide: A gpu-based sub-liness deep learning engine via Ish sparsification. IEEE Transactions on Parallel and Distributed Systems, 33(1 1):3015-3022 2022. doi: 10.1 2..1199/PDDS.222.3313 33133999
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Lev- skaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. arXiv preprint arXiv:2211.051022 2022.
Siva Reddy, Danqi Chen, and Christopher D Manning. Coqa: A conversational question answering challenge. Transactions of the Association for Computational Linguistics, 7:249-266 266, 2019.
Baptiste Rozière, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Jingyu Liu, Romain Sauvestre, Tal Remez, Jérémy Rapin, Artyom Kozhevnikov, Ivan Evtimov, Joanna Bitton, Manish Bhatt, Cristian Canton Ferrer, Aaron Grattafiori, Wenhan Xiong, Alexandre Défosse7, Jade Copet, Faisal Azhar, Hugo Touvron, Louis Martin, Nicolas Usunier, Thomas Scialom, and Gabriel Synnaeve. Code Ilama: Open foundation models for code, 2024. URL
httpss:/aaxxt org/ abs///b///// 2308 12950
Caitlin Sadowski. Simhash : Hash-based similarity detection. 2007. URL https://ppp
semant iettantisscholce org
:::.::la......rrooo::::::::::::: org./oorpus1D:11949711 199497199999
Anshumali Shrivastava and Ping Li. Asymmetric Ish (alsh) for sublinear time maximum inner product search (mips). In Advances in Neural Information Processing Systems (NeurIPS), pp. 2014.
Prajwal Singhania, Siddharth Singh, Shwai He, Soheil Feizi, and Abhinav Bhatele. Loki: Low-rank keys for efficient sparse attention. arXiv preprint 22X5v:2406.02542, 2024.
14

---

Published as a conference paper at ICLR 2025
Malcolm Slaney, Yury Lifshits, and Junfeng He. Optimal parameters for lecalityseensiiieee hashing. Proceedings of the IEEE, 100(9))2604-200) 200044--6,,,,000000 doi: 10.1 10.1099/PRRC222022229999999
Ryan Spring and Anshumali Shrivastava. A new unbiased and efficient class Ish-based samplers and estimators for partition function computation in log-linear models. arXiv preprint a 05. 2017.
Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: Query- aware sparsity for efficient long-context IIm inference. arXiv preprint arXiv:2406.1077 107744 2024.
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312. /805, 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cris- tian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, lliyan Zarov, Yuchen Zhang. Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic. Sergey Edunov, and Thomas Scialom. Ilama 2: Open foundation and fine-tuned chat models, 2023.
Minzheng Wang, Longze Chen, Cheng Fu, Shengyi Liao, Xinghua Zhang, Bingli Wu, Haiyang Yu, Nan Xu, Lei Zhang, Run Luo, et al. Leave document behind: Benchmarking long- ang-context Ilms with extended multi-doc qa. arXiv preprint arXiv:2406. 17419, 2024.
Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mechanistically explains cong-context factuality, arXiv preprint 1554444 2024.
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, Song Han, and Maosong Sun. Infllm: Unveiling the intrinsic capacity of lims for understanding extremely long sequences with training-free memory. arXiv, 2024.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv: 7753,, 2023.
Zihao Ye, Ruihang Lai, Bo-Ru Lu, Chien Yu Lin, Size Zheng. Lequn Chen, Tianqi Chen, and Luis Ceze. Cascade inference: Memory bandwidth efficient shared prefix batch decoding, February 2024.
htlltttps://ffaaann aicieerrrriii/i wwtpaa///////a/ cascade leaiiddiiiifeeee inference httml
Amir Zandieh, Insu Han, Majid Daliri, and Amin Karbasi. Kdeformer: Accelerating transformers via kernel density estimation. In International Conference on Machine Learning, pp. 40605-40623. PMLR, 2023.
Hailin Zhang, Xiaodong Ji, Yilin Chen, Fangeheng Fu, Xupeng Miao, Xiaonan Nie, Weipeng Chen, and Bin Cui. Pqcache: Product quantization- kvcache for long context 11m inference. arXiv preprint arXiv:2407. 12820. 2024.
Jun Zhang, Jue Wang, Huan Li, Lidan Shou, Ke Chen, Gang Chen, and Sharad Mehrotra. Draft & verify: Lossless large language model acceleration via self-specclatiiv decodiin CoRR, abs/2309 68, 2023a. doi: 10. 0044550/ARXIV.230.000 168. URL https://doi. 1 48550/arXiv 2309. 08169]
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang "Atlas" Wang, and Beidi Chen. H2o: Hitte- oracle for efficient generative inference of large language models. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neu- ral Information Processing Systems, volume 36, pp. 34661-34710. Curran Associates, Inc.,
15

---

Published as a conference paper at ICLR 2025
2023b.
URL
http:///pooeeddingnnnnrroooo/ooo 2
A PROOFS FOR THEOREMS
A.1
PROOF FOR THEOREM
3.2
Proof.
8
1
6
(12)
B
B
1=1
1=1
Assume Ej is the covariance matrix of 2g is the covariance matrix of
Tr(E,) - - TT
[Bi]l;:
(13)
-|||||1 constant, so the trace covariance matrix monotonically decreases with B.
A.2
PROOF FOR THEOREM 3.3
Proof.
B||
E
E
1 iEs Σ (1 - w;) 8)
(1 - w)
(14)
1=1
1=1
==1
Without loss of generality, let 02 = WA and 01
-mmmin
Siss 04 = then
n
1
E||
=n
8
B
(15)
B
a
8
(16)
(()==0* 5 is convex function with B>
1 and +20
Then, with Jensen's incquality, we have
Σ 8 2 (n 1)
2i=200 (n
(Σ a:) 01
B
(17)
n-1
= (n 1)(
11 )8=
(n - 1)
)
(18)
n1
1
Bx 1, We can prove g(x) for any I (0,1),B2 1. Then we have
eB 1
B 2 - 1)(1
-
(19)
2
Then we finally have
1
E
B||||= B
8
<1+
(20)
B OBSERVATIONS ON ATTENTION DISTRIBUTION
To better understand the attention distribution, study the geometry of q.1 and make the following three observations. (1) Key states of the initial token (also known as attention sink, denoted by ) remain almost the same for arbitrary input. In Figure Da 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 similarity 0.99. (2) The orientation of the center of key states (i.c. Kave = 1 НΣ= k2) remains stable for different input sentences. Figure 9b) we measure the mutual cosine similarity of kavg of different input sentences. Although variance exists, the similarity of kang is over 0.9. (3) The orientations of kaog and kiink are almost opposite. In Figure we find that for each head, kisnkk and kang has cosine similarity between -0.9~
These observations shape the geometry 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 TopK will place even more weight the sink token, thus
16

---

Published as a conference paper at ICLR 2025
Cosine Similarity
1.000
Cosine ne Simliarity
1,000
-1.0
Head-o Head
Head-4
0.975
Head-6
80
0.30
20
0.925
15
-0.997
-0.900
-0.85
-0.8
-0.955
24
40
0..7
40
16
24
40
Sontece
(a) Beonk
(b) kang
(c) Cosine similarity of Ksonk and kuvg
Figure 9: Geometric information of attention. Left: With arbitrary input, the orientation of &sink almost remains the same, with minimum similarity 0.99 across sampled inputs. Mid: The orientation of Ravg is stable across various input sentences with a similarity >> 0.9 observed. Right: sint kang are almost opposite with similarity between 0.9 ~ -.8..
Value
-2
logllv, 011
qk//vd
-6
position ids
Figure 10: The range fluctuation of log/us of and in single decoding step. Compared stable, hence we do consider log va ol in our proposed sampling probability.
gk vd
ok? vd log/o-
losing contextual information. In addition, misaligning and k also causes difficulty in search LLu al. 2024a)
C
ORACLE SAMPLING
The optimal sampling probability to guarantee estimation unbiased in terms of lowest variance is not directly using attention score distribution but u' x wi w://e:]|. However, this sampling probability not optimal regarding 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 [T], given its lower probability of being sampled, which may influence the functionality of attention. Second, due to the same reason, uf w/ [e: is flatter than w,, resulting larger computation cost (as analyzed by Theorem 3.3
D SUPPLEMENTARY ANALYSIS
Figure 10 shows that compared to gk?"
loglo -o is stable in a decoding step. Vd
Figure 11 shows that the norm of the value states of attention sink is smaller than others.
ADDITIONAL EVALUATION
In this section, provide additional experimental results demonstrate that
17

---

Published as a conference paper at ICLR 2025
10
V norm
13
position
Figure 11: The y-axis is the norm values states ||e,] for token (on the x-axis). We observe that the value norm ||oo|| of the attention sink is significantly smaller than others.
MAGICPIG support longer context lengths and wide range of LLMs (Appendix E.1 MAGICPIG can scale up with 70B level LLM (Appendix(E.2)
MAGICPIG improves decoding throughput with various hyper- (K, L). (Appendix E.3
E.1 LONGER CONTEXTS
Following the setups of Table
we
evaluate
two
additional models.
aws httpa/awsppotttiinnnooon
MegaBeam Mi stral-lll--. 73 512k) and
(Gao et al. 2024) with
context lengths extended to 256K. The results are shown in Table
Table 4: Synthesized tasks RULER (Hsieh et al. 2024 MAGICPIG preserves high accuracy with extended context lengths and different models. Config and cost are defined as in Table
Methods
Config
16K
32K
64K
96K
128K
256K
Avg.
Cost:
Costz
Costsotal-
Full
91.7 88.1
83.5
83.7
83.5
82.5
85.5
0.00
1.00
1.00
MAGIPPIG
(10,150)
89.8 86.5
81.7
80.7
81.6
79.0
83.2
0.00
0.02
0.02
MAGICPIG
(9.120)
90.7 88.5
82.9
82.4
82.3
80.1
84.5
0.00
0.04
MAGICPIG
(8,75)
90.6
86.4
82.8
81.6
82.3
80.8
84.1
0.00
0.05
0.05
Quest
(16,0.04)
83.3
83.2
79.3
78.6
78.5
78.5
80.2
0.06
0.04
0.10
Full
93.5
90.8
85.1
83.5
81.7
78.4
85.55
0.00
1.00
1.00
MAGICPIG
(10,150)
88.0
86.4
81.3
78.8
77.3
71.1
80.5
0.00
0.02
0.02
MAGICPIG
(10,170)
89.0
88.7
82.8
80.0
77.7
73.7
82.0
0.00
0.025
0.025
MAGICPIG
(9.120)
91.4
88.2
82.4
80.4
79.2
75.2
82.8
0.00
0.04
0.04
MAGICPIG
(8,75)
91.4
88.6
83.1
80.5
79.1
73.9
82.8
0.00
0.05
0.05
Quest
(16,0.04)
84,9
83.7
78.7
78.6
76.3
72.1
79.2
0.06
0.04
0.10
E.2 SCALING UP TO LARGER MODELS
We evaluate MAGICPIG for meta- meta-Ilama/L Ialaam//llamm3..... 1--....1-----stmrrrmm (Dubey at. 2024) to demonstrate that our approach can work well with larger LLMs in Table
Table Synthesized tasks from RULER HHiihh
2024
MAGICPIG preserves high accuracy with low are preserved, Contig and cost are defined as in Table]
computation for 70B level models. 4 layers (0.16.32,48)
Methods
Config
16K
32K
64K
Avg.
Cost,
Costa
Costcotal-
Llama-3.1-70 1-.0-------nrte
Full
96.4
94.6
89.2
93.6
0.00
1.00
1.00
MAGICPIG
(10,135)
95.1
92.3
86.7
91.7
0.00
0.016
0.016
MAGICPIG
(10,150)
94.8
93.4
88.2
92.1
0.00
0.02
0.02
MAGICPIG
(9.110)
95.7
93.5
88.4
92.5
0.00
0.034
0.034
MAGICPIG
(9,120)
96.0
94.3
89.1
93.1
0.00
0.04
0.04
18

---

Published as a conference paper at ICLR 2025
E.3 SYSTEM PERFORMANCE
In this section, we evaluate the system performance (latency, throughput) of MAGICPIG under different hyper-parameter configurations. We use Llama-3.1-8B- 1-8--Instruct (Dubey et 2024 with 96K contexts as an example.
Table 6: System performance for MAGICPIG using Llama-3. 1aama-3.1-8B5 1-8B-Instruct with a 96K context length under varying hyper-parameter configurations. We report the decoding latency (time between tokens, TBT) when the batch size is I, the maximum throughput, and the throughput with latency constraint of 200ms (Throughput, gooss the table). Config and cost are defined T
The number with means memory limit CPU.
Config
TBT (ms)
Max Throughput (tokens/sec)
Throughpul c0as (tokens/sec)
Costcotal-
(11,300)
17.38
41.68"
40.84
0.02
(10,220)
14.07
32.29*
26.66
0.04
(10,170)
16.79
46.52
39,90
0.025
(10,150)
18.31
53.78
48.89
0.02
(9.120)
13,93
32.50
26.60
0.04
(8,75)
12.47
27.43
21.17
0.05
F SELECTION OF
(K, L)
In this section, we discuss the impact of the LSH hyper-parameter (K, L) and how to select it. First, we briefly explain what aper-paaaeeter (K, L) does for LSH sampling. Then, explain the relations between (K, L) and attention computation cost and accuracy. Finally, we show how we decide the parameters by ablation studies.
F.1 (K, L) IN LSH
In each hash table, we K hash functions to compute the hash code of ie and q. In Simhash (Charikar 2002), the hashing we use in MAGICPIG, the hash functions are random projections. With K random projections, we able partition the space (in problem, the space is R2228) into 2 subspace. If and only k and fall in the same subspace, we say they collide in this hash table. We have I. hash tables in total. In MAGICPIG, if and only & and q collide least two hash tables, is sampled by Here are some intuitions about how (K, L) will influence the LSH sampling in MAGICPIG.
K is too small, then we cannot partition the space well; we will sample too many ks, which might be far away from g (in the attention problem, this means their inner production is small), increasing computation cost.
On the other hand, if K is too large, although the quality of sampled ks will be better, the collision probability in each table will small; thus, the number of the sampled ks will be reduced. We need to increase L to ensure that a certain number of keys are sampled and involved in the computation. However, increasing (K, L) too much will bring more memory overhead on CPU DRAM since we build L hash tables for each key-value head.
Thus, (K, L) is important because it balances computation cost, overhead, and sampling quality (which determines accuracy). Tuning (K, L) necessary in LSH (Lvet al. 2017 Slaney et al. 2012).
F.2 (K, L.) AND MEMORY OVERHEAD
(K, L) will change overheads brought by MAGICPIG: the memory occupied by hash tables the CPU and extra computation for random projections (hash functions) on the GPU (as shown in Table
LLM decoding is a haamorr-bandwidtteee bound process and the majority of time is spent loading the data (parameters/KY cache) to GPU cores rather than actually doing the computation (Miao et al. 2023 Zhang et al. 2023a| Chen et al. 2024). Besides, the consumnnn part, i.e., the long-context attention computation, moved to the CPU Thus, the 1.8% ~ 8.5% extra computation on GPU will only make a minor difference in execution time. However, the enlarged size of hash tables prevents us from always increasing (K, L) to get more accurate results.
As shown in Table 7 under the same (K, L), the memory overhead of hash tables grows linearly with context length and the total number of key-value heads in models (which is determined by model sizes).
19

---

Published as a conference paper at ICLR 2025
Table 7: The overhead of Locality sensitive hashing during decoding. We report the size of random projectors (on GPU) hash tables (on CPU), the computation overhead CO (refers by random projections LSH and the computation of the original model'' linear projections (e.g.. Wq. W Wy. and MLP)). Notice that when the context length exceeds 64K, we need to use 32-bit integers to store the indices for the KV cache in hash tables. Llama-3. 1ana-3.1-8BB70B Instruttt (Dubey et al. (2024) and Code-Llama-344 (4b-166K Rozière
ratio between computation introduced
et al. 2024) use group query attention, thus the sizes of hash tables reduced.
Models
(K, L)
Context length
Projectors
Hash tables
CO
Llama-3.
1-8B--8-888 1.ama-3.1-8B
--Stttrrruttt Instruct
(10, 150)
96K
384KB
14GB
3.8%
Llama-3, 1-1-BB- ---nsttruct
(11,300)
96K
825KB
28GB
8.5%
Llama-3. 1-8B-Instruct
(10, 150)
64K
384KB
4.7GB
3.8%
Llama-3. 1-suu---nssruc
(10, 150)
64K
384KB
11.8GB
1.8%
13-KKKK
(10, 150)
16K
384KB
7.3GB
5.2%
Code-Llama-34b-101
(10, 150)
16K
384KB
1.8GB
2.2%
F.3 (K, L) AND COMPUTATION COST/BUDGET
In summary, increasing K will make the budge smaller, and increasing L. will increase the budget.
Theoretically, as introduced in Section 4.3 in our approach, the key k; is sampled only at least two hash tables exist where k; shares the hash value with query q- With the assumption that k; is well-distributed (In each hash table out of L, each hash value corresponds to roughly the same number of k,s), the ratio of retrieved kys be estimated with
B/n= 1- (1-0.5k] - - (1-0..4))
(21)
where n is the context length. Here, we estimate the collision probability of k; and q in a single hash table 0.5*
Empirically, the ratio of retrieved keys and values (B/n) might differ from the above estimation since the data is not perfectly distributed. We present the empirically measured budget in Table
Table 8: Empirical measured budget/cost for different (K, I.).
K/L
75
100
120
150
200
300
14%
21%
27%
35%
48%
66%
5%
8%
11%
15%
22%
36%
1.6%
2.7%
4%
5.4%
8.5%
15.4%
0.5%
0.9%
1.5%
2%
3%
6%
0.15%
0.3%
0.5%
0.6%
1%
2%
F.4 (K, L) AND ACCURACY
There are no naive relations between (K, L) and downstream accuracies since (K, L) not only influences sampling quality but also the computation budget. One safe way to discuss the relation between (K, L) and accuracy is: Fixing the computation budget, larger (K, L) will potentially produce higher accuracy, since the sampling quality higher. Our experimental results show that,
Increasing (K, L) can significantly improve accuracy in relatively longer contexts Table
Table 9: We show the effectiveness of larger hash tables for longer contexts by evaluating MegaBeam-Mistrale 512K on RULER (Hsich et al. 2024). With the same computation cost (~ 2%), config (11, 300) achieves higher accuracy compared to (10, 150).
(K, L)
16K
128K
256K
Full
91.7
83.7
82.5
(10, 150)
89.8
80.7
79.0
(11,
90.6
83.3
81.9
Same set of (K, L) can generalize to larger LLMs Table
*Costa in Tables 1o
20

---

Published as a conference paper at ICLR 2025
Table 10: 8B and 70B models RULER (Hsich et al.
2024 64K.
Models/Config
Full
(10, 150)
(10, 135)
(9, 120)
(9, 110)
Llama-3 3.1-8B-Innste -B--Instruct
86.1
84.8
83.6
84.7
84.7
Llama-3. 3.1-70B- Instruct
89.2
88.2
86.7
89.1
88.4
F.5
HOW TO SELECT (K,
Finding the optimal (K, L) for high accuracy well as efficiency is long-standing problem in LSH. Similar to the traditional hyper-parameter tuning process machine learning, K, and are configured offline based data subsets. In LSH, K is more sensitive nyper-parameter than L. A slight change of can drastically influence the number retrieved items (i.e., budget/cost) and quality. In MAGICPIG, K=8-10 manually determined by ablations small-scale tasks and found to be effective across various models and tasks; then, we adjust L to obtain the wanted computation cost/budget.
Here, we present two ablations to demonstrate the selection of K in Tables 11 and 12
Table 11: Fixing the budget/cost to 4%, we ablation the performance of different (K, L) on RULER (Hsich et. 2024 16K.
Models/Config
Full
(10, 240)
(9, 120)
(8, 65)
(7,
LJama-3. 1-0mma-3--BB-Insttucc
94.2
94.2
92.8
92.3
88.5
Table 12: Fixing L as 120, we ablation the performance of different K on RULER (Hsich et al. Llama-3. 3.1-8B-Instruce
2024)
16K for
(K. L)
Full
(10, 120)
(9, 120)
(8, 120)
(7, 120)
Cost
1.0
0.012
0.04
0.11
0,27
Accuracy
94.2
92.8
92.8
94.1
94.3
If we want the computation cost to be below 5% and I. below 200 (to reduce memory overhead in the CPU), then K=8-10 is a reasonable choice. Unlike K, L is not that sensitive. We select L based on the following principle after determining K: we can allow the computation cost to be smaller for larger K since the sampling is more precise. This is why we choose to use (8, 75), (9, 120), and (10, 150).
It's worth pointing out that tuning (K, L) is challenging and long-standing problem in LSH, and we only give an example of practice MAGICPIG. More advanced hashing algorithms (such as Cross-polytope (Andoni et al. 2015) or data-dependent ones (Andoni & Razenshteyn 2013)) can improve the trade-off between memory overhead and accuracy. We leave as a future direction.
G TOPK VS. SAMPLING
In this section, provide an intuitive understanding of how sampling can work better than TopK. TopK only captures the ranking information when estimating attention output. In contrast, sampling considers the entire data distribution (i.e., the attention score after Softmax).
Here an example. Imagine a zoo with 100 animals: 10 elephants, 10 pigs, 10 tigers, and other unique animals. The daily food consumption for each group is as follows:
Elephants: 50 Ib/day each
Pigs: 20 Ib/day each
Tigers: 10 Ib/day each
Other unique animals: Ib/day each
To compute the true average daily food consumption per animal in the zoo:
True Average= (10x50) +10 20) +10 (10x -871b.
100
we use a Top-K approach (e.g., selecting the top 10 animals based on the numbers of animals), we include elephants, pigs, tigers, and 7 randomly selected animals from the unique ones. The estimated
21

---

Published as a conference paper at ICLR 2025
average is:
TopK Average=
=221b.
37
This overestimates the average because it disproportionattly weights high-consumption animals.
Instead, we perform sampling with replacement from the animal distribution, proportional to their numbers. The probabilities for each group are:
Sampling Probabilities (0.1,0.I.0.1.0.000 x7],
70),
where 0.1 represents the probabilities for elephants, pigs, and tigers (10/100 each), and 0.01 corresponds to each unique animal (1/100).
Perform 10 random draws. A possible sampling outcome could be: [elephant, pig, tiger, other, other, other, other, other, other, other]. The corresponding daily food estimate is:
Sample Estimate 50+200+ 10++10++711] =8.7lb.
10 This estimate is unbiased, meaning the expected value the estimates equals the true average (8.7 lb). While there is variance across individual trials, the standard deviation (std) can be calculated 4.7 lb for 10-sample budget.
Increasing the sampling budget reduces variance. For example, with 20 samples, the std decreases to 3.4 lb. Meanwhile, Top-K with a budget of 20 adds 17 unique animals, yielding:
TopK Average (K=20)= (10 x 50) +(10 20) (10x 10) +17 x 17Ib.
47
Again, the Top-K K estimate remains biased, significantly overestimating the average.
Note that this is intended as an intuitive example. For a detailed and formal derivation of the sampling methodology, please refer Kloek & Van Dijk 1978) Owen (2013) Lohr (2021).
H LIMITATIONS AND FUTURE WORK
MAGICPIG stores the offloaded KV cache and hash tables on CPU DRAM. which is unsuitable for serving scenarios with insufficient DRAM. KV cache quantization methods like QServe (Lin et al. 2024) and KIVI (Liu et al. 2024b) can help reduce the KV cache memory. Currently, another limitation is that, we have not implemented MAGICPIG in prefilling stage, which is also an important direction in long context LLM serving. Applying more advanced LSH algorithms, such as Cross- polytope hash (Andoni et al. 2015), can reduce the size of hash tables while improving estimation accuracy. Building CPU-GPU pipelines (He Zhai 2024) and leveraging the new avx5 12 12.bf16 features of CPUs will improve efficiency. For higher- end GPUs with sufficient HBM, leveraging LSH to accelerate GPU attention computation is also an interesting topic, as GPU-friendly LSH algorithms and efficient GPU kernels (Nguyen Mau & Inoguchi 2020 Pan et al. 2022 are required to do sampling. Besides, how to automatically tune the LSH hyper-parameter also an interesting future work.
parameter (K, L) (L et al. 2017) is
22


Reviews:
[{"review_id": "OyHSYBG7sP", "reviewer": "Reviewer_KBnM", "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.", "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. \n\nQ2. 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. \n\nQ3. 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. \n\nQ4. 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). \n\nQ5. 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.\n\nQ6. 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. \n\nQ7. 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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 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. \n\nS2. Experiments are highly promising, outperforming Quest in both accuracy and efficiency. \n\nS3. 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.\n\nW2. Some parameter evaluations are missing in the experiments. \n\nW3. Only Llama series models are evaluated.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wk7QAvHuT4", "reviewer": "Reviewer_UYfg", "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.", "questions": "See the weakness part", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tThe observation that top-k attention does not perform well is interesting.\n\n2.\tThe idea of using LSH to conduct importance sampling is novel.\n\n3.\tMAGICPIG 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.\tLacks 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)?\n\n2.\tThe 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.     \n\n3.\tWhat 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RHLYg6njbl", "reviewer": "Reviewer_ePKx", "summary": "This paper introduce a novel approach that leverages LSH sampling to approximate the oracle sampling. Empirical evaluation shows improvement over the baseline.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "S1. The problem and the solution is well-motivated in the paper.\nS2. The CPU-GPU co-design enables the storage of large LSH index table.\nS3. The empirical results outperforms baselines.", "weaknesses": "Weak Points\n----\nW1. 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.\n\nW2. 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.\n\nW3. Another direction of accelerating the inference is to quantize the model. How does the proposed method work on quantized LLM is not discussed.\n\nW3. No code is provided. It might be hard for readers to reproduce the results.\n\nPresentation\n----\nP1. 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.\nP2. Although it might be obvious for readers with retrieval and word extraction background, the acronym niah, cwe, and fwe and not explained before usage. \nP3. The numbers in Figure 6 might be a bit outdated. In addition, the connection between CPU and GPU could be faster SXM.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FWH0u6rgDB", "reviewer": "Reviewer_Q8fb", "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.", "questions": "1.\tIs the size of the hash table related to model size and sequence length? How does the size of the hash table affect the performance?\n\n2.\tWhat is the time overhead of constructing hash tables, and which factors influence this overhead?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.\tMAGICPIG 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.\n\n2.\tMAGICPIG 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.\tWhile 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.\n\n2.\tThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "c9V0yEbY1C", "reviewer": "Reviewer_qU5o", "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.", "questions": "See weaknesses", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "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:\n1. what is the intuition for selecting (K, L) for the hash table size? \n2. For Table 1/2/3, why is latency not included in the comparison? \n3. Does the author believe further improvement can be made by combining this approach with PEFT?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)