KVLink: Accelerating Large Language Models via Efficient KV Cache Reuse

We sincerely thank Reviewer 5Ajo for the valuable recognition of our work and insightful questions. Your comments have been a great help in improving and improving our work. Below is a detailed response to the key points you raised.

Q1: Storing documents as KV Cache remains too costly at present.

A1: Thank you for raising this important point. We agree that storing the KV cache is more costly than storing the text. We would like to discuss this trade-off between storage overhead and efficiency, as well as the strategies to reduce the cost, as outlined below.

• First, we can reduce storage overhead by combining KVLink with existing KV cache compression techniques. In our paper, we explore two different strategies and show that the performance loss is minimal with a well-designed compression scheme. We believe the compression method can be further improved to reduce the storage cost.

• Second, GPU cost is more significant than storage cost. As shown in our experiment, serving a Llama3.1-8B model on one A100 80GB GPU, prefilling a 5000-token context with KVLink can reduce latency by about 96%. Therefore, for requests of this length, KVLink can serve 25 times more requests than standard decoding using the same amount of GPU time. For one million such requests, KVLink uses about 9 GPU hours for prefilling, which costs around 16 USD, while standard decoding takes about 246 GPU hours, costing around 440 USD (based on the current market price of ~1.79 USD/hour for an A100 80GB GPU). Meanwhile, storage is much cheaper than GPU resources. For example, on cloud storage like Amazon S3 standard tier, storing 1 GB per month costs about 0.023 USD.

• Finally, we can reduce storage cost through system-level design. Specifically, two strategies can be used:

  • Store the KV cache only for documents with high hit rates. Strategies like Least Recently Used (LRU) or Least Frequently Used (LFU) can help decide which documents to keep in KV cache. For less frequently used documents, we can store them as plain text.

  • Use hierarchical KV cache storage. Less frequently accessed caches can be stored on cheaper storage (e.g., SSD), while more critical ones can be kept in faster memory (e.g., CPU RAM).

In summary, our work focuses on speeding up LLM inference. We leave the full system design that balances efficiency and storage cost as future work.

Q2: Based on the results shown in the paper, most of the experimental improvements are observed in 1B and 3B models, with only marginal gains for the 8B model. Can larger models (30B/70B) be evaluated to demonstrate the general applicability of link tokens?

A2: Thank you for the question. Firstly, we observed that when using the larger models, it is harder to achieve further performance improvements because these models already have high absolute performance, leaving less room for gains compared to smaller models. We agree that it is important to show that the improvements hold across different model sizes. To address the concern, we have trained the KVLink5 model and our strongest baseline, BlockAttention, using Qwen-2.5-32B-instruct as the base model. The result is given below.

We highlight that our method also outperforms the best baseline, BlockAttention, on the 32B model with a 1.3%–2.7% absolute improvement. The only exception is the TriviaQA dataset. We believe the reason is that the training set of TriviaQA is used to train both methods, so their performance on in-distribution data becomes similar. These improvement numbers are consistent with the results on 1B and 3B models, showing that our method works well across different model sizes.

Q3: In the experiments presented in the paper, many experimental designs for certain task types, such as the NQ task, do not fully demonstrate the effectiveness of link tokens.

A3: The tasks that we tested in the paper for the first experiment are QA and summarization tasks. The three reasons that we are testing these tasks are:

1. The prefill-heavy nature of these tasks. These tasks usually have much longer context and short model response, which makes them suitable to benefit from KVLink through reusing context cache.
2. To provide a comprehensive evaluation and comparison, these tasks are all the tasks tested by our baselines.
3. The multi-hop reasoning QA tasks rely more on cross attention across documents, and QA is a task that has strong demand in real-world applications.

Thank you for the response. While my concerns have been addressed, my score will remain unchanged, as the impact of link tokens on large models appears to be limited.

Best,

We sincerely thank Reviewer vcib for the valuable recognition of our work and insightful questions. Your comments have been a great help in improving and improving our work. Below is a detailed response to the key points you raised.

Q1: Storing the entire KV cache demands an overwhelmingly large storage footprint.

