ZipCache: Accurate and Efficient KV Cache Quantization with Salient Token Identification

Authors: yefei  he
luoming  zhang
weijia  wu
jing  liu
hong  zhou
bohan  zhuang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 07:50:20

Abstract: KV cache stores key and value states from previous tokens to avoid re-computation, yet it demands substantial storage space, especially for long sequences. 
  Adaptive KV cache compression seeks to discern the saliency of tokens, preserving vital information while aggressively compressing those of less importance.
  However, previous methods of this approach exhibit significant performance degradation at high compression ratios due to inaccuracies in identifying salient tokens. 
  Additionally, the compression process introduces excessive overhead, substantially increasing memory burdens and the generation latency.
  In this paper, we present ZipCache, an accurate and efficient KV cache quantization method for large language models (LLMs). 
  First, we construct a strong baseline for quantizing KV cache. Through the proposed channel-separable tokenwise quantization scheme, the memory overhead of quantization parameters are substantially reduced compared to fine-grained groupwise quantization.
  To enhance the compression ratio, we propose normalized attention score as an effective metric for identifying salient tokens by considering the lower triangle characteristics of the attention matrix. The quantization bit-width for each token is then adaptively assigned based on their saliency.
  Moreover, we develop an efficient approximation method that decouples the saliency metric from full attention scores, enabling compatibility with fast attention implementations like FlashAttention.
  Extensive experiments demonstrate that ZipCache achieves superior compression ratios, fast generation speed and minimal performance losses compared with previous KV cache compression methods. For instance, when evaluating Mistral-7B model on GSM8k dataset, ZipCache is capable of compressing the KV cache by $4.98\times$, with only a 0.38% drop in accuracy. In terms of efficiency, ZipCache also showcases a 37.3% reduction in prefill-phase latency, a 56.9% reduction in decoding-phase latency, and a 19.8% reduction in GPU memory usage when evaluating LLaMA3-8B model with a input length of 4096. Code is available at https://github.com/ThisisBillhe/ZipCache/.

Main paper: 
 introduction: 
 LLMs with the next-token-prediction scheme have achieved remarkable advancements in various text-related tasks, such as language understanding [13, 34, 10] , content creation [1, 5, 36] , coding [3, 29, 42] and mathematics [33, 23, 35] . In this generation scheme, the forthcoming token interacts with all previous tokens via the attention mechanism [38] , where the query, key and value states will be calculated for each token. As the past tokens will not be altered, previously computed key and value states can be stored as KV cache to prevent re-computations, significantly improving the generation speed. However, as the batch size and the input context length grows, the stored KV cache emerges as a new memory bottleneck for LLMs. For example, when serving a 175B-parameter LLM [1] with a batch size of 64 and a context length of 4096, the KV cache can occupy 1.2TB of memory space, while the model weights only require 350GB. Meanwhile, the size of KV cache will continue to increase as decoding progresses. Therefore, the compression of KV cache is crucial for the efficient deployment of LLMs.Recent compression methods for KV cache can be broadly categorized into two types. The first type of methods compresses the KV cache uniformly, without considering the significance of individual tokens. To preserve performance, these methods often rely on either high-precision quantization [21] or maintaining recent tokens in full-precision [32] , which undoubtedly compromise the compression ratio. Additionally, if salient tokens are not among the most recent ones, such as in information retrieval tasks, it may result in degraded performance. The other type of methods [46, 43, 16] compress KV cache adaptively by identifying salient tokens and compresses them separately. This approach aligns with the observation that a minority of tokens contribute the majority of attention scores [41] , potentially achieving higher compression ratios than non-adaptive methods. However, current adaptive KV cache compression methods [46, 43] use accumulated attention scores as a metric of token saliency, which is insufficient in two aspects. First, accumulated attention scores is inaccurate in identifying important tokens. Due to the presence of attention masks, the attention matrix is a lower triangular matrix. Earlier tokens tend to have larger softmax attention values and more attention scores to be accumulated, as illustrated in Figure 3 . Under this metric, the saliency of the most recent tokens can never surpass that of the first token, thereby introducing a bias in determining token saliency. Additionally, to obtain accumulated attention scores, full attention matrices must be explicitly computed and stored, which can be inefficient for serving LLMs. Given an input context length of l, fast attention implementations such as FlashAttention [8, 7] only require O(l) memory by computing attention output in blocks without retaining complete attention matrices. By contrast, storing full attention matrices requires O(l 2 ) memory, and the large number of memory accesses significantly slows down the inference speed, as depicted in Figure 4 . Figure 1 : Accuracy and efficiency comparisons across various KV cache compression methods. Data is collected with LLaMA3-8B model on Line Retrieval dataset. Among these methods, ZipCache achieves the highest accuracy, generation speed and compression ratio. Details can be found in the supplementary material.To address these challenges, we introduce ZipCache, an efficient KV cache compression method that attains exceptionally high compression ratios by accurate salient token identification. Figure 1 presents an overview of latency-accuracy comparisons among ZipCache and diverse KV cache compression methods. We start by designing an efficient quantization baseline for compressing the KV cache. To preserve performance, predecessor methods [32, 21] employ fine-grained groupwise quantization, which involves independent quantization for a small channel group within each token. However, this method necessitates storing extensive quantization parameters and results in significant memory overhead. By contrast, we introduce a channel-separable quantization scheme that decouples the quantization along channel and token dimensions. This method significantly reduces the quantization overhead without compromising performance. To accurately recognize salient tokens, we introduce a new token saliency metric based on normalized attention scores, which alleviates the bias towards earlier tokens that accumulate more values. All tokens, without exception, will be quantized to the target bit-width based on their estimated saliency, boosting the overall compression ratio. Moreover, to ease integration with fast attention implementations, we introduce an efficient approximation of the token saliency metric. This approximation only relies on computing and storing attention scores from a few number of tokens, which we refer to as probe tokens. An effective probe token selection strategy is then introduced to minimize performance loss. As a result, the majority of tokens can benefit from fast attention implementations, significantly enhancing the generation speed.In summary, our contributions are as follows:• We establish an efficient channel-separable quantization scheme for KV cache, which significantly reduces the overhead of quantization parameters without compromising performance compared to fine-grained groupwise quantization approach. • We propose an accurate metric for assessing token saliency based on normalized attention scores.All tokens are adaptively quantized according to their assessed saliency, thereby improving the overall compression ratio. • We further develop an efficient approximation method for the token saliency metric that integrates seamlessly with fast attention implementations, enhancing generation speed. [26, 17, 14] and quantization-aware training (QAT) [28, 31] . For large language models (LLMs), where fine-tuning can be data-and computation-intensive, PTQ is often the preferred method [40, 11, 45, 27] .In this paper, we also quantize KV cache in a post-training manner. For both approaches, quantization can be implemented at various levels of granularity, including channelwise, tokenwise, and groupwise approach. Typically, a finer quantization granularity involves the independent quantization of smaller parameter groups, which often results in improved performance albeit at the cost of more quantization parameters and increased memory overhead. In the context of LLMs, fine-grained quantization is frequently utilized due to the presence of outliers [22, 45] . However, for KV cache compression, this will greatly reduce the overall compression ratio.Mixed precision quantization [39, 44, 12, 2] allocates varying bit-widths to distinct parts of a model or tensor, enabling a more compact compression. This approach originates from the observation that model components exhibit differing sensitivities to quantization. Consequently, components with low sensitivity can utilize reduced bit-widths without impairing performance. For LLMs, previous studies [46, 43, 30, 18] have shown significant disparities in the importance of tokens, indicating that heavy compression of non-critical tokens has minimal impact on overall performance. This insight highlights the applicability of mixed precision quantization for compressing the KV cache. 
