Source-linked AI summary

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

Ruoling Qi, Yirui Liu, Xuaner Wu, Yuxin Jin, Jian Chen, Jiayu Qin, Yin Chen, Jiawei Shao

arXiv:2609.33503v1cs.AI

TL;DR

Cross-request KV reuse can lower RAG prefill cost, but prefix caching limits reuse and independently cached chunks omit cross-chunk interactions. RelaxKV jointly selects repair targets and recomputation context using query attention. It improves aggregate LongBench performance over ProphetKV at a 15% anchor ratio across four decoder models, with controlled results highlighting context selection.

  • Problem

    Prefix caching limits cross-request reuse, while independently prefetched chunks omit cross-chunk interactions; existing selective repair methods choose targets but retain the full causal prefix as context.

    Source: §1 · p. 1
  • Method

    RelaxKV uses layer-wise query attention to jointly select repair targets and a query-relevant recomputation context.

    Source: §1 · pp. 1–2
  • Results

    At a 15% anchor ratio, RelaxKV improves aggregate LongBench performance over ProphetKV on all four models.

    Source: §1 · p. 2
  • Takeaways & Limitations

    Matched-target and context-construction controls support recomputation context selection as an important factor in the quality–latency trade-off.

    Source: §1 · p. 2, §4.4 · p. 7, §6 · p. 9

    The context-construction ablation is not a strict equal-work comparison because causal truncation can change the number of context positions and measured recomputation latency.

    Source: §4.4 · p. 7

Abstract

from arXiv · show

Cross-request KV caching reduces the prefill cost of Retrieval-Augmented Generation (RAG), but conventional prefix caching severely limits cache reuse across requests. Position-Independent Caching (PIC) removes this constraint by reusing independent chunks, but their KV states miss cross-chunk interactions. Existing methods selectively recompute token states to recover these missing interactions, but primarily allocate the recomputation budget to selecting which states to recompute, while fixing the recomputation context to the full causal prefix. We introduce RelaxKV, which formulates selective cache repair as a joint allocation problem over repair targets and recomputation context. Guided by the user query, RelaxKV identifies layer-specific repair targets and restricts their recomputation to a query-relevant context, reducing attention computation. Across four decoder models, RelaxKV at a 15% anchor ratio improves aggregate LongBench performance over ProphetKV on all models. On Qwen3-14B, RelaxKV provides a stronger quality-TTFT trade-off than ProphetKV across a 5%-30% anchor-ratio sweep, and achieves the best selective results on RULER-MV and LV-Eval at 16K and 32K context lengths. Controlled ablations further demonstrate the importance of recomputation context selection.

1 INTRODUCTION

Cross-request KV reuse can reduce RAG prefill cost, but prefix caching limits reuse and independently prefetched chunks miss cross-chunk interactions. RelaxKV jointly allocates repair targets and recomputation context, improving reported quality–latency results across long-context evaluations.

1 INTRODUCTION

  • Selective repair addresses the trade-off between restoring missing cross-chunk interactions and retaining the prefill-efficiency benefit of cache reuse.Prefix caching limits reuse to requests sharing a prefix, while Position-Independent Caching allows independent chunks but their KV states omit cross-chunk interactions.Source: §1 · p. 1
  • Existing selective repair methods vary in target selection, but selected tokens still attend to their full causal prefixes during recomputation.Thus, even a small repair-target set can incur substantial attention computation; context selection is a second allocation dimension.Source: §1 · p. 1
  • RelaxKV jointly allocates repair targets and recomputation context, extending selective cache repair beyond choosing which token states to recompute.Layer-wise query attention derives layer-specific repair targets and a shared context; active targets use this context to update their KV entries.Source: §1 · pp. 1–2
  • At a 15% anchor ratio, RelaxKV improves aggregate LongBench performance over ProphetKV on all four decoder models.The reported results also show a stronger Qwen3-14B quality–TTFT trade-off over a 5%–30% sweep and best selective results on RULER-MV and LV-Eval at 16K and 32K.Source: §1 · p. 2

2 BACKGROUND AND MOTIVATION

Independent chunk prefetching makes reusable caches inexpensive but approximate because each chunk omits preceding-chunk context. Query attention varies across layers, motivating separate choices for which states to repair and which KV states they can access.

2.1 APPROXIMATE PREFIX KV REUSE

  • Independent chunk prefetching is inexpensive, but its concatenated KV cache generally differs from full prefill because chunks omit preceding-chunk hidden states.The paper distinguishes this contextual staleness from RoPE positional misalignment and evaluates chunks using final global position IDs.Source: §2.1 · p. 2
  • RelaxKV addresses missing context from preceding chunks; positional alignment alone does not eliminate the cache’s approximation.When a cached chunk is placed at a different offset, RoPE keys can be realigned before repair.Source: §2.1 · p. 2

2.2 QUERY-AWARE CACHE REPAIR

  • Query-to-context attention identifies highly attended positions that can shift across layers, while selective repair separates target selection from context selection.Repair targets determine which token states are recomputed; recomputation context determines which cached KV states those tokens can attend to.Source: §2.2 · pp. 2–3
  • RelaxKV distinguishes repair anchors, layer-specific active targets, and context tokens to make both dimensions of selective repair explicit.A token selected at a deeper layer must be recomputed through preceding layers to reach that layer.Source: §2.2 · p. 3

3 METHOD

RelaxKV uses layer-wise query attention to select repair anchors, derives active targets and a shared context, and recomputes targets using only causally valid context positions. Updated cache entries combine fresh and retained cached KV states.

3.2 LAYER-WISE QUERY ATTENTION AND REPAIR TARGETS

  • Layer-wise query-attention scores over prefix tokens show that high-attention positions shift across layers.RelaxKV selects the highest-scoring k = max(1, round(rTc)) positions independently at each layer, where r is the anchor ratio.Source: Figure 1 · p. 3, §3.2 · p. 4

3.1 OVERVIEW

  • RelaxKV selects layer-wise anchors from query attention, derives active repair targets and a shared context, then updates each target’s KV entry using causally valid context positions.Only active targets execute the Transformer block at each layer.Source: §3.1 · p. 3

3.2 LAYER-WISE QUERY ATTENTION AND REPAIR TARGETS

  • Deeper-layer selections require recomputation through preceding layers, so active target executions can exceed the nominal anchor count.A token remains active through its highest selected layer and then becomes inactive.Source: §3.2 · p. 4

3.3 RECOMPUTATION CONTEXT SELECTION

  • At each layer, RelaxKV restricts an active target’s attention to the causally valid positions in the shared anchor union rather than its full causal prefix.The retained targets keep their original global position IDs, so the restriction changes available KV positions without changing positional coordinates.Source: §3.3 · p. 4

3.4 CACHE UPDATE SEMANTICS

  • Recomputed active targets overwrite their corresponding cache entries; all other entries retain their naive values.The update set follows directly from the active target set and requires no additional selection step.Source: §3.4 · p. 4

3.3 RECOMPUTATION CONTEXT SELECTION

  • With the same active targets, RelaxKV uses only the causal subset of U as recomputation context, unlike full-context recomputation’s complete causal prefix.Figure 3 marks excluded context positions with dashed entries.Source: §3.3 · p. 4, Figure 3 · p. 5

3.4 CACHE UPDATE SEMANTICS

  • An active target attends to fresh KV states for positions also active at that layer and cached KV states for the remaining context positions.Tokens that leave the active set may still contribute cached KV states in later layers, making the mixed fresh–cached context part of RelaxKV’s approximation.Source: §3.4 · p. 5

4 EVALUATION

The evaluation tests quality across decoder models and long-context workloads, then isolates the effects of recomputation context and budget allocation. Results compare RelaxKV with selective-repair baselines and full recomputation.

4 EVALUATION

  • The evaluation asks whether quality persists across models and longer contexts, how context selection affects quality and cost, and whether joint allocation improves quality–TTFT trade-offs.RelaxKVG holds ProphetKV’s repair targets and cache-update mask fixed while changing the recomputation context.Source: §4 · p. 5, §4.1 · p. 6

4.1 EXPERIMENTAL SETUP

  • The evaluation spans four decoder models from 3B to 14B, LongBench, RULER-MV, and LV-Eval, with context-length tests from 4K to 32K tokens.LongBench uses seven QA, multihop QA, and passage-retrieval tasks; RULER-MV and LV-Eval test long-context performance.Source: §4.1 · p. 5
  • Comparisons use shared retrieved chunks and matched prompts, tokenization, model weights, 512-token chunks, and evaluation examples; retrieval latency is excluded from measured request time.The evaluation includes full recomputation, full cache reuse, and four selective-repair methods.Source: §4.1 · p. 5

4.2 ACCURACY ON STANDARD TASKS

  • At a 15% anchor ratio, RelaxKV improves aggregate LongBench performance over ProphetKV on all four models while remaining close to full recomputation.The directly comparable selective baselines and RelaxKVG use a nominal 20% repair ratio.Source: §4.2 · p. 6, Table 1 · p. 6

4.3 CONTEXT LENGTH SCALING

  • At 32K on Qwen3-14B RULER-MV, RelaxKV reaches 95.00 accuracy versus ProphetKV’s 91.17 and RelaxKVG’s 92.00, remaining close to full recomputation.On LV-Eval, RelaxKV records the best selective performance at 16K and 32K: 28.71 and 25.40 token F1 versus ProphetKV’s 24.78 and 19.32.Source: §4.3 · p. 6, Table 2 · p. 7

4.4 ABLATION STUDY OF RECOMPUTATION CONTEXT CONSTRUCTION

  • With repair targets and nominal context budget fixed, context construction produces substantially different task quality despite similar recomputation latency.The comparison is a controlled context-construction ablation, not a strict equal-work comparison: causal truncation can change effective attention work.Source: §4.4 · p. 7, Table 3 · p. 7

4.5 BUDGET ACCOUNTING AND ONLINE SERVING LATENCY

  • At a 15% anchor ratio on Qwen3-14B’s 8K RULER-MV workload, the anchor union covers 63.05% of context positions while active targets account for 42.85% of layer-token states.The remaining 57.15% of layer-token states are not recomputed; across evaluated budgets, RelaxKV trades additional recomputation for higher task quality.Source: §4.5 · pp. 7–8, Figure 4 · p. 8

5 RELATED WORK

Related work addresses reusable KV caches, selective repair, query-guided selection, and long-context evaluation. RelaxKV focuses on online repair of reusable caches by selecting both repair targets and the context available during recomputation.

5.1 KV CACHE REUSE AND SELECTIVE RECOMPUTATION

  • PIC enables reuse of independently encoded chunk KV states across requests, while selective-repair methods address missing cross-chunk interactions after cache composition.RelaxKV leaves the base model and cache construction unchanged, allocating online recomputation across repair targets and context.Source: §5.1 · p. 8

5.2 QUERY-AWARE SELECTION AND SPARSE ATTENTION

  • Query-aware methods select salient cache positions, decoding pages, prefill keys, or repair targets, but address different stages or cache representations than RelaxKV’s joint repair allocation.SnapKV and KVzip compress contextualized caches rather than repairing missing cross-chunk interactions in composed independent caches.Source: §5.2 · p. 9

5.3 LONG-CONTEXT EVALUATION

  • LongBench, RULER, and LV-Eval cover natural task diversity, controlled length scaling, and multi-document settings with natural distractors.The paper uses LongBench for overall quality and RULER-MV and LV-Eval to examine performance as context length and irrelevant context increase.Source: §5.3 · p. 9

6 CONCLUSION

RelaxKV jointly allocates cache repair across repair targets and recomputation context, using layer-wise query attention to derive targets and a shared query-relevant context. Its context restricts each target to causally preceding positions in the anchor union, combining recomputed and cached KV states.

6 CONCLUSION

  • RelaxKV extends selective cache repair to jointly allocate repair targets and recomputation context, with controls validating the effectiveness of context selection.The conclusion reports competitive quality–latency trade-offs across models and particularly strong performance in long-context settings.Source: §6 · p. 9

A.1 REPAIR SCHEDULE

  • Layer-wise query attention selects anchors independently; a token selected at a deeper layer remains active through that layer, then becomes inactive.Active repair targets at each layer are determined by their highest selected layer.Source: §A.1 · p. 12

A.2 RECOMPUTATION CONTEXT AND CACHE UPDATE

  • Each active target attends only to causally preceding positions in the shared anchor union, using recomputed KV states for active positions and cached states for the rest.The method retains original global position IDs for RoPE.Source: §A.2 · pp. 12–13, §A.1 · p. 12
  • The unmasked attention-edge count is a structural measure, not an exact FLOP count for the eager GQA implementation.The paper uses measured recomputation latency and TTFT as primary efficiency metrics.Source: §A.2 · p. 13

A.3 CONTROLLED CONFIGURATIONS

  • Table 4’s controlled configurations vary recomputation context or active repair targets, with global repair targets shared across layers.Context-construction variants hold layer-wise anchors and active targets fixed unless otherwise noted.Source: §A.3 · p. 13, Table 4 · p. 13

A.4 FULL-PREFIX QUALITY–COST CONTROL

A full-prefix control compares context restriction against complete causal-prefix attention while fixing repair targets. The appendix also describes cross-model quality–TTFT curves and reports that nominal anchor ratios can differ from realized context and active-target ratios.

A.4 FULL-PREFIX QUALITY–COST CONTROL

  • At a 20% anchor ratio, mean context-set ratios are 65.1% on RULER-MV and 70.4% on MuSiQue in the Llama-3.1 sweep.For Qwen3-14B at 15%, context-set/active-target ratios are 64.1%/44.8% on MuSiQue, 78.1%/55.2% on RULER-MV 32K, and 80.3%/57.8% on LV-Eval.Source: §A.4 · p. 13
  • 9.9–23.2% lower mean recomputation latency, with absolute quality differences within 0.45–1.00 points, for RelaxKV versus the full-prefix control across four settings.The control fixes active repair targets and varies only the context available to them.Source: §A.4 · p. 14, Table 5 · p. 14

A.5 CROSS-MODEL QUALITY–TTFT TRADE-OFF CURVES

  • Figure 5 plots RULER-MV accuracy and MuSiQue F1 against cache-ready TTFT across four models, with curves tracing nominal repair or anchor ratios.The models are Qwen3-14B, Llama-3.1-8B, Phi-4-14B, and Llama-3.2-3B.Source: §A.5 · p. 14, Figure 5 · p. 14

A.6 CHUNK SIZE SENSITIVITY

The chunk-size sensitivity evaluation varies reusable chunk size while keeping the model, prompts, tokenization, examples, and nominal repair ratios fixed. The reported pattern is that all reuse-based methods improve as chunk size increases.

A.6 CHUNK SIZE SENSITIVITY

  • All evaluated reuse-based methods improve as reusable chunk size increases on Qwen3-14B MuSiQue without truncation.Table 6 compares evaluated chunk sizes; selective baselines and RelaxKVG use a nominal 20% repair target ratio, while RelaxKV uses a 15% anchor ratio.Source: §A.6 · p. 14, Table 6 · p. 15

B.1 CONTROLLED SETUP

The controlled sweep tests whether a larger globally shared repair-target set can replace layer-specific targets. It finds sharply higher TTFT for limited aggregate quality gains, with setup and evaluation details defining the scope.

B.1 CONTROLLED SETUP

  • The diagnostic fixes context U to the union of layer-wise Top-20% anchors and recomputes the same globally ranked target set at every layer, each attending to its causal subset of U.Targets are ranked over the full context, and the target-count cap does not require the target set to be a subset of U.Source: §B.1 · p. 15
  • Requested target ratios can be capped at the per-example context-set size, so realized ratios may differ across datasets.Some examples are already capped at a requested ratio of 0.7.Source: §B.1 · p. 15, Table 8 · p. 16
  • The diagnostic indicates target count alone does not explain the benefit of a layer-varying repair schedule, and does not establish a universally optimal selector.End-to-end quality–latency comparisons remain those in the main paper.Source: §B.1 · p. 15
  • From 20% targets to Target-Union, RULER-MV TTFT rises 361.58 to 914.47 ms and MuSiQue TTFT 267.86 to 606.73 ms, while LongBench average shifts 57.84 to 58.10.The seven-task average varies non-monotonically at intermediate ratios.Source: §B.1 · p. 15, Table 8 · p. 16

C.1 MODEL, CACHE, AND PRECISION

  • Experiments use four decoder models in BF16 with batch size one and greedy decoding; reusable contexts are independently encoded in 512-token chunks before their KV caches are concatenated.All selective methods run on one H100 80GB GPU; only the 16K and 32K full-recomputation references use two-GPU model parallelism.Source: §C.1 · p. 16

C.2 TASK DEFINITIONS AND EVALUATION SPLITS

The evaluation covers long-context benchmarks and task-specific input settings. Accuracy comparisons control prompts, tokenization, model weights, and examples, while cache-ready TTFT includes selective-repair work through first-token decoding.

C.2 TASK DEFINITIONS AND EVALUATION SPLITS

  • RULER-MV uses a nominal 8K context for standard comparisons and 4K–32K for length stress tests; LongBench uses an 8,192-token input limit with head–tail truncation.The context ablation uses untruncated MuSiQue inputs, while LV-Eval is evaluated at 16K and 32K.Source: §C.2 · p. 16

C.3 ACCURACY PROTOCOL

  • Within each comparison, methods share prompts, tokenization, model weights, and evaluation examples; the reported LongBench aggregate averages seven selected task scores.The global target-capacity analysis uses Llama-3.1-8B-Instruct.Source: §C.3 · p. 16

C.4 CACHE-READY TTFT PROTOCOL

  • Cache-ready TTFT uses three timed repeats after warm-up and includes selective scoring, selection, KV gathering, recomputation, cache update, query prefill, and first-token decoding.Reusable-cache construction is excluded, and the reusable naive cache is resident on GPU before timing.Source: §C.4 · p. 16
Loading 2609.33503v1…