A1: Thank you for highlighting this important concern. We acknowledge that KV cache storage is more expensive than storing plain text. Below, we address this trade-off between storage and efficiency, along with possible ways to reduce the cost.

• First, storage overhead can be minimized by combining KVLink with existing KV cache compression methods. In our paper, we study two such strategies and show that with proper compression design, the performance drop is minimal. These compression methods can be further refined to lower the storage burden.

• Second, the cost of GPU usage outweighs that of storage. In our experiments, serving a Llama3.1-8B model on an A100 80GB GPU, using KVLink for prefilling a 5000-token context reduces latency by around 96%. This means that for such requests, KVLink enables 25 times more requests to be served using the same GPU time compared to standard decoding. For one million such requests, KVLink requires roughly 9 GPU hours (costing 16 USD), while standard decoding needs about 246 GPU hours (440 USD), based on the current price (1.79 USD/hour for A100 80GB). Meanwhile, storage costs remain low. For instance, Amazon S3 standard tier charges about 0.023 USD per GB per month.

• Lastly, storage can be further optimized at the system level. Two strategies are applicable:

  • Store KV cache only for documents with high access frequency. Techniques like Least Recently Used (LRU) or Least Frequently Used (LFU) can help choose which documents to cache. Infrequently used documents can be stored as plain text.

  • Use a hierarchical storage system for KV cache. Caches accessed less often can be offloaded to lower-cost storage like SSDs, while high-priority caches can be kept in fast memory such as CPU RAM.

In a nutshell, this work focuses on improving LLM inference speed. Designing a complete system to optimize the efficiency–storage trade-off is an important direction for our future work.

Q2: While extending into the domain of RAG and KV compression is meaningful, the methodology appears similar to AnLLM.

A2: Thank you for the question. Here we want to highlight the differences between our compression method and the AnLLM method from three dimensions:

• We include the KV compression technique to help reduce the storage overhead of pre-computed KV cache. To this extent, we explore different strategies to compress the KV cache, including prompt compression [1], our method, and AnLLM. However, we empirically found that AnLLM does not give good performance after compression. Please find the experiment results below using AnLLM for compression. The AnLLM model here is first trained with the pre-training data as described in the AnLLM paper, and then supervised fine-tuned using QA tasks for 2 epochs, the same as our method.

The sub-optimal performance of AnLLM motivates us to make a few modifications to it. We acknowledge that our method is directly built upon AnLLM with slight changes to improve performance. We also mention this in the main paper (lines 160–164).

• Second, the main difference between our revised compression method and the original AnLLM is the attention mechanism. In AnLLM, the compressed tokens (i.e., the anchor tokens) in the previous chunk can be attended by the following raw tokens and anchor tokens. However, in our method, to follow the cache reuse setting where each chunk is encoded separately, the compressed tokens can only be attended by the link tokens and query tokens, which follow the standard self-attention. We empirically find this attention mechanism more suitable for our application setting.

• Third, as you have mentioned, the intended use cases of the two methods are also different. AnLLM focuses on dynamically dropping the generated tokens with compression during the decoding phase for long context tasks, instead of pre-computing and compressing the KV cache for reuse. In contrast, our method compresses the document in the context into a reusable KV cache for efficient prefilling.

[1] Pan, Zhuoshi, et al. "LLMLingua-2: Data Distillation for Efficient and Faithful Task-Agnostic Prompt Compression." ACL (Findings). 2024.

Q3: The on-the-fly adjustment of RoPE positions is well known as “pre-RoPE”, in KV cache compression and quantization papers.

A3: Thank you for pointing this out. We acknowledge there are multiple ways of restoring the cache position information, and we want to highlight that our key contribution in this paper is proposing the link tokens for restoring the cross attention between the reused context. As discussed in the related work section, current KV reuse works also handle the position information using similar methods. For example, BlockAttention removes the position information inside the caches by rotating the key cache with a negative angle. Nevertheless, we will definitely discuss and credit the pre-RoPE techniques in our final manuscript.

Q4: During training, “link” tokens take positional indices between documents; is this policy maintained at inference time, too? In other words, do link tokens possess non-consecutive positional indices in the implementation?

