WakeKV improves memory use for changing AI attention heads
WakeKV: Reactive, Reversible KV Residency for Heads That Change Their Minds
Computation and LanguageHardware ArchitectureDistributed, Parallel, and Cluster ComputingMachine Learning
Summary
Many AI models use attention heads that help decide what parts of information to focus on, but these heads often change how they work during a task. Existing methods label these heads once and keep that label fixed, which can waste memory or slow down the model. The authors created WakeKV, a system that moves attention heads to a special memory area when they 'cool down' instead of freezing or discarding them permanently. This approach improves memory efficiency and keeps model quality high during tasks like long reasoning or recall.
What this means in practice
- •For ai infrastructure engineers: Optimize memory management in inference servers to reduce cache misses and improve throughput for large language models during complex tasks.
- •For cloud service operators: Deploy WakeKV in multi-tenant systems to handle dynamic memory usage efficiently and maintain quality in multi-turn conversational AI services.
Authors
Utkarsh Ranjan
Abstract
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.