PaperlayerView exact revision on arXiv

Source-linked AI summary

LookaheadKV: Fast and Accurate KV Cache Eviction by Glimpsing into the Future without Generation

Jinwoo Ahn, Ingyu Seong, Akhil Kedia, Junhan Kim, Hyemi Jang, Kangwook Lee, Yongkweon Jeon

arXiv:2603.10899v1cs.LGcs.AI

TL;DR

Long-context LLM inference is constrained by KV-cache growth and the overhead of draft-based future-response estimation. LookaheadKV predicts true token-importance scores with lightweight trainable modules, outperforming prior methods while reducing eviction overhead.

  • Problem

    KV-cache size grows linearly with sequence length, while draft-based importance estimation adds substantial compute and latency during prefilling.

  • Method

    LookaheadKV augments a frozen LLM with learnable lookahead tokens and LoRA modules trained to predict true importance distributions without explicit draft generation.

  • Results

    Across models and long-context benchmarks, LookaheadKV consistently outperforms prior methods, with eviction overhead below 2.16% at 32K context length and up to 14.5× lower than draft-based approaches.

  • Takeaways & Limitations

    LookaheadKV provides a resource-efficient KV-cache eviction solution with less than 0.5% additional parameters and marginal prefill-latency increase.

  • Takeaways & Limitations

    Experiments did not include larger-sized models, and the method currently focuses on prefill KV-cache eviction rather than decoding-stage eviction.

Abstract

from arXiv · show

Transformer-based large language models (LLMs) rely on key-value (KV) caching to avoid redundant computation during autoregressive inference. While this mechanism greatly improves efficiency, the cache size grows linearly with the input sequence length, quickly becoming a bottleneck for long-context tasks. Existing solutions mitigate this problem by evicting prompt KV that are deemed unimportant, guided by estimated importance scores. Notably, a recent line of work proposes to improve eviction quality by "glimpsing into the future", in which a draft generator produces a surrogate future response approximating the target model's true response, and this surrogate is subsequently used to estimate the importance of cached KV more accurately. However, these approaches rely on computationally expensive draft generation, which introduces substantial prefilling overhead and limits their practicality in real-world deployment. To address this challenge, we propose LookaheadKV, a lightweight eviction framework that leverages the strength of surrogate future response without requiring explicit draft generation. LookaheadKV augments transformer layers with parameter-efficient modules trained to predict true importance scores with high accuracy. Our design ensures negligible runtime overhead comparable to existing inexpensive heuristics, while achieving accuracy superior to more costly approximation methods. Extensive experiments on long-context understanding benchmarks, across a wide range of models, demonstrate that our method not only outperforms recent competitive baselines in various long-context understanding tasks, but also reduces the eviction cost by up to 14.5x, leading to significantly faster time-to-first-token. Our code is available at https://github.com/SamsungLabs/LookaheadKV.

1 INTRODUCTION

LOOKAHEADKV addresses the KV-cache bottleneck in long-context inference by predicting token importance without explicit draft generation. It aims to preserve eviction quality while substantially reducing latency overhead.

  • A 128K-token sequence requires 40GB to store the KV cache for LLaMA3.1–70B in half-precision.
  • Draft-based eviction improves quality but incurs expensive token-generation overhead, while simple heuristics lose performance under constrained budgets.
  • LOOKAHEADKV uses learnable tokens and special LoRA modules to predict true importance scores without generating an approximate response.
  • LOOKAHEADKV is evaluated across diverse long-context benchmarks and multiple models, outperforming strong baselines across budgets and context lengths.
  • Less than 2.16% eviction overhead at 32K context length is up to 14.5× lower than draft-based approaches.

2 BACKGROUND

KV cache eviction estimates prompt-token importance to retain a smaller cache, but accurate future-response-based scoring is unavailable during prefilling. Existing methods therefore trade eviction quality against latency: heuristics are inexpensive but less accurate, while draft-based approaches improve performance with substantial compute.

  • KV Cache Eviction Using Importance Scores: KV cache eviction retains the most important prompt key-value pairs to reduce attention-output perturbation.Importance scores are derived from attention weights, and TopK pairs can be retained.
  • KV Cache Eviction Using Importance Scores: Ground-truth importance scores average cross-attention between generated-response queries and prompt keys.The scores quantify each prompt token’s relative contribution to response generation.
  • KV Cache Eviction Using Importance Scores: Because the true future response is unknown during prefilling, its importance scores cannot be computed directly.Prior methods construct a surrogate response sequence to predict the future attention pattern.
  • Prompt-based Approaches: SnapKV estimates future importance from the prompt suffix while reusing prefill attention weights with marginal extra computation.It exemplifies an inexpensive prompt-based heuristic.
  • Draft-based Approaches: Draft-based methods generate an approximate response to estimate future importance, improving performance over simple heuristics but increasing latency.The explicit draft-generation step creates the central accuracy–latency trade-off.

3 PROPOSED METHOD: LOOKAHEADKV

