ICML 2026poster0 citations

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

Shihao Wang, Jiahao Chen, Yanqi Pan, Hao Huang, Yichen Hao, Xiangyu Zou, Wen Xia, Chongyang Qiu

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).

TransformerRetrievalBenchmark
BibTeX
@inproceedings{
wang2026prophetkv,
title={Prophet{KV}: User-Query-Driven Selective Recomputation for Efficient {KV} Cache Reuse in Retrieval-Augmented Generation},
author={Shihao Wang and Jiahao Chen and Yanqi Pan and Hao Huang and Yichen Hao and Xiangyu Zou and Wen Xia and Chongyang Qiu and Wentao Zhang and PengFei Wang},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=SwGwmKIgUb}
}