KV Admission: Learning What to Write for Efficient Long-Context Inference

2025/12/19 by Yen-Chieh Huang, Huang, Yen-Chieh, Hsiu, Pi-Cheng +2 · 1 citation
Computer Science · #Advanced Data Storage Technologies #Advanced Neural Network Applications #Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Parallel Computing and Optimization Techniques

paper · doi:10.48550/arxiv.2512.17452

openalex publication_date 2025/12/19 · openalex created_date 2025/12/23 · openalex updated_date 2026/07/28

Abstract

Long-context LLM inference is bottlenecked by the quadratic attention complexity and linear KV cache growth. Prior approaches mitigate this via post-hoc selection or eviction but overlook the root inefficiency: indiscriminate writing to memory. In this paper, we formalize KV cache management as a causal system of three primitives: KV Admission, Selection, and Eviction. We instantiate KV Admission via Write-Gated KV (WG-KV), a lightweight mechanism that learns to predict token utility before cache entry. By filtering out low-utility states early to maintain a compact global cache alongside a sliding local cache, WG-KV reduces memory usage by 46-68% and delivers 3.03-3.70x prefill and 1.85-2.56x decode speedups on Llama and Qwen models, while maintaining compatibility with FlashAttention and Paged-KV systems. These results demonstrate that learning what to write is a principled and practical recipe for efficient long-context inference. Code is available at https://github.com/EMCLab-Sinica/WG-KV.

Citations

• Nemotron-Math: Efficient Long-Context Distillation of Mathematical Reasoning from Multi-Mode Supervision
• Cache What Lasts: Token Retention for Memory-Bounded KV Cache in LLMs
• Expected Attention: KV Cache Compression by Estimating Attention from Future Queries Distribution
• Cache Me If You Can: How Many KVs Do You Need for Effective Long-Context LMs?
• Inference-Time Hyper-Scaling with KV Cache Compression
• KVzip: Query-Agnostic KV Cache Compression with Context Reconstruction
• XAttention: Block Sparse Attention with Antidiagonal Scoring
• Plan-and-Act: Improving Planning of Agents for Long-Horizon Tasks
• Q-Filters: Leveraging QK Geometry for Efficient KV Cache Compression
• FlexPrefill: A Context-Aware Sparse Attention Mechanism for Efficient Long-Sequence Inference
• SpargeAttention: Accurate and Training-free Sparse Attention Accelerating Any Model Inference
• LServe: Efficient Long-sequence LLM Serving with Unified Sparse Attention
• MoBA: Mixture of Block Attention for Long-Context LLMs
• Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention
• Unveiling Simplicities of Attention: Adaptive Long-Context Head Identification
• AttentionPredictor: Temporal Patterns Matter for KV Cache Compression
• LongProc: Benchmarking Long-Context Language Models on Long Procedural Generation
• LongBench v2: Towards Deeper Understanding and Reasoning on Realistic Long-context Multitasks
• ClusterKV: Manipulating LLM KV Cache in Semantic Space for Recallable Compression
• Squeezed Attention: Accelerating Long Context Length LLM Inference
• ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference
• MagicPIG: LSH Sampling for Efficient LLM Generation
• MatryoshkaKV: Adaptive KV Compression via Trainable Orthogonal Projection
• MoH: Multi-Head Attention as Mixture-of-Head Attention
• In-context KV-Cache Eviction for LLMs via Attention-Gate
• DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads
• RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval
• Eigen Attention: Attention in Low-Rank Space for KV Cache Compression
• NACL: A General and Effective KV Cache Eviction Framework for LLMs at Inference Time
• Palu: Compressing KV-Cache with Low-Rank Projection
• ThinK: Thinner Key Cache by Query-Driven Pruning
• RazorAttention: Efficient KV Cache Compression Through Retrieval Heads
• Model Tells You Where to Merge: Adaptive KV Cache Merging for LLMs on Long-Context Tasks
• MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention
• PQCache: Product Quantization-based KVCache for Long Context LLM Inference
• The FineWeb Datasets: Decanting the Web for the Finest Text Data at Scale
• A Simple and Effective L₂ Norm-Based Strategy for KV Cache Compression
• Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference
• A Training-free Sub-quadratic Cost Transformer Model Serving Framework With Hierarchically Pruned Attention
• Loki: Low-rank Keys for Efficient Sparse Attention
• MiniCache: KV Cache Compression in Depth Dimension for Large Language Models
• PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference
• Reducing Transformer Key-Value Cache Size with Cross-Layer Attention
• SirLLM: Streaming Infinite Retentive LLM
• Layer-Condensed KV Cache for Efficient Inference of Large Language Models
• Retrieval Head Mechanistically Explains Long-Context Factuality
• SnapKV: LLM Knows What You are Looking for Before Generation
• Get More with LESS: Synthesizing Recurrence with KV Cache Compression for Efficient LLM Inference
• KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization
• Transformers are Multi-State RNNs
• SGLang: Efficient Execution of Structured Language Model Programs
• SparQ Attention: Bandwidth-Efficient LLM Inference
• Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs
• SEA: Sparse Linear Attention with Estimated Attention Mask
• Efficient Streaming Language Models with Attention Sinks
• Efficient Memory Management for Large Language Model Serving with PagedAttention
• H₂O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
• Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time
• Dynamic Context Pruning for Efficient and Interpretable Autoregressive Transformers
• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
• When Not to Trust Language Models: Investigating Effectiveness of Parametric and Non-Parametric Memories
• The Stack: 3 TB of permissively licensed source code
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• Efficient Intent Detection with Dual Sentence Encoders
• Root Mean Square Layer Normalization
• An Evaluation Dataset for Intent Classification and Out-of-Scope Prediction
• Benchmarking Natural Language Understanding Services for building Conversational Agents
• HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering
• The NarrativeQA Reading Comprehension Challenge
• Attention Is All You Need
• TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for\n Reading Comprehension
• RetroInfer: A Vector-Storage Approach for Scalable Long-Context LLM Inference

Cited by

• Keyless Attention: Value-Space Routing and Value-Only Caching for Efficient Transformers

Related

• LLMs Know What to Drop: Self-Attention Guided KV Cache Eviction for Efficient Long-Context Inference
• Cache Me If You Can: How Many KVs Do You Need for Effective Long-Context LMs?
• Compressing KV Cache for Long-Context LLM Inference with Inter-Layer Attention Similarity
• Cache What Lasts: Token Retention for Memory-Bounded KV Cache in LLMs
• C²KV: Compressed and Composable KV Cache Reuse for Efficient LLM Inference