SKVQ: Sliding-window Key and Value Cache Quantization for Large Language Models

Q2: novelty

First, we want to clarify that the problem we are addressing is different from the issues tackled by existing methods. Our primary focus is on ultra-low-bit quantization of KV-cache. Our method is designed based on the key insight

1. we identified outlier channels in the key cache and innovatively recognized that the locality of the attention module can be combined with mixed-precision quantization.

2. We hypothesize that preserving a small portion of the KV-cache (designed as a sliding window) in high precision can effectively reduce quantization error, and our detailed experimental analysis supports this hypothesis.

While some techniques have been employed in various domains before, integrating them to minimize KV-cache quantization error is not trivial. For example, our experiments compared smoothquant and RPTQ, which utilize smoothing and channel-reordering, respectively, and demonstrated that both perform poorly in ultra-low-bit quantization of KV-cache.

Regarding the three techniques you mentioned, we would like to highlight the differences from previous work:

1. Designs similar to the sliding-window have been extensively explored in approximate linear attention models like Longformer[1], and the structural property of attention, particularly locality, has been well-studied in previous works. These studies inspired us and provided a solid theoretical foundation for our method. However, no prior work has combined quantization with the inherent structured nature of attention. We are the first to integrate these two aspects, thoroughly exploring the potential of attention locality in KV-cache quantization and designing an algorithm that seamlessly fuses it with KV-cache quantization.

2. While auto-clip techniques have been applied in works like AWQ[2], they were used for weight quantization. We are the first to apply clipping techniques to KV-cache quantization. There is a significant difference between the two: since weight distributions are known, precomputing clip values offline is straightforward. In contrast, KV-cache is generated dynamically, making online auto-clipping in dynamic quantization impractical due to unacceptable overhead. Our research discovered that approximating and precomputing clip thresholds on a small calibration set effectively addresses the clipping challenge for KV-cache and effectively improves accuracy.

3. KVQuant does not employ channel reordering; it performs per-channel non-uniform quantization by clustering within each channel. In contrast, SKVQ conducts per-token uniform quantization by extracting value distribution features from each channel, grouping similar channels, and applying uniform quantization to each group. The concept of channel reordering has indeed been used in RPTQ[3], but [3] aims to reduce weight-activation quantization loss through reorder-indexing. We firstly apply channel-reordering to KV-cache quantization.

We will further emphasize our insights and clarify the differences between our approach and previous work in the revised version. Thank you for pointing it out!

[1] Beltagy, Iz, Matthew E. Peters, and Arman Cohan. "Longformer: The long-document transformer." arXiv preprint arXiv:2004.05150 (2020).

[2] Lin, Ji, et al. "Awq: Activation-aware weight quantization for llm compression and acceleration." arXiv preprint arXiv:2306.00978 (2023).

[3] Yuan, Zhihang, et al. "Rptq: Reorder-based post-training quantization for large language models." arXiv preprint arXiv:2304.01089 (2023).

Q6: RoPE and Quantization?

Because we perform pre-RoPE quantization, we store pre-RoPE KVCache, requiring RoPE to be applied at each decoding step. Experiments in FlashInfer[1] has shown Fused-RoPE Attention kernel can apply RoPE on the fly with negligible overhead.

[1] Ye, Z. FlashInfer: Kernel Library for LLM Serving.

Q8: Is fake quantization employed? Quantization Kernel and Memory Storage Details

Regarding the details of FP8: We simulate the use of FP8 by first computing the quantization parameters (zero-point and scale) in FP16, then converting the values to FP8 using round-to-nearest-even rounding, and finally converting them back to FP16.

kernel implementation details: We have currently implemented the quant/dequant CUDA kernel and are continuing to optimize it. Specifically, as mentioned in the answer to Q5, the kernel details for the quantization process are as follows:

1. We adjust the output channels of the weight (QKV_projection) to avoid explicit reordering.
2. Using the pre-computed reorder index, we determine the range of each group and calculate the group-shared scale and zero-point in FP16, then convert them to FP8, store in sequence by group.
3. Next, we perform packing: for the 2-bits K/V cache, we pack every four adjacent numbers within the same group into one uint8 for storage. Although the group size might not align perfectly with 4, leading to some waste, this is negligible (at most $0.8 \times NumGroup$ bytes).

After quantization and packing, the original KV-cache shape $[bs, SeqLen, NumHeads, HeadDim]\Rightarrow [bs, SeqLen, PackDim]$, where $PackDim \approx NumHeads \times HeadDim / 8$. This reduces the memory footprint of the KV-cache by around 800% compared to FP16.

During the decoding step:

1. Based on the starting position index of the packed groups (obtained from the reorder index during the offline pre-computation), we determine the starting position of each group, find the corresponding scale and zero-point, and dequantize to FP16. Then, we reorder the results back to the original sequence using the inverse reorder index and write them back to HBM.
2. Execute RoPE and the Attention Kernel with FP16 KV-cache.

