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, Wentao Zhang, Pengfei Wang
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:125415-125434, 2026.

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

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-wang26y, title = {{P}rophet{KV}: User-Query-Driven Selective Recomputation for Efficient {KV} Cache Reuse in Retrieval-Augmented Generation}, author = {Wang, Shihao and Chen, Jiahao and Pan, Yanqi and Huang, Hao and Hao, Yichen and Zou, Xiangyu and Xia, Wen and Qiu, Chongyang and Zhang, Wentao and Wang, Pengfei}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {125415--125434}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/wang26y/wang26y.pdf}, url = {https://proceedings.mlr.press/v306/wang26y.html}, 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).} }
Endnote
%0 Conference Paper %T ProphetKV: User-Query-Driven Selective Recomputation for Efficient KV Cache Reuse in Retrieval-Augmented Generation %A Shihao Wang %A Jiahao Chen %A Yanqi Pan %A Hao Huang %A Yichen Hao %A Xiangyu Zou %A Wen Xia %A Chongyang Qiu %A Wentao Zhang %A Pengfei Wang %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-wang26y %I PMLR %P 125415--125434 %U https://proceedings.mlr.press/v306/wang26y.html %V 306 %X 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).
APA
Wang, S., Chen, J., Pan, Y., Huang, H., Hao, Y., Zou, X., Xia, W., Qiu, C., Zhang, W. & Wang, P.. (2026). ProphetKV: User-Query-Driven Selective Recomputation for Efficient KV Cache Reuse in Retrieval-Augmented Generation. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:125415-125434 Available from https://proceedings.mlr.press/v306/wang26y.html.

Related Material