SpeCache: Speculative Key-Value Caching for Efficient Generation of LLMs

We thank the reviewer for their thoughtful feedback! We address specific concerns and questions below.

Q1. Unclear Overhead of Parallel Prefetching and Computation

Since our method is built upon KIVI, we can consider KIVI as the baseline when parallel pre-fetching and computation are not used. On Mistral-7B-Instruct-v0.2, we test the latency and VRAM usage of KIVI and KIVI+SpeCache with batch size of 1 and context length of 64k.

In addition, we also provide the throughput of Mistral-7B-Instruct-v0.2 with context length of 32k when the batch size is increased to the maximum capacity that can be accommodated by 48GB.

We can find that SpeCache only slightly increase the latency of KIVI.

Q2. Questionable Relevance of the KV Cache Bottleneck Example

Thank you for pointing that out. Our example is indeed somewhat outdated. We will revise the description to specify the KV cache size for long sequences with a batch size of 1, such as for Mistral-v0.2-7B with a 32k context length, where the KV cache size exceeds 4B, and for Llama3.1 with a 128k context length, where the KV cache size exceeds 16B.

Q3. Peak Memory Usage Concern in Prefill Phase

In fact, during the prefilling phase, we do not need to keep the entire KV cache in GPU memory. This is because the computation, quantization, and offloading of the KV cache are done layer-by-layer. In other words, at any given moment, only the KV cache of a single layer needs to be fully stored in GPU memory. Before the KV cache of one layer is computed, the KV cache from the previous layer has already been quantized and offloaded. For example, Mistral-7B has 32 layers, so at any given time, we only need 1/32 of the 16-bit KV cache in GPU memory. This also allows SpeCache to handle context lengths that full cache cannot accommodate.

Q4. Experiments on LLaMA-3.1-8B-Instruct

We implemented SpeCache on Llama-3.1-8B-instruct and evaluated it on longbench.

As you expected, the performance of Full KV significantly improves compared to the 8K context at 128K context. Our method, with just 10% KV size, has a 3% gap compared to Full KV, but still shows a 17.8% improvement over the purely quantization method, KIVI.

Q5. Comparision with SnapKV

We added SnapKV as a baseline on longbench. Since SnapKV compresses the KV cache to a fixed size, to ensure a fair comparison, we dynamically calculated the SnapKV budget for each sample length, aligning its KV cache size with that of our method.

Under the same KV cache size budget, SnapKV outperforms all other baselines. However, SpeCache still delivers superior performance. Our method is capable of working with an extremely small KV cache size (about 0.1x), whereas SnapKV tends to experience performance degradation under such conditions. This is because SnapKV compresses the KV cache all at once, which may result in the loss of information needed for subsequent tokens.

Q6. Compare throughput gains at a fixed batch size

Thank you for your suggestion. We have provided some relevant results in our response to Q1.

We thank the reviewer for their thoughtful feedback! We address specific concerns and questions below.

Q1. Why the 1-bit improvement is so large? Performance of KVQuant + SpeCache.

• The improvement of SpeCache on 2-bit KIVI is relatively small because the performance of 2-bit KIVI on longbench is already close to that of 16-bit KV Cache, which can be considered as the performance upper bound after KV Cache compression. As a result, there is limited optimization space left. On the more complex RULER benchmark, SpeCache can provide a significant improvement to 2-bit KIVI (see Q5).
• In our comparison, all instances of "1-bit KIVI" refer to our improved KIVI implementation described in Section 3.4, since original KIVI can hardly work with 1-bit (see Table 5). The larger improvement of SpeCache on the improved 1-bit KIVI is due to the fact that the improved 1-bit KIVI is still far from the upper bound, leaving more room for optimization.
• Since the official KVQuant code does not include 1-bit quantization, we compared the performance of SpeCache with 2-bit KVQuant. We conducted experiments on LLaMA-2-7B and measured the PPL on Wikitext-2. The results show that SpeCache also improves the performance of KVQuant.

Q2. Does this speculative token work layer by layer? What is pre-generated?

The pre-decoding phase is a single inference step we introduce between prefilling and decoding. Its purpose is to generate the first speculative token $T’_2$. During pre-decoding phase, $T_1$ is the model input, and the entire model output (i.e., the predicted next token) is used as $T’_2$. In the first decoding step, $T’_2$ is used alongside the $T_1$ as the query, and $T’_2$ will record the top-k KV pairs from each attention layer by layer during decoding, in order to prefetch them for use in the next step.