kv cache compression: 
 While KV cache effectively prevents re-computation and significantly enhances generation speed, its memory footprint is notably substantial with long-context input. To alleviate this, many efforts have been made to reduce the KV cache size. Based on the compression method, these methods can be categorized into two groups: token dropping [46, 16, 30] and KV cache quantization [43, 21, 32] . The former identifies and drops unimportant tokens in the KV cache. For example, H2O [46] only maintain 20% heavy-hitted tokens and 20% recent tokens while evicting the rest. However, discarding tokens permanently erases their information, which proves to be suboptimal for tasks such as retrieval [43] .Conversely, the latter category employs quantization on the cached key and value states, and mixed precision quantization can further be applied once token importance is identified [43] . To tackle the outliers present in the KV cache, these methods extract the outlier as full precision [21] or use finer-grained quantization scheme [32] , which increases the quantization overhead. In this study, we propose an efficient channel-separable quantization scheme with reduced quantization overhead and strong performance. Additionally, both categories of methods commonly adopt accumulated attention scores as the metric for token importance [46, 43] . However, we observe that this criterion is inaccurate and can result in significant performance deterioration at low bit-widths. In contrast, we achieve superior compression performance by utilizing a more accurate metric for identifying salient tokens.3 Preliminary 
attention block in llms: 
 Given an input prompt, the generation process of LLMs can be broadly categorized into two distinct phases: the prefill phase, which computes and stores the KV cache for input tokens, and the decoding phase, where new tokens are generated through a next-token-prediction scheme. Given input data X and an attention block with its weight matrices W Q , W K and W V , the prefill phase can be formulated as:EQUATIONEQUATIONHere, d k is the dimension of the key, and A refers to the attention scores. K and V will be stored as KV cache. For clarity, we have omitted the output projection.For the decoding phase, given x as the embedding vector of the current token, the query q becomes a vector and the KV cache matrices will be updated as follow:EQUATIONThe attention output are then computed as follows:EQUATIONTo ensure clarity and consistency, we introduce notation to define the hyper-parameters used in the paper. Specifically, we denote the batch size as b, the number of attention heads as h, the sequence length as l, and the head dimension as d.Uniform quantization is adopted in our study and all experiments. Given a floating-point vector x, it can be uniformly quantized to k-bit as follows:EQUATIONHere, ⌊•⌉ denotes the round operation, s = max(x)-min(x)2 k -1and z = -⌊ min(x) s ⌉ are quantization parameters. It should be noted that the quantization parameters are stored in full-precision, which can lead to significant overhead if the quantization is fine-grained.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [NA] 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). 
