SAFE-KV: Causal Suffix Repair for Behavior-Preserving Non-Prefix KV Cache Reuse in RAG

Xingjian Zhang

cs.CLOct 10, 2026

SAFE-KV: Causal Suffix Repair for Behavior-Preserving Non-Prefix KV Cache Reuse in RAG

Authors: Xingjian Zhang

Organizations: Carnegie Mellon University, Pittsburgh, PA, USA

Abstract

Retrieval-augmented generation (RAG) reuses the same documents across many requests, making non-prefix KV-cache reuse attractive. Yet a cached document state is not reusable text: it is a context-conditioned causal state whose keys and values depend on the prefix, position, and preceding documents under which it was computed. We formulate non-prefix KV reuse as a behavior-preservation problem with full prefill on the exact prompt as the reference, and introduce SAFE-KV, which repairs a composed cache by re-encoding the causal suffix from a selected start document. On Qwen2.5-7B-Instruct, prefix-conditioned composition preserves only 47.2% of HotpotQA-500 and 37.8% of 2Wiki-500 full-prefill answers; SAFE-KV with risk-guided K2 selection restores 99.8%/98.4% (donor-copy protocol) and 100.0%/98.4% with actual suffix forwarding. Document-only and fixed-window repairs recover far less, because the document states after the repaired region keep histories that omit it. Reuse without repair is exact up to four documents and degrades at sixteen, where suffix repair recovers 98-100%. Faithful repair re-encodes most of the context, so in a single-request L40S study SAFE-KV's operating rule serves every workload by full prefill. SAFE-KV thus establishes the causal-state contract for non-prefix reuse and a validated repair mechanism that meets it.

Figures & tables

Figure 1: SAFE-KV treats a non-prefix cache hit as a causal-state contract. (a) Full prefill encodes every document under its true history. (b) Prefix-conditioned composition encodes each document after the prefix alone, so later slices and the query miss their predecessors and the answer can drift. (c) Re-encoding only the selected document leaves every later slice with a history that omits it. (d) SAFE-KV re-encodes the causally closed suffix from the selected start

s

, restoring a consistent history from

s

onward.

Figure 2: Local repair versus suffix repair on 2Wiki-100 (values of Table 3 ; donor-copy protocol; risk selector fixed). Document-only (Doc K1) and window (w2–w4) re-encoding recover answer preservation gradually; the suffix from the same start (Suffix K1) recovers 97.0% at a budget close to w4’s. Plotted values: figure_local_vs_suffix.data.csv (artifact).

Figure 3: Position sensitivity (values of Table 5 ; donor-copy protocol; 50 examples per position). With the evidence document first, patching only that document is no better than prefix composition, because every later slice keeps its mismatched history; suffix repair removes most of the position dependence. Plotted values: figure_position_sweep.data.csv (artifact).

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Robustness checks (donor-copy protocol; plotted values in figure5_external_sanity.data.csv ). Left: on examples where full prefill is exact-match correct (HotpotQA 272, 2Wiki 226), risk suffix K2 preserves nearly all correct full-prefill behavior. Right: the Mistral-7B noisy diagnostic (250 examples) shows the same unsafe-composition failure and suffix-repair recovery pattern; its query and hybrid rows are largely decided by the id tie-break (Section 6.5 ).

Path	Ans. eq. [95% CI]	Token exact	Gold EM	Changed-wrong
Full prefill	1.000	1.000	0.920	0
Donor-copy Risk K1 (diagnostic)	0.640 [0.50, 0.76]	0.640	0.600	18
Actual-forwarding Risk K1	0.440 [0.30, 0.58]	0.440	0.440	27

Appendix

Table 20: Noise64 quality gate on the exact 50 examples and suffixes timed in Table 7 (mean repair ratio 0.519 for both paths; identical selected documents). Donor-copy and actual forwarding differ by +0.20 [0.06, 0.34] in normalized answer preservation (paired bootstrap). The mechanical median speed-up of Risk K1 on this workload therefore does not correspond to a behavior-preserving operating point.

Figure 5: Quality–budget and latency view. Left: 2Wiki-500 operating points under the donor-copy protocol (values of Table 1 ; inset: high-fidelity cluster). Right: median time to first logit for full prefill and risk suffix K1 with actual forwarding and selector planning on one L40S (values of Table 7 ; percentages are the mean re-encoded context fraction). Risk K1 is slower on four workloads, and the Noise64 median crossover preserves only 44% of full-prefill answers with worse p95 latency, so the operating rule serves every workload by full prefill. Every plotted value is listed in figures/figure4_budget_gate_pareto.data.csv .

