Search papers, labs, and topics across Lattice.
To resolve the quality degradation that occurs when reusing precomputed KV caches across ultra-long retrieval-augmented generation (RAG) contexts, the authors combine KV-concatenation-aware model fine-tuning with selective KV cache recomputation during inference. While naive cache reuse substantially reduces prefill latency, it disrupts attention representations over long contexts, causing severe retrieval accuracy degradation. Evaluated on the 124k-token RULER benchmark, this dual strategy yields a 9.7-point score improvement over recompute-only baselines while achieving an 80% reduction in time to first token (TTFT) compared to standard full attention.
Prefill latency can be slashed by 80% without sacrificing long-context retrieval accuracy by pairing concatenation-aware fine-tuning with selective KV recomputation.
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.