method 4.1 a strong baseline for kv cache quantization: 
 Tokenwise quantization, as depicted in Figure 2(b ) is prevalent in quantizing large language models (LLMs) due to the distinct representations of individual tokens. However, it has been widely observed, as illustrated in Figure 2 (a), that outliers emerge within the channel dimensions of key and value matrices [43, 32] , posing challenges for tokenwise quantization. To address this, recent work [32] resorts to groupwise quantization, where outlier channels are processed in distinct groups, as illustrated in Figure 2(c ). However, this fine-grained quantization approach introduces excessive memory overhead, thereby significantly impacting the compression ratio. For instance, considering X ∈ R b×h×l×d as the data to be quantized and a group size of n, tokenwise quantization only results in 2bl quantization parameters, while groupwise quantization would yield 2bhld n quantization parameters. Since these parameters are usually stored in full precision, this overhead would constitute a substantial portion of the storage cost for quantized data.Motivated by depthwise separable convolution [19] , we introduce an efficient channel-separable tokenwise quantization scheme, which disentangles the channel and token dimensions. As shown in For the i-th channel in X, the normalization process can be formulated as:EQUATIONAfter normalization, each channel is scaled to a closed magnitude, mitigating the influence of outliers during tokenwise quantization. Subsequently, tokenwise quantization can be reliably applied and the scales c are multiplied back to restore the magnitude of each channel. The process of channel-separable tokenwise quantization is summarized in the supplementary material. Within this quantization scheme, the total number of quantization parameters amounts to hd + 2bl, representing a notable reduction compared to groupwise quantization, while effectively balancing the outlier channels and the representation of each token. As referred to Figure 2 (a), since the differences in token representations are small in key cache, we employ channelwise quantization for the key cache to further reduce overhead and employ channelseparable tokenwise quantization for the value cache. As depicted in Table 1 , this configuration yields superior performance with reduced quantization overhead compared with groupwise quantization, thereby establishing a robust baseline for KV cache quantization.Adaptive KV cache compression [46, 43, 16] aims to discern the saliency of each token, keeping the information of salient tokens while evicting or aggressively compressing the rest, to achieve a higher compression ratio. These salient tokens, also referred to as "Heavy Hitters" [46] , are often identified based on accumulated attention scores. Given attention score matrix A ∈ R l×l , the saliency of token i is estimated by:EQUATIONSo there are five eggs Tokens with large saliency values are then considered salient tokens. However, this approach has inherent limitations due to the lower triangular nature of the attention score matrix, as illustrated in Figure 3 (a). There are two primary issues. Firstly, earlier tokens benefit from having more values accumulated since the elements above the diagonal are all zero. For instance, in a sequence of length l, the initial token accumulates l positive values, whereas the final token only accumulates one. Secondly, Softmax function converts real numbers into probabilities, so that the earlier rows of the attention matrix tending to have higher values, as fewer numbers are involved in the Softmax calculation. Consequently, the accumulated attention score of the final token will always be smaller than that of the first, which exceeds 1. To address this, previous works, such as H2O [46] , always maintain recent caches in full precision. Nevertheless, this solution is suboptimal since recent tokens are not necessarily the most significant ones.To enhance the evaluation of each token's saliency, we introduce an accurate token saliency metric based on normalized attention scores pi :EQUATIONHere, nnz(A :,i ) denotes the number of non-zero elements in the i-th column of A. As evidenced in Figure 3 (a), normalizing the accumulated attention scores mitigates the influence of excessively large values in the initial rows of the attention score matrix, thereby delivering a more precise assessment.To validate the efficacy of our new metric, we input a sample from GSM8k dataset with chain-ofthoughts (CoT) prompting to the LLaMA3-8B model and identify saliency of each token by Eq. 7 and Eq. 8, respectively. As depicted in Figure 3 (b) and (c), the salient tokens are at the end of the prompt, which correspond to the question for LLM to answer. However, these tokens are identified as low saliency by accumulated attention scores. Under the KV cache compression framework, these tokens would either be discarded or quantized to extremely low bit-width, resulting in a significant performance deterioration. In contrast, our method accurately identifies the salient tokens. Additional experimental results regarding the accuracy of our method will be detailed in Section 5.2.As analyzed in Section 4.2, adaptive KV cache compression requires the explicit computation of full attention scores, as referred to Figure 4 (b), which clashes with fast attention implementations like FlashAttention [8, 7, 9] . As shown in Figure 4 (c), FlashAttention computes attention outputs in tiles without storing the intermediate attention scores. To reconcile the efficiency of FlashAttention with the substantial compression offered by adaptive KV caching, we devise an effective approximation for Eq. 8 as a measure of token saliency. Specifically, we sample a small group of tokens, designated as probe tokens, and compute their attention scores A probe as follows:EQUATION)By substituting A probe into Eq. 8, we can approximate the saliency of all tokens. For the remaining non-probe tokens, their attention scores do not have to be computed explicitly, enabling the integration of fast attention implementations to expedite the generation process, as illustrated in Figure 4 (a). However, the positions of the probe tokens will undoubtedly affects the accuracy of the approximated token saliency and the selection of probe tokens is under explored.In this study, we suggest four strategies for sampling probe tokens:• Random tokens. The probe tokens are randomly sampled from all positions.• Special tokens. The special tokens and punctuation tokens will be treated as probe tokens.• Recent tokens. The most recent tokens are selected as probe tokens.• Random+recent tokens. The probe tokens will be divided into two parts, one using recent tokens and the other randomly selecting from the remaining tokens.It should be emphasized that our study diverges from prior research [16] in that, rather than directly choosing special or recent tokens as salient tokens, we opt to sample a subset of tokens as "probes" to detect the salient ones. As depicted in Table 2 , we present a comprehensive comparison of the performance among four distinct sampling strategies. Among the four strategies examined, a hybrid approach that combines recent tokens with randomly selected tokens emerges as the most effective. Unless otherwise specified, this hybrid strategy with 5% recent tokens and 5% random tokens will be employed in our method. 
