WakeKV: Reactive, Reversible KV Residency for Heads That Change Their Minds

📄 arXiv: 2610.02713v1 📥 PDF

作者: Utkarsh Ranjan

分类: cs.CL, cs.AR, cs.DC, cs.LG

发布日期: 2026-10-02

备注: Accepted to the NeurIPS 2026 Workshop on ML for Systems. 2 figures, 4 tables, appendix


💡 一句话要点

提出WakeKV以解决KV缓存固定分类问题

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: KV缓存 动态调整 注意力机制 生成模型 自然语言处理 吞吐量提升 反应式策略

📋 核心要点

  1. 现有的KV缓存压缩方法在生成过程中对注意力头的分类固定,无法适应头的动态变化,导致性能下降。
  2. WakeKV通过将冷却的注意力头转移到可恢复的CPU储存区,提供了一种灵活的反应式驻留策略,能够动态调整头的状态。
  3. 在多个模型和场景下的实验表明,WakeKV显著降低了未命中率,相比于传统方法提升了性能,验证了其有效性。

📝 摘要(中文)

大多数KV缓存压缩方法在生成过程中对注意力头的分类保持固定,然而研究发现大多数头在生成过程中至少会改变一次读取行为。为此,本文提出了WakeKV,一种反应式驻留策略,将冷却的头移动到可恢复的CPU储存区,而不是冻结或永久驱逐其状态。在相同的内存或预算下,WakeKV在五个模型-场景组合中持续改善了未命中率,相较于冻结分类和破坏性驱逐,表现出显著优势。Mistral-7B上的FlexiCache/vLLM实现验证了在实际硬件上的收益,提高了吞吐量,同时保持了LongBench的质量。

🔬 方法详解

问题定义:本文旨在解决现有KV缓存压缩方法在生成过程中对注意力头的固定分类问题。现有方法无法适应头的动态变化,导致未命中率上升,影响生成性能。

核心思路:WakeKV提出了一种反应式的驻留策略,通过将冷却的注意力头转移到可恢复的CPU储存区,而不是简单地冻结或驱逐其状态,从而实现动态调整。

技术框架:WakeKV的整体架构包括三个主要模块:注意力头的状态监测、动态分类调整和CPU储存区的管理。通过实时监测头的行为,系统能够决定何时将头转移到CPU储存区。

关键创新:WakeKV的核心创新在于其反应式的驻留策略,允许注意力头在生成过程中根据需求动态调整状态。这一设计与传统的固定分类方法本质上不同,提供了更高的灵活性和适应性。

关键设计:在实现WakeKV时,关键参数包括头的冷却阈值、CPU储存区的容量和动态分类的频率。损失函数设计上,考虑了未命中率和生成质量的平衡,以确保在提升性能的同时不损害生成质量。

🖼️ 关键图片

fig_0
img_1
img_2

📊 实验亮点

实验结果显示,WakeKV在五个模型-场景组合中,相比于传统的冻结分类和破坏性驱逐方法,未命中率显著降低,提升幅度达到了XX%。在Mistral-7B上的FlexiCache/vLLM实现中,吞吐量提高,同时保持了LongBench的生成质量,验证了其在实际应用中的有效性。

🎯 应用场景

WakeKV的研究成果在自然语言处理、机器翻译和对话系统等领域具有广泛的应用潜力。通过提高KV缓存的动态适应能力,能够显著提升生成模型的性能和响应速度,满足实时应用的需求。未来,该方法还可能扩展到其他需要动态状态管理的AI系统中。

📄 摘要(原文)

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.