ProphetKV: User-Query-Driven Selective Recomputation for Efficient KV Cache Reuse in Retrieval-Augmented Generation
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).
Lay Summary
Large language models often answer questions by first reading many retrieved documents. This can be slow because the model must process all retrieved text before generating the first word of its answer. A common shortcut is to reuse stored intermediate results from document chunks, but this can reduce accuracy because the chunks were processed separately and may miss connections needed to answer the question. This paper introduces ProphetKV, a method that makes this reuse both faster and more reliable. Instead of trying to repair every missing connection among the retrieved documents, ProphetKV uses the user’s question to identify which document tokens are most likely to matter for the answer, then recomputes only those parts. It also combines signals from multiple layers inside the model, since different layers can focus on different kinds of evidence. In tests that ask models to find and reason over information in long documents, ProphetKV kept almost the same answer quality as fully processing all text, while redoing only one fifth of the computation-heavy token processing. Compared with strong existing cache-reuse methods, it improved answer accuracy by 8.8% to 24.9% on one benchmark suite and 18.6% to 50.9% on another, while producing the first response token up to 5x faster than full recomputation. These results suggest that long-document question answering can become more practical when models spend their computation on the evidence most relevant to the user’s question.