TokenSelect: Efficient Long-Context Inference and Length Extrapolation for LLMs via Dynamic Token-Level KV Cache Selection

2024/11/05 by Wei Wu, Wu, Wei, Zhuoshi Pan +14 · 17 citations
Computer Science · Mathematics · #Advanced Data Storage Technologies #Artificial intelligence #Biology #Cache #Computer science #Computer security #Context (archaeology) #Extrapolation #Inference #Mathematics #Network Packet Processing and Optimization #Parallel computing #Security token #Selection (genetic algorithm) #Statistics

paper · pdf · doi:10.48550/arxiv.2411.02886

published in arXiv (Cornell University) (Cornell University)

openalex publication_date 2024/11/05 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Rapid advances in Large Language Models (LLMs) have spurred demand for processing extended context sequences in contemporary applications. However, this progress faces two challenges: performance degradation due to sequence lengths out-of-distribution, and excessively long inference times caused by the quadratic computational complexity of attention. These issues limit LLMs in long-context scenarios. In this paper, we propose Dynamic Token-Level KV Cache Selection (TokenSelect), a training-free method for efficient and accurate long-context inference. TokenSelect builds upon the observation of non-contiguous attention sparsity, using QK dot products to measure per-head KV Cache criticality at token-level. By per-head soft voting mechanism, TokenSelect selectively involves a few critical KV cache tokens in attention calculation without sacrificing accuracy. To further accelerate TokenSelect, we design the Selection Cache based on observations of consecutive Query similarity and implemented the efficient Paged Dot Product Kernel, significantly reducing the selection overhead. A comprehensive evaluation of TokenSelect demonstrates up to 23.84× speedup in attention computation and up to 2.28× acceleration in end-to-end latency, while providing superior performance compared to state-of-the-art long-context inference methods.

Cited by

• LOCKS: Page-Local Compact Key Summaries for Efficient Long-Context Decoding
• CoSA: Accelerating Long-Context Inference via Proxy-Kernel Co-Designed Sparse Attention
• CTkvr: KV Cache Retrieval for Long-Context LLMs via Centroid then Token Indexing
• Accelerating Large-Scale Reasoning Model Inference with Sparse Self-Speculative Decoding
• KVSwap: Disk-aware KV Cache Offloading for Long-Context On-device Inference
• ContextPilot: Fast Long-Context Inference via Context Reuse
• AssoMem: Scalable Memory QA with Multi-Signal Associative Retrieval
• PatternKV: Flattening KV Representation Expands Quantization Headroom
• ProxyAttn: Guided Sparse Attention via Representative Heads
• SparseServe: Unlocking Parallelism for Dynamic Sparse Attention in Long-Context LLM Serving
• Recall Before You Rank: Similarity-Guided Top-K Reuse for Efficient Long-Context Attention
• UniGist: Towards General and Hardware-aligned Sequence-level Long Context Compression
• Retrospective Sparse Attention for Efficient Long-Context Generation
• Beyond Exponential Decay: Rethinking Error Accumulation in Large Language Models
• RetroInfer: A Vector-Storage Approach for Scalable Long-Context LLM Inference
• Token Reduction Should Go Beyond Efficiency in Generative Models -- From Vision, Language to Multimodality
• TALSC: Timeliness-Aware Large-Small VLM Collaboration for Infrastructure-Assisted Autonomous Driving

Related

• LLMs Know What to Drop: Self-Attention Guided KV Cache Eviction for Efficient Long-Context Inference
• Compressing KV Cache for Long-Context LLM Inference with Inter-Layer Attention Similarity
• Unshackling Context Length: An Efficient Selective Attention Approach through Query-Key Compression
• Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference
• KVCompose: Efficient Structured KV Cache Compression with Composite Tokens