LOOKAHEADKV predicts future-response importance during prefilling using learnable lookahead tokens and selectively activated LoRA modules, avoiding explicit draft generation. It trains these modules against true importance scores while preserving normal-token behavior and reports strong performance across evaluated settings.

  • Proposed Method: LOOKAHEADKV: LOOKAHEADKV uses lightweight learnable modules to predict ground-truth importance scores without explicit draft generation.The framework is designed to provide faster eviction while retaining future-response information.
  • Learnable Lookahead Tokens: Learnable lookahead tokens estimate true-response attention patterns during prefilling and guide KV cache eviction.Their queries are used to estimate prompt-token importance before autoregressive decoding.
  • Lookahead LoRA: Lookahead LoRA selectively activates for lookahead tokens, improving their representations while leaving normal input-token outputs unchanged.The adapters provide complementary accuracy gains without altering the original model behavior for ordinary tokens.
  • Importance Estimation: LOOKAHEADKV retains the Top-K KV pairs with the highest importance scores from attention involving lookahead queries.The complete-sequence queries and keys incorporate the lookahead embeddings and LoRA projections.
  • LOOKAHEADKV Training: Training minimizes average KL divergence between normalized ground-truth and predicted importance scores across layers and heads.Only lookahead embeddings and LoRA modules are updated; the other LLM layers remain frozen.
  • Efficiency and Training: Applying lookahead LoRA to every linear layer increases latency by less than 1.3% while improving performance over omitting LoRA.FlashAttention is used in the forward pass, with eager attention for importance-score computation and backpropagation.
  • Evaluation: LOOKAHEADKV consistently outperforms baselines across tested LongBench budgets, RULER context lengths, and models.Figure 4 summarizes LongBench results across budgets and models and RULER results at a fixed budget of 128.

4 EXPERIMENTS

LOOKAHEADKV is evaluated across diverse long-context benchmarks, model families, context lengths, budgets, and output settings. It consistently outperforms prior eviction methods, including costly draft-based approaches, especially under constrained cache budgets.

  • Evaluation setup: The evaluation covers LongBench, RULER, LongProc, and MT-Bench across multiple model families and sizes.Experiments include LLaMA and Qwen architectures and context lengths from 4K to 32K.
  • LongBench evaluation: LOOKAHEADKV consistently outperforms all tested methods across LongBench models and cache budgets.The advantage is particularly pronounced at lower cache budgets, where draft-based approaches also outperform simple baselines but remain behind LOOKAHEADKV.
  • RULER evaluation: LOOKAHEADKV outperforms other baselines across evaluated RULER context lengths and generalizes from 16K training to 32K evaluation contexts.The results use a fixed cache budget of 128.
  • Long-form output evaluation: LOOKAHEADKV outperforms prior approaches on HTML-to-TSV generation at both 12K–0.5K and 23K–2K input–output settings.Both settings use a fixed cache budget ratio of 30% with LLaMA-3.1-8B.
  • Multi-turn evaluation: LOOKAHEADKV is on par or superior on MT-Bench across all tested models and budgets, with particular robustness at C = [64, 128].Responses are evaluated using Qwen3-235B-A22B as the LLM judge.

5 ANALYSIS

The analyses examine efficiency, temperature sensitivity, trainable-module choices, and context-length generalization. LOOKAHEADKV maintains strong performance while limiting overhead, with performance gains saturating beyond a lookahead size of 32.

  • Efficiency comparison: 14.5× lower eviction overhead than LAQ is achieved by LOOKAHEADKV at 32K sequence length.LOOKAHEADKV requires marginal additional cost across tested context lengths, whereas draft-based methods incur computation or memory-traffic overhead.
  • Temperature analysis: LOOKAHEADKV remains superior across temperature settings, while all methods degrade by 3–4% at T = 0.8.The degradation is also observed for FullKV, indicating that stochasticity affects all evaluated approaches similarly.
  • Ablation on trainable modules: Performance gains generally improve with larger lookahead windows and broader LoRA coverage, but gains saturate at n_lookahead = 32.The main configuration therefore uses n_lookahead = 32 with LoRA applied to all linear modules.
  • Context-length generalization: Training on shorter contexts still supports LOOKAHEADKV evaluation beyond the training window, although longer training contexts perform better.The study trains LLaMA-3B with 2K, 4K, and 8K contexts and evaluates on RULER.

6 RELATED WORK

Prior KV-cache eviction work estimates token importance from attention patterns, prompt suffixes, or approximate future responses. Related parameter-efficient methods provide precedents for learned prompts, adapters, and selective activation.

  • KV cache eviction: Sparse attention and stable token importance motivated methods that evict unimportant KV entries while preserving performance.H2O, NACL, and TOVA use attention scores to estimate token importance.
  • Prefill KV cache eviction: Prefill eviction methods range from SnapKV’s prompt observation window to SpecKV’s approximate-response observation window.The passage identifies a shift from prompt-based to response-informed importance estimation.
  • Prompt tuning: Prompt Tuning, Prefix-Tuning, and P-Tuning v2 establish learned embeddings or vectors for parameter-efficient adaptation.These methods are related to LOOKAHEADKV’s use of trainable prompt-like components.
  • Related objectives and adapters: Related training objectives appear in distillation and ranking or retrieval, while contemporaneous work selectively activates LoRA modules for some tokens.These connections situate LOOKAHEADKV’s objective and selective adapters within existing techniques.