implementation details: 
 Models and datasets. To validate the efficacy of our proposed method, we conduct experiments with three open-source LLMs: Mistral [20] , LLaMA2 [37] and LLaMA3. These models are evaluated on three challenging benchmarks: GSM8k [6] for math problem solving, HumanEval [4] for code generation, and Line Retrieval [25] for data retrieval. To ensure reproducibility, the reported results are obtained using the Language Model Evaluation Harness [15] and LongEval [24] .Quantization and generation settings. We employ mixed precision quantization for KV cache where salient tokens will be quantized to 4-bit while the remaining will be quantized to 2-bit. For both subsets, we apply channelwise quantization for the key cache and channel-separable tokenwise quantization for the value cache. The proportion of salient tokens will be denoted by "Saliency Ratio" in the experimental results. During the decoding process, ZipCache adopts a streaming strategy [21] and repeats the compression process for the KV cache whenever 100 new tokens are generated.We begin our evaluation on GSM8k dataset with chain-of-thoughts (CoT) prompting, and the results are presented in Table 3 . This task requires LLM to solve mathematical problems and return the final answer without multiple options. This task poses considerable challenges and previous KV cache compression methods manifest notable declines in accuracy. For instance, KIVI [32] shows an accuracy drop of 7.89% on LLaMA3-8B model, indicating the suboptimality of preserving recent tokens in full precision instead of identifying salient ones. Moreover, there is a substantial decrease in accuracy, amounting to 20.4%, for MiKV [43] under the high compression ratio. This suggests that accumulated attention scores mistakenly identify salient tokens, resulting in the loss of vital information during compression. By contrast, the proposed normalized attention scores can accurately measure token saliency, leading to a substantial enhancement in accuracy by 18.27% for LLaMA3-8B models in comparison to MiKV. In comparison to GEAR [21] , which quantizes the entire KV cache to 4-bit, our approach additionally quantizes 40% tokens to 2-bit with enhanced performance on Mistral-7B model. This underscores the superiority of accurate adaptive compression of KV cache.Table 3 : Performance comparisons on GSM8k with CoT prompts. Here, "H/L" denotes the bit-width for salient tokens (high-precision) and regular tokens (low-precision), respectively. The compression ratio is calculated with an average input length of l = 840. In this subsection, we assess the performance of code generation across various KV cache compression methods, as summarized in Table 4 . Remarkably, ZipCache attains a compression ratio of 4.94× without sacrificing performance when tested with the Mistral-7B model, outperforming predecessor methods. Moreover, when evaluating on LLaMA3-8B model, our approach outperforms KIVI-2 [32] by 7.32% with a significantly higher compression ratio (4.39× vs. 2.55×). It should be noted that the average input length for this task is only 119, while KIVI retains the recent 32 tokens in full-precision, thereby considerably diminishing its overall compression ratio. This underscores the advantage of ZipCache over methods that consistently retain information of recent tokens.Table 4 : Performance comparisons on HumanEval for code generation. Here, "H/L" denotes the bit-width for salient tokens (high-precision) and regular tokens (low-precision), respectively. 0-bit denotes the tokens are evicted. The compression ratio is calculated with an average input length of l = 120. We further evaluate the data retrieval performance of various KV cache compression methods on Line Retrieval [25] dataset, where LLMs are required to retrieve specific content from a record of lines using a corresponding line index. The accuracy results under various number of lines are depicted in Figure 5 . Notably, all quantization-based compression methods exhibit superior performance compared to the eviction-based approach H2O [46] . For eviction-based methods, information is permanently discarded upon eviction, whereas quantization introduces only minor errors while preserving the integrity of the data. Additionally, in comparison to KIVI [32] , which always maintains recent caches at full precision, our approach consistently achieves better retrieval accuracy. This can be attributed to the nature of retrieval tasks, where salient tokens may appear at any position within the context, rather than being confined to the most recent caches. Moreover, when compared to MiKV [43] , which employs accumulated attention scores as a saliency metric, our method yields a remarkable 42% accuracy improvement when evaluated using 200 lines on the Mistral-7b model. This substantial enhancement once more highlights the effectiveness of normalized attention scores in identifying salient tokens.Additional experimental results on HumanEval [4] can be found in the supplementary material. In this subsection, we compare the latency and memory consumption of ZipCache and MiKV [43] under various input lengths, as depicted in Figure 6 . Data is collected by serving LLaMA3-8B model on a Nvidia A100 GPU. MiKV employs accumulated attention scores to estimate token saliency, necessitating the use of standard attention for both prefill and decoding phases. Conversely, through an efficient approximate saliency metric, ZipCache requires only the calculation of the attention matrix for 10% of the tokens, while the remaining 90% tokens can be computed using either FlashAttention [7] or FlashDecoding [9] . Consequently, ZipCache achieves faster inference speed and lower memory usage, boasting a 37.3% reduction in prefill-phase latency, a 56.9% reduction in decoding-phase latency, and a 19.8% reduction in GPU memory usage when the input length scales to 4096.  
