SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Implementation Details. We provide a general framework that can support any KV cache eviction algorithm and enhance throughput and reduce the latency of LLM generation with careful implementation. For example, to ensure I/O efficiency, we do not swap memory when stored KV is evicted, but directly fill with newly-added KV. More details are included in Appendix A.
Figure 4: Comparison results between the baseline model with full cache, our  $\mathsf{H}_2\mathsf{O}$ , and the "Local" strategy that utilizes the most recent KV embeddings.
In this section, our goal is to demonstrate that  $\mathsf{H}_2\mathsf{O}$ , a remarkably simple KV cache eviction policy is capable of enhancing end-to-end throughput and reducing latency in wall-clock while maintaining generation quality across a broad spectrum of domains and tasks.
- In Section 5.1, we show that  $\mathrm{H}_2\mathrm{O}$  can reduce the memory footprint of KV cache by up to  $5\times$  without accuracy degradation on a wide range of model architectures (OPT, LLaMA, GPT-NeoX), sizes (from 6.7B to 175B) and evaluation benchmarks (HELM and lm-eval-harness). More importantly, can enhance the performance of existing KV cache sparsification techniques.
- In Section 5.2, we demonstrate that  $\mathrm{H}_2\mathrm{O}$  can increase the inference throughput by up to  $3\times$ ,  $29\times$ ,  $29\times$  compared to the state-of-the-art inference engine FlexGen, DeepSpeed and the widely used Hugging Face Accelerate without compromising model quality.
- In Section 5.3, we present extensive ablation studies to show the effectiveness of  $\mathsf{H}_2\mathsf{O}$  under different sequence lengths, especially the input with infinite sequence length and its compatibility with quantization.
All details (hyperparameters, data splits, etc.), along with additional experiments, are in Appendix A.
5.1 End-to-End Results
We demonstrate that  $\mathsf{H}_2\mathsf{O}$  can reduce KV cache memory footprint by  $5 - 10\times$  while achieving comparable accuracy on a majority of tasks.
Setup. Our experiments are based on three representative model families of LLMs, including the OPT [39] with model sizes, LLaMA [40], and GPT-NeoX-20B [41]. We sample eight tasks from two popular evaluation frameworks (HELM [16] and lm-eval-harness [15]): COPA [42], MathQA [43], OpenBookQA [44], PiQA [45], RTE [46], Winogrande [47], XSUM [48], CNN/Daily Mail [49]. Also, we evaluate our approach on recent generation benchmarks, AlpaceEval [50] and MT-bench [51], and the details are included in Appendix. We use NVIDIA A100 80GB GPU.
Baselines. Since  $\mathsf{H}_2\mathsf{O}$  evenly assigns the caching budget to  $\mathsf{H}_2$  and the most recent KV, except for full KV cache, we consider the "Local" strategy as a baseline method. In addition, we also provide two different variants of Sparse Transformers (strided and fixed) as strong baselines. Also, the full
KV cache with fewer shots (0/1-shot) prompts are considered as the baseline, which has a similar sequence length of the 5-shot tasks with  $20\%$  KV cache budget.
Main Results. We evaluate LLMs with KV cache budget ranging from  $4\%$  to  $100\%$  on 5-shot downstream tasks. Results are summarized in Figure 4 and Table 1& 2. The following observations can be drawn: (1) With different KV cache budgets, our  $\mathsf{H}_2\mathsf{O}$  demonstrates consistent and significant improvements against the "Local" strategy across various model sizes, model types, and downstream tasks. We can draw sim
Table 1: Quantatively comparison between  $\mathsf{H}_2\mathsf{O}$  with Full methods of different number of shots.
Methods	PiQA	COPA	OpenbookQA	Winogrande
Full	80.09	81.00	44.80	71.51
0-shot Full	78.89	76.00	41.40	70.00
1-shot Full	79.11	76.00	43.60	70.24
Local	57.94	56.00	28.40	51.30
H2O	79.22	85.00	43.80	71.67
ilar conclusions comparing  $\mathsf{H}_2\mathsf{O}$  with other baselines like Sparse Transformer; (2) Meanwhile, with less than  $20\%$  KV cache budget(i.e., more than  $5\times$  memory reduction),  $\mathsf{H}_2\mathsf{O}$  achieves comparable performance as the model with full KV embeddings; (3)  $\mathsf{H}_2\mathsf{O}$  with  $20\%$  KV cache budget approximately uses 1.2 samples per input and show consistent improvement over zero-shot and one-shot full model that use 1 and 2 samples, respectively. (4) Our  $\mathsf{H}_2\mathsf{O}$  shows consistent effectiveness in the more challenging long sequence generation tasks, XSUM, and CNN/Daily Mail.
Analysis. Since the evicted KV will not be seen in the future steps, dropping certain critical KV embeddings can cause a severe functional collapse, resulting in significant performance degradation, e.g., in {LLaMA-13B, XSUM} {LLaMA-7B, CNN/Daily Mail}, the "Local" strategy collapses at  $60\%$  budgets while our  $\mathsf{H}_2\mathsf{O}$  can still match the full cache performance with
Table 2: Results of different sparsification methods w. or w.o. H₂. Experiments are conducted with OPT-30B with 20% KV cache budget.
Models	COPA	OpenBookQA	PiQA	Winogrande
Full	85.00	43.20	78.51	70.24
Local w.o. H2	48.00	25.20	55.82	49.17
Local w. H2	84.00	43.00	78.45	69.06
Sparse Transformer (strided) w.o. H2	50.00	24.60	56.20	47.59
Sparse Transformer (strided) w. H2	83.00	42.60	78.24	69.61
Sparse Transformer (fixed) w.o. H2	61.00	23.80	58.60	49.88
Sparse Transformer (fixed) w. H2	76.00	41.40	77.80	64.96
20% budgets. In some tasks, our methods even surpass the baseline models, which demonstrates a regularization effect of our  $\mathsf{H}_2\mathsf{O}$ . For example, in {OPT-66B, RTE}, {OPT-30B, MathQA} and {GPT-NeoX-20B, XSUM}, our  $\mathsf{H}_2\mathsf{O}$  achieves an extra performance improvement of 0.73%, 0.64% and 0.18 with 20% KV cache budget, respectively. These consistent results validate the effectiveness of our  $\mathsf{H}_2\mathsf{O}$  framework.
Enhancing Baseline Techniques. Importantly, we observe other sparsification baselines fail under an extremely low cache budget while combining the most recent KV embeddings with the ones of heavy hitters successfully achieves comparable performance as using full KV embeddings. From Table 2, we can observe that both "strided" and "fixed" sparse attention fail under  $20\%$  KV cache budgets, encountering a significant performance drop (up to  $35\%$  compared with the full cache). After combining with  $\mathsf{H}_2$ , both approaches reach a similar performance as using full KV embeddings.
5.2 Heavy Hitter for High-Throughput Generative Inference
Table 3: Generation throughput (token/s) on a T4 GPU with different systems. In the sequence length row, we use “512 + 32” to denote a prompt length of 512 and a generation length of 32. “OOM” means out-of-memory. The gray text in the bracket denotes the effective batch size and the lowest level of the memory hierarchy that the system needs for offloading, where “C” means CPU and “G” means GPU.
Seq. length	512+32	512+512	512+1024
Model size	6.7B	30B	6.7B	30B	6.7B	30B
Accelerate	20.4 (2, G)	0.6 (8, C)	15.5 (1, G)	0.6 (8, C)	5.6 (16, C)	0.6 (8, C)
DeepSpeed	10.2 (16, C)	0.6 (4, C)	9.6 (16, C)	0.6 (4, C)	10.1 (16, C)	0.6 (4, C)
FlexGen	20.2 (2, G)	8.1 (144, C)	16.8 (1, G)	8.5 (80, C)	16.9 (1, G)	7.1 (48, C)
H2O (20%)	35.1 (4, G)	12.7 (728, C)	51.7 (4, G)	18.83 (416, C)	52.1 (4, G)	13.82 (264, C)
Table 5: Generation throughput and latency on an A100 GPU. In the sequence length row, we use “ $7000 + 1024$ ” to denote a prompt length of 7000 and a generation length of 1024. “OOM” means out-of-memory.
Seq. length	Model size	Batch size	Metric	FlexGen	H2O (20%)
7000+1024	30B	1	latency (s)	57.0	50.4
5000+5000	13B	4	latency (s)	214.2	155.4
2048+2048	6.7B	24	latency (s)	99.5	53.5
2048+2048	6.7B	24	throughput (token/s)	494.1	918.9
2048+2048	6.7B	64	throughput (token/s)	OOM	1161.0
We implement our KV cache eviction policy in a state-of-the-art inference engine, FlexGen [19], and report the throughput and latency improvements.  $\mathrm{H}_2\mathrm{O}$  is orthogonal to existing optimizations in FlexGen, such as offloading and quantization, so they can be combined to achieve better performance.
Setup We conducted experiments on two GPUs: an NVIDIA T4 (16GB) GPU and an NVIDIA A100 (80GB) GPU. On the T4 GPU, we evaluate the generation throughput following the settings in the FlexGen paper. The evaluated models are OPT-6.7B and
Table 4: Results of generation throughput (token/s) on a T4 GPU with different systems on real-world datasets, XSUM.
Model size	6.7B	30B
Accelerate	11.98 (1, G)	0.23 (2, C)
DeepSpeed	3.52 (6, C)	0.31 (2, C)
FlexGen	10.80 (1, G)	3.29 (44, C)
H2O (20%)	30.40 (1, G)	6.70 (180, C)
OPT-30B. When the model and KV cache do not fit into a single GPU, we turn on CPU offloading. The results of both pure GPU and GPU with CPU offloading are reported. All the speedup results are tested in an end-to-end setting, including both the pre-filling and generation phases. And it includes the time for constructing the  $\mathrm{H}_2\mathrm{O}$  KV cache. We use synthetic datasets where all prompts are padded to the same length. The system is then required to generate the same number of tokens for each prompt. We test different combinations of prompt and generation lengths. We also test our method on real-world datasets (XSUM) for further assessment. The evaluation metric is generation throughput, which is the number of generated tokens / (prompt time + decoding time). We use DeepSpeed ZeRO-Inference [17], Hugging Face Accelerate [18], and FlexGen [19] as baselines. On the A100 GPU, with more GPU memory, we evaluate the performance of the systems with sequence lengths up to 10K. Although OPT is only trained on 2K sequence length, we benchmark the throughput and latency performance to show the potential of  $\mathrm{H}_2\mathrm{O}$  for better models in the future.
Results. Table 3& 4 shows the generation throughput of all systems on the T4 GPU. With our KV cache eviction policy, the memory usage is reduced, which brings two advantages: 1) we can use a much larger batch size; 2) we can make a setting from requiring offloading to not requiring offloading. As shown in Table 3& 4,  $\mathsf{H}_2\mathsf{O}$  with a  $20\%$  budget improves the generation throughput over FlexGen, DeepSpeed, and Accelerate by up to  $3\times$ ,  $29\times$ , and  $29\times$ , respectively, across both synthetic and real-world dataset. The results on the A100 GPU with sequence lengths from 4K to 10K are listed in Table 5. With the same batch size,  $\mathsf{H}_2\mathsf{O}$  can reduce the latency by  $1.1 - 1.9\times$  compared to FlexGen. Additionally,  $\mathsf{H}_2\mathsf{O}$  saves memory so it allows a larger batch size, which brings  $2.3\times$  improvement on generation throughput for OPT-6.7B.
5.3 Ablation Results
We present extensive ablation studies of  $\mathsf{H}_2\mathsf{O}$  on (1) infinite-length input, (2) different number of shots, (3) compatibility with quantization methods on KV cache, and (4) dissecting the effectiveness of different components. We find a surprising property of  $\mathsf{H}_2\mathsf{O}$  - it not only improves the efficiency of LLMs, but also increases the diversity of the generated text.
Q1: Can  $\mathsf{H}_2\mathsf{O}$  empower LLMs to process infinite-length inputs? A1: Effective generation with sequence length up to four million tokens. Some recent works [52, 53] demonstrate the possibility of handling infinite-length inputs, a notable challenge in current LLMs. These methods employ an attention sink that retains the first few tokens and applies position rolling in the KV cache, empowering LLMs to process infinite-length inputs. Inspired by this progress, we further implement our  $\mathsf{H}_2\mathsf{O}$  for infinite-length inputs. Figure 5 showcases the positive results of  $\mathsf{H}_2\mathsf{O}$ , i.e.,  $\mathsf{H}_2\mathsf{O}$  can empower LLMs to tackle input with length up to four mil
lion tokens, achieving a better performance (lower perplexity) than the original StreamLLM method [52] across various cache size. Further comparisons are reported in Appendix C.4.
Q2: Does the number of shots during inference effects the effectiveness of  $\mathsf{H}_2\mathsf{O}$ ? A2: Effective across zero-shot to ten-shots inference. We further examine  $\mathsf{H}_2\mathsf{O}$  under different numbers of shots during inference, and the results are reported in Table 10 and Figure 8. With different shots inference, our  $\mathsf{H}_2\mathsf{O}$  achieves matching performance (difference less than  $1.00\%$ ) as the full model across different downstream tasks. The "Local" strategy encounters significant performance degradation (up to  $37.00\%$ ). Such results demonstrate the effectiveness of our  $\mathsf{H}_2\mathsf{O}$  under different inference scenarios. More details about zero-shot and one-shot inference are reported in Appendix C.3.
Q3: Compatible with Quantization? A3: Yes. To pursue further efficiency, we show the compatibility of  $\mathsf{H}_2\mathsf{O}$  with another orthogonal approach, i.e., quantization in Table 6. We use OPT-30B as our base model and COPA, OpenBookWA, and PiQA as evaluation tasks. Intuitively sparsity and quantization are highly related so combining them might introduce larger errors. Surprisingly the combination almost always achieves better accuracy than  $\mathsf{H}_2\mathsf{O}$  or quantization alone. Experiments about throughput improvement are detailed in Appendix C.2.
Figure 5: (Upper) streaming with  $\mathsf{H}_2\mathsf{O}$  to handle inputs with sequence lengths of four million tokens. (Bottom) Perplexity comparison between the original StreamLLM method and our  $\mathsf{H}_2\mathsf{O}$ , results are collected on the first text sample of PG-19 [54].
Q4: When does  $\mathrm{H}_2\mathrm{O}$  match the baseline with full KV
embeddings? A4: With both  $\mathsf{H}_2$  and the recent tokens. We investigate the separate effects of KV embeddings of  $\mathsf{H}_2$  and the local tokens. We conduct experiments on 4 tasks with OPT-13B and OPT-30B. For each task, we compare the performance of three KV cache eviction policies, including only the KV embeddings of  $\mathsf{H}_2$ , only the ones of local tokens, and our  $\mathsf{H}_2\mathsf{O}$  that keep both. As shown in Table 9, only retaining the embeddings of  $\mathsf{H}_2$  or local tokens can't maintain a similar performance as the model using full embeddings, with a performance degradation from  $2.85\%$  to  $22.75\%$ . Incorporating both components, our  $\mathsf{H}_2\mathsf{O}$  successfully retains the baseline performance with full embeddings. Besides, the model with only  $\mathsf{H}_2$  shows a consistent improvement against the one with only local tokens, which indicates  $\mathsf{H}_2$  might contribute more to maintaining the performance.
Q5: Extra benefits from  $\mathsf{H}_2\mathsf{O}$ ? A5: Increased diversity of generated text. Besides all the benefits of our  $\mathsf{H}_2\mathsf{O}$ , we also observe an bonus introduced by  $\mathsf{H}_2\mathsf{O}$ , i.e., the improved diversity of generated content. The results are reported in Appendix C.1. Given the same prompts, we visualize the generated text of the models with different KV cache budgets. Compared with the model of full KV cache, our  $\mathsf{H}_2\mathsf{O}$  can generate sentences with fewer repeated words and more creativity.
In this paper, we study one of the key bottlenecks of LLM deployment, KV cache, particularly for long-content and large-batch generation applications. We propose  $\mathsf{H}_2\mathsf{O}$ , a simple KV cache eviction policy for significantly reducing its memory footprint. The main insight of our approach is the recognition of a subset of tokens, known as Heavy Hitters, which contribute the most value when computing attention scores. We formulate the KV cache eviction as a dynamic submodular problem and provide the theoretical guarantees for our algorithm. Through extensive evaluations, we demonstrate that  $\mathsf{H}_2\mathsf{O}$  can significantly improve end-to-end throughput and decrease latency in wall-clock time, without compromising the generation quality of LLMs across a variety of tasks.
Ying Sheng and Clark Barrett are partly supported by NSF-2110397 and the Stanford Center for Automated Reasoning. Z. Wang is in part supported by a Google Research Scholar Award and the NSF AI Institute for Foundations of Machine Learning (IFML).
[141] Yuzhou Gu, Zhao Song, and Lichen Zhang. A nearly-linear time algorithm for structured support vector machines. arXiv preprint arXiv:2307.07735, 2023.
[142] Lianke Qin, Zhao Song, Lichen Zhang, and Danyang Zhuo. An online and unified algorithm for projection matrix vector multiplication with application to empirical risk minimization. In AISTATS, 2023.
[143] Zhao Song, Mingquan Ye, and Lichen Zhang. Streaming semidefinite programs:  $O(\sqrt{n})$  passes, small space and fast runtime. arXiv preprint arXiv:2309.05135, 2023.
[144] Haotian Jiang, Yin Tat Lee, Zhao Song, and Lichen Zhang. Convex minimization with integer minima in  $\widetilde{O}(n^4)$  time. In ACM-SIAM Symposium on Discrete Algorithms (SODA), 2024.
[145] S Cliff Liu, Zhao Song, Hengjie Zhang, Lichen Zhang, and Tianyi Zhou. Space-efficient interior point method, with applications to linear programming and maximum weight bipartite matching. In ICALP. arXiv preprint arXiv:2009.06106, 2020.
A More Implementation Details 21
B Extended Related Works, Discussions, and Limitations 22
B.1 Extended Related Works 22
B.2 Discussions and Limitations 22
C Extended Experiments 24
C.1 Increased Diversity of Generated Text 24
C.2 Throughput Improvement of  $\mathsf{H}_2\mathsf{O}$  Combining with Quantization 25
C.3 Effectiveness on Zero-shot and One-shot Inference 26
C.4 Comparison with StreamingLLM 26
C.5 Enhancing the "Top-K" Baseline 27
C.6 Separate Effects of Heavy-Hitter and Local Tokens 27
C.7 Performance of Inference with Different Number of Shots. 27
C.8 Heavy-Hitter in Attention Blocks 28
C.9 Comparison with SpAtten 28
C.10 Heavy-Hitter in MLP Blocks 28
D Theoretical Analysis 33
D.1 Notations 33
D.2 Submodular 34
D.3 Dynamic Submodular 35
D.4 Static Attention 35
D.5 Recursive Attention Definition 35
D.6Eviction Policy 37
D.7 Explaining Submodular Diminishing Return Property in Attention Scheme 38
D.8 Submodular: High Level Ideas 39
D.9 Robust Greedy with Error Propagation 39
D.10 Robust Submodular and Adding Items 40
D.11 Universal Conditions 41
D.12 Induction Lemma for Exact Function 41
D.13 Induction Lemma for Approximate Function 42
D.14 Theoretical Result 43
D.15 Extended Related Work for Theoretical Attention Problems 44
D.16 Sparsity Preserving 44
D.17 Definition of Loss Function 45
D.18 Gradient 46
D.19 Hessian 47
D.20 Hessian is Positive Definite 47
D.21 Hessian is Lipschitz 48
D.22 Greedy Type Algorithm 49
In this section, our goal is to provide the details of system implementation (mentioned in Section 4.2) and experiment settings (mentioned in Section 5), as well as the pseudocode.
System Details. We implement  $\mathsf{H}_2\mathsf{O}$  on top of FlexGen. FlexGen is a white-box implementation of OPT models, and we have done some surgery on handling the KV cache. Specifically, for the given parameter  $K$ , we always maintain a list of KV cache with the first  $K$  entries as heavy hitters, and the last  $K$  entries as most recent tokens. In order to avoid data movement in memory, the memory for KV cache is preallocated. We use a circular queue to update the last  $K$  entries efficiently.
Experiment Details. Our study involves the evaluation with varying sizes of KV cache that encompass  $4\%$ ,  $10\%$ ,  $20\%$ , and  $60\%$  of the prompt's length. We select two tasks (XSUM and CNN/Daily Mail) from the HELM framework [16] and present the performance based on 1000 test samples. Additionally, we employ the lm-eval-harness framework [15] for six other tasks (COPA, MathQA, OpenBookQA, PiQA, RTE, and Winogrande). For the tasks derived from the HELM framework, we report the performance for zero-shot inference, while for the tasks from lm-eval-harness, we default to conduct five-shot inference.
Pseudocode. We show a simplified pseudocode below to demonstrate our implementation skeleton. The function generation_loop() is the base loop in FlexGen that controls prefetch and overlap the I/O streams and computation. Then in function compute_layer(), the function attention_forward() will be called for attention layers. During prefill iterations, the function computeattention() will return  $K$  heavy hitters and  $K$  recent tokens if the prompt has a length greater than or equal to  $2K$ , otherwise return all the KV cache. The function compute attention () will return new KV cache and the indices for evicted entries during decoding iterations, which would be the place to implement a customized eviction strategy. (If the current number of stored KV cache is less than  $2K$ , the eviction will be ignored.) During each decoding iteration, the oldest one among the last  $K$  tokens (if we have no less than  $K$  tokens' KV cache stored) will be removed from the reserved last  $K$  entries, one heavy hitter will be evicted for each head, and the newest token will be added to the KV cache with a position in the last  $K$  entries. This happens in the function store_cache().
def generation_loop(..):
    # Prologue
    ...
    # Generate
    for i in range(gen_len):
        for j in range(num_layers):
            for k in range(numgpu_batches):
                load_weight(i, j+1, k)
                load_cache(i, j, k+1)
                store Hidden(i, j, k-1)
                load Hidden(i, j, k+1)
                compute_layer(i, j, k)
                store_cache(i, j, k-1)
            sync()
    # Epilogue
    ...
# h is the hidden states (activations)
def attention_forward(h, ...):
    # the read/write buffer are intermediate stops for prefetching
    if prefetch:
        h, new_k_cache, new_v_cache = computeattention(h, ...) 
        # select K heavy hitters and K recent tokens
        new_k_cache, new_v_cache = select(new_k_cache, new_v_cache, K)
        cache_write_bufferstore(new_k_cache, new_v_cache)
    else:
        k_cache, v_cache = cache_read_buf.pop()
        # evict_ids track the entries that will be evicted
        h, new_k_cache, new_v_cache, evict_ids =
            computeattention(h, k_cache, v_cache, ...)
cache_write_buffer.store(new_k_cache，new_v_cache，evict_ids) return h   
def store_cache(）: ifprefill: #store cache directly else:k_new，v_new，evict_ids  $=$  cache_write_buffer.pop() #circular queue for the last K entries #extract the index for the oldest token at i-th iteration oldest  $= ((i - 1)\% K) - K$  #update the KV cache (k_home and v_home) cache_replace(k_home，evict_ids，k_new，K,oldest) cache_replace(v_home，evict_ids，v_new，K,oldest)
B Extended Related Works, Discussions, and Limitations
The goal of this section is to first introduce more background and related works for Section 2, then describe some previous attempts in our experiments as well as discuss the social impact and limitations of this work.
B.1 Extended Related Works
Quantization, Pruning, Distillation for Inference. Previously, model compression algorithms have been extensively investigated as a viable approach for mitigating the computational resource requirements of model inference. These algorithms can be broadly categorized into three groups: (1) quantization [55, 56, 57, 58], which involves mapping model parameters or activations from high-precision data types to low-precision counterparts, such as using 8-bit integers instead of the commonly employed 32-bit floating point format; (2) pruning or sparsity [59, 60, 61, 62], which aims to eliminate unnecessary neurons or weights within the models; (3) and distillation [63, 64, 65, 66] where predictions from larger models are utilized as supervised information to train smaller models.