7 CONCLUSION AND LIMITATION

LOOKAHEADKV predicts KV-token importance without explicit draft generation by combining learnable lookahead tokens with selectively activated LoRA modules. The method adds little parameter or prefill cost, but its evaluation is limited in model scale and currently targets prefill eviction.

  • Conclusion: LOOKAHEADKV trains a frozen LLM with learnable lookahead tokens and lookahead LoRA modules to match ground-truth importance distributions.The training procedure collects ground-truth and lookahead estimates across layers and heads before optimizing their difference.
  • Efficiency: LOOKAHEADKV introduces less than 0.5% additional parameters and marginal prefill latency.The efficiency analysis estimates TTFT using FLOPs and memory traffic under a fixed LLaMA3.1-8B configuration.
  • Limitation: Experiments could not include larger-sized models because of limited compute resources.The authors report improvements across tested model sizes but identify larger-scale evaluation as unaddressed.
  • Limitation: The current method focuses on prefill KV-cache eviction; extending it to decoding-stage eviction remains future work.This defines the present scope boundary of the framework.
  • Eviction procedure: During eviction, lookahead-token attention to prompt tokens is pooled into scores, and TopK retains the corresponding KV pairs.The method appends learned tokens, performs a prefill pass, and caches selected keys and values under budget k.

C IMPLEMENTATION OPTIMIZATION

LOOKAHEADKV combines flash attention for forward computation with eager attention for importance-score computation and backpropagation. This reduces the memory burden of eager attention when prompts are much longer than responses.

  • Naive eager attention can run out of memory because the full attention matrix grows quadratically with sequence length.
  • Flash attention handles the attention layer’s forward pass, while eager attention computes importance scores and enables backpropagation.
  • O(|X| · |Y| + |Y|2) replaces O((|X| + |Y|)2) for importance computation when |X| ≫ |Y|.Here, |X| and |Y| are the prompt and response lengths, respectively.

D NEED FOR DATA GENERATION

LOOKAHEADKV normally requires target-model responses as training data, but the authors test whether source-dataset responses can substitute for them. The substitution causes only a relatively minor average LongBench drop in lower-budget regimes, suggesting a practical alternative when response generation is impractical.

  • LOOKAHEADKV training requires generated responses from the target model, whose production can be costly across multiple models.
  • Source-dataset responses produce a relatively minor drop in average LongBench performance in lower-budget regimes.
  • The authors suggest source responses may suffice when their attention distributions are moderately similar to those of model-generated responses.
  • LOOKAHEADKV achieves the best RULER performance at 64K and 128K context lengths with a cache budget of 128.The evaluation uses LLaMA3.1-8B-Instruct and 50 randomly sampled examples per task.

E.3 DISCUSSION OF GENERATION STOCHASTICITY IN LOOKAHEADKV TRAINING

The authors examine whether LOOKAHEADKV trained with greedy responses remains applicable under stochastic decoding and report strong persistence of importance patterns. Additional results cover longer contexts, multiple models, and theoretical versus practical latency measurements.

  • Generation stochasticity: Even at temperature T = 0.8, importance patterns show strong persistence between greedy and stochastic decoding responses.Table 8 compares recall@512 and Kendall rank correlation across temperatures, averaged over 30 samples and all layers and heads.
  • Generation stochasticity: LOOKAHEADKV’s stochastic-decoding deviation is smaller than the deviation induced by responses from a speculative model.
  • Additional evaluation: RULER results are reported across six models at cache budgets of 64 and 128.
  • Efficiency analysis: The supplementary latency results note that empirical TTFT overhead can exceed theoretical estimates because of measurement noise and implementation inefficiencies.Better implementations may reduce these overheads toward theoretical costs.
  • Configurations: Training and eviction configurations include searched learning rates, standard baseline settings, and hyperparameters summarized in Tables 16 and 15.

G DATASETS, BENCHMARKS, AND SOFTWARE

The implementation uses publicly available training datasets, established long-context benchmarks, and KVCache-Factory-based software. LLMs assisted with selected wording and literature searches, while the majority of the paper was written and inspected by humans.

  • Software: The source code is provided in the supplementary materials, and the implementation is built on KVCache-Factory.
  • Datasets: The training mixture includes ChatQA2-Long-SFT-data, tulu-3-sft-olmo-2-mixture, The Stack, MetaMathFewshot, HellaSwag DPO Fewshot, and ARC DPO Fewshot.
  • Benchmarks: Evaluation uses LongBench, RULER, and LongProc from their stated official or project sources.
  • Writing process: LLMs assisted with selected wording and related-literature searches, while most text was human-written and generated text was inspected for harmful or controversial content.
Loading 2603.10899v1…