conclusion and future work: 
 In this paper, we have proposed ZipCache, an accurate and efficient mixed-precision quantization framework for compressing KV cache. To commence, we introduce a channel-separable quantization scheme for KV cache, effectively reducing the overhead of storing quantization parameters compared to traditional fine-grained quantization schemes without performance degradation. Additionally, we present a novel metric for accurately assessing token saliency based on normalized attention scores. This metric enables adaptive quantization of all tokens according to their saliency, leading to improved compression ratios without sacrificing model performance. Moreover, we introduce an efficient approximation method for the token saliency metric, seamlessly integrating with fast attention implementations such as FlashAttention and FlashDecoding. This enhancement significantly boosts generation speed and reduces GPU memory requirements. Our extensive experiments have demonstrated that ZipCache achieves state-of-the-art compression performance in terms of compression ratio, accuracy and generation speed. We believe that ZipCache will pave the way for more practical and scalable deployment of LLMs in various real-world applications.Limitations and Broader Impacts. While ZipCache presents promising advancements in KV cache mixed-quantization frameworks for LLMs, the saliency ratio is manually specified before evaluation and cannot be automatically adjusted based on task datasets. Moreover, similar to other generative models, ZipCache can potentially be used to generate malicious content.Algorithm 2: ZipCache for Prefill Phase procedure ZipCachePrefill: Input: Query states Q, key states K, value states V, saliency ratio r%, bit-width for salient tokens k h , bit-width for regular tokens k l // Salient Token Identification Select probe tokens and compute their attention scores A probe by Eq. 9 Measure the token saliency p with A probe by Eq. 8 // Computing Attention Output with FlashAttention O = FlashAttention(Q, K, V) // Compressing KV Cache Partition key states: K salient , K regular = Split(K, p, r%) Partition value states: V salient , V regular = Split(V, p, r%)K salient = ChannelQuant(K salient , k h ), V salient = CSTQuant(V salient , k h ) K regular = ChannelQuant(K regular , k l ), V regular = CSTQuant(V regular , k l ) K = Concat(K salient , K regular ) V = Concat(V salient , V regular ) // Return Attention Output and Compressed KV Cache return O, ( K, V)Algorithm 3: ZipCache for Decoding Phase procedure ZipCacheDecoding:Input: Query vector q, key vector k, value vector v, KV cache ( K, V), saliency ratio r%, bit-width for salient tokens k h , bit-width for regular tokens k l , decoding token index i, probe attention scoreA probe K = Concat(k, K) // Concatenate key cache V = Concat(v, V) // Concatenate value cache o = FlashAttention(q, K, V) // Compute attention output i = i + 1 if i == 100 then// Re-compress every 100 tokens Extract K[: -100] and V[: -100] and adaptively compress them with A probe Reset i = 0, A probe = None else if i > 95 or randint(0, 100) < 5 then // probe tokens consists of 5% recent and 5% random tokens.Compute attention scores a of current token by Eq. 4 A probe = Concat(a, A probe ) // Return Attention Output, KV Cache and Attention Scores from Probe Tokens return o, (K, V), A probe retrieval task with a Nvidia A100 GPU. We use a batch size of 8 and an average input length of 3072. To ensure a fair comparison, we implement GEAR and KIVI with FlashAttention integration. Notably, the latency of ZipCache is lower than that of GEAR, which can be attributed to our efficient quantization scheme, whereas GEAR has a high overhead due to outlier extraction. Compared to KIVI, ZipCache's latency is slightly higher (2584.01 ms vs. 2482.26 ms), but ZipCache achieves a higher compression ratio and better performance. This difference is due to KIVI's fixed compression strategy, while we adaptively compress the KV cache based on the saliency. In comparison to MiKV [43] , which identifies salient tokens through accumulated attention scores, our method achieves a notable 10.0% accuracy improvement by accurately pinpointing salient tokens and a substantial 38.0% decrease in prefill latency by integrating FlashAttention [7] .Moreover, to demonstrate the efficacy of the proposed efficient quantization scheme, we also implement ZipCache with groupwise quantization. The results show that using groupwise quantization Table A: The effect of various saliency metric on GSM8k with CoT prompts. Here, "H/L" denotes the bit-width for salient tokens (high-precision) and regular tokens (low-precision), respectively. "Locality" means the recent tokens are identified as salient tokens. The compression ratio is calculated with an average input length of l = 840.  
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: The paper has 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: The authors have cited the related papers. 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. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? 
 
 Appendix: In this subsection, we evaluate the performance of ZipCache on LongBench using the longchat-7b-v1.5-32k model, as shown in Table C . The results show that ZipCache outperforms the previous state-of-the-art method, KIVI on long context scenario. 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. 


Decision: This paper presents ZipCache for KV cache quantization with salient token identification. After the rebuttal, it receives one strong accept, one accept, and two borderline accept. Its merits, including clear motivation, good organization, interesting idea, and good results, are well recognized by the reviewers. The authors' response well resolves the concerns of unsatisfying writing, insufficient implementation details, and incorrect formulation. Based on the enclosed reviews, I think this work should be accepted. Please also incorporate these details in the revised manuscript.

Reviews: review: paper_summary: paper_summary: The paper introduces a channel-separable quantization scheme that decouples the quantization along channel and token dimensions. This method significantly reduces the quantization overhead without compromising performance. To accurately recognize salient tokens, the paper introduces a new token saliency metric based on normalized attention scores, which alleviates the bias towards earlier tokens that accumulate more values. The authors have demonstrated on different LLM models with different tasks.
strength_weakness: strength_weakness: The paper needs to be proof read by a native English speaker, to improve its readability.

The contribution is incremental, and the idea of group quantization (different for different K and V) are not new as already noted by the authors.

The idea of speeding up through flashattention and flashdecoding while the baselines not leveraging that is a bit of unfair comparison. Both GEAR and KIVI can exploit similar speed up benefit, thus Fig 1 comparison and later in table speed comparison is not fair.