A4: The use of “link” tokens is the same for both training and inference time, where the positional indices of the “link” tokens are consecutive, ranging from the position of the last token in the preceding document to the position of the first token in the next following document.

Q5: In TTFT experiments (Section 3.3), on which model and GPU were the experiments conducted, and how does TTFT vary with model size?

A5: In TTFT experiments, the Llama3.1-8B-instruct model is used with an A100 80GB GPU. As the model size increases, the reduction in TTFT latency will also increase given the same length of context and same device. This is because larger models take more time to prefill the contexts, whereas loading a larger KV to GPU memory adds only a small extra time cost. We will include this information in our final manuscript.

Q6: In Table 1, it would be helpful to indicate which methods require fine-tuning (BlockAttention, KVLink) and which do not, for clarity.

A6: Thank you for pointing this out. We will include this information in Table 1 in our final manuscript. Specifically, for the KV reuse methods, BlockAttention and KVLink require training, while PromptCache and CacheBlend are training-free methods.

We sincerely thank Reviewer 94Ka for the valuable recognition of our work and insightful questions. Your comments have been a great help in improving and improving our work. Below is a detailed response to the key points you raised.

Q1: Missing TurboRAG as related work.

A1: Thank you for pointing out this important work. We will include the discussion of TurboRAG in our experiments and related work section in the final manuscript. We would like to compare our method with TurboRAG as below.

• First, we want to clarify that our key contribution, the link token mechanism, is different from TurboRAG. Instead, TurboRAG is very similar to our baseline, BlockAttention, where the reused caches are directly concatenated. More specifically,

  • TurboRAG introduces two extra tokens: prepending doc_start to the document and appending doc_end to the document. Similar to BlockAttention, TurboRAG computes and stores the KV cache of the document and the two tokens with local self-attention. That means the two tokens are precomputed offline as part of the document and only maintain local attention within the document to mark the document boundaries.

  • Therefore, these two tokens are mainly used to indicate the boundaries of documents. In contrast, our method introduces link tokens and recomputes their KV cache at testing time to reconnect the separately encoded documents.

  • During the training phase, we also introduce the link tokens and train them with the objective of connecting the separately computed contexts, which brings better performance compared to TurboRAG and BlockAttention.

• Second, we also conduct experiments to empirically compare our method to TurboRAG. We replicate TurboRAG and train it using the same training data and training setup as our method. Below is the result:

We observe a consistent performance improvement of KVLink over all tasks, which shows the importance of our link tokens in restoring cross attention.

Q2: Scalability issue of storing pre-computed KV cache.

A2: We appreciate your thoughtful feedback. We agree that storing KV caches comes with a higher cost than storing the original text. Below, we discuss the trade-off between storage and efficiency and how it can be managed.

• First, combining KVLink with existing KV cache compression methods helps reduce storage usage. Our paper presents two such approaches, demonstrating that a well-designed compression strategy causes minimal performance loss. Ongoing improvements in compression can further reduce the storage demand.

• Second, GPU usage cost is typically higher than storage cost. As shown in our experiments, when running a Llama3.1-8B model on an A100 80GB GPU, KVLink can cut latency by 96% for a 5000-token input. Under fixed GPU hours, this allows KVLink to handle 25 times more requests than standard decoding. For every million of these requests, KVLink takes 9 GPU hours (16 USD), while standard decoding uses 246 GPU hours (440 USD), based on the current rate (1.79 USD/hour for A100 80GB). On the other hand, storage remains inexpensive. For example, Amazon S3’s standard plan charges just 0.023 USD per GB each month.

• Finally, we can lower storage cost through system-level solutions. Two design strategies can be used:

  • Cache only high-hit-rate documents. Strategies like LRU or LFU can help identify which documents to keep in cache. Others can be stored in plain text as usual.

  • We can use tiered KV cache storage to further save storage cost. For example, although a 1,000-token document stored as UTF-8 text occupies about 5KB, whereas its Llama3-8B KV cache requires roughly 131MB. We can put this cache in cheaper storage if it is seldom used. In general, low-access caches should be saved on cheaper storage like SSDs, while important ones stay in faster memory like CPU RAM.

