SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$\mathbf{H}_2\mathbf{O}$ : Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
Zhenyu Zhang $^{1}$ , Ying Sheng $^{2}$ , Tianyi Zhou $^{3}$ , Tianlong Chen $^{1}$ , Lianmin Zheng $^{4}$ , Ruisi Cai $^{1}$ , Zhao Song $^{5}$ , Yuandong Tian $^{6}$ , Christopher Ré $^{2}$ , Clark Barrett $^{2}$ , Zhangyang Wang $^{1}$ , Beidi Chen $^{6,7}$
<sup>1</sup>University of Texas at Austin, <sup>2</sup>Stanford University, <sup>3</sup>University of California, San Diego, <sup>4</sup>University of California, Berkeley, <sup>5</sup>Adobe Research, <sup>6</sup>Meta AI (FAIR), <sup>7</sup>Carnegie Mellon University {zhenyu.zhang, tianlong.chen, ruisi.cai, atlaswang}@utexas.edu, ying1123@stanford.edu,
{chrismre,barrett}@cs.stanford.edu,t8zhou@ucsd.edu,lianminzheng@gmail.com, zsong@adobe.com, yuandong@meta.com, beidic@andrew.cmu.edu
Large Language Models (LLMs), despite their recent impressive accomplishments, are notably cost-prohibitive to deploy, particularly for applications involving long-content generation, such as dialogue systems and story writing. Often, a large amount of transient state information, referred to as the KV cache, is stored in GPU memory in addition to model parameters, scaling linearly with the sequence length and batch size. In this paper, we introduce a novel approach for implementing the KV cache which significantly reduces its memory footprint. Our approach is based on the noteworthy observation that a small portion of tokens contributes most of the value when computing attention scores. We call these tokens Heavy Hitters  $(\mathsf{H}_2)$ . Through a comprehensive investigation, we find that  $(i)$  the emergence of  $\mathsf{H}_2$  is natural and strongly correlates with the frequent co-occurrence of tokens in the text, and  $(ii)$  removing them results in significant performance degradation. Based on these insights, we propose Heavy Hitter Oracle  $(\mathsf{H}_2\mathsf{O})$ , a KV cache eviction policy that dynamically retains a balance of recent and  $\mathsf{H}_2$  tokens. We formulate the KV cache eviction as a dynamic submodular problem and prove (under mild assumptions) a theoretical guarantee for our novel eviction algorithm which could help guide future work. We validate the accuracy of our algorithm with OPT, LLaMA, and GPT-NeoX across a wide range of tasks. Our implementation of  $\mathsf{H}_2\mathsf{O}$  with  $20\%$  heavy hitters improves the throughput over three leading inference systems DeepSpeed Zero-Inference, Hugging Face Accelerate, and FlexGen by up to  $29\times$ ,  $29\times$ , and  $3\times$  on OPT-6.7B and OPT-30B. With the same batch size,  $\mathsf{H}_2\mathsf{O}$  can reduce the latency by up to  $1.9\times$ . The code is available at https://github.com/FMInference/H20.
Large Language Models (LLMs) have demonstrated remarkable proficiency in a wide range of natural language processing applications such as content creation, summarization, and dialogue systems [1, 2, 3, 4]. However, their deployment is very costly. In addition to the widely-studied bottlenecks of model size and the quadratic cost of attention layers, the problem of the size of the KV cache, which stores the intermediate attention key and values during generation to avoid re-computation, is becoming increasingly prominent [5]. For instance, a 30 billion-parameter model with an input batch size of 128 and a sequence length of 1024 results in 180GB of KV cache. A natural approach is to limit its maximum size as is done in classical software or hardware caches [6]. However, it is challenging to reduce KV cache memory footprints in LLMs without accuracy drops.
Figure 1: Upper plots illustrate symbolic plots of an attention map deploying different KV cache policies in LLM generation. Lower right: contrasts their accuracy-memory trade-off. Left: the overview of  $\mathsf{H}_2\mathsf{O}$  framework.
While there exists substantial literature on sparse attention approximation in training, they have not seen wide adoption for alleviating KV cache bottleneck. First, most existing methods, e.g., Reformer [7] and Flash Attention [8], are designed to overcome the quadratic memory required by attention mechanisms when modeling long sequences but still require a large cache size. Second, variants like sparse transformer [9], low-rank based transformers [10, 11] or multi-query attention [12, 13, 5] can reduce the cache size, but directly applying them on pre-trained LLMs for generation results in high miss rates and degrades the accuracy as shown in Figure 1. Finally, some recent advances such as gisting tokens [14] can learn to compress the KV cache for documents, but their expensive eviction policies are difficult to deploy during generation.
Therefore, an ideal KV cache should have (i) a small cache size to reduce memory footprint, (ii) a low miss rate to maintain the performance and long-content generation ability of LLMs, and (iii) a low-cost eviction policy to reduce the wall-clock time during generation. However, there are three technical challenges. First, it is not immediately clear whether the size of the KV cache can be restricted—each decoding step might, in principle, require access to all previous attention keys and values. Second, identifying an optimal eviction policy that maintains generation accuracy is a combinatorial problem<sup>1</sup>. Finally, even if an optimal policy can be brute-forced, it is infeasible for deployment on real-world applications.
Fortunately, our preliminary exploration has yielded intriguing observations about the empirical properties of LLMs. These findings pave the way for the potential design of an efficient KV cache.
Sparsity for small cache size: We observe that even when trained densely, the attention matrices of LLMs are over  $95\%$  sparse at inference time (shown in Figure 2). This holds for a wide range of pre-trained LLMs. Therefore, only  $5\%$  of the KV cache is sufficient for decoding the same output token at each generation step, which suggests it may be possible to have up to a  $20\times$  reduction in KV cache size without an accuracy drop.
Heavy-Hitters for low miss rate: We discover that the accumulated attention scores of all tokens in attention blocks adhere to a power-law distribution. It suggests that there exists a small set of influential tokens that are critical during generation, named heavy-hitters  $(\mathsf{H}_2)$ .  $\mathsf{H}_2$  provides an opportunity to step away from the combinatorial search problem and identify an eviction policy that maintains accuracy.
Greedy algorithm for low-cost policy: we surprisingly find that retaining the  $\mathsf{H}_2$  based on local statistics at each decoding step—summing the attention scores of only the preceding tokens—is as effective as considering the attention of future tokens (shown in Figure 2).
Based on the above, we first rigorously define the generative process of LLMs operating with a size-constrained KV cache in Section 2.1. Then we propose Heavy-Hitter Oracle  $(\mathsf{H}_2\mathsf{O})$ , a framework
<sup>1</sup> Belady's Algorithm is optimal for standard cache, but not necessarily for KV cache.
that exploits the properties of LLMs and uses simple, low-cost eviction policies that retrain the quality of LLMs throughout the generation process. Specifically,
- In Section 3, we explore the emergence of  $\mathsf{H}_2$  in attention, revealing their fundamental and critical roles: (i)  $\mathsf{H}_2$  exhibit a strong correlation of frequently co-occurring words in textual data; and (ii) removing  $\mathsf{H}_2$  completely damages the model's functionality. We demonstrate that  $\mathsf{H}_2$  can largely lower the cache miss rate of the existing policies mentioned above. Theoretically, assuming the attention scheme is submodular,  $\mathsf{H}_2$  corresponds to a greedy algorithm and is therefore near-optimal.
- In Section 4, we present a greedy but low-cost variant of  $\mathsf{H}_2$  which is dynamically determined by the accumulated attention score at each decoding step. We formulate the eviction policy with greedy  $\mathsf{H}_2$  as a variant of dynamic submodular maximization. The analysis shows that it results in a similar generative process as the one using the  $\mathsf{H}_2$  eviction policy.
We perform extensive experiments on OPT, LLaMA, and GPT-NeoX on a single NVIDIA A100 (80GB) GPU to evaluate  $\mathsf{H}_2\mathsf{O}$  across a range of tasks from lm-eval-harness [15] and HELM [16]. We implement  $\mathsf{H}_2\mathsf{O}$  on top of FlexGen that can easily adapt different cache eviction techniques to produce a system with high-throughput inference. Performance experiments show our framework achieves  $29\times$ ,  $29\times$ ,  $3\times$  higher throughputs compared to three leading inference systems, DeepSpeed Zero-Inference [17], Hugging Face Accelerate [18], and FlexGen [19] respectively. With the same batch size,  $\mathsf{H}_2\mathsf{O}$  achieves up to  $1.9\times$  lower latency compare to FlexGen.
2 Related Work and Problem Setting
Efficient Inference of LLMs. The substantial parameter counts of large language models (LLMs) present significant challenges for inference. To overcome this limitation, previous efforts have employed model compression techniques with specific designs to achieve efficient LLM inference, such as the method described in [20, 21, 22], which employs one-shot pruning on LLMs, resulting in negligible performance degradation even without retraining. Additionally, alternative approaches explore quantization methods specifically tailored to LLMs, as discussed in [23, 24, 25, 26, 27, 28]. Also, CoLT5 [29] employs a token-wise conditional computation strategy to reduce the overall computation cost. These methods address efficient inference from orthogonal perspectives and can be organically integrated. The techniques investigated in this study are closely associated with pruning or sparsity but focus on a distinct inference bottleneck, namely, KV cache. One closely related work[30] utilizes a learnable mechanism that determines necessary tokens during inference but requires an extra fine-tuning process, which makes it less practical.
Sparse, Low-rank Attention Approx. The quadratic computational complexity of attention modules is one of the major bottlenecks of transformer inference [31]. Various efforts are devoted to addressing this challenge [7, 9, 10]. For example, Reformer [7] reduces the computational cost from quadratic to superlinear complexity via locality-sensitive hashing. Performer [10] employs positive orthogonal random features to approximate attention kernels. One relevant work, Sparse Transformer [9], introduces sparsity to reduce KV cache memory footprint and achieve an efficient attention mechanism, considered as our baseline in this paper. Moreover, SpAtten [32] utilizes accumulated attention scores to select important tokens for efficient attention inference while they don't consider the variance of token importance across attention heads and layers. Comparison with SpAtten is detailed in Appendix C.9.
Caching. Caching, which plays a pivotal role in optimizing system performance, entails the development of effective eviction policies to handle frequently accessed data. Conventional approaches such as Least Recently Used and Least Frequently Used [33, 34] prioritize the recency and frequency of data access. And the design of KV cache encounters many similar challenges as traditional caching.
LLM Inference Breakdown. The generative procedure of LLMs encompasses two distinct phases: (i) the prompt phase, in which an input sequence is utilized to produce the KV cache (consisting of the key and value embeddings), similar to the forward pass employed during LLM training; and (ii) the token generation phase, which leverages and updates the KV cache to generate new tokens incrementally. Each generation step relies on the previously generated tokens. The primary focus of this paper is to enhance the efficiency of the KV cache in attention during the token generation phase, thereby accelerating LLM inference.
(b)
Figure 2: (a) Attention Sparsity in pre-trained LLMs. (b) The distribution of accumulated attention scores with respect to the corresponding word (red scatter) and the co-occurrence times of words in the data (gray curve). The x-axis represents the word index in the vocabulary. (c) The performance comparison between the baseline model with full KV and the model w.o. heavy hitter. (d) Comparison between the baseline model with full KV,  $\mathsf{H}_2\mathsf{O}$  with the local statistic,  $\mathsf{H}_2\mathsf{O}$  with the global statistic, and the model with only the most recent KV (Local). Apart from the baseline model, each model is evaluated with  $20\%$  KV cache budget.
We formally define the generative process with limited KV cache size. Denote attention query matrix as  $Q \in \mathbb{R}^{n \times d}$  and key matrix as  $K \in \mathbb{R}^{n \times d}$ .  $Q_{i,*}$  represents the  $i$ -th row of  $Q$  and  $K_{\leq i,*}$  represents the first  $i$  rows of  $K$ . Let  $k$  denote the budget of space and  $k < n$ . For simplicity,  $K_{S_i,*} (\in \mathbb{R}^{i \times d})$  denotes a sub-matrix of  $K$  which selects  $S_i$  rows from  $K$ . (For the non-selected rows  $[i] \backslash S_i$ , we put all zeros in that row) Eviction policy is defined as:
Definition 2.1 (Eviction Policy, informal). Let  $S_{i-1}$  denote the source set. Let  $S_i$  denote the target set. We defined the eviction policy  $g: S_{i-1} \to S_i$  such that
-  $|S_{i}| = k$  (KV cache size is not changing over the time)
-  $|S_i \backslash S_{i-1}| \leq 1$  or equivalently  $|S_i \cap S_{i-1}| \geq k - 1$  (we can avail at most  $1 \, \text{KV}$  in the KV cache)
Then, we define the generative process with our eviction policy.
Definition 2.2 (The generative process with eviction policy, informal). Let  $k$  denote the size of the KV cache. For each  $i \in [n]$ , for the  $i$ -th token, we have
- Let  $S_{i} \subset [n]$  denote the tokens in KV cache when predicting the  $i$ -th token.
- The information we have is a length- $i$  vector  $o_i \coloneqq D_i^{-1} \cdot \exp(Q_{i,*}(K_{S_i,*})^\top)$  (normalized attention)
- scalar  $D_{i} \coloneqq (\exp(Q_{i,*}(K_{S_{i},*})^{\top}) - 1_{[i] \setminus S_{i}}) \cdot \mathbf{1}_{i}$  (the evicted KV is set to 0, and we need to subtract them when computing the normalization)
- Replacing  $S_{i}$  by  $[i]$  in the above definition of  $o_{i}$  and  $D_{i}$  leads to standard generative process.
- The eviction policy (Definition 2.1) updates  $S_{i}$  based on  $S_{i - 1}$  and their corresponding information.
Remark 2.3. Our goal is to find a KV cache eviction policy such that the output of the generative process is similar or comparable to the original one without limiting the cache size.
3 Observations
We present two key empirical insights of LLMs that inspire the design of  $\mathsf{H}_2\mathsf{O}$ , as follows.
3.1 Sparsity for Small Cache Size
Inspired by previous literature, which reveals the existence of attention sparsity in DistillBERT [35] and bounded-norm self-attention heads [36]. We first show an observation on the sparsity of attention in pre-trained LLMs. Then we discuss how it can potentially unlock the possibility of reducing KV cache size without an accuracy drop. Given the normalized attention score  $\mathrm{Softmax}(QK^{\top})$  matrix that is calculated by the query matrix  $Q$  and the key matrix  $K$ , we set the threshold as one percent of the maximum value in each row and calculates the corresponding sparsity.
Observation. We conduct zero-shot inference with the pre-trained OPT model on the validation set of Wiki-Text-103. We plot the layer-wise sparsity within attention blocks and visualize the normalized attention score matrix. The results are presented in Figure 2 (a). We observe that although the LLMs are densely trained, the resulting attention score matrices are highly sparse, with a sparsity over  $95\%$  in almost all layers.
Insights. The attention blocks' sparsity suggests that access to all previous key and value embeddings is unnecessary for generating the next token. This suggests it is possible to evict unessential KV embeddings and reduce the requirement of KV cache during generation.
3.2 Heavy-Hitters for Low Miss Rate
The previous section showed the sparse nature of attention blocks in pre-trained LLMs, which provides the opportunity for designing small KV cache size while still maintaining the performance of LLMs. However, determining the best eviction policy that preserves generation accuracy presents a combinatorial challenge. Although Belady's Algorithm [37] is optimal and easy to compute for standard cache (offline), it is not applicable for KV cache design. Because once evicting important KVs, it could destroy the performance of LLMs due to the sequential dependency of LLM generation.
Observation. Fortunately, in the early stage of our exploration, we find that the accumulated attention scores of all the tokens within attention blocks follow a power-law distribution, as shown in Figure 2. This suggests the existence of a small set of tokens that are critical during generation. We denote those tokens as heavy-hitters  $(\mathsf{H}_2)$ . In order to verify the importance of these tokens, we compare the quality of LLM generation after masking heavy hitters with that of the original model. Not surprisingly, as shown in Figure 2, the accuracy drops drastically, confirming the importance of those tokens. Additionally, we can see the accumulated attention score of each word (in red dots) have a high correlation with their co-occurrences in the data (gray curve).
Analysis. First, based on  $\mathsf{H}_2$ , we see an opportunity to side-step from the combinatorial search problem and design a KV cache eviction policy that preserves the LLM generation quality. We conduct an empirical study implementing a KV cache eviction policy that retains only the  $\mathsf{H}_2$  and the recent KV embeddings in the cache. The intuition is that recent words typically exhibit stronger correlations with current tokens. We assess the effectiveness of this eviction policy through pretrained OPT-30B and six downstream tasks. The outcomes of these evaluations are illustrated in Figure 2. It is obvious that the  $\mathsf{H}_2$  based eviction policy can largely reduce the KV cache size without degrading the performance of OPT-30B.
Moreover, during the post analysis, inspired by [38], we find that  $\mathsf{H}_2$  based policy is related to the classical greedy algorithm (a polynomial-time algorithm with provable guarantees) under the assumption that the attention schema is submodular. We present details in Appendix D.
Lemma 3.1 (informal). Assuming the attention scheme is submodular, then greedily constructing the set  $S_{i}$  (without cache size limitation) satisfies the near-optimal property in terms of submodular.
4 Heavy-Hitter Oracle
The goal of this section is to propose the greedy algorithm using the  $\mathsf{H}_2$ -based policy and to show the provable guarantees. We first present the  $\mathsf{H}_2$ -based policy called  $\mathsf{H}_2\mathsf{O}$  cache eviction policy and formulate its deployment in LLM generation as a variant of submodular maximization problem, named dynamic submodular. Then we present  $\mathsf{H}_2\mathsf{O}$  in the generative process, followed by a practical example of deploying our proposal. Finally, we provide theoretical guarantees for  $\mathsf{H}_2\mathsf{O}$  and show our efficient system implementation.
4.1 Greedy Algorithm for Low-Cost Policy
We have shown a simple yet effective KV cache policy based on  $\mathsf{H}_2$ . However, it is impractical to deploy such an algorithm because we do not have access to the future-generated tokens. Fortunately, we empirically observe that local  $\mathsf{H}_2$ , which is calculated using local statistics at every decoding step by summing up the attention scores of the previous tokens, is equally effective as taking into account the attention of future tokens (Figure 2). In the following, we formally define this dynamic attention score computation (with space limitation) as a novel dynamic submodular type problem.
Definition 4.1 (Dynamic submodular framework, informal). Define function  $F:2^{[n]} \times 2^{[n]} \to \mathbb{R}$ , then for any set  $Z \subset [n]$ , we assume that  $F(Z,\cdot):2^{[n]} \to \mathbb{R}$  is a submodular function w.r.t. to  $Z$ , i.e.,
- For all sets  $X, Y \subset [n]$  satisfy that  $Z \subset X \subset Y$ ,
- For all element  $x \in [n]$  satisfy that  $x \in [n] \backslash Y$ ,
we have  $f(X \cup \{x\}) - f(X) \geq f(Y \cup \{x\}) - f(Y)$ , where  $f(\cdot) \coloneqq F(Z, \cdot)$ .
Remark 4.2. We provide practical insights of Definition 4.1.  $X$  denotes the existing words in the KV cache.  $Y$  is any superset of  $X$ .  $x$  can be viewed as a "word" which is either newly added to KV cache or existing deleted from KV cache. An example  $f$  can be attention score, i.e., see Algorithm 1.
If we load the sequence of  $S_{1}, S_{2}, \dots, S_{n}$  (we promise that  $|S_{i}| \leq k$  and  $|S_{i} \setminus S_{i-1}| \leq 1$ ) into Definition 4.1, i.e., for each  $i \in [n]$ , we choose  $Z = S_{i}$ , then it becomes a particular instance of the dynamic submodular problem.
Next, we provide a formal description of our algorithm, followed by an example.
Definition 4.3 ( $\mathsf{H}_2\mathsf{O}$  Eviction Policy). Let  $F_{\mathrm{score}}:2^{[n]}\to \mathbb{R}$  denote certain score function. Let  $S_{i - 1}$  denote the source set. Let  $S_{i}$  denote the target set. We defined the eviction policy  $g:S_{i - 1}\rightarrow S_i$  s.t.
-  $|S_{i}| = k$  (KV cache size is not changing over the time)
-  $|S_i \backslash S_{i-1}| \leq 1$  or equivalently  $|S_i \cap S_{i-1}| \geq k - 1$  (we can avail at most  $1 \, \text{KV}$  in the KV cache)
- We construct  $S_{i} \gets (S_{i - 1} \cup \{i\}) \backslash \{u\}$  as  $u \gets \arg \max_{v \in (S_{i - 1} \cup \{i\})} F_{\mathrm{score}}(S_{i - 1} \cup \{i\} \backslash \{v\})$ .
To describe our algorithm (Algorithm 1), we choose a particular instantiation of the function  $F_{\mathrm{score}}$ , i.e., the summation of that sets in the attention matrix.
Algorithm 1  $\mathsf{H}_2$  Eviction Algorithm
1: procedure  $\mathrm{H}_2\text{-EVICTION}(Q,K\in \mathbb{R}^{n\times d},k\in \mathbb{N})$   
2: Let  $k$  denote the budget size of cache  
3:  $S_0\gets \emptyset$   
4: for  $i = 1\rightarrow n$  do  
5: if  $i\leq k$  then  
6:  $S_{i}\leftarrow S_{i - 1}\cup \{i\}$   
7: else  
8:  $D_{i}\gets (\exp (Q_{i,*}(K_{S_{i - 1},*})^{\top}) - 1_{[i]\backslash S_{i - 1}})\cdot \mathbf{1}_{i}$   
9:  $o_i\gets D_i^{-1}\cdot (\exp (Q_{i,*}(K_{S_{i - 1},*})^{\top}) - 1_{[i]\backslash S_{i - 1}})$   
10:  $F_{\mathrm{score}}(T):= \sum_{s\in T}o_s$   
11:  $G_{i}\gets S_{i - 1}\cup \{i\}$   
12:  $u\gets \operatorname *{argmax}_{v\in G_i}F_{\mathrm{score}}(S_{i - 1}\cup \{i\} \backslash \{v\})$   
13:  $S_{i}\gets (S_{i - 1}\cup \{i\})\backslash \{u\}$   
14: end if  
15: end for  
16: end procedure
Figure 3: Illustration of Algorithm 1 during two consecutive decoding steps.
Figure 3 presents an illustrative example of our  $\mathsf{H}_2$  Eviction Algorithm. We assume that the budget size of KV cache is 3. Following the completion of the fourth decoding step, the KV embeddings associated with the third token are evicted based on the accumulated attention score. Consequently, these evicted KV embeddings become inaccessible in the subsequent decoding steps.
4.2 Theoretical Guarantee and System Implementation
We state a theoretical result as follows. The proofs and more details are provided in Appendix D.
Theorem 4.4 (informal). Under the mild assumption, let  $k$  denote the budget of space limitation. If for each token, we greedily compute the attention score based on top- $k$  choice, then we can show the set  $\widetilde{S}_i$  we generate each for token  $i$  satisfy that  $f(\widetilde{S}_i) \geq (1 - \alpha)(1 - 1/e) \max_{|S|=k} f(S) - \beta$ , where  $\alpha, \beta > 0$  are parameters.
Remark 4.5. We remark the above theorem provides a theoretical explanation of why can we hope our greedy algorithm (with cache limitation) can provide a good solution to the problem.