Top-K Retrieval with Fixed-Size Linear-Attention Completion: Backbone- and KV-Format-Preserving Attention for KV-Cache Read Reduction
- URL: http://arxiv.org/abs/2604.05438v1
- Date: Tue, 07 Apr 2026 05:19:54 GMT
- Title: Top-K Retrieval with Fixed-Size Linear-Attention Completion: Backbone- and KV-Format-Preserving Attention for KV-Cache Read Reduction
- Abstract summary: We propose a retrieval-completion attention module that keeps backbone weights and the KV-cache format unchanged.<n>For each query, we compute exact attention over sink/tail anchors and the query-dependent retrieved Top-K tokens.<n>We add the exact and estimated contributions in the unnormalized domain and apply a single normalization, recovering the missing softmax mass without additional attention-side KV reads.
- Score: 3.1897721055763024
- License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
- Abstract: Long-context generation is increasingly limited by decode-time key-value (KV) cache traffic, particularly when KV is offloaded beyond GPU memory. Query-aware retrieval (e.g., Top-K selection) reduces this traffic by loading only a subset of KV pairs, but renormalizing the softmax over the subset introduces bias when attention mass is spread over unretrieved tokens. We propose a retrieval-completion attention module that keeps backbone weights and the KV-cache format unchanged. For each query, we compute exact attention over sink/tail anchors and the query-dependent retrieved Top-K tokens, and estimate the remaining mid-region numerator and denominator using a fixed-size feature-map summary computed at prefill time. We add the exact and estimated contributions in the unnormalized domain and apply a single normalization, recovering the missing softmax mass without additional attention-side KV reads. Across long-context benchmarks, the proposed method improves over selection-only Top-K at matched token-equivalent read budgets, with the largest gains in high-entropy heads.
Related papers
- SGD-KV: Summarization Guided KV Cache Compression [21.816110519918876]
Large language models (LLMs) face severe memory bottlenecks in long-context inference due to the linearly growing key-value ( KV) caches.<n>We present SGD-KV, a head-aware framework that leverages a novel chunk-summarization diagnostic task to systematically identify and prioritize attention heads specialized in hierarchical information aggregation.<n>Our findings show that strategically allocating the KV cache budget based on the summarization score distribution of attention heads yields a superior efficiency-accuracy trade-off for long-context inference.
arXiv Detail & Related papers (2026-09-03T00:31:11Z) - CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration [67.17709877327232]
Cate KV is a hybrid KV cache method that retains only critical token information for consistent heads.<n>Cate KV reduces memory usage by up to $2.72times$ and accelerates decoding by $2.18times$ in single-sample inputs.
arXiv Detail & Related papers (2026-08-31T06:02:37Z) - ResKV: Reconstructing Omitted Attention Contributions for Fixed-Budget KV Cache Compression [25.24197102648478]
KV cache compression is essential for efficient long-context inference.<n>Merging-based alternatives preserve more information but can perturb retained keys and values that should remain exact.<n>We propose ResKV, which divides a fixed KV budget into an exact main cache and a compact residual cache.
arXiv Detail & Related papers (2026-07-31T16:16:45Z) - Back from the Future: Key-Value Cache Management by Counter-Causal Surprise [67.1056509495879]
Key-value (KV) cache management through compression and eviction strategies has emerged as an important research direction in recent years.<n>We propose a simple yet effective KV eviction scheme motivated by the insight that past tokens which can be well-predicted from more recent tokens are redundant.<n>We evaluate our strategy on various open-source LLMs and benchmark datasets showing competitive or improved performance over other state-of-the-art methods.
arXiv Detail & Related papers (2026-07-30T02:42:51Z) - Make Each Token Count: Towards Improving Long-Context Performance with KV Cache Eviction [65.710271475739]
We introduce a global retention-based KV eviction method that learns each token's future utility under a unified memory budget.<n>We show that our method substantially reduces KV memory while matching or surpassing full-cache inference.<n>These results suggest that learned, globally calibrated KV eviction is not only a compression technique, but also a mechanism for improving long-context reasoning.
arXiv Detail & Related papers (2026-05-10T16:47:50Z) - Kwai Summary Attention Technical Report [69.40814939510126]
Long-context ability, has become one of the most important direction of next-generation Large Language Models.<n>Standard softmax attention exhibits quadratic time complexity with respect to sequence length.<n>We propose Kwai Summary Attention (KSA), a novel attention mechanism that reduces sequence modeling cost by compressing historical contexts.
arXiv Detail & Related papers (2026-04-27T12:59:53Z) - SpecAttn: Co-Designing Sparse Attention with Self-Speculative Decoding [3.47505166101801]
SpecAttn is a self-speculative decoding method with verification-guided sparse attention.<n>It achieves 2.81$times$ higher throughput over vanilla auto-regressive decoding and 1.29$times$ improvement over state-of-the-art sparsity-based decoding methods.
arXiv Detail & Related papers (2026-02-06T22:12:52Z) - ProphetKV: User-Query-Driven Selective Recomputation for Efficient KV Cache Reuse in Retrieval-Augmented Generation [22.835149054167122]
We propose Prophet KV, a user-query-driven KV Cache reuse method for RAG scenarios.<n> Prophet KV prioritizes tokens based on their semantic relevance to the user query.<n>Our evaluation results show that Prophet KV retains 96%-101% of full-prefill accuracy with only a 20% recomputation ratio.
arXiv Detail & Related papers (2026-01-31T09:53:31Z) - Crystal-KV: Efficient KV Cache Management for Chain-of-Thought LLMs via Answer-First Principle [22.39292106468653]
Chain-of-Thought (CoT) reasoning significantly improves accuracy on complex tasks.<n>CoT emphasizes the final answer, rendering conventional KV compression strategies ineffective.<n>We present Crystal-KV, an efficient KV cache management framework tailored for CoT reasoning.
arXiv Detail & Related papers (2026-01-05T07:42:58Z) - G-KV: Decoding-Time KV Cache Eviction with Global Attention [57.47409249054187]
Large language models (LLMs) excel in complex tasks but encounter significant computational and memory challenges due to long sequence lengths.<n> KV cache compression has emerged as an effective approach to greatly enhance the efficiency of reasoning.<n>We propose G-KV, a KV cache eviction method that employs a global scoring mechanism, combining local and historical attention scores to more accurately assess token importance.
arXiv Detail & Related papers (2025-11-29T14:21:33Z) - More Tokens, Lower Precision: Towards the Optimal Token-Precision Trade-off in KV Cache Compression [71.42818367729573]
In large language models (LLMs), the memory usage of KV cache has become a critical bottleneck during inference.<n>The mainstream KV compression methods, including KV pruning and KV quantization, primarily focus on either token or precision dimension separately.<n>In this paper, we comprehensively investigate the token-precision trade-off in KV cache compression.
arXiv Detail & Related papers (2024-12-17T09:20:31Z) - Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs [82.08922896531618]
We introduce adaptive KV cache compression, a plug-and-play method that reduces the memory footprint of generative inference for Large Language Models (LLMs)
We conduct targeted profiling to discern the intrinsic structure of attention modules.
Based on the recognized structure, we then construct the KV cache in an adaptive manner: evicting long-range contexts on attention heads emphasizing local contexts, discarding non-special tokens on attention heads centered on special tokens, and only employing the standard KV cache for attention heads that broadly attend to all tokens.
arXiv Detail & Related papers (2023-10-03T05:17:08Z)
This list is automatically generated from the titles and abstracts of the papers in this site.
This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.