SinkQ: Accurate 2-bit KV Cache Quantization with Dynamic Sink Tracking

Thank you for your valuable comments. We will address your concerns as follows.

Q:

Writing: Authors did not explain enough of the baseline method of KIVI, especially the part of group and residual in KIVI. The method is partly based on KIVI but authors mainly explained how they improved upon KIVI. Adding background on how KIVI separates between grouping and residual will be nice.

A:

We have changed the related context in Section 3.2.

Q:

Lacks comparison with other non-quantization works .

A:

Thanks for your suggestion. We have added the results of some non-quantization methods and some other quantization-based methods (ZipCache, H2O, Streamingllm and SnapKV). The results are shown below and the experimental details can be found in Appendix B.

We are also reproducing the results of KVQuant, but there are some issues with the results. We hope we can attach the results of KVQuant before the deadline.

Q:

Lacks experiment to larger models (ex. 70B)

A:

For the scenario of larger models, we have added experiments of LLaMA-2-70b-chat-hf, the results are shown below and the experiment details can be found in Appendix B.

Thank you very much for your feedback. We believe that you have two main concerns:

1. How do throughput and memory change when hyperparameters change?
2. Can SinkQ maintain its performance in other scenarios, such as long sequences and larger models with a small sink num?

Q:

How do throughput and memory change when hyper-parameters change?

A:

We will provide a more detailed explanation of the memory usage during inference. Firstly, KIVI will have the following full precision tokens:

1. When the group is not full, the Keys in the group will be retained.
2. Recent Values.

SinkQ will have the following full precision tokens:

1. When the group is not full, the Keys and Values in the group will be retained
2. Recent Keys and Values.
3. A very small number of sinks

So, the main difference is that there are some differences between groups and recent in full precision tokens. We choose this processing method mainly to preserve the quantization consistency of k and v for ease of engineering implementation. And there seem to be some differences of KIVI in the paper and code so we choose to fulfill it by ourself. To achieve fair comparison, we use G=128, R=32 for SinkQ and G=128, R=128 for KIVI in the main experiments.

So, come back to your question:

1. When the choice of group size (G) / residual length (R) varies, would the throughput increase as compared to KIVI be consistent?

The answer is yes.

2. Would the additional memory overhead of SinkQ scale even more when sink_num goes up?

The answer is yes.

Q:

Can SinkQ maintain its performance in other scenarios, such as long sequences and larger models?

A:

For the scenario of long sequences, we give a brief description to justify that a small sink num is enough. The sum of attention scores is always 1 regardless of the length of the sequence. Tokens with extreme high attention scores in a long sequence cannot be more than that in a short sequence from a probability perspective .

For the scenario of larger models, we have added experiments of LLaMA-2-70b-chat-hf, the results are shown below and the experiment details can be found in Appendix B. We have also included comparisons with other baselines in Appendix B.

So, come back to your concrens:

1. How SinkQ would succeed (or fail) when the number of model parameters or the context length continuously go up?

Answer: SinkQ can succeed in these scenarios.

2. If the memory overhead is too large and you have to use a small, fixed number of sinks (which means that you cannot store all outlier tokens), would SinkQ still be superior to KIVI in this case?I hope these can address your concerns.

Answer: SinkQ is definitely superior to KIVI because it preserves more full-precision tokens than KIVI.

Thank you for your valuable comments. We will address your concerns as follows.

Q:

The paper lacks an in-depth discussion of the computational overhead introduced by steps such as slicing, concatenation, and the calculation of outlier tokens. A profiling of these operations would help to understand the real-world efficiency impact.

A:

The computational overhead comes from two aspects.

1. In the compression stage, we calculate the magnitude of each token's key, perform threshold comparison, select the index, and quantify it. The cost of the sink operation is relatively high compared to quantization, but the compression is only performed every G steps (G is the group size in the paper and we set it to 128 in the main experiment), so this time cost can be almost negligible compared to the whole decoding process.
2. In the attention calculation stage, we need to calculate the qkv in the sink pool and cover the attention score according to the sink id, which has a certain cost.

We take an example for illustration. We use SinkQ to infer an example which is randomly selected from GSM8K with an input length of 269 and an output length of 20. It has 2 compress steps and 20 decode steps. We record the time consumption in the attention block and create a pie chart in appendix C. The sink operation accounts for about 18% of the time. However, considering the pre-processing, post-processing and FFN calculation in the entire forward step, the time proportion of sink operations is also very small overall.

We have also retained the profiler and added it to the supplementary material, which you can open with TensorBoard.

Q:

The paper compares SinkQ with KIVI but does not include comparisons with other KV cache compression methods, such as KVQuant, [1], GEAR [2] and SKVQ [3].

A:

Thanks for your suggestion. We did not select these baselines mainly because there were some differences in settings which could lead to unfair comparisons. For example, KVQuant uses many complex operations to improve accuracy and is not tuning-free; GRAR requires computation of low-rank matrix, which leads to additional latency and memory. To address your concerns, we have added the results of ZipCache, H2O, Streamingllm and SnapKV. The results are shown below and the experimental details can be found in Appendix B.

We are also reproducing the results of KVQuant, but there are some issues with the results. We hope we can attach the results of KVQuant before the deadline.

Thank you for your valuable comments. We will address your concerns as follows.

Q:

The authors' observation of sink tokens is not new or novel. Previous works have observed the existence of heavy-hitter tokens [1] or salient tokens [2] in LLMs, which may occur at any token position. Hence, the claim on line 211 is inaccurate: "Previous work suggesting that attention sinks occur only in the initial tokens..."

A:

For sink and heavy hitters, we provide further explanation: We believe that attention sinks and heavy hitters are not completely identical concepts, and their differences are as follows:

1. Attention sinks refer to tokens with consistently high attention scores during the decoding process, with only a few tokens (usually less than 4).
2. Heavy hitters are relatively important tokens in the decoding process compared to other tokens, typically accounting for over 10%-20% of all tokens.
3. Heavy hitters contain attention sinks, but there are differences in the distribution within tokens and the number of tokens. We have plotted the accumulative attention scores (sorted) in Appendix D, indicating that attention sinks are more unique than heavy hitters.

Moreover, SinkQ reveals the outlier tokens in channel-wise key quantization, which was not mentioned in previous works. This is almost the most important observation in our work, and it plays a crucial role in the effectiveness of SinkQ (see Figure 1 (f)).

Q:

The proposed method lacks novelty. The proposed mixed-precision approach for KV cache quantization is similar to previous works [2,3], which also identify outlier tokens in the KV cache and preserve in higher precision or full precision.

A:

SinkQ can be seen as a mixed-precision quantization method, but our innovation lies in other aspects:

1. SinkQ reveals the outlier token phenomenon in channel-wise key quantization, which was not mentioned in previous work
2. The identification method of the outlier token in SinkQ is different from other methods. It directly uses the magnitude of the key as the evaluation metric, which does not require the calculation and storage of the accumulated attention score at each step, but only needs to be calculated at the compress step. This greatly speeds up efficiency, and this method is perfectly adapted to our preliminary exploration.
3. SinkQ only needs to retain a few outlier tokens to significantly increase quantization accuracy, while methods such as ZipCache require retaining more tokens which it considers important.

We further explain why our recognition strategy, namely the magnitude of the key, is better than using accumulative attention score:

1. Selecting outlier tokens based on the magnitude of the key can accurately identify them in the channel-wise key quantization, while the accumulative attention score may not (the top k of the attention score in Figure 1 (c) does not completely overlap with the top k of the magnitude of the key).
2. The calculation of accumulative attention score may not be accurate under this setting. Calculating based on the global attention score may result in subsequent tokens having a small proportion of attention score because of the causal mask; Calculating based on the sliding window may be difficult to obtain accurate global information.
3. The calculation of attention scores need to be done at every step, which requires additional computational and memory costs.

Q:

The "Method" section is incomplete, missing some important details. It is not clear how SinkQ identifies the sink tokens, which is a critical part of the proposed method. During inference, the tokens are ingested in a streaming manner, while the authors only describe how the outlier tokens can be identified after saving all keys in an offline manner on line 152.

A:

We further explain the quantization process here:

We follow the quantization process of KIVI, performing quantization every G steps. At this point, we have G KV caches, and we directly calculate the magnitude of these key caches as the metric. Then, we select the outlier tokens based on this metric and quantize the remaining tokens.

We have changed the related context in our paper (see Section 3.2).

Thank you for your valuable comments. We will address your concerns as follows.

Q:

The major claim of this paper lacks sufficient justification. The foundation of the work is based on the statement, "Previous work suggests that attention sinks occur only in the initial tokens; however, our observations reveal that they can appear at any position within a sentence." However, there is no concrete evidence or examples provided to support this strong claim. To strengthen the paper, it would be beneficial if the authors included specific evidence, such as visualizations or quantitative analysis, to demonstrate instances where attention sinks appear at various positions within a sentence.

A:

We have provided an example of key visualization in Figure 1 (b). Additionally, we have randomly selected 10 examples in GSM8K and provided the positions of the four tokens with the highest attention scores:

Generally , [bos] token is attention sink, while attention sinks do not have obvious distribution patterns, but it is true that initial tokens are more likely to obtain higher attention scores.

Q:

In the experiments, Figure 3(b) and Figure 3(c) suggest that the proposed method may not be as memory-efficient as Kiwi, despite claims of improved efficiency. It would be helpful for the authors to directly address this observation and clarify the reasons behind this apparent discrepancy. A more detailed analysis of the trade-offs between the proposed method and Kiwi in terms of memory efficiency would strengthen the discussion.

A:

We will provide a more detailed explanation of the memory usage during inference. Firstly, KIVI will have the following full precision tokens:

1. When the group is not full, the Keys in the group will be retained.
2. Recent Values.

SinkQ will have the following full precision tokens:

1. When the group is not full, the Keys and Values in the group will be retained
2. Recent Keys and Values.
3. A very small number of sinks

So, the main difference is that there are some differences between groups and recent in full precision tokens. We choose this processing method mainly to preserve the quantization consistency of k and v for ease of engineering implementation. And there seem to be some differences of KIVI in the paper and code so we choose to fulfill it by ourself. To achieve fair comparison, we use G=128, R=32 for SinkQ and G=128, R=128 for KIVI in the main experiments. Due to the above reasons, there may be some differences between SinkQ and KIVI when the batch size is large. However, when the batch size is 1 and the seq length is long, the number of these full precision tokens is very small compared to the quantized tokens, and can be almost ignored. At this time, the memory of the two methods is relatively close.