In conclusion, our paper focuses on accelerating LLM inference. We leave the challenge of designing an optimal system that balances performance and storage for future work.

Q3: The performance of the proposed method with different numbers of retrieved documents.

A3: We agree that it is crucial to test the generalizability of KVLink under different numbers of retrieved documents. To address your concern, we evaluate the performance of our method and the best baseline, BlockAttention, on the NaturalQuestions benchmark by varying the number of retrieved documents. The experiment results are as below.

Here we can observe that for almost all the numbers of retrieved documents, the KVLink model outperforms the BlockAttention model with up to 10% performance gain. This result indicates that KVLink can generalize to different numbers of retrieved documents. We will add this experiment in our final manuscript.

We sincerely thank Reviewer 534y for the valuable recognition of our work and insightful questions. Your comments have been a great help in improving and improving our work. Below is a detailed response to the key points you raised.

Q1: Are link tokens just a special set of markers and NOT special tokens (OR) there is an additional mid-training phase where the tokens are baked in?

A1: Thank you for your question. These link tokens are indeed special tokens, designed to restore the cross attention between the documents at test time. Please find our explanations below.

• First, we highlight that the hidden representations of these link tokens are recomputed at inference time using standard self-attention. Instead of treating them as markers to separate each document and saving their KV cache together with the documents, the KV cache of these tokens is recomputed to restore the cross attention between documents. We do not have a separate mid-training phase for link tokens. Instead, the embeddings of link tokens and all other tokens are trained together during the fine-tuning stage. Through this stage, the link tokens learn to compensate for the missing cross-attention between the separately computed KV caches.

• Second, we empirically verify whether the link tokens act as special tokens and not just as markers to separate documents. We conduct another experiment where we insert two special tokens, doc_start and doc_end, as markers to indicate the boundaries of each reused document. Unlike link tokens, these two tokens are precomputed offline with the document and only maintain local attention inside the document. This method uses the same training process and computation as our method. Below is our experimental result:

We observe a consistent performance improvement with link tokens. This performance gap shows the necessity of using link tokens to restore cross attention.

Q2: How does KVLINK perform in scenarios where the overlap between context segments across queries is low? Is there a break-even point where the overhead of precomputing and managing caches outweighs the benefits?

A2: Thank you for this insightful question! We agree that KVLink speeds up pre-filling by reusing the cache across different queries, but when the overlap between the queries becomes low, the cost of managing the cache may outweigh the benefits it brings. One extreme example is when all queries retrieve different documents.

However, we believe that the overhead in this case can be reduced in two ways:

• Tier-based caching system. We believe the cache management cost can be further optimized through system-level design. Specifically,

  • Some documents are frequently retrieved, while others are rarely reused. So it is more storage-efficient to store only the cache of documents with high hit rates. Here, strategies like Least Recently Used (LRU) or Least Frequently Used (LFU) can be applied.

  • It is helpful to offload less-used caches to cheaper storage tiers and reserve fast memory for the most important ones. In practice, we can set up a hierarchical cache across different hardware: keep the most-used KV caches in GPU VRAM (for fastest reuse), move moderately-used caches to CPU RAM, and store rarely-used caches on SSD or disk.

• On-the-fly KV computation. Instead of precomputing caches for all documents offline in advance, we can compute the cache for a document only when it is retrieved. This way, only documents that have been retrieved before are cached, which avoids overhead from managing rarely used documents. This strategy can also be combined with the tier-based caching system above as an engineering solution when deploying the RAG system.

We also want to highlight that the benefit of KVLink is to save GPU cost by serving more requests with the same GPU hours, and GPU cost is usually much higher than the storage cost of the cache.

Q3: Is there any degradation in model performance for tasks that require strong cross-document reasoning?

A3: Many of our experiment tasks already require strong cross-document reasoning. In particular, 2WikiMultiHopQA, HotpotQA, MuSiQue, and TriviaQA all require the model to combine knowledge from multiple documents to produce the correct answer. Although the absolute accuracy on MuSiQue is lower because it is a harder task, we observed that for all of these tasks, our KVLink method gives a consistent performance improvement over all the baseline methods.