It worth noting that the dequant, subsequent RoPE, and Decoding Kernel can be fused together, eliminating the need to write the intermediate dequantized KV-cache back to HBM. We are continuously optimizing this fused kernel, and will open source it.

Q5: Effectiveness of H2O

Firstly, we'd like to clarify that while H2O serves as a cache eviction policy, it isn't directly applicable to quantization in KV-cache. Achieving the goal of using accumulative attention scores to determine the important token for mixed-precision quantization requires careful design.

There are a lot of details, let's begin from vanilla dynamic per-token RTN quantization during decoding. We face a dilemma: on one hand, we need to generate more tokens to assess the importance of the current token; on the other hand, we need to make an immediate decision whether to quantize this token to a low-precision.

A sliding-window strategy resolves this effectively! Upon generating a token pair cache_t at decoding step $t$, we defer quantization until the it slides out of the window (i.e., at decoding step $t+W$). Then, based on the accumulative attention scores of the subsequent $W$ tokens on token_t, we decide whether to retain it at high precision.

Our strategy is implemented as follows: for a token $t_i$ that slides out of the window, we accumulate it's attention score in $step_{i+1},\cdots, step_{i+W}$. If this score exceeds the cumulative score of the previous $96$(or other threshold) of tokens, we consider this token as a heavy hitter and protect it from quantizing to low-precision.

The score in the table is the average score in needle in haystack test, same with Figure5 in our paper. We set the ratio of full precision to 4%. we use RTN, without channel reorder and clipping.

The results indicate that enabling heavy hitter in filter rules marginally improves accuracy. It can't outperform the random baseline, indicating that our current design approach does not yield significant benefits. At the same time, it comes with 2 limitations:

1. Heavy hitter determination relies on attention scores, which are inaccessible during inference with existing fusion operators like flashattention, flashinfer. This necessitates a custom operator for accessing cumulative scores, leading to significant overhead.
2. The dynamic nature of heavy hitters requires an additional index at runtime, introducing substantial overhead, particularly in ultra-low-bit quantization.

These drawbacks hinder the feasibility of this approach. Thus, while we haven't further adopted this strategy in our experiments, we believe there are more efficient ways to implement heavy hitters detection and more effective filter rules to enhance SKVQ performance.

Q1: The full-precision prefill stage can bottleneck when the prompt length is long.

Since we are not sure whether your concern about the "bottleneck" at prefill phase refers to computation or memory, in addition to the above explanation of memory, we will add an additional explanation of computation in prefill phase:

For each layer, we use the full-precision activation generated from the QKV_projection for attention computation, which is identical to the full-precision model. Additionally, we quantize the full-precision KV embedding into low-precision KV-cache for storage, which introduce a lightweight overhead.

Using full-precision KV embeddings for attention at prefill phase is a common practice in many quantization algorithms, such as FlexGen and KIVI. Overall, SKVQ provides no speedup compared to full-precision model at prefill phase, but also no significant additional overhead.

Q2: Evaluation on long generation scenarios

More explanation about the experiments:

1. The average generated length was approximately 1500 tokens.

2. Although there are some significant shortcomings in the above experiment, such as the limited sample size, we believe the results to som extent indicate that SKVQ has a good performance in long generation scenarios. In some cases, the diversity of the generated content due to quantization even results in higher scores compared to the full precision model.

Should you have any further inquiries, please let us know and we are more than delighted to discuss with you and run more experiments for any pieces of your interests in our work.

Q3: Evaluation on larger models?

We evaluated SKVQ on Llama3-70B-Instruct on LongBench, which adopt GQA with $\text{NumHeads}=64, \text{NumKVHeads=8}$ (i.e. 8 queries share one KV-cache group). Our results show that for most tasks, SKVQ's accuracy loss is within 2%, and for some tasks, it is lossless. However, for a few tasks like Qasper, SKVQ experiences more significant accuracy loss. We believe this is largely because the full-precision model itself performs poorly on this task. We believe this experiment further demonstrates the generalization ability of SKVQ on models of larger size.

We sincerely appreciate the time and efforts you have dedicated to reviewing our paper. Should you have any further inquiries, please let us know and we are more than delighted to discuss with you and run more experiments for any pieces of your interests in our work.

I am satisfied with author replies.

About 2 bits and 1.5 bits. I understand that in this paper, those are just statistical averages. Yet I want to emphasize that with 2 bits, one can distinguish just 4 values. And with 1.5 bits - only '3' (of course practically one cannot encode half a bit). But it all means that after quantization you end up with a very small repertoire of values. It looks like those values are close to acting as yes/no gates on the attention. This is speculative, but maybe the authors want to consider this.