Source-linked AI summary
WakeKV: Reactive, Reversible KV Residency for Heads That Change Their Minds
Utkarsh Ranjan
TL;DR
Most approaches fix head treatment offline or at prefill, although head behavior can change during generation. WakeKV keeps demoted state recoverable and fetches it on demand. Its miss rate is no worse than frozen classification at 17 matched-memory levels and it beats destructive eviction in 19 of 20 matched-budget settings.
Problem
Method
WakeKV caps each head’s GPU-resident set with an LRU budget, demotes pages to a CPU reservoir, and fetches their exact state when demanded again.
Source: §3 · p. 3Results
At 17 matched-memory levels, reactive miss rate is ≤ frozen’s, and at matched budget it beats evict in 19/20 settings.
Source: §5.1 · p. 4Takeaways & Limitations
Reactive, reversible residency is a promising alternative to fixed per-head classification, with gains reported from simulation through real hardware.
Source: §6 · p. 4Evaluation is limited to 1.5B–8B models with uneven regime coverage; the real-system result covers one model/regime on one A30, and the simulator counts misses rather than stall time.
Source: §6 · p. 4
Abstract
from arXiv · showhide
Most KV-cache compression methods classify attention heads once, either offline or during prefill, and keep this classification fixed throughout generation. Across three models (1.5B-8B) and three regimes (needle retrieval, long chain-of-thought, and multi-turn recall), we measure head behavior on four model-regime combinations and find that most heads change their reading behavior at least once during generation. We introduce WakeKV, a reactive residency policy that moves cooling heads to a recoverable CPU reservoir rather than freezing or permanently evicting their state. At matched memory or budget, WakeKV consistently improves miss rate over frozen classification and destructive eviction, evaluated across five model-regime combinations and over three cited baselines (SnapKV, uniform R-KV, and ReasonAlloc) across four eligible combinations. A FlexiCache/vLLM implementation on Mistral-7B confirms the benefit on real hardware, improving throughput while retaining LongBench quality.
1 Introduction
KV-cache growth makes residency decisions central to long-context serving, yet prevailing methods fix head treatment offline or at prefill. The paper introduces reactive, reversible WakeKV and reports gains over fixed and destructive alternatives, including a real-hardware evaluation.
1 Introduction
- KV-cache size grows linearly with sequence length and can exceed GPU memory, while most approaches hold each head’s treatment fixed after offline selection or prefill.Serving systems must decide which cache state to keep resident, compress, or evict.Source: §1 · p. 1
- The paper identifies an open gap: surveyed systems do not combine runtime head-priority updates with reversible demotion, despite prior evidence of within-generation non-stationarity [1].WakeKV addresses the gap with reactive residency and is validated across five model/regime combinations at matched memory and budget.Source: §1 · p. 1
- Across three models and three regimes, head behavior changes during generation, while inexpensive online signals do not reliably predict those changes.Two of the regimes are ones competing frozen-role methods do not evaluate.Source: §1 · p. 1
- On A30/Mistral-7B, WakeKV reactive improves throughput over stock at every rerank interval (1.34–2.33×) while retaining 95–99% of LongBench quality.Across matched memory or budget, the paper also reports gains over frozen classification, destructive eviction, and three cited baselines.Source: §1 · p. 2
2 Motivation
The motivation is that head behavior shifts during generation, while inexpensive signals do not reliably predict when to fetch state. Measurements quantify both the shifts and the limits of prediction.
2.1 Head Churn
- Adjacent-step Jaccard is 0.53 for continuous-score NIAH, within the 0.28–0.51 range reported in [1].Jaccard measures overlap of attended sets across consecutive decode steps; 1 means unchanged and 0 means complete turnover.Source: §2.1 · p. 2
- Across regimes, 63–85% of heads experience at least one drift event, although only 2–9% are persistently unstable; the 8B CoT drift-event fraction is 0.79.Drift events use windowed-median overlap below τ=0.5; persistent instability uses mean RCO below threshold throughout.Source: §2.1 · p. 2
2.2 Wake-up Prediction
- The best attention-derived predictor reaches only 0.08–0.20 precision at 0.41–0.70 recall, and neither tested proactive route was reliable enough to prefetch.Token-based wake-burst prediction reaches F1 0.53–0.62 on CoT but yields no surviving markers on NIAH or multi-turn; a coarse trigger catches 54% of wake-ups at a 2%-of-steps rate.Source: §2.2 · p. 2
2 Motivation
- Because heads shift but are not predictably enough anticipated for prefetch, the findings motivate reacting to demand instead.The motivating question is whether head behavior changes over time and whether those changes can be predicted well enough to prefetch KV state.Source: §2.2 · p. 2, §2 · p. 2
3 Method
WakeKV replaces permanent per-head classification with capped GPU residency and reversible demotion to a CPU reservoir. Demand-driven fetching restores a page’s exact state when it is wanted again.
3 Method
- Each head maintains an LRU-ordered GPU-resident set capped by B and a CPU reservoir holding the full, uncompressed state of demoted pages.The runtime rule demotes under pressure, fetches when needed, and never permanently discards a head’s state.Source: §3 · p. 2
- Figure 1 contrasts prior methods by decision timing and reversibility, then contrasts WakeKV’s reservoir with an eviction foil using the same LRU schedule but destroying state.The figure marks the dynamic-and-reversible quadrant as previously empty among surveyed approaches.Source: Figure 1 · p. 3
- On each decode step, demanded pages move forward in an LRU order; when residency exceeds budget B, the least-recently-used page is copied to the CPU reservoir and fetched back on demand.A miss fetches the page’s exact, uncompressed key/value state before that step’s attention completes, at the cost of one PCIe transfer.Source: §3 · p. 3
4 Evaluation Setup
The evaluation logs per-head attended positions and weights across three models and three regimes, with an additional Mistral-7B/NIAH combination for results.
4 Evaluation Setup
- The setup covers three models and NIAH, long chain-of-thought, and multi-turn recall; each decode step logs the top 64 attended positions and weights per head.The 1.5B and 3B models were measured on an RTX 2080 Ti, while the 8B model and real-system follow-up ran on an A30; results add Mistral-7B/NIAH.Source: §4 · p. 3
5 Results
At matched memory or budget, WakeKV’s reactive residency outperforms frozen classification and destructive eviction in the reported comparisons, and the simulator gains extend to a real-hardware deployment. The section reports results against additional baselines and a throughput–quality trade-off on Mistral-7B.
5 Results
- The evaluation compares policies at matched memory or budget because equal nominal budget labels can yield different realized memory.The section compares WakeKV with frozen classification, an evict policy using the same LRU schedule, and three additional baselines.Source: §5 · p. 3
5.1 Reactive vs. Frozen and Evict
- 17/17 matched-memory points favor reactive over frozen, while reactive beats destructive evict in 19/20 matched-budget settings.The sole evict exception is 8B/CoT at the tightest budget; Table 1 reports comparisons made at matched memory or budget.Source: Table 1 · p. 3, §5.1 · p. 4
5 Results
- 32/32 matched-memory comparisons favor reactive over SnapKV, uniform R-KV, and ReasonAlloc across four eligible combinations.At ∼28k pages on 8B/CoT, R-KV’s miss rate is 0.132 against reactive’s 0.078.Source: Figure 2 · p. 4
5.2 Real-Hardware Deployment
- At R=1, WakeKV improves stock throughput by 1.34× at 97.4% LongBench quality; at R=16, throughput reaches 2.33× while quality stays above 95%.The Table 2 sweep is on an A30 with Mistral-7B; part of the 2.33× gain comes from less attention work.Source: §5.2 · p. 4, Table 2 · p. 4
6 Conclusion and Future Work
WakeKV replaces fixed head classification with a reactive rule that demotes under pressure, fetches on demand, and never discards state. The conclusion describes it as a promising alternative while noting limits in model coverage, simulator metrics, and real-system scope.
6 Conclusion and Future Work
- WakeKV’s single reactive rule consistently beats frozen classification and destructive eviction, with gains carrying from simulation to real hardware.The rule demotes under pressure, fetches on demand, and never discards state.Source: §6 · p. 4
- Evaluation is limited to 1.5B–8B models with uneven regime coverage, simulator miss counts, and one real-system model/regime on one A30.The stated next step is a real-system comparison of reversible and destructive eviction using stall time.Source: §6 · p. 4
- The conclusion presents reactive, reversible residency as a promising alternative to fixed per-head classification.This conclusion is bounded by the stated evaluation limitations.Source: §6 · p. 4
A Full Per-Combo, Per-Budget Sweep
The appendix reports simulator-compatible reimplementations of three baselines and a raw miss-rate sweep across five model/regime combinations. It distinguishes direct same-row comparisons from matched-memory comparisons and notes a non-default pruning interval excluded from headline win counts.
A Full Per-Combo, Per-Budget Sweep
- SnapKV, uniform R-KV, and ReasonAlloc are simulator-compatible reimplementations, each a declared simplification of its cited method.SnapKV [5] uses prefill-only top-token selection; uniform R-KV [6] is reduced to importance-only pruning; ReasonAlloc [4] uses an equal split rather than offline Reasoning-Wave allocation.Source: A Full Per-Combo, Per-Budget Sweep · p. 7
- Table 3 reports raw miss rates at nominal budgets B ∈{8, 16, 32, 64} across all five model/regime combinations.Reactive and Evict share realized memory within a row; the other policies realize different memory at the same nominal B and require matched-memory comparisons.Source: A Full Per-Combo, Per-Budget Sweep · p. 7, Table 3 · p. 7
- Mistral-7B/NIAH uses a non-default R-KV/ReasonAlloc pruning interval of buffer/delta=8 rather than 128, and is excluded from Table 1 win counts.The interval avoids a degenerate no-pruning result on this model’s short NIAH answers.Source: A Full Per-Combo, Per-Budget Sweep · p. 7
B Head Churn Metrics by Regime
The head-churn breakdown reports adjacent-step Jaccard, persistent instability, and discrete drift events by regime. At 8B, binary Jaccard tightens on math CoT while page-level drift remains high.
B Head Churn Metrics by Regime
- Table 4 breaks out adjacent-step Jaccard, persistently unstable heads, and heads with at least one drift event by regime.Adjacent-step Jaccard is continuous for NIAH and binary for CoT and multi-turn.Source: B Head Churn Metrics by Regime · p. 8, Table 4 · p. 8
- At 8B on CoT, page-level drift remains 0.79 despite tighter binary Jaccard as copy behavior stabilizes.The passage identifies page-level drift as the metric relevant to a residency controller.Source: Table 4 · p. 8
- The 8B CoT case retains high page-level drift even as copy behavior stabilizes, so the caption describes the case as undiminished at scale.The reported page-level drift value is 0.79.Source: Table 4 · p. 8