The details of system level implementation is missing, it will be good to help the reviewer provide how the speed up can be demonstrated on a single GPU inference system.
questions: questions: refer to weakness.
limitations: limitations: The system benefits are not clear or thorough for the reviewer to clearly appreciate. Also more fair system eval is needed.
paper_summary: paper_summary: This paper proposes a post training quantization framework named ZipCache for quantizing the key and value cache of LLMs. The authors introduce a channel-separable token-wise quantization scheme, which consumes less memory than the group quantization in terms of the quantization parameters. They also select saliency tokens based on a normalized attention matrix to avoid the accumulated bias. Additionally, the authors introduce an approximation method to apply it to FlashAttention.
strength_weakness: strength_weakness: - The idea of channel-separable quantization is not novel. It is a common practice to smooth the outliers of activations in the channel dimension, e.g., SmoothQuant [1], OmniQuant [2].
- As shown in Table 3 (a), the values on the diagonal of the matrix are relatively large, while the values at other positions are relatively small. Based on this circumstance,  if the mean operation is performed for none-zero positions, earlier tokens may have smaller $p_i$ for they have more tiny values counted. Therefore, it may cause a prefer to the latest tokens. I think this is a issue for the proposed technique in Section 4.2, which needs more analysis.
- Ablation studies are not provided. 
- As stated in Line 268 in the paper, new tokens remain un-quantized unless they reach 100. As far as I know, keeping part of tokens in full precision can improve the accuracy [3, 4]. Therefore, it is better to distinguish the contributions of this implementation and the proposed techniques. 
- I think the value of 'H/L' should be carefully considered. This is because both KIVI and ZipCache keep new tokens in FP16 at first and then perform quantization. Since the 'H' value of KIVI is 16, why it becomes 4 for ZipCache?
- The equation (5) is not accurate for not subtracting the zero-point 'z' in the de-quantization.

[1]. SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models.

[2]. Omniquant: Omnidirectionally calibrated quantization for large language models.

[3]. KIVI: ATuning-Free Asymmetric 2bit Quantization for KV Cache.

[4]. SKVQ: Sliding-window Key and Value Cache Quantization for Large Language Models.
questions: questions: - The authors have compared the channel-separable quantization with many dynamic methods, but lacks comparison of accuracy with static methods, e.g., WKVQuant [1]. It is better to further explore whether the $c_i$ in equation (6) can outperform well-trained parameters, since they consumes the same memory space.
- How does the calculation of $c_i$ in equation (6) effect the inference latency? It is better to compare the channel-separable quantization with group quantization and static methods in terms of the inference latency.

[1]. WKVQuant: Quantizing weight and key/value cache for large language models gains more.
limitations: limitations: This paper contains the describtion of limitations. For suggestions, please refer to the Weakness and Questions.
paper_summary: paper_summary: The paper introduces ZipCache, an adaptive KV cache compression method for LLMs by accurately identifying salient tokens. It first presents a channel-separable tokenwise quantization scheme that reduces the overhead of quantization parameters. Next, it proposes a metric for identifying salient tokens based on normalized attention scores. By efficiently approximating the saliency metric, the method integrates seamlessly with fast attention mechanisms like FlashAttention. The efficacy of ZipCache is demonstrated through extensive experiments, showing superior performance in terms of compression ratio, speed, and minimal accuracy loss compared to existing methods.
strength_weakness: strength_weakness: 1.The experiments should encompass a broader diversity of tasks to comprehensively assess the method's effectiveness, such as LongBench.

2.The font size in Figures 5 and 6 is too small, making them difficult to read. It is recommended to increase the font size for better clarity and accessibility.
questions: questions: See weaknesses
limitations: limitations: See weaknesses
paper_summary: paper_summary: The paper presents an adaptive, mixed-precision quantization method for compressing KV cache in LLMs. It proposes a channel-separable tokenwise quantization scheme to establish a robust quantization baseline, reducing the memory overhead of quantization parameters. A saliency metric, based on normalized attention scores, is employed to identify salient tokens accurately. This approach preserves essential information while aggressively quantizing less salient data, allowing for higher compression ratios with minimal impact on model accuracy. The method can be integrated with FlashAttention with an efficient approximation of the proposed metric. Experimental results validate the superiority of ZipCache on speed and accuracy over existing methods.
strength_weakness: strength_weakness: 1. Evaluating ZipCache on other generation or comprehending tasks will make the findings more robust.
questions: questions: 1. Have you considered using only the attention scores from the most recent token to determine token saliency? What impact might this have on the overall performance?
2. What does saliency mean in Figure 3c? What is the probability of each token being selected as a salient token, as mentioned in the caption of Figure 3?
limitations: limitations: Yes

Comments: title: Appreciation for Your Valuable Feedback
Final remark
Appreciation for Your Valuable Feedback
Appreciation for Your Valuable Feedback
Response to additional questions (Part 2)
Response to additional questions (Part 1)
Review for Rebuttal
Response.
Follow-Up on Rebuttal
Follow-Up on Rebuttal
Follow-Up on Rebuttal
Follow-Up on Rebuttal
Appreciation for Your Valuable Feedback
Review for rebuttal
comment: Dear Reviewer ozSB,

Thank you for your feedback. We truly appreciate your careful consideration of our responses and will carefully revise the paper based on your and other reviewers' suggestions.

Best regards,

