PaperlayerView exact revision on arXiv

Source-linked AI summary

KVLink: Accelerating Large Language Models via Efficient KV Cache Reuse

Jingbo Yang, Bairu Hou, Wei Wei, Yujia Bao, Shiyu Chang

arXiv:2502.16002v4cs.CL

TL;DR

LLMs redundantly re-encode overlapping context segments across queries, motivating KVLink’s independent document-level KV caching and reuse. KVLink corrects positional shifts and restores cross-document attention, improving QA accuracy while reducing time-to-first-token latency and supporting cache compression.

  • Problem

    LLMs redundantly re-encode shared context segments for different queries, despite unchanged document content and potential for reuse.

  • Method

    KVLink precomputes each document’s KV cache independently, concatenates retrieved caches at inference, re-encodes positions globally, and reconnects documents with link tokens.

  • Results

    Across diverse datasets, KVLink outperforms state-of-the-art baselines, with up to 96% lower time-to-first-token latency and reported gains of 6.6% on Natural Question and 7.3% on HotpotQA.

  • Takeaways & Limitations

    KVLink provides an efficient context-reuse strategy that can also be combined with KV-cache compression for large-scale deployment.

  • Takeaways & Limitations

    Certain benchmarks, including ARC-C and Winogrande, show minor drops versus the original models, typically below 3%.

Abstract

from arXiv · show

We describe KVLink, an approach for efficient key-value (KV) cache reuse in large language models (LLMs). In many LLM applications, different inputs can share overlapping context, such as the same retrieved document appearing in multiple queries. However, the LLMs still need to encode the entire context for each query, leading to redundant computation. In this paper, we investigate a new strategy to eliminate such inefficiency, where the KV cache of each document is precomputed independently. During inference, the KV caches of retrieved documents are concatenated, allowing the model to reuse cached representations instead of recomputing them. To mitigate the performance degradation when using KV caches computed independently for each document, KVLink introduces two key techniques: adjusting positional embeddings of the KV cache at inference to match the global position after concatenation, and using trainable special tokens to restore self-attention across independently encoded documents. Experiments across 7 datasets demonstrate that KVLink improves question answering accuracy by an average of 4% over state-of-the-art methods. Furthermore, by leveraging precomputed KV caches, our approach reduces time-to-first-token by up to 96% compared to standard LLM inference, making it a scalable and efficient solution for context reuse. Additionally, KVLink can be combined with KV cache compression to further save cache loading and storage overhead while outperforming the baselines.

1 Introduction

KVLink targets redundant context encoding by precomputing document-level KV states for reuse, then addresses the accuracy and attention gaps caused by independent encoding. Experiments report stronger accuracy than baselines, sharply reduced time-to-first-token latency, and compatibility with KV-cache compression.

  • Motivation: Conventional LLM inference re-encodes concatenated contexts for every query, redundantly processing shared documents and incurring high prefilling costs.This prevents separately encoding and reusing precomputed key-value states for recurring context segments.
  • Approach: KVLink precomputes KV states for each document independently and concatenates the cached states during inference to eliminate redundant computation.The approach is designed for contexts containing reusable documents or segments.
  • Challenge and solution: Up to 35% relative QA accuracy decreases have been reported for naive independent document encoding because of positional discrepancies and missing cross-document attention.KVLink introduces positional re-encoding and cross-document reconnection to address these issues.
  • Results: 6.6% on Natural Question and 7.3% on HotpotQA are reported improvements over the best baseline, while time-to-first-token latency falls by up to 96%.The method sacrifices only minimal performance relative to conventional full encoding.
  • Results: KVLink remains effective when combined with LLMLINGUA and ANLLMS KV-cache compression, supporting lower cache loading and storage overhead.The paper presents this combination as relevant to practical large-scale deployment.

2 Methodology

KVLink independently caches retrieved documents, concatenates their KV states at inference, and restores the positional and cross-document interactions lost by context-free encoding. The method also explores link-token-based reconnection and compressed KV caches to reduce reuse overhead.

  • Problem Formulation: Current pipelines encode each query’s full concatenated context, so KV states for shared documents remain conditioned on query-specific preceding documents and cannot be directly reused.This creates redundant computation across queries with overlapping retrieved segments.
  • Problem Formulation: KVLink precomputes every document’s KV cache independently, then concatenates retrieved caches during inference to reuse overlapping documents across queries.Only the retrieved cache states are loaded and combined rather than re-encoding the documents.
  • Challenges: Independent encoding can reduce QA accuracy by up to 35% because later documents lose cross-document dependencies and separately cached positions mismatch full-sequence positions.KVLink addresses both sources of degradation.
  • KV Cache Positional Re-encoding: KV cache positional re-encoding stores key-value states without positional transformations and applies each token’s global rotary embedding after cache concatenation.This aligns cached positions with their locations in the full inference sequence with negligible added time.
  • Cross-Document Reconnection with Link Tokens: KVLINK’s trainable link tokens reconnect documents by attending to preceding documents and current-document tokens through a customized attention map.The mechanism preserves local causal attention within documents while transmitting information across segments.
  • Cross-Document Reconnection with Link Tokens: Inference loads and concatenates precomputed document caches, then processes newly appended link tokens with the customized attention map.This avoids recomputing retrieved-document KV caches while matching the intended training-time attention behavior.
  • Compressed KV Cache Linking: A 1,000-token document occupies about 5KB as UTF-8 text but roughly 131MB as an Llama3-8B KV cache, motivating cache compression.The paper considers token dropping with LLMLINGUA and anchor-token compression with ANLLMS.
  • Compressed KV Cache Linking: Compressed KV linking divides documents into fixed-size chunks and stores multiple anchor-token KV caches so remaining tokens attend only to compressed representations.The default chunk size is 100 tokens, and increasing anchor tokens can better preserve performance.

3 Experiments

KVLINK is evaluated for separately encoded KV caches, inference efficiency, capability preservation, and compressed-cache reuse across QA, summarization, and general-ability settings. It consistently improves benchmark performance, reduces TTFT by 85%–96%, preserves general capabilities, and benefits from link tokens and cache compression.

  • 3.2 Main Results: KVLINK consistently outperforms baselines across evaluated QA and summarization datasets, surpassing BLOCKATTENTION by up to 5% on most QA tasks.It also achieves accuracy close to fine-tuned Llama evaluated with fully concatenated contexts.
  • 3.2 Main Results: Link tokens bridge separately encoded documents by restoring inter-document connections and improve performance with negligible computational cost.Performance consistently increases as the number of link tokens grows; without them, BLOCKATTENTION performs worse despite using the same training data.
  • 3.3 Inference Efficiency Evaluation: 85%–96% lower TTFT is achieved than standard decoding when precomputed KV caches are loaded, positionally re-encoded, and linked.These operations introduce only negligible latency, and the efficiency gap widens as reused context length increases.
  • 3.4 General Capability Preservation: KVLINK maintains performance comparable to fine-tuned Llama models across general reasoning and instruction-following tasks.The method does not significantly degrade general capabilities despite restructuring context encoding.
  • 3.4 General Capability Preservation: Less than 3% drops occur on certain benchmarks, including ARC-C and Winogrande, compared with original models.The authors identify data-mixture refinement and additional reasoning and instruction-following data as potential future improvements.
  • 3.5 KVLINK with Cache Compression: Adding link tokens consistently mitigates accuracy losses introduced by KV cache compression across QA benchmarks.The modified ANLLMS compression retains more information than LLMLINGUA, producing stronger performance across QA benchmarks.

4 Related Work

Prior work improves LLM efficiency through model and decoding optimizations, KV-cache reuse, retrieval integration, and cache compression. KVLINK combines cache reuse with complementary compression methods to address both computation and storage overhead.

  • Efficiency research includes pruning, quantization, and optimized decoding methods such as speculative decoding and early exiting.
  • KV-cache reuse methods mainly target prefix reuse, while arbitrary-position approaches can suffer from cross-chunk attention loss or positional mismatches.
  • KV compression reduces memory through cache eviction or quantization, and KVLINK can be integrated with these techniques.
  • Retrieval-augmented generation commonly retrieves relevant documents using dense retrievers or BM25 before integrating them into the model context.

5 Conclusion

The paper presents KVLINK as a method for improving LLM efficiency by reusing precomputed KV caches for retrieved documents or context segments. It identifies future work in data mixture, fine-tuning, and real-world deployment.

  • KVLINK reuses precomputed KV caches to avoid redundant computation for overlapping contexts across queries.
  • Future work includes refining the data mixture, investigating fine-tuning strategies, and exploring real-world deployment scenarios.

NeurIPS Paper Checklist

The checklist assessment states that the paper’s main claims match its stated KVLINK contribution and experimental validation. It also notes that claims should reflect assumptions, limitations, and generalization scope.

  • The checklist judges the abstract and introduction’s claims accurate representations of KVLINK’s contribution and experimental support.
  • The checklist emphasizes that paper claims should clearly state contributions, assumptions, limitations, and expected generalization scope.

2. Limitations

The checklist reports that limitations are discussed in Appendix A.9, while experimental, reproducibility, resource, ethics, and societal-impact information is addressed through the paper or supplemental material. The paper reports no theoretical results or human-subject research.

  • The paper’s limitations are reported as discussed in Appendix A.9.
  • The paper reports no theoretical results, so theorem assumptions and proof requirements are not applicable.
  • The authors state that code, datasets, KVLINK models, and reproduction details will be released in supplemental material.
  • The paper states that training and test details are provided in Section 3.1 and Appendix A.2.
  • The paper reports compliance with the NeurIPS Code of Ethics and states that it does not involve crowdsourcing or human-subject research.
  • The checklist indicates that societal-impact discussion should address potential harms, including malicious use, fairness, privacy, and security concerns.

A.1 Data Mixture

KVLINK is fine-tuned on a mixed dataset spanning retrieval-augmented QA, multi-turn conversation, summarization, standard QA, instruction following, and pre-training data. This mixture trains the model to integrate independently encoded contexts while preserving broader capabilities.

  • Retrieval-augmented QA: Retrieval-augmented QA uses separately encoded retrieved passages from TriviaQA and 2WikiMQA with GPT-4-generated reference answers.Ten Contriever-retrieved Wikipedia passages are independently encoded for each question.
  • Multi-turn conversation: Multi-turn conversation training independently encodes earlier turns so the model integrates segmented context while producing later responses.The model is trained to generate appropriate responses in subsequent turns.
  • Summarization: Summarization training splits XSum documents into consecutive segments, independently encodes them, and trains the model to generate coherent summaries.Each segment is independently encoded into KV cache before summary generation.
  • Training format: Training examples are truncated to 4096 tokens, with 100-token reused segments for summarization and loss computed only on the final assistant response for conversations.Each prior user–assistant conversation is independently encoded as reused context.

A.2 Implementation Details

KVLINK fine-tuning uses a mixed dataset and task-specific prompts, with explicit markers and indexed link tokens organizing reused contexts. Evaluation measures QA accuracy using exact substring matching, alongside few-shot general-capability tests.

  • Training: Fine-tuning uses a mixed dataset to teach integration of disjoint context segments while preserving original capabilities.The training mixture is drawn from different sources.
  • Prompts: GPT-4-generated system prompts are randomly selected during training for SFT, QA, and summarization task categories.QA training also shuffles retrieved reused documents in the context.
  • Context reuse: KVLINK marks reused-context boundaries with KV-START and KV-END special tokens and assigns link tokens according to document index.The data-preprocessing pipeline uses these markers to organize reused documents or context segments.
  • Context reuse: Multi-turn examples place independently encoded earlier turns between KV-START and KV-END before the current user request and assistant response.The sequence includes indexed link tokens between prior conversation turns.
  • Context reuse: Link tokens for the n-th reused document follow the indexed pattern linkn-1, linkn-2, ..., linkn-K when K tokens are used.The token pattern depends on the document’s index among reused documents.
  • Evaluation: QA evaluation uses accuracy, counting a prediction as correct when any predicted substring exactly matches any gold answer.General capability preservation is evaluated with few-shot GSM8K and MMLU examples.

A.3 More Implementation Details of Baselines

The baseline implementations differ in how they concatenate and position reused KV caches. BLOCKATTENTION concatenates independently encoded caches without special tokens, while PROMPTCACHE requires positional handling for accurate reuse.

  • BLOCKATTENTION: BLOCKATTENTION concatenates independently encoded reused-document caches directly without inserting special tokens between them.Examples without context reuse use the same data processing as KVLINK.
  • PROMPTCACHE: PROMPTCACHE originally leaves positional encoding unchanged when caches are reused, creating discontinuous position information.The evaluation maximizes PROMPTCACHE performance by supplying gold position information.

A.4 Ablation Studies on Data Mixure

The data-mixture ablation evaluates how removing summarization, multi-turn conversation, or all but QA affects KVLINK5. The reported results indicate that cross-chunk reasoning tasks and diverse cache-reuse data support robust cache reuse.

  • Ablation setup: Three KVLINK5 configurations remove summarization, remove multi-turn conversation, or retain QA only.The first two configurations preserve relative proportions among remaining tasks, while QA Only trains for two epochs to limit overfitting.
  • Ablation findings: Tasks requiring cross-chunk reasoning are essential for robust cache reuse, and multiple cache-reuse task types improve the training mixture.These conclusions follow from the reported data-mixture experiments.

A.5 Impact of Answer Document Position

KVLink’s QA training data must prevent models from relying on answer-document position. Without shuffling, performance drops when the correct document appears farther back in the context.

  • Contriever-based QA data tends to place answer-containing documents near the front, encouraging models to search early context rather than reason across reused documents.
  • Performance drops significantly when the ground-truth document is located farther back in the context.
  • Shuffling retrieved documents during QA training is necessary to reduce the positional bias created by relevance-ranked retrieval.
Loading 2502.16002v4…