PaperlayerView exact revision on arXiv

Source-linked AI summary

Fine-Tuning a KV Cache Concatenation-Aware Model or Recomputing KV Caches? Why Not Both?

Fumihiko Tachibana, Daisuke Miyashita, Jun Deguchi

arXiv:2609.09768v1cs.LGcs.AIcs.CL

TL;DR

Long RAG contexts increase prefill workload and TTFT, while it remains unclear whether KV-cache reuse preserves response quality at very long lengths. The paper combines concatenation-aware fine-tuning with selective KV-cache recomputation and reports higher long-context accuracy alongside lower TTFT. At 124k tokens, the method improves the RULER score by 9.7 points over CacheBlend-only and reduces TTFT by 80% versus full attention.

  • Problem

    KV-cache reuse reduces TTFT for long contexts, but response-quality preservation as context length becomes very long remains underexplored.

  • Method

    The method combines fine-tuning that accounts for KV-cache concatenation with selective recomputation of a subset of KV caches.

  • Results

    9.7 points higher RULER score than CacheBlend-only at 124k tokens, with TTFT reduced by 80% versus full attention.

  • Takeaways & Limitations

    Combining fine-tuning and selective recomputation improves long-context generation quality while retaining KV-cache loading's TTFT benefits.

  • Takeaways & Limitations

    The evaluation did not fully compare attention-score-based recomputation methods, and selective recomputation's benefit depends on the attention implementation used.

Abstract

from arXiv · show

In Retrieval-Augmented Generation (RAG) systems, a large number of retrieved chunks are concatenated to form the input context so that users can receive high-quality responses based on external knowledge. As a result, the input context length increases substantially, leading to a larger prefill workload and, in turn, a longer time to first token (TTFT). While previous works that reuse precomputed key-value (KV) caches effectively reduce TTFT for long-context inputs, it remains unclear whether response quality is preserved when the input context becomes very long. In this paper, we propose a combined approach that (i) fine-tunes the model while taking KV cache concatenation into account and (ii) selectively recomputes a subset of the KV caches. By applying both techniques, we demonstrate improved accuracy for long-context inputs. Experiments on the RULER benchmark show that, for a 124k-token input, our method improves the RULER score by 9.7 point over the baseline that recomputes KV caches only. Moreover, TTFT is reduced by 80% compared with full attention.

1 Introduction

Longer RAG contexts increase prefill cost and TTFT, while concatenated KV-cache reuse can degrade generation quality because reused non-prefix caches lack preceding cross-attention. The paper combines concatenation-aware fine-tuning with selective KV-cache recomputation to improve long-context quality and latency.

  • RAG retrieval increases context length, making quadratic prefill computation and KV-cache construction costly and increasing TTFT.
  • KV-cache concatenation improves reuse beyond prefix caching but introduces deviations from full-attention caches because preceding cross-attention is missing.
  • Existing recomputation and fine-tuning methods improve quality on some benchmarks, but accuracy under very long contexts remains underexplored and declines as context length increases.
  • The proposed method combines concatenation-aware model fine-tuning with selective recomputation of a small fraction of KV caches.
  • 9.7 points higher RULER score over CacheBlend-only at 124k tokens, while TTFT is reduced by 80% when caches are loaded from SSD.

2 Related Works

Prior work addresses concatenated-cache quality degradation through cache recomputation, fine-tuning, or attention-score and positional adjustments. However, long-context accuracy evidence is limited, and combining fine-tuning with other approaches had not been examined.

  • Prompt Cache and PIE improve cache merging through prompt markup or positional alignment, but quality degradation remains because retrieved chunks lack cross-attention.
  • Recomputation methods select KV-cache positions using deviation rules, auxiliary models, or fixed chunk regions.
  • Fine-tuning methods train models for absent cross-attention between chunks so concatenated KV caches can be reused with less degradation.
  • Other methods mitigate attention-sink or attention-score differences by adding prefix tokens, changing softmax temperature, or scaling attention scores.
  • Combining CacheBlend, Link0, and APE had been explored, but combining a fine-tuned model with other approaches had not been examined.
  • Few studies examine how response accuracy depends on context length, despite prefill recomputation scaling as O(n^2) and reported speed gains for long contexts.

3 Methodology

The methodology evaluates how fine-tuning and selective KV-cache recomputation affect long-context accuracy and key-entry deviations. It combines Block-attention with CacheBlend and overlaps cache loading, transfer, merging, and recomputation to improve efficiency.

  • 3.1 Issues in existing approaches: response quality at long context inputs: At 4k tokens, accuracy is close to full attention, but the gap widens as context length increases through 124k tokens.The evaluation compares 4k, 8k, 16k, 32k, 64k, and 124k-token inputs.
  • 3.2 Motivation of proposed method: effect of fine-tuning/recomputing on the deviations in the key entries: Block-attention fine-tunes the model for absent cross-attention between documents, while CacheBlend selectively recomputes KV caches at positions with large deviations.The two approaches address KV-cache concatenation through model adaptation and targeted recomputation, respectively.
  • 3.2 Motivation of proposed method: effect of fine-tuning/recomputing on the deviations in the key entries: Combining Block-attention and CacheBlend achieves the smallest key-entry deviations because they reduce deviations through different mechanisms.CacheBlend checks large first-layer deviations, whereas Block-attention changes model parameters across all layers.
  • 3.2 Motivation of proposed method: effect of fine-tuning/recomputing on the deviations in the key entries: CacheBlend reduces large key-entry deviations more broadly than EPIC because EPIC recomputes only around each chunk’s attention sink.The comparison uses no recomputation, CacheBlend at 15%, and EPIC at the first layer.
  • 3.3 Other Candidates for KV Cache Recomputation and Performance Comparison: The Block-attention and CacheBlend combination improves accuracy while keeping inference latency essentially the same as CacheBlend alone.This combination was selected because it achieved the best accuracy among the evaluated strategies.
  • 3.3 Other Candidates for KV Cache Recomputation and Performance Comparison: The implementation overlaps SSD cache loading, GPU transfer with merging and position adjustment, and transformer-layer recomputation.Layer-wise execution uses ThreadPoolExecutor for loading and CUDA streams for later stages.

4 Experiments

Experiments evaluate the combined fine-tuning and selective KV-cache recomputation approach across long-context RULER, RAG, and general tasks. The approach generally preserves accuracy better than using either technique alone while reducing TTFT.

  • Generation Quality on the RULER Benchmark: The proposed approach shows the smallest degradation from full attention across RULER input context lengths for Llama3.1-8B-Instruct.The same performance pattern is reported with Qwen2.5-7B-Instruct.
  • Generation Quality on the RULER Benchmark: At 124k tokens, reducing the number of chunks recovers the RULER score toward the full-attention result.The proposed approach’s improvement over CacheBlend-only becomes larger as the number of documents per chunk decreases.
  • Generation Quality on the RULER Benchmark: As input context grows, the score gap relative to 100% recomputation widens, while Block-attention with 15% recomputation has the smallest accuracy degradation.This indicates that combining fine-tuning with selective recomputation improves scores when KV caches are merged.
  • Performance in RAG and General Tasks: The proposed approach achieves the best accuracy among non-full-attention methods on four RAG benchmarks.It also improves performance on the evaluated few-shot ICL tasks, with selective recomputation benefiting general tasks.
  • Token-Length Dependence of TTFT with the Proposed Approach: KV-cache loading becomes advantageous over full prefill from 8k tokens onward, reducing TTFT by 80% at 124k tokens.The Block-attention fine-tuned model gains accuracy over the Normal model without additional TTFT cost.

5 Conclusions

The paper combines KV-cache-concatenation-aware fine-tuning with selective KV-cache recomputation to address quality degradation in long-context inputs. This combination reduces generation-quality degradation relative to either technique alone while reducing TTFT by 80% at 124k tokens when caches are loaded from SSD.

  • Conclusion: Figure 10 relates averaged TTFT over 20 tasks to RULER HQA accuracy at a 124k-token input.The experiment uses Llama3.1-8B-Instruct with 16 documents per cache file.
  • Conclusion: The proposed method combines KV-cache-concatenation-aware fine-tuning with recomputation of selected KV caches for long-context inputs.The combination targets degradation that occurs when precomputed KV caches are concatenated.
  • Conclusion: The combined technique reduces generation-quality degradation compared with either existing technique alone.The paper also reports an 80% TTFT reduction at a 124k-token input when KV caches are loaded from SSD.

Limitations

The paper identifies limitations in selective-recomputation efficiency, evaluation coverage, untested method combinations, and fine-tuning of thinking models.

  • Limitation 1: Selective recomputation requires a faster attention mechanism that supports selected query indices to fully realize its benefits.Using an attention mask is necessary to avoid attention to future KV positions and associated quality degradation.
  • Limitation 2: Attention-score-based recomputation methods such as A3 and KVShare were not evaluated in detail against the proposed approach.The authors note that these methods could achieve better performance than CacheBlend.
  • Limitation 3: The combination of Link0 and APE was not evaluated because the study focused on KV-cache recomputation with model fine-tuning.The authors report that this combination could achieve better performance.
  • Limitation 4: The study did not fine-tune thinking models such as Qwen3 and gpt-oss.A new strategy is required to incorporate Block-attention-style design without reducing chain-of-thought capability.

A Detailed Results for NIAH tasks

The detailed NIAH results report Single and Multi-key NIAH evaluations within the RULER benchmark at 4k-token input using Llama3.1-8B-Instruct.

  • The detailed results cover the Single NIAH 1–3 and Multi-key NIAH 1–3 tasks.
  • Table 5 reports Single NIAH evaluation results at a 4k-token input.
  • Table 6 reports Multi-key NIAH evaluation results at a 4k-token input.

B Data Proportions

The appendix documents training and evaluation data quantities, including separate NIAH results at 124k-token input and chunk counts for RULER configurations.

  • Samples with code-related context are used only for full attention, making the Block-attention training set smaller than the Normal/Block-attention set.
  • Tables 7 and 8 report Single and Multi-key NIAH evaluation results at a 124k-token input using Llama3.1-8B-Instruct.
  • Tables 11 and 12 report average RULER chunk counts for Llama3.1-8B-Instruct and Qwen2.5-7B-Instruct, respectively.
  • Table 9 reports the quantities of data used during model fine-tuning.

C Comparison of TTFT between Flashinfer and FlashAttention-2

The TTFT comparison examines full attention with FlashAttention-2 alongside recomputation and KV-cache-loading approaches across input context lengths.

  • FlashAttention-2 reduces full-attention computation time, making its TTFT comparable to 15% recomputation with Flashinfer.
  • With one document per cache file, Block-attention has TTFT comparable to the proposed approach because many cache files add overhead beyond I/O-limited reading.
  • Loading KV caches becomes advantageous when context length reaches 32k tokens or more under the described condition.

E RULER benchmark results using KVShare

The KVShare evaluation examines attention-score-based recomputation and compares it with the proposed fine-tuned-model combination on RULER configurations.

  • KVShare is an attention-score-based recomputation method evaluated only during the prefill stage.
  • The combination of the fine-tuned model and recomputation remains effective when KVShare is used.
  • At 124k-token input, CacheBlend achieves a better RULER score than KVShare.
  • Figure 11 compares averaged TTFT over 20 RULER HQA tasks with 16 documents per cache file.

F Evaluation in the RULER HQA and four RAG benchmarks

The evaluation covers the RULER HQA and four RAG benchmarks, using prompt examples and measurements of TTFT, accuracy, and RULER scores across context lengths. Experiments include Llama3.1-8B-Instruct and Qwen2.5-7B-Instruct, with comparisons involving KVShare and its combination with a fine-tuned model.

  • Benchmark setup: The evaluation uses the RULER HQA and four RAG benchmarks with a shared input-context prompt format.The prompt example is provided in Figure 16, and the RULER HQA chunk count differs slightly from QA2 because of task handling.
  • Implementation: The generated answers used for the training dataset were produced with GPT-4o mini and compiled into a dataset.The manuscript also states that AI assistants were used for paraphrasing support.
  • Latency evaluation: Figure 12 reports averaged TTFT over 20 RULER HQA tasks for Llama3.1-8B-Instruct.The number of documents per cache file is set to one.
  • Latency and accuracy: Figure 13 relates averaged TTFT over 20 RULER HQA tasks to accuracy at a 124k-token input.The experiment uses Llama3.1-8B-Instruct and one document per cache file.
  • RULER results: For Llama3.1-8B-Instruct, combining the fine-tuned model with KVShare achieves a better RULER score than KVShare alone across several context lengths.These results are reported in Figure 14.
  • RULER results: For Qwen2.5-7B-Instruct, combining the fine-tuned model with KVShare also achieves a better RULER score than KVShare alone across several context lengths.These results are reported in Figure 15.
Loading 2609.09768v1…