Authors of #5969.
Thanks to authors for their comprehensive response! Overall majority of my concerns seems to be addressed by the authors. However, regarding Q3, I would recommend the authors to change their wordings as it is not true that majority of the quantization can not support flashattention. Flashattention is an orthgonal method that can be merged with most of the quantization schemes. Please tone down this claim about others don't supporting it. For Q4, if so, then please explicitly mention in L307 of the manuscript that flashattention during prefill and decoding during decode, despite this being intuitive (as flashattn does not provide much benefit during decode phase). I thus increase my score.
Dear Reviewer S3gk,

Thank you for your feedback. We truly appreciate your careful consideration of our responses.

Best regards,

Authors of #5969.
Dear Reviewer Nc8R,

Thank you for your feedback. We truly appreciate your careful consideration of our responses and will carefully revise the paper based on your and other reviewers' suggestions.

Best regards,

Authors of #5969.
**Q9: The probe tokens during decoding phase are kept at higher precision.**
This is not the case. All KV caches, including those for probe tokens, are quantized after every 100 new tokens are generated. As referred to Q5, during the decoding phase, the primary difference between probe tokens and other tokens is that we use the queries of probe tokens to compute attention scores explicitly, while other tokens are processed using FlashDecoding. 

**Q10: Keeping probe tokens at higher precision will impact accuracy as well as throughput.**
Please refer to Q9. Probe tokens are not kept in higher precision.

**Q11: What is the computation and memory overhead for probe tokens?**
As shown below, computing attention scores with the queries of probe tokens introduces limited computation and memory overhead. 

Table: Computation and memory overhead for probe tokens. Here, "ZipCache w/o Probe Tokens" denotes the token saliency is randomly generated rather than approximated with probe tokens. Data is collected by serving LLaMA3-8B on a NVIDIA A100 GPU with a batch size of 8 and sequence length of 3072.

| Model                  | Method  | Prefill-phase Latency (ms)  | Max GPU Memory (MB)|
|------------------------|---------|-------|--------|
| **LLaMA3-8B**          | ZipCache w/o Probe Tokens| 2503.06 | 34990 |
|                        | ZipCache w/ Probe Tokens| 2584.01 | 34992  |