Q3. Is this throughput run with flash decoding?

For Full KV Cache, we consistently use Flash-Attention2. For SpeCache and KIVI, during the prefilling phase, we also use Flash-Attention2. However, in the decoding phase, due to the involvement of low-bit operations, we use the low-bit CUDA operators provided by KIVI to compute attention.

Q4. Throughput comparison between SpeCache and KIVI.

Thank you for your suggestion. Here, we first provide a throughput comparison between SpeCache and KIVI with context length of 32k on Mistral-7B-Instruct-v0.2.

SpeCache only slightly reduces the throughput of KIVI. Since the computation of speculative tokens is parallelized with the output tokens, the simultaneous decoding of both tokens does not significantly increase the latency.

Q5. Add more experiment results on RULER.

We evaluate LLaMA-3.1-8B-instruct on RULER with a 32k context length, and the results are as follows:

We find that 2-bit KIVI experiences significant performance degradation on RULER compared to Full KV, especially on the N-MK3 and VT tasks. However, SpeCache effectively mitigates this issue, keeping the performance gap from Full KV within 1%. While 1-bit KIVI completely fails on all tasks in RULER, SpeCache provides a significant performance boost and even reaches the Full KV cache level on certain tasks, such as N-S1.

Q6. Add a cache hit rate for speculative tokens.

We evaluate 2-bit SpeCache (g=64) and Mistral-7B-Instruct-v0.2 on two long sequence generation tasks, MultiNews and GovReport in Longbench. We track two metrics:

• Speculative Token Exact Hit Rate: The proportion of speculative tokens that perfectly match the next output token.
• Speculative Token Top-K KV Cache Hit Rate: The proportion of the top-k KV cache needed for the next output token that is hit by the top-k KV cache of the speculative token.

Although the exact hit rate is only around 57%, SpeCache does not directly output speculative tokens. Instead, it uses them as a medium to guess the top-k KV cache. The top-k KV cache hit rate is above 90%, which ensures the effectiveness of the speculative tokens.

We thank the reviewer for their thoughtful feedback.

Q1. How sensitive is SpeCache to changes in the distribution of attention patterns?

We add the phrase "How many paragraph are there in the article? Translate the first sentence into German." to the end of a 15k-token text composed of several of Paul Graham's essays, then input it into Mistral-7B-Instruct-v0.2, and observe the output token and speculative token generated via 1-bit SpeCache (black bold text represents the erroneous prediction of speculative tokens):

• Output token:

There are 13 paragraphs in the article above. The first sentence in German translates to "Es gibt zwei deutlich unterschiedliche Arten, politisch mittler zu sein: vorsätzlich und zufällig. Absichtslose Mitte stehen dafür, dass die Extreme etwa gleich weit entfernt sind." (There are two distinct ways to be politically moderate: intentionally and by accident. Accidental moderates are...

• Speculative token:

are 13 paragraphs in the article above. The first sentence in German translates to "Ichlieberhnte Wegeliche Arten, politisch mittellangigeben: zufzlich und zufällig." (There are "Intentionenander aus eigentlich Menschen, wenn sie beideehentenfert sind." (There are two distinct ways to be politically moderate: intentionally and by accident. Accidental moderates are...

We find that when switching languages and topics, speculative tokens diverge from the output tokens within some decoding step. However, as explained in our response to Q2, SpeCache does not rely on the exact match of speculative tokens. It only requires that speculative tokens have a high top-k KV cache hit rate for the next output token.

Q2. Accuracy of the speculation mechanism.

We evaluate 2-bit SpeCache (g=64) and Mistral-7B-Instruct-v0.2 on two long sequence generation tasks, MultiNews and GovReport. We track two metrics:

• Speculative Token Exact Hit Rate: The proportion of speculative tokens that perfectly match the next output token.
• Speculative Token Top-K KV Cache Hit Rate: The proportion of the top-k KV cache needed for the next output token that is hit by the top-k KV cache of the speculative token.

Although the exact hit rate is around 57%, SpeCache does not directly output speculative tokens. Instead, it uses them as a medium to guess the top-k KV cache. The top-k KV cache hit rate is above 90%, which ensures the effectiveness of the speculative tokens.

Q3. SpeCache requires running two forward passes per token generation. Have you explored distilling a smaller model?

• Throughout the inference process, we only add one step of pre-decoding, and in the subsequent decoding phase, SpeCache runs only one forward pass for each token generation. This is because SpeCache uses both the Output token (e.g., $T_1$) and the Speculative token (e.g., $T'_2$) as input (i.e., [$T_1$, $T'_2$]) for decoding, which generates two tokens (i.e., $T_1$ generates $T_2$, $T'_2$ generates $T'_3$) in a single forward pass. These tokens are then used as input for the next step, with $T_2$ serving as the model's official output.
• This is also the reason why the generation of the speculative token is a “free lunch” attached to the forward pass of the output token. Therefore, there is almost no additional overhead, and no need to distill an extra model to generate the speculative token.

Q4. How would the performance change with different hardware setups.

Due to the varying computational capabilities of different GPUs and the differences in the optimization levels of their operators, it is challenging to conduct a direct ablation study on bandwidth. Therefore, we perform some theoretical analysis to address this. For example, on Mistral-7B-Instruct-v0.2, when the context length is 2K, the maximum batch size for 1-bit SpeCache supported by a 48GB GPU memory is 410. This requires transmitting 3.4GB of KV cache during each decoding step. Theoretically, since SpeCache can fully parallelize data transfer and computation, the decoding latency is 775ms/step. In other words, under optimal conditions, as long as the bandwidth is above 4.4GB/s, the decoding latency of SpeCache will not significantly change.

Q5. Theoretical or empirical upper bound on the batch size where "decoding latency < fetching latency" no longer holds?

As the batch size increases, the fetching latency grows linearly, while the decoding latency increases sub-linearly due to improved computation parallelism. This results in the overall decoding latency being lower than the latency for fetching the entire KV cache when the batch size is large. Empirically, we have observed the following behavior: For Mistral-7B-Instruct-v0.2 on an A6000 GPU, with a context length of 2K, when batch size >= 4, fetching latency > decoding latency. Additionally, when the context lengths are 8K or 32K, the boundary conditions shift to 2 and 1, respectively.

We thank the reviewer for their thoughtful feedback! We address specific concerns and questions below.

Q1. Discuss the difference between SpeCache and QUEST and the contribution of SpeCache.

Although both QUEST and SpeCache focus on how to accurately select sparse KV pairs, their emphasis differs due to the distinct scenarios they address.

• QUEST aims to reduce computation and the loading of KV cache from GPU HBM by reducing the number of KV pages involved in the attention computation, thereby accelerating the model. In contrast, SpeCache focuses on how to predict and load important KV pairs in parallel when the KV cache is offloaded to lower-level storage devices (e.g., CPU memory) for GPU memory saving. Since the CPU-GPU bandwidth (e.g., 16GB/s) is much smaller than the GPU HBM bandwidth (e.g., 1.5TB/s), SpeCache selects a much smaller number of KV pairs to load (64) compared to QUEST (over 1024), and focuses on parallel loading and computation.

• From the experimental results, SpeCache demonstrates a more efficient use of sparse KV pairs. According to the results in QUEST, when the KV budget is less than 1024, model performance significantly declines on longbench (Fig. 7 of the QUEST paper). However, SpeCache performs well even when only 32 KV pairs are loaded (Fig. 5 in our paper). This is because the low-bit KV cache copy we employ can distinguish important KV pairs at a finer granularity, and the low-bit KV cache itself provides some coarse-grained information during decoding.

Q2. As the sequence becomes longer, will the speculative token be less accurate?

We evaluate 2-bit SpeCache (g=64) and Mistral-7B-Instruct-v0.2 on two long sequence generation tasks, MultiNews and GovReport in Longbench. We track the Speculative Token Top-K KV Cache Hit Rate -- the proportion of the top-k KV cache needed for the next output token that is hit by the top-k KV cache of the speculative token. We separate calculated the hit rate at different decoding step intervals.

We find that as the output grows, the hit rate gradually decreases, but after dropping to around 87%, it stabilizes and maintains a relatively high hit rate.

Q3. Does the speculative decoding introduce more latency and GPU memory usage?

Since both tokens share the same model weight matrices and KV cache during decoding, speculative decoding neither introduces additional latency from loading pages into the GPU HBM nor increases VRAM usage. Furthermore, as the decoding process is IO-bound and GPU units are not fully utilized, the added computational load does not significantly increase the latency. We verify this through experiments. On Mistral-7B-Instruct-v0.2, we test the latency and VRAM usage of KIVI and KIVI+SpeCache with a batch size of 1 and a context length of 64k.

We can find that SpeCache only slightly increases the latency of KIVI.