Experiment	Key result	Interpretation	Design decision
External 7B composition	Prefix composition preserves only 47.2% HotpotQA and 37.8% 2Wiki full-prefill answers.	Non-prefix document reuse is a behavioral problem even after matching the instruction prefix.	Make full-prefill behavior preservation the primary contract.
Document-only repair	On Noise16, oracle evidence-document patching reaches only 54.4% equivalence.	Gold evidence labels are not causally closed cache repair units.	Repair causally closed units rather than evidence documents.
Downstream windows	On 2Wiki-100, w4 reaches 79.0% equivalence at 72.9% patch ratio; suffix K1 reaches 97.0% at 79.3%.	Repair quality changes sharply when the recomputed region crosses the causal frontier.	Use suffix closure rather than a fixed local window.
Linked intervals	On Noise16, oracle linked intervals reach 80.0%, below oracle suffix closure at 93.6%.	Human-plausible reasoning chains do not always match transformer cache dependencies.	Keep linked intervals as a negative ablation.
Adaptive suffix	On Noise16, risk adaptive suffix keeps 86.8% equivalence at 57.9% patch ratio versus risk suffix K2 at 88.8% and 69.9%.	Budget can be traded for fidelity, but aggressive trimming increases residual drift.	Present SAFE-KV as a quality–budget family.
Learned selector	On external selector diagnostics, synthetic-trained learned selection has near-zero evidence recall.	Selector learning is sensitive to supervision shift; the stable contribution is the repair contract.	Use lightweight query/risk selectors in the main paper and report the learned-selector failure.

Appendix

Table 27: Design ledger for SAFE-KV (donor-copy protocol). Each row records an experimental route that shaped the algorithm; suffix repair is the simplest repair unit that remained consistent across these positive and negative results.

Work	Reuse setting	Main goal	Full-prefill safety oracle?	Repair granularity?	Matched random controls?
vLLM/PagedAttention	serving cache management	throughput/memory efficiency	no	no	no
LMCache	cache storage/transfer/reuse layer	system reuse	no	no	no
Prompt Cache	modular prompt reuse	latency reduction	limited	no	no
CacheBlend	non-prefix chunk fusion	serving efficiency + quality	partially, systems-focused	selective token recompute	not this framing
Cache-Craft	chunk-cache management for RAG	TTFT/throughput with quality control	no (task accuracy)	selective token recompute	no
KVLink	independent document cache reuse	speed + QA with adjustment/training	limited	no	no

Appendix

Table 32: Positioning against adjacent cache-reuse work. SAFE-KV’s contribution is the evaluation framing: full-prefill behavior preservation on the exact prompt, comparison of repair units, answer-level drift, and matched random controls. The CacheBlend and Cache-Craft rows summarize their published descriptions.

KV-Cache Reuse

Key-value (KV) cache reuse lets language models reuse previously computed attention states to cut repeated prefill work and serving latency.

See all 73 papers

CacheRepair: Learning to Repair Cross-Chunk Context in RAG for KV Cache Fusion
Sep 28, 2026Genglin Wang, Wangsong Yin, Yeerzhati Abudunuer +3Retrieval-Augmented Generation KV-Cache Management
The Chinese University of Hong Kong · Peking University · The Hong Kong University of Science and Technology
KV-Repurpose: A Testbed for Evaluating Non-Prefix KV-Cache Reuse
Oct 9, 2026Taiga Ikeda, Daisuke Miyashita, Jun DeguchiKV-Cache Reuse KV-Cache Management
Kioxia Corporation, Japan
RelaxKV: Recomputation Guided by the Query with Sparse Context Attention for Efficient KV Cache Reuse
Sep 27, 2026Ruoling Qi, Yirui Liu, Xuaner Wu +5Retrieval-Augmented Generation KV-Cache Management
Shanghai Jiao Tong University · Institute of Artificial Intelligence, China Telecom (TeleAI) · State University of New York at Buffalo

Sep 28, 2026J/K move · Enter open · S save

CacheRepair: Learning to Repair Cross-Chunk Context in RAG for KV Cache Fusion

Genglin Wang, Wangsong Yin, Yeerzhati Abudunuer, Haoxuan Xu+2

The Chinese University of Hong Kong · Peking University · The Hong Kong University of Science and Technology

Retrieval-Augmented Generation KV-Cache Management LLM Inference Acceleration

Open paper AskPDF arXiv

Multi-document retrieval-augmented generation (RAG) requires a language model to process multiple retrieved text chunks before answering a question. Precomputing each chunk's KV cache independently and concatenating the caches when the chunks are retrieved can accelerate this step. However, the assembled cache lacks cross-chunk attention information, reducing answer quality. Selective recomputation methods recover the missing cross-chunk context by rerunning the target LLM on selected tokens, incurring substantial online computation. We introduce CacheRepair, a lightweight network that learns the difference between independently computed KV caches and those produced by processing the chunks together. The network combines compressed KV features with token embeddings and uses attention that is bidirectional within each chunk and flows from earlier to later chunks. Each repair block receives the compressed cache features, and the predicted residual is added to every document token's cache. Each repair network is trained for a specific frozen target LLM on a generic retrieval corpus and reused across downstream datasets. Our analysis shows that repair reduces KV errors both near chunk boundaries and throughout chunk interiors. Evaluation across three target LLMs and four downstream datasets places CacheRepair on the measured answer-quality-latency Pareto frontier in eleven of twelve model-dataset combinations. Reported time to first token (TTFT) includes online cache transfer and repair. Across all twelve combinations, the largest repairers achieve 1.69-4.61

\times

speedups in median TTFT over full prefill and improve mean F1 by 2.1-26.1 percentage points over direct cache reuse.

Abstract

KV-Cache Reuse

Explore similar work

CacheRepair: Learning to Repair Cross-Chunk Context in RAG for KV Cache Fusion

KV-Repurpose: A Testbed for Evaluating Non-Prefix KV-Cache Reuse

RelaxKV: Recomputation Guided by the Query with Sparse Context Attention for Efficient KV Cache Reuse