[1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models.

[2] No Token Left Behind: Reliable KV Cache Compression via Importance-Aware Mixed Precision Quantization.

[3] GEAR: An Efficient KV Cache Compression Recipe for Near-Lossless Generative Inference of LLM.

[4] KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache.

[5] On the Efficacy of Eviction Policy for Key-Value Constrained Generative Language Model Inference.
Thanks to the reviewer for the valuable comments.

**Q1: The KV compression can be useful for longer context based evaluation.**
Firstly, in alignment with prior literature [1-4], we have evaluated our method on three challenging and widely-recognized benchmarks, including the long context Line Retrieval task. As highlighted in Table 3, our approach consistently surpasses previous state-of-the-art methods, underscoring its effectiveness. Secondly, as mentioned in lines 36-38 of our paper, during **batch inference**, the KV cache with an input length of approximately 4K is already a significant bottleneck in the storage system. For example, the KV cache can occupy 1.2TB of memory space with a batch size of 64 and an input length of 4096. Moreover, to further address your point, we have also evaluated the performance of ZipCache on LongBench using the longchat-7b-v1.5-32k model, as shown below. Due to the limited time slot of the rebuttal, we will conduct experiments on more long-context tasks and add the results to the revised version.

Table: Performance comparisons on LongBench with LongChat-v1.5-7B-32k. The saliency ratio is 60% for ZipCache.
| Model                  | Method  | TREC↑  | SAMSum↑ | LCC↑   | RepoBench-P↑ |
|------------------------|---------|-------|--------|-------|-------------|
| **LongChat-v1.5-7B-32k**| FP16    | 66.06 | 41.19  | 52.96 | 56.8        |
|                        | KIVI-2  | **66.0** | 40.57 | 47.99 | 52.6      |
|                        | ZipCache| **66.0** | **40.64**  | **51.25** | **52.87** |

**Q2: Similar idea of normalized attention scores has been proposed.**
Thank you for sharing this literature [5] and we will include it in the references in the revised version. 
However, we would like to emphasize that our proposed saliency metric can be applied universally to all tokens without exception, whereas the approach in [5] excludes certain tokens from the eviction scope based on their standard deviation. Additionally, our work provides a comprehensive analysis of the limitations of using accumulated attention scores as a saliency metric, as discussed in lines 200-210 and illustrated in Figure 3 of the paper. 

**Q3: The generation efficiency is compared with only MiKV and the comparison is unfair.**
As referred to Figure 1 in the paper, Table A in the Appendix and Table A in the rebuttal PDF, we have compared the generation efficiency with H2O, GEAR and KIVI. Moreover, the previous adaptive KV cache compression methods [1-2] are not compatible with FlashAttention, which is a major drawback of them. By contrast, our method integrates seamlessly with FlashAttention, enhancing generation speed.

**Q4: How do probe tokens help perform flashattention or flashdecoding efficiently?**
Firstly, as referred to lines 230-232 and Equation (9) of the paper, we select a small set of probe tokens and compute attention scores using their queries, which allows us to approximate the saliency for all tokens. Once we have determined the token saliency, we can efficiently compute the attention output using FlashAttention. For a more detailed explanation of our method, please refer to Algorithm 2 and Algorithm 3 in the Appendix.

**Q5: Do you implement flashattention or flashdecoding during decoding phase?**
Yes. As mentioned in lines 267-268, 304-307, and detailed in Algorithm 3 in the Appendix, we implement FlashDecoding during the decoding phase. Similar to the prefill phase, we explicitly compute attention scores for a small set of probe tokens to approximate the saliency of all tokens, enabling the majority of tokens to be computed using FlashDecoding. 

**Q6: How does the selection of probe tokens help guide the bit width assignment of salient tokens if a token does not have any representative in the probe token set.**
As mentioned in lines 254-255 of the paper, our approach to selecting probe tokens involves a **hybrid** strategy, where 5% of the tokens are the **latest** ones, and 5% are randomly selected. All previous tokens are attended when computing attention scores for the latest token.

**Q7: The algorithm of decode stage seems ambiguous. For example, what is the precision of the KV cache for new tokens.**
As referred to lines 267-268 of the paper, we implement a streaming strategy during the decoding phase, where the new KV cache is quantized every 100 new tokens generated. This approach is consistent with the strategy used in GEAR [3] and is designed to enhance decoding speed. Consequently, before reaching the 100-token threshold, the new KV cache is maintained in full precision. After 100 new tokens are generated, all KV cache will be quantized based on their estimated saliency. 

**Q8: Attention and softmax need to be computed at high precision to have numerical stability.**
Indeed. The KV cache will be dequantized to full-precision when calculating attention. This is consistent with previous KV cache quantization work [2-4].
Thanks to the authors for the detailed response. I understand the motivation of the channel-seperable token-wise quantization, as well as the effectiveness of the token saliency metric after the rebutall. The integration of this kind of methods to flash-attention is also a valuable contribution. Therefore, considering the answers from Q1 to Q5, I decide to raise my score to 5. Although, I think the description of the Section 4.1 can be improved, especially from the aspect of accuracy and latency. As claimed to be a more advanced baseline, it should process better generalization ability compared to group quantization and well-trained static parameters, rather than solely performing well on a single dataset. Additionally, I think the answers to Q6 and Q8 should be included into the main text.
Thanks to the authors for the new results with attached pdf. Particularly, I appreciate the fair pre-fill stage comparison with all quantization schemes leveraging flashattention.

1. While I appreciate the author's effort, I think further clarifications are needed. In specific, the KV compression can be useful for longer context based evaluation, the current manuscript including that in the  one in rebuttal pdf. The LLaMA2 model has pretraining context length of ~4k which is not sufficient enough for the long context demonstration. Beyond long context evaluation the benefits of KV cache compression is not that usefully demonstrable. And the main manuscript demonstrates results with around 900 and 200 tokens on average as said by the authors.

2. The idea of normalized attention score is not new either, so I believe their is a bit of overclaim here. Please refer to this paper, that proposed something similar, namely MAS [1].

3. Why the generation efficiency is compared with only MiKV, instead of other SoTA quantization methods like KIVI? It is understandable that MiKV falls short as it cant leverage flash-decoing/attention, this is again unfair comparison.

4. It is also not clear if the probe tokens are selected randomly how it will still help manage perform flashattention or flashdecoding efficiently. Also, during decode phase do you implement flashattention or flashdecoding? please provide more details.

5. Overall, I believe the paper is built on top of the main findings of probe tokens, however the details in that front is yet to be sufficient. Additionally, how the selection of probe tokens help guide the bit width assignment of salient tokens if a token does not have any representative in the randomly selected probe token set?

6. The algorithm of decode stage seems ambiguous, or not fully informative. For example, the new tokens getting added on top of previous KV, is it the hybrid H/L quantized KV or or some high precision KV? As we may need to compute attention and softmax etc at high precision to have numerical stability. This needs clarification.


7. It looks like the probe tokens during decode phase are kept at higher precision, did the author include that in their total compression scheme? Also, the probe token computation and memory overhead is not clearly discussed in the paper's results. This should impact accuracy as well as throughput compared to schemes like KIVI.  



[1] On the Efficacy of Eviction Policy for Key-Value Constrained Generative Language Model Inference, 2024.
Dear Reviewer Nc8R,

Thank you for dedicating your time to reviewing our paper. As the discussion period deadline is approaching, we kindly invite any further comments or concerns you might have. Your feedback has been immensely valuable to us for refining the paper.

Best regards,

Authors of #5969.
Dear Reviewer ozSB,

Thank you for dedicating your time to reviewing our paper. As the discussion period deadline is approaching, we kindly invite any further comments or concerns you might have. Your feedback has been immensely valuable to us for refining the paper.

Best regards,

Authors of #5969.
Thank the authors' detailed response. After reviewing all comments and responses, my concerns are addressed well. The proposed method is novel and effective, and I believe this paper is ready to be published. So I would like to raise my score to 8.
Dear Reviewer Nc8R,

We greatly appreciate the time and effort in reviewing our work. We have carefully considered your comments and suggestions and made significant revisions to address the concerns you raised. We are eager to ensure that our paper meets the high standards of our respected reviewers.

Please don’t hesitate to let us know if there is any additional feedback you might have at this stage.

Best regards,

Authors of #5969.
Dear Reviewer ozSB,

We greatly appreciate the time and effort in reviewing our work. We have carefully considered your comments and suggestions and made significant revisions to address the concerns you raised. We are eager to ensure that our paper meets the high standards of our respected reviewers.

Please don’t hesitate to let us know if there is any additional feedback you might have at this stage.

Best regards,

Authors of #5969.
Dear Reviewer E63Z,

Thank you for your feedback. We truly appreciate your careful consideration of our responses.

Best regards,

Authors of #5969.
Thanks for the efforts from the authors. My concerns are addressed in the rebuttal. I think this paper is well-motivated and has its merit to the community. It demonstrates strong performance in the experiments. Therefore, I raise my score after rebuttal.