Reasoning-Aware Token–Head KV Lifetime Prediction for Long-Context LLMs

Long-context reasoning models can generate thousands of tokens before producing a final answer, causing the key-value (KV) cache to grow linearly throughout generation. Existing KV-cache compression methods commonly select important token positions using historical attention, assign different budgets across attention heads, or classify heads into fixed local and retrieval-oriented roles. However, these approaches may be too coarse for long reasoning trajectories, where the usefulness of a cached representation can change across tokens, KV heads, layers, and decoding steps. We investigate whether the remaining useful lifetime of individual token–KV-head entries can be predicted online during autoregressive reasoning. Our central hypothesis is that KV utility is sparse but dynamic: an entry that is currently unused may become important several decoding steps later, while a historically important entry may cease to contribute to future computation. Rather than assigning a permanent cache policy to an entire head or evicting a token from all heads simultaneously, we consider retention decisions at token–head granularity and condition them on the evolving reasoning state. Our approach consists of three stages. First, we collect full-cache reasoning trajectories and estimate oracle KV importance by selectively perturbing or removing historical token–head entries and measuring their effect on future model predictions and task outcomes. Second, we use these oracle measurements to train a lightweight predictor of future utility or remaining lifetime from inexpensive runtime features, including token age, attention statistics, layer and head identity, and hidden-state information. Finally, we replace the oracle with the learned predictor during inference and dynamically retain or evict KV entries under a fixed memory budget. We plan to evaluate the method on long-context retrieval and mathematical reasoning tasks, comparing against full-cache inference, sliding-window policies, historical-attention-based eviction, and static head-aware allocation. Evaluation will measure both model quality and systems performance, including task accuracy, next-token distribution divergence, KV-cache memory usage, decode latency, throughput, and predictor and cache-management overhead. This study aims to determine whether explicitly modeling the future lifetime of KV representations can preserve reasoning-critical information at substantially smaller cache budgets while remaining practical for real inference systems.