ProphetKV: User-Query-Driven Selective Recomputation for Efficient KV Cache Reuse in Retrieval-Augmented Generation

TL;DR

ProphetKV tackles the prefill bottleneck in long-context retrieval-augmented generation by performing user-query-driven selective recomputation of KV caches. It introduces a dual-stage pipeline: a lightweight, query-guided token-importance assessment and a layer-fusion step that yields a unified set of tokens for recomputation, avoiding the deadlock of per-layer decisions. The method preserves cross-attention relevant to the user query, achieving near full recomputation accuracy with only 20% token recomputation, and substantially outperforms state-of-the-art baselines on RULER and LongBench. The approach demonstrates consistent accuracy and efficiency gains across multiple models and long-context benchmarks, highlighting its practical impact for production-grade RAG systems.

Abstract

The prefill stage of long-context Retrieval-Augmented Generation (RAG) is severely bottlenecked by computational overhead. To mitigate this, recent methods assemble pre-calculated KV caches of retrieved RAG documents (by a user query) and reprocess selected tokens to recover cross-attention between these pre-calculated KV caches. However, we identify a fundamental "crowding-out effect" in current token selection criteria: globally salient but user-query-irrelevant tokens saturate the limited recomputation budget, displacing the tokens truly essential for answering the user query and degrading inference accuracy. We propose ProphetKV, a user-query-driven KV Cache reuse method for RAG scenarios. ProphetKV dynamically prioritizes tokens based on their semantic relevance to the user query and employs a dual-stage recomputation pipeline to fuse layer-wise attention metrics into a high-utility set. By ensuring the recomputation budget is dedicated to bridging the informational gap between retrieved context and the user query, ProphetKV achieves high-fidelity attention recovery with minimal overhead. Our extensive evaluation results show that ProphetKV retains 96%-101% of full-prefill accuracy with only a 20% recomputation ratio, while achieving accuracy improvements of 8.8%-24.9% on RULER and 18.6%-50.9% on LongBench over the state-of-the-art approaches (e.g., CacheBlend, EPIC, and KVShare).

Figures (11)

• Figure 1: A high-level overview of different methods of KV Cache reuse in RAG scenarios.
• Figure 2: An example illustrating the selected tokens of existing approaches under a 20% recomputation ratio. Text irrelevant to the user query is colored in gray, and tokens selected by each method are colored in red. See Appendix \ref{['sec:appendix-example-prompt']} for the full prompt.
• Figure 3: Overlap ratio ($|S \cap G| / |G|$) between selected tokens ($S$) and query-attended tokens ($G$), where $G$ is the average query-to-context attention under full-prefill, across selection ratios ($p$).
• Figure 4: The overlap ratio between query-attended tokens and actual critical tokens in the decoding stage. Across various model families, query attention consistently predicts the tokens that will be attended during the decoding stage.
• Figure 5: Overview of the proposed method: utilizing the query as a prophet to guide selective KV cache recomputation.