ArkVale: Efficient Generative LLM Inference with Recallable Key-Value Eviction

Authors: renze  chen
zhuofeng  wang
beiquan  cao
tong  wu
size  zheng
xiuhong  li
xuechao  wei
shengen  yan
meng  li
yun  liang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 14:06:49

Abstract: Large Language Models (LLMs) are widely used in today's tasks of natural language processing. 
To support applications like multi-turn chats, document understanding, and content generation, models with long context lengths are growing in importance.
However, managing long contexts brings substantial challenges due to the expansion of key-value cache (KV cache). Longer KV cache requires larger memory, limiting the batch-size thus decreasing throughput. Also, computing attention over long KV cache incurs more memory access, hurting the end-to-end latency.
Prior works find that it is sufficient to use only the recent and high-impact tokens for attention computation, allowing the eviction of less vital tokens to shrink cache size.
Nonetheless, we observe a dynamic shift in token importance across different decoding steps. Tokens initially evicted might regain importance after certain decoding steps.
To address this, we propose ArkVale, a page-based KV cache manager that can recognize and recall currently important tokens evicted before. We asynchronously copy the filled page into external memory (e.g., CPU memory) as backup and summarize it into a much smaller digest by constructing the bounding-volume of its keys. Before attention computation, we measure all pages' importance based on their digests, recall the important ones, evict the unimportant ones, and select the top-ranked pages for attention computation. 
Experiment results show that ArkVale performs well on various long context tasks with negligible accuracy loss under 2k$\sim$4k cache budget and can improve decoding latency to $2.2\times$ and batching throughput to $4.6\times$ because it applies attention on only a small subset of pages and reduce per-sample memory usage of KV cache.

Main paper: 
 introduction: 
 Large Language Models (LLMs) are rapidly gaining universal presence, underpinning a myriad of natural language processing applications, including dialogue systems [2, 52, 16] , document summarization [62, 24] , code completion [13, 46] , and question answering [31] . The context length supported by large models is also growing progressively to support more applications like multi-turn chat [2, 52, 16 ] and text summarization [28, 68, 22] . This shift has seen an expansion from an initial range of up to 16k, advancing steadily towards lengths of 32k [1] , 128k [43] , and even 2048k [21] .However, handling long contexts poses significant challenges. The key-value cache (KV cache) of LLMs scales with both the batch size and the length of historical context. For each token generated by an LLM, the query must attend to the entire set of context key-values. On one hand, longer KV caches occupy more memory spaces, limiting batch-size as well as throughput. On the other hand, computing attention with long KV cache incurs more memory accesses, hurting the inference latency.Previous works [58, 64, 37, 23, 40, 6] reveal sparsity in KV cache, which means that only a subset of tokens in the KV cache significantly influence the accuracy of attention computation. They propose two characteristics of token importance: locality and persistence. Locality implies that recent tokens tend to be important for current attention, while persistence suggests that tokens previously deemed important are likely to retain their importance over time. They devise algorithms to retain recent tokens [58, 37, 64] (Figure 1 (b)) or vital tokens [58, 37, 64, 40, 23, 6] (Figure 1 (b)) while evicting older or less crucial ones to reduce memory usage of KV cache and memory accesses of attention.In this work, we delve deeper into the dynamism of token importance, observing that tokens previously deemed unimportant may regain importance over time and vice versa. Previous works ignore this and may risk permanently discarding the tokens that are vital later on, inadvertently decreasing model accuracy. To address these challenges, we propose ARKVALE, inspired by vLLM [36] to organize tokens into pages for fine-grained management of KV cache. Upon filling a KV-page, we asynchronously copy it to CPU memory as a backup, and summarize/compress it into a much smaller digest by constructing a bounding-volume of the keys of the KV-page. Before each attention computation, we dynamically estimate the importance of each page based on their digests and current query, and selectively recall & evict some pages based on their importance scores (Figure 1 (d) ).Our contribution can be summarized as follows:• We characterize the dynamism of the importance of tokens in the KV cache of LLM, and find that some unimportant tokens may regain importance over time.• We propose ARKVALE, a KV cache manager which organize tokens into pages and dynamically evict & recall them based on their importance.• We propose a method based on bounding-volume to summarize the pages and estimate the importance of them with a given query.We evaluate ARKVALE against many state-of-the-art KV cache eviction works on various longcontext benchmarks. Experimental results show ARKVALE performs well on all the long context tasks with negligible accuracy loss under a cache budget of 2k∼4k and speedup model decoding latency up to 2.2× (1.7× in average) and batching throughput up to 4.6× (3.5× in average) in long context scenarios. Our code is now available at https://github.com/pku-liang/ArkVale.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
related work: 
 Many works have been proposed to extend context window of pre-trained models. One prevalent method is the integration of Rotary Position Embeddings (RoPE) [49] . Fine-tuning RoPE's scaling enables the Llama-2 model [55] , initially handling 4k tokens, to be expanded to support 32k tokens in LongChat [1] and 128k tokens in Yarn-Llama-2 [43] . Recently, LongRoPE [21] pushes this boundary even further to accommodate up to 2048k tokens. However, as models become more capable of dealing with long context, they also face new challenges in terms of memory usage and inference efficiency. ARKVALE is designed to address such issues.Some training-aware methods have been proposed to handle these problems. Multi-Query Attention (MQA) [47] and Group-Query Attention (GQA) [7] aim to train LLMs with fewer attention heads in KV cache. Works like RWKV [42] , RetNet [50] , and Mamba [25] propose Linear Attention/RNN to limit KV cache size. Sparse attention architectures [61, 10, 56, 17, 45, 33, 53] design special sparse pattern of attention during training to control the KV cache size during inference. But these works require much training effort for pre-training or fine-tuning while ARKVALE is a training-free method.There are also some post-training methods, which evict tokens in KV cache to reduce memory usage and accelerate attention computation. StreamingLLM [58] and LM-Infinite [26] only keep initial tokens and recent tokens to maintain a fixed-size KV cache. Methods like H2O [64] , Scissorhands [37] and TOVA [40] keep important tokens in KV cache based on the historical or current attention scores. FastGen [23] categorizes tokens and employs a more nuanced approach for choosing which KV cache tokens to preserve. Keyformer [6] improves the eviction score function by leveraging the Gumbel distribution. Q-Hitter [63] combines KV-cache quantization with KV-cache eviction. These methods rely on historical data to dictate cache eviction and may risk discarding the tokens that are important in the future. Works like IceFormer [38] , SparQ [5] , ALISA [65] , and Quest [51] use post-training sparse attention to alleviate such issue but they require all KV cache residing in memory and thus cannot control KV cache size to optimize memory usage. Many scheduling-based methods [19, 18, 60, 15, 66, 14, 36, 48, 32, 27, 67] are also employed to optimize memory efficiency for LLMs or other DNNs, aiming to enhance memory usage and execution latency through proper computation partitioning and scheduling.3 BackgroundQuestion: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We discuss it in Section 7.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. " • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
attention computation and generative inference of llm: 
 The attention computation in a typical LLM involves mixing token-level information of queries Q ∈ R sq×d and keys K ∈ R s kv ×d with values V ∈ R s kv ×d , where d is the hidden-dimension and s q (s kv ) means the number of query (key/value) tokens. The attention computation can be formulated as:S := (Q • K ⊤ / √ d) ∈ R sq×s kv , P := softmax(S) ∈ R sq×s kv , O := (P • V) ∈ R sq×d, where P is often called "attention map/scores" and P i,j reflects the importance of key/value token j to query token i. Then, the generative inference process of LLM mainly consists of two stages: the prefill/prompt stage and the decoding/generation stage. The prefill stage takes a prompt sequence with length s in as input and caches the keys & values computed by each layer of the LLM. The decoding stage uses and updates the KV cache to generate tokens step-by-step, where the current token depends on past tokens' keys & values stored in the KV cache. For the attention in prefill stage, s in = s q = s kv , while s q = 1 and s kv = #past-tokens + 1 in each decoding step.For a decoding step with batch-size as b, number of transformer layers as l, history sequence length as s, hidden-dimension as h, and data type as float16, the KV cache requires 4bsh bytes memory accesses for each attention computation and occupies 4blsh bytes of memory in total. The latency & memory breakdown of a decoding step of model LongChat-7b-v1.5-32k [1] (l = 32, h = 4096) with b = 8 and s = 2 9 , 2 10 , 2 11 , 2 12 , 2 13 , 2 14 , 2 15 is shown in Figure 2 . As the sequence length increases, accessing the KV cache incurs significant overhead, thereby impacting the end-to-end latency of inference. Besides, as the sequence length increases, the KV cache itself consumes more memory space; for instance, with just a context length of 16k, a batch size of 8 nearly saturates the 0.5k 1k 2k 4k 8k 16k 32k Sequence Length 80GB memory of an A100 GPU. This not only hampers the use of larger batch sizes but also poses challenges for deploying models with more parameters.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Justification: Due to space constraints, we don't delve into the broader impacts within the main text. In a nutshell, our work positively contributes to the wider application of LLMs by reducing inference latency and memory usage while maintaining accuracy. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
token-level sparsity of kv cache: 
 To tackle the aforementioned issues, we aim to exploit the sparsity of the KV cache. We collect data by running LongChat-7b-v1.5-32k [1] on LongBench [9] . Figure 3 illustrates the sparsity of the KV caches of different layers. We organize the tokens in KV cache into pages (with page-size=32), and rank the pages of each layer based on the highest attention scores of the tokens within each page for each decoding step. We can observe that, in each layer except the initial one, less than 10 pages (320 tokens) contribute more than 99% scores in sum, achieving over 95% sparsity. Thus we can retain only a subset of KV tokens for attention. 3 : Token-level sparsity of KV cache. We organize the tokens in KV cache into pages (with page-size=32) and rank the pages of each layer based on the highest attention (softmax) scores of the tokens within each page.Previous efforts leverage the sparsity of KV cache to control its size by retaining recent tokens while discarding older ones or selectively evicting less important tokens based on their historical attention scores, shown in Figure 1 (b) (c ). However, we find that the importance of tokens in the KV cache can dynamically change over time, as shown in Figure 4a . For example, page 256 initially holds little importance but later becomes crucial, with its significant role (at position 12620) occurring nearly 4500 tokens after its sequence position (at position 8192). Previous methods risk prematurely discarding such pages to save KV cache space. To address this, we propose ARKVALE, a shown in Figure 1 (d), a method designed to recognize shifts in token importance and properly recall vital tokens as well as evict unimportant ones, thereby preventing substantial drops in accuracy.5 Techniques of ARKVALE 
system design: 
 Figure 5 shows the design of ARKVALE. Inspired by vLLM [36] , ARKVALE arranges KV cache tokens into pages, enabling coarse-grained management of tokens. As shown in Figure 5 (a), once a page is filled, its keys & values are asynchronously moved from GPU to external memory (CPU memory typically) for backup, with keys summarized into a digest kept on the GPU (explained in § 5.2). Prior to attention computation, ARKVALE gauges the importance of each page (whether cached or evicted), using the query and page digests ( ). If any top-k pages were evicted before, they will be recalled from backup, and bottom-ranked pages in GPU will be evicted for exchange (Figure 5 (d) ).Depending on the cache-size and the number of pages selected for attention, the overhead of page recall may be different. To validate the feasibility of our approach and ensure that recalls do not incur excessive overhead, we uniformly sample four instances from each dataset included in LongBench [9] and gather data from the inference process of LongChat-7b-v1.5-32k [1] to simulate eviction and recall scenarios. We configure the page-size p = 32 and vary both the cache capacity c (pages) and the attention size k (the number of top-ranked pages selected for attention). The results, depicted in Figure 4b , reveal that as long as k < min(40, c/2), we can limit the number of recalls per decoding step to under 10, amounting to a total of 5 MB of data transfer. This transfer takes merely hundreds of microseconds over contemporary CPU-GPU communication interfaces, which is one or two orders of magnitude faster than the end-to-end latency of each decoding step (usually a few (a) Bounding-sphere need to record the center c and the radius r.(b) Bounding-cuboid need to record the max-vector b (2) = max 5 i=1 k (i) and the min-vector b (4) = min 5 i=1 k (i) .Figure 6 : Summarize page keys {k (i) } 5 i=1 into their bounding-volume (sphere/cuboid). We can estimate the max-dot-product between query q and keys {k (i) } 5 i=1 using the bounding-volume.For every filled page, we maintain a digest as a compression of it. The digest is much smaller than the page itself, and can be used to estimate the importance of each page regardless of whether the page has been evicted before attention computation. This estimation enables us to rank pages to selectively recall or evict pages based on their importance.Given a page with n keys K = {k (i) } n i=1 (k (i) ∈ R d ), and a query q ∈ R d , we measure its importance using the maximum dot product between the query and the keys:I(q, K) = max k∈K q•k. For two sets of keys, K 1 and K 2 , I(q, K 1 ) > I(q, K 2 ) implies that there's at least one token in K 1 with higher attention (softmax) scores compared to every token in K 2 .It's evident that argmax k∈K q • k must lie on a vertex of the convex hull of the point set K. Hence, we can approximate I(q, K) by constructing a convex covering set of K. We leverage the concept of bounding-volume [3, 34] to summarize the keys K of a page. A bounding-volume for a set of points is a closed region that completely contains all of them, which is widely used in computer graphics [34, 29, 11, 12] .Bounding-sphere. We can cover the keys K with a sphere, needing to store only a center c ∈ R d and a radius r ∈ R as the digest. Given a query q ∈ R d , we can estimate the importance I(q, K) leveraging c and r. Among the points k ′ on the surface of the sphere, those for which k ′ -c is parallel to q yield the maximum dot product with q. Thus we have k ′ = r |q| q, leading to:I(q, K) = max k∈K q • k ≈ q • (c + r |q| q) = q • c + r|q| (1)An example is shown in Figure 6a . Finding the minimal bounding sphere is complicated; hence, we employ an efficient approximation method: we take the center of the Axis-Aligned Bounding Box (AABB) [12, 11] of K as our sphere center: c = 1 2 (min k∈K k + max k∈K k) (min and max are element-wise operations in this paper). For the radius r, three alternatives are considered:EQUATIONwhere |...| means L2-norm. Strictly speaking, only r max guarantees to enclose all points in K, but the latter two can avoid overestimating I(q, K) in practice, as shown in Figure 7 and discussed in §6.2.Bounding-cuboid. We can also cover the keys K directly using an Axis-Aligned Bounding Box (AABB) [12, 11] . It needs to store the max-vector b max = max k∈K k and min-vector b min = min k∈K k as the digest. Given a query vector q ∈ R d , we can estimate the importance I(q, K) of K with the boundaries of the box. This is achieved by summing the maximum product by the boundary values and q across all dimensions:I(q, K) = max k∈K q • k ≈ d i=1 max k∈K q i k i = d i=1 max(q i b max i , q i b min i ) = 1 • max(q ⊙ b max , q ⊙ b min )(3) Here, 1 ∈ R d denotes a vector of ones, and ⊙ means element-wise multiplication. An example is shown in Figure 6b . Similar to the bounding-sphere approach, we can employ different "radius vectors" to define cuboids with varying sizes:EQUATIONwhere abs(...) computes element-wise absolute values. For each r, we can derive b max = c + r and b min = cr, where c = 1 2 (min k∈K k + max k∈K k) is the center of the AABB of K.6 Evaluation 
experimental setup: 
 We apply our method to LongChat-7b-v1.5-32k [1] and use 6 datasets from LongBench [9] for benchmarking: HotpotQA [59] , NarrativeQA [35] , Qasper [20] , GovReport [28] , TriviaQA [30] ,and PassageRetrieval [9] , along with the passkey-retrieval tasks. For comparison, we choose the state-of-the-art KV cache eviction methods including StreamingLLM [58] , H2O [64] , and TOVA [40] as baselines. As Figure 3 illustrates, the initial layers exhibit relatively low sparsity; hence, we do not apply ARKVALE and baselines to the first two layers of model.Our experiment platform comprises an Intel(R) Xeon(R) Gold 6348 CPU @ 2.60GHz and an NVIDIA A100 80GB PCIe GPU. The software stack includes CUDA version 12.3, PyTorch [41, 8] version 2.3.0, and HuggingFace Transformers [57] version 4.40.0. We implement ARKVALE on top of Huggingface Transformers, with CUTLASS [54] , FlashInfer [60] , and RAFT [44] for certain kernels. Recall Accuracy centroid sphere-max sphere-center sphere-mean cuboid-max cuboid-center cuboid-meanFigure 7 : Recall accuracy (the proportion of pages predicted to be among the top-k that indeed belong to the top-k) of different estimation methods. "Centroid" is the baseline which just uses the centroid of keys to estimate the max-dot-product with given query.We begin by evaluating the accuracy of the page summarization & importance estimation methods discussed in §5.2. Using data from the LongBench [9] datasets employed in our experiments, we simulate each decoding step to collect both the actual page rankings and the estimated rankings derived using the digests. For varying values of k, we define recall accuracy as the proportion of overlap between the estimated top-k page set E k and the true top-k page set R k , that is:|E k ∩ R k |/|R k |.Figure 7 illustrates the recall accuracy of different estimation techniques across various k values. The "centroid" method, which uses the centroid [4] (element-wise average of keys) of page keys as the digest and estimates page importance by directly taking the dot product of the query and the centroid, serves as a straightforward baseline. Other methods are the six introduced in §5.2: spheres with r max , r center , and r mean as well as cuboids with r max , r center , and r mean . As shown in Figure 7 , the centroid method cannot recall the top-4 pages with even 50% accuracy and achieves less than 5% accuracy in recalling the top-1 page, substantially undermining inference accuracy. Conversely, our six proposed methods guarantee at least a 60% recall accuracy, with the standout cuboid-mean method ensuring a 95% accuracy for top-1 recall and consistently over 80% for other k values. The cuboid-based methods generally outperform their sphere-based counterparts due to more retained information (two vectors for cuboids versus one vector and a scalar for spheres). Furthermore, the "mean" variants tend to perform better than others, potentially because they provide a more balanced boundary estimation, avoiding overestimation of page importance.Next, we evaluate the accuracy of ARKVALE in general long-context tasks. We employ 6 datasets from LongBench [9] as benchmarks, covering tasks such as multi-document QA with HotpotQA [59] , single-document QA with NarrativeQA [35] , and Qasper [20] , text summarization with GovReport [28] , few-shot learning with TriviaQA [30] , and synthetic tasks with PassageRetrieval [9] . The target model is LongChat-7b-v1.5-32k [1] . Our comparison baselines include the original model and the state-of-the-art KV cache eviction methods: StreamingLLM [58] , H2O [64] , and TOVA [40] . We configure four cache budget settings: 4096, 2048, 1024, and 512. Notably, H2O and TOVA's original implementations require the full attention scores matrix to maintain states, preventing the use of optimizations like Flash-Attention [19] in prefill stage and causing Out-of-Memory (OOM) for many benchmark samples. To address this, we split each input into context and question/instruction sections and adopt a two-phase prefill: the first phase uses Flash-Attention [19] to handle the lengthy context without OOM (while we also separately compute the last row of attention scores for H2O We further examine the accuracy of ARKVALE under long-range dependency scenarios in detail. We employ the passkey retrieval task [39] as our benchmark. We establish three context lengths for our tests: 10k, 20k, and 30k. For each context length, we generate 20 unique test cases, each with the passkey inserted at depths corresponding to 0%, 5%, ..., up to 95% of the text's length. Following the setup in Section §6.3, we compare ARKVALE against StreamingLLM [58] , H2O [64] , and TOVA [40] . Similarly, we adopt two-phase prefill for input texts to prevent OOM incidents of H2O and TOVA and more effectively evaluate each method's long-range dependencies handling.The results are shown in Table 1 . Baselines, which permanently evict some tokens, risk discarding passkey-related information during inference, leading to accuracy falling short of 50% and declining further as the context length increases or the cache budget decreases. In contrast, ARKVALE examines the importance of evicted tokens (pages) and promptly recalls vital ones, thereby maintaining a stable passkey retrieval accuracy above 95% across all tested context lengths and cache budgets. We set up test-cases with lengths close to 10k, 20k, and 30k from GovReport [28] in LongBench [9] and examined ARKVALE's inference latency with settings of a batch-size=4, page-size=32, and KV cache budgets of 512, 1024, 2048, and 4096. Figure 9a shows the results. Here, "Full" refers to the baseline without KV cache eviction, whereas "Estimate," "Select," and "Recall" denote the overhead for estimating page importance, choosing pages for computations and evictions, and recalling pages from CPU memory, respectively. ARKVALE outperforms the baseline across all the tested text lengths, reaching up to 2.2× speedup (1.7× in average). The main bottleneck of the baseline is the full-attention, whose latency grows linearly with seq-len, whereas our attention's latency is mainly governed by the fixed cache budget. Our extra latency overhead mostly arises from estimating importance, which grows linearly with the #pages = seq-len page-size . Recalling, conversely, adds small overhead (especially when seq-len is large and cache-budget is small), consistent with the discussion in §5.1.We further compare the achievable maximum throughput in serving scenarios of ARKVALE versus the baseline under a 40 GB memory limit for KV cache. Figure 9b shows the results. Our throughput reaches up to 4.6× (3.5× in average) to that of baseline. This achievement is not only due to the reduced inference latency but also to a decreased per sample memory usage of the KV cache, enabling larger batch-size to enhance weight sharing to lessen the average latency of Linear computations per sample. 
limitation: 
 The limitation of ARKVALE mainly lies on storing a backup for each KV cache page in external memory (typically the CPU memory). Although the latency of data transfer for backup during decoding phase can be hided by asynchronous copy, the backup latency during prefilling phase is hard to completely overlapped. Also, the backups may occupy much CPU memory. When the CPU's memory capacity is insufficient, these backups may need to be offloaded to disk storage. 
conclusion: 
 In this paper, we propose ARKVALE, a method designed for managing KV cache eviction and recall. It organizes KV cache tokens into pages, backs them up, and creates summaries, enabling the recognition of page importance to recall vital pages evicted before. Our method performs well on various long context tasks with few accuracy loss under a cache budget of 2k∼4k and speeds up decoding latency up to 2.2× and boosts decoding throughput up to 4.6× in long-context scenarios. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: This paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: In §6.1, we cite the methods we compare to and the datasets we use and specify versions of the software tools employed.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix:  


Decision: This paper looks into the problem of KV cache compression for more memory-efficient LLM inference. Existing eviction-based KV cache compression methods permanently evict tokens which may later turn out to be important as generation moves forward. The paper introduces a method to offload a subset of tokens that are not important for the time being but can recall those tokens when needed if token importance changes dynamically. The paper also takes into account of page attention, which supports easy swap-in and swap-out of tokens to second memory during inference. Overall, the reviewers are positive to the paper. The main reasons for accepting the paper are: 
1) It seems to be novel to add the mechanism of dynamically retrieving tokens back from secondary memory based on the change of token importance. 
2) Promising empirical speedups.
3) Relatively comprehensive evaluation. 

There were some concerns from reviewers about the generalizability of the approach to other models and comparison with additional baselines. However, those were addressed via the rebuttal. The paper can be improved if it can incorporate the discussion of concurrent work, as suggested by reviewer YdSG.

Reviews: review: paper_summary: paper_summary: The paper introduces ARKVALE, a novel page-based key-value (KV) cache management approach designed to optimize the performance of Large Language Models (LLMs) when dealing with long context lengths. As the demand for higher context lengths in tasks such as multi-turn chats and content generation increases, the management of extended key-value caches becomes crucial due to memory constraints and the impact on computation latency and throughput. ARKVALE addresses these challenges by dynamically managing the cache, selectively evicting less important tokens while recalling those that become relevant again at different decoding steps. This strategy leverages a mechanism that organizes tokens into pages and uses digests to estimate the importance of these pages. By summarizing pages into smaller, more manageable units, ARKVALE efficiently decides which pages to recall from external memory and which to evict, enabling focused attention computations on only the most relevant subsets of data.

The paper's experiments demonstrate that ARKVALE effectively handles various long-context tasks with minimal loss in accuracy, even under constrained cache sizes between 2K and 4K tokens. The results indicate significant improvements in model decoding latency and batching throughput, specifically up to 2.2x faster latency and 4.6x higher throughput. These enhancements are achieved by applying attention to a reduced subset of pages, thus decreasing the per-sample memory usage of the KV cache. This system not only improves the operational efficiency of LLMs but also maintains high accuracy levels, suggesting a scalable solution for managing extensive data in real-time language processing applications.
strength_weakness: strength_weakness: - Insufficient Model Evaluation:
The paper evaluates the proposed method using only the LongChat-7b-v1.5-32k model, which is relatively outdated in the current landscape. It does not demonstrate the method's generality and robustness across different model architectures (e.g., MoE, GQA) and scales (70B models).
questions: questions: Clarification on Model Choice: The paper primarily utilizes the LongChat-7b-v1.5-32k model for evaluating ARKVALE. Can the authors provide specific reasons for choosing this model over others? Additionally, how do the authors anticipate ARKVALE would perform with other, potentially larger or more recent LLM architectures?
limitations: limitations: The authors have not discussed the limitations of their work in the submission. It would be great to add sections discussing the limitations and societal impacts of ARKVALE.
paper_summary: paper_summary: The paper proposed a method to minimize the risk of KV cache eviction by efficiently and soundly offloading some of them into external memory, which is realized by page organization, page digest, and digest ranking/scoring. The method gets much better performance in context retrieval tasks compared to other KV eviction methods.
strength_weakness: strength_weakness: It's not very clear if the dynamics of importance will change along the decoding in other tasks, so the observations may be limited.
It would be interesting to see how the page size affect the methods, as vLLM's default page size is 16. Is there any reason to scale to 32?
Would the method still work when TP>1?
questions: questions: I'm wondering if there will be an open-sourced implementation/PR to vLLM?
limitations: limitations: It's desirable to see normal length tasks, e.g., GSM8K.
paper_summary: paper_summary: ARKVALE presents a page-based key-value (KV) cache manager designed to address the challenges associated with long-context processing in large language models. The main contribution is its ability to recognize and recall important tokens that were previously evicted, thereby optimizing memory usage and improving throughput and latency without significant accuracy loss. The method involves asynchronously backing up filled pages in external memory and using bounding-volume techniques to summarize and estimate the importance of these pages for efficient recall and eviction. Experiments show strong performance in terms of both accuracy and efficiency.
strength_weakness: strength_weakness: - The benchmark is only conducted on LongChat, lacking evaluations on different models such as Mistral and LLaMA-3.
- There is a lack of comparison with other methods focusing on efficiency, and some other baselines like [1,2].
- The paper lacks discussion and ablation studies on some hyper-parameters, such as top-k, page-size and the relationship between top-k, page-size and the cache budget.

[1]  Infllm: Unveiling the intrinsic capacity of llms for understanding extremely long sequences with training-free memory

[2] Snapkv: Llm knows what you are looking for before generation
questions: questions: - What is the performance of ARKVALE on Mistral-7B-Instruct-v0.2 and LLaMA-3-8B-Instruct?
- How does the performance (accuracy, latency/throughput) of ARKVALE compare to other methods [1, 2]?
- What is the process for selecting the hyperparameters, and what their influence on performance?
- Does the prefill phase use all tokens? For example, for the first token generated, is full attention used? Additionally, is this method compatible with FlashAttention/FlashDecoding?

[1]  Infllm: Unveiling the intrinsic capacity of llms for understanding extremely long sequences with training-free memory

[2] Snapkv: Llm knows what you are looking for before generation
limitations: limitations: See weaknesses and questions.
paper_summary: paper_summary: This paper proposes a paged-based KV cache manager that identifies and recalls important tokens for LLM inference, termed ArkVale. Results show that ArkVale achieves 2.2x latency and 4.6x throughput improvement on various long context tasks.
strength_weakness: strength_weakness: 1. How does the page size affect the memory consumption for the KV cache? Would smaller page size lead to potential fragmentation issues?
2. In the related work section, it would be nice if the authors could the relationship between ArkVale and some concurrent works [1,2,3].

[1] Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference, MLSys 2024.

[2] Q-Hitter: A Better Token Oracle for Efficient LLM Inference via Sparse-Quantized KV Cache, MLSys 2024.

[3] ALISA: Accelerating Large Language Model Inference via Sparsity-Aware KV Caching, ISCA 2024.
questions: questions: Please see the weaknesses above.
limitations: limitations: Please see the weaknesses above.

Comments: title: Comparison with Quest
Thank you for the response
comment: Quest [1] shares similarities with ArkVale (and some other works like InfLLM [2]) in estimating and performing topk-filtering on kv-cache at the page/block granularity. The main differences are as follows: 
- Their approach aligns more closely with post-training dynamic sparse attention methods like SparQ [3] and IceFormer [4]. These methods do not involve eviction of kv-cache and cannot save GPU memory.  
- The estimation method we employ is based on bounding-volume, whereas their estimation method is similar to the "cuboid-max" (a subclass of our approach) introduced in our paper.

[1] Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference 

[2] InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory 

[3] SparQ Attention: Bandwidth-Efficient LLM Inference 

[4] IceFormer: Accelerated Inference with Long-Sequence Transformers on CPUs
I appreciate the authors' responses; most concerns have been addressed. I will keep my evaluation for acceptance.

I also noticed that this paper is quite similar to a paper [1] presented at ICML 2024, published after the NeurIPS submission deadline. While not obligatory, discussing and comparing this work with that paper would be beneficial.

[1] Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference
Thanks for the reply, this answers my questions.
Thank you for your detailed response. I don't have any further questions. 

My only suggestion would be to incorporate the discussion of concurrent works into the final version if the paper is accepted.