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

Introduction

With the rapid development of large language models (LLMs), the number of parameters is no longer the sole factor significantly affecting model performance. The ability to effectively process longer context information has become one of the key metrics for evaluating LLMs’ capabilities. The latest Web applications such as cross-document understanding (Bai et al., 2024), LLM-powered search systems (Sharma et al., 2024), repository-level code completion (Zhang et al., 2023; Di et al., 2024), and complex reasoning (OpenAI, [n. d.]) have all placed higher demands on the long-context abilities of LLMs. There are two main difficulties in using pre-trained LLMs for long-context inference. On one hand, LLMs are limited by their context length during pre-training (e.g. Llama 3 only has 8192 tokens). Directly inferencing on longer sequences can lead to severe performance degradation due to reasons including sequence lengths out-of-distribution (Xiao et al., 2024; Han et al., 2024). On the other hand, even if LLMs possess sufficiently large context lengths, the quadratic computational complexity of attention with respect to sequence length makes the response time for long-context inference unbearable.

Previous works have made numerous attempts to address these difficulties. To extend the context length of LLMs, the current common practice is to perform post-training on long texts (Team et al., 2024; Yang et al., 2024; GLM et al., 2024). However, this approach comes with significant computational costs, particularly in two aspects: the synthesis of high-quality long-text data and the training process on extended sequences. To accelerate long-context inference, many studies focus on the sparsity of attention, attempting to reduce the scale of KV Cache involved in computation. The key to this type of method lies in designing sparse patterns for attention, which can be mainly divided into two categories: one uses predefined sparse patterns (Wang et al., 2019; Zaheer et al., 2020; Xiao et al., 2024; Han et al., 2024), while the other estimates the potential importance of KV Cache during the inference process (Zhang et al., 2024; Oren et al., 2024; Xiao et al., 2024; Tang et al., 2024; Jiang et al., 2024), attempting to select relevant KV Cache tokens into attention calculations. However, the design of these sparse patterns is often heuristically based on historical criticality or coarse-grained criticality estimation of tokens, making it difficult to ensure that the selected tokens are truly critical, thus resulting in sub-optimal performance, as shown in Fig. 1.

In this paper, we further observe the non-contiguous sparsity of attention, revealing the importance of designing more fine-grained dynamic sparse patterns. To this end, we propose TokenSelect, a model-agnostic and training-free approach that utilizes token-level selective sparse attention for efficient long-context inference and length extrapolation. Specifically, for each Query, TokenSelect dynamically calculates token-level per-head criticality for the past KV Cache and selects the $k$ most critical tokens through our head soft vote mechanism, involving them in the attention calculation. This reduces the scale of attention calculation to a constant length familiar to the model, while maintaining almost all of the long-context information, thereby simultaneously addressing the two main difficulties for long-context inference. To reduce the overhead of token selection, TokenSelect manages the KV Cache in token-level pages (Zheng et al., 2024) and design efficient kernel for token selection based on Paged KV Cache management through Triton (Tillet et al., 2019). Furthermore, based on our observation of high similarity between consecutive queries, we have designed the Selection Cache, which allows consecutive similar queries to share token selection results, thereby reducing the selection frequency while ensuring its effectiveness.

We evaluate the performance and efficiency of TokenSelect on three representative long-context benchmarks (Zhang et al., 2024; Hsieh et al., 2024; Bai et al., 2024) using three open-source LLMs (Yang et al., 2024; Dubey et al., 2024; AI et al., 2024). The experimental results demonstrate that our TokenSelect can achieve up to $23.84\times$ speedup in attention computation compared to FlashInfer (flashinfer ai, [n. d.]), and up to $2.28\times$ acceleration in end-to-end inference latency compared to state-of-the-art long-context inference method (Xiao et al., 2024). Simultaneously, it provides superior performance on three long-text benchmarks. In summary, we make the following contributions:

An observation on the non-contiguous sparsity of attention that highlights the importance of token-level selection.

TokenSelect, a model-agnostic and training-free method that achieves accurate and efficient long-context inference and length extrapolation, which is compatible with mainstream LLM serving systems and ready for Web applications.

A comprehensive evaluation of TokenSelect, demonstrating up to $23.84\times$ speedup in attention computation and up to $2.28\times$ acceleration in end-to-end latency while exhibiting superior performance.

Related Works

Due to computational complexity constraints, current LLMs based on Transformers often utilize limited context lengths during pre-training (Touvron et al., 2023; Dubey et al., 2024; Jiang et al., 2023; Yang et al., 2024; GLM et al., 2024; AI et al., 2024). To extend the long-context capabilities of LLMs, current methods can be broadly categorized into three approaches (Huang et al., 2024; Zhou et al., 2024; Zhao et al., 2024): 1) Modifying positional encodings: A widely adopted method is positional interpolation (Chen et al., 2023). Chen et al. first proposed linear scaling of RoPE (Su et al., 2024) to map longer positional ranges within the original training window. Subsequent works (bloc97, 2023; emozilla, 2023) further improved this method using Neural Tangent Kernel (NTK) theory (Jacot et al., 2018), achieving longer context windows while maintaining model performance. Methods like YaRN (Peng et al., 2024) and Giraffe (Pal et al., 2023) optimize interpolation effects by adjusting frequency components or introducing temperature parameters. 2) Long-context post-training: This approach extends the model’s context length through additional training steps on longer documents after pre-training (Yang et al., 2024; Tian et al., 2024). It has been widely adopted by leading LLMs (Team et al., 2024; Yang et al., 2024; GLM et al., 2024) with the support of sequence parallelism techniques (Shoeybi et al., 2020; Jacobs et al., 2023; Liu et al., 2024). 3) Incorporating additional memory modules: Notable examples include Transformer-XL (Dai* et al., 2019), Compressive Transformer (Rae et al., 2020), RMT (Bulatov et al., 2022) and Infini-attention (Munkhdalai et al., 2024). Although these methods have expanded the context length of LLMs, long-context inference still faces the challenge of high computational costs.

Efficient Long-context Inference

In state-of-the-art LLMs serving systems (Kwon et al., 2023; Huggingface, 2024; NVIDIA, 2024; Zheng et al., 2024), technologies such as Flash Attention (Dao et al., 2022; Dao, 2024) and Paged Attention (Kwon et al., 2023) have greatly optimized LLMs inference efficiency by improving GPU I/O bottlenecks. However, in long-context inference scenarios, the quadratic computational complexity of attention with respect to sequence length poses new challenges for LLMs inference. Numerous studies focus on the sparsity of attention, selecting partial KV Cache for attention calculations to improve long-context inference efficiency. Sliding window (Wang et al., 2019; Zaheer et al., 2020) is one of the most widely used sparse patterns, reducing complexity to linear by executing attention computations within localized windows. Recent works like StreamingLLM (Xiao et al., 2024) and LM-infinite (Han et al., 2024) retain the initial tokens of the sequence in addition to sliding windows, effectively maintaining LLMs’ performance when processing long sequences. While these approaches are simple to implement, they cannot retain information from long contexts. Another approach focuses on dynamic KV Cache selection during inference. Methods like H2O (Zhang et al., 2024), TOVA (Oren et al., 2024), FastGen (Ge et al., 2024), Scissorhands (Liu et al., 2023), and SnapKV (Li et al., 2024) evaluate token criticality based on historical attention scores, selecting tokens within a limited budget. However, these methods permanently discard parts of the KV Cache, causing information loss from long contexts. To address this, InfLLM (Xiao et al., 2024) introduces Block Memory Units for KV Cache management, retrieving information from long contexts and offloading less-used blocks to CPU. Similarly, QUEST (Tang et al., 2024) proposes query-aware sparsity at page granularity, while MInference (Jiang et al., 2024) optimizes long-context inference using three sparse patterns. Apart from considering all attention heads, some other works (Ribar et al., 2024; Lee et al., 2024; Tang et al., 2024) attempt to focus on only a subset of attention heads. Beyond selection, some other research focuses on KV Cache quantization (Liu et al., 2024; Yang et al., 2024; He et al., 2024; Kang et al., 2024) and merging (Liu et al., 2024; Wan et al., 2024; Zhang et al., 2024; Wang et al., 2024). However, existing methods struggle to be applied in real-world Web applications, both in terms of accuracy and efficiency.

Preliminaries

In this section, we first introduce the inference process of LLMs, and then define the Selective Sparse Attention Problem.

Nowadays, mainstream LLMs are primarily based on the Decoder-only Transformer architecture, consisting sequentially of a word embedding layer, a series of transformer layers, and a token prediction head. Each transformer layer includes a multi-head attention (MHA) module and a feed-forward networks (FFN) module. The inference process of LLMs can be divided into two stages: the Prefill Stage and the Decode Stage.

where $\mathbf{W}_{q},\mathbf{W}_{k},\mathbf{W}_{v}$ are linear projections, $[\cdot]$ represents tensor concatenation operation, and Eq.(2) is also known as Scaled Dot-Product Attention (SDPA). After these computation, $\mathbf{K}_{\text{prefill}}$ and $\mathbf{V}_{\text{prefill}}$ are stored as the KV Cache for current layer $\mathbf{K}_{\text{cache}}$ and $\mathbf{V}_{\text{cache}}$, and $\mathbf{O}_{\text{prefill}}$ is used for subsequent calculations.

where $\mathbf{K}_{\text{decode}},\mathbf{V}_{\text{decode}}$ are composed of the KV Cache and the KV corresponding to the current input, which are then used to update the KV Cache of the current layer for use in the next forward pass.

LLMs inference, unlike training, is memory-bound, necessitating frequent GPU I/O operations between HBM and SRAM while underutilizing processing units. This bottleneck is particularly evident in SDPA computation. Optimizing for I/O is crucial for enhancing LLMs inference efficiency, especially in long-context scenarios.

2. Selective Sparse Attention

As discussed in the Sec. 1, the high attention sparsity in LLMs suggests sparse attention as a promising solution for long-context inference challenges. Sparse attention can keep the number of tokens participating in attention computations at a constant scale, rather than increasing with sequence length. Given that predefined sparse patterns are detrimental to performance, we aim to dynamically select crucial tokens for attention computation at each step during the inference process. Therefore, we formalize this problem according to the following definition.

where $\mathcal{I}$ is the set of selected indices. The objective is to find an appropriate selection function $\mathcal{S}$ that minimizes the difference between the outputs of the SDPA and the selective sparse attention:

Existing works on long-context inference (Wang et al., 2019; Zaheer et al., 2020; Xiao et al., 2024; Han et al., 2024; Zhang et al., 2024; Oren et al., 2024; Li et al., 2024; Xiao et al., 2024; Tang et al., 2024; Jiang et al., 2024) can be categorized under the Selective Sparse Attention Problem, with variations primarily in the design of the selection function $\mathcal{S}$. (Wang et al., 2019; Zaheer et al., 2020; Xiao et al., 2024) have developed input-independent selection functions $\mathcal{S}()$, while (Zhang et al., 2024; Oren et al., 2024; Li et al., 2024) propose query-independent functions $\mathcal{S}(\mathbf{K}_{\text{cache}})$ for improved performance. Current state-of-the-art methods (Xiao et al., 2024; Tang et al., 2024; Jiang et al., 2024) utilize query-aware selection functions $\mathcal{S}(\mathbf{Q},\mathbf{K}_{\text{cache}})$. However, these approaches typically operate at a block-level, which limits their effectiveness and overall performance.

Motivations and Observations

Previous works (Wang et al., 2019; Zaheer et al., 2020; Xiao et al., 2024; Han et al., 2024; Zhang et al., 2024; Oren et al., 2024; Li et al., 2024; Xiao et al., 2024; Tang et al., 2024; Jiang et al., 2024) have demonstrated the sparsity of attention scores in LLMs, particularly when processing long texts. Recent approaches (Xiao et al., 2024; Jiang et al., 2024; Tang et al., 2024) partition the KV Cache into non-overlapping blocks, estimating block criticality for sparse attention calculations. These methods assume that tokens with higher attention scores tend to be contiguous. However, our further observations reveal that this assumption does not always hold true in practice. As illustrated in Fig. 2a, attention scores are sparsely distributed at the token-level, with critical tokens not necessarily contiguous. This non-contiguity leads to significant omissions in block-level token selection. Fig. 2b demonstrates that finer selection granularity improves recall of critical tokens, motivating us to perform token-level selection. For token-level selection, an intuitive approach would be to directly select the top-$k$ tokens with the highest attention logits. However, observation in Fig. 2c reveals considerable disparity in the $L_{1}$ norm of attention logits across attention heads. As a result, the selection result tends to be dominated by a few heads with disproportionately large attention logits, driving us to design a more robust selection function that maintains the independence of heads.

Consecutive Queries are similar

As sparsity of attention is dynamic (Zhang et al., 2024; Xiao et al., 2024; Tang et al., 2024; Jiang et al., 2024), token selection should be performed for every Query, which inevitably increases the computational overhead of selective sparse attention. Fortunately, we observe that consecutive Queries exhibit high similarity, as shown in Fig. 3a. Intuitively, when two consecutive Queries are highly similar, their dot products with the Keys will also be similar, leading to substantial overlap in the token selection results. Due to space constraints, we provide an informal lemma about this below. The formal version and corresponding proof can be found in the Appendix A.

Fig. 3b illustrates this lemma experimentally. It can be seen that the overlap rate of token selection tends to increase with query similarity. This key insight motivates us to reuse selection results for similar queries, improving computational efficiency. Moreover, the similarity distribution of consecutive Queries remains consistent across different tasks, as demonstrated in Fig. 3a, allowing us to apply a global similarity threshold across all scenarios.

Designs of TokenSelect

In this section, we will introduce the design details of TokenSelect, primarily encompassing the Selection Function, the Selection Cache, and efficient implementation of TokenSelect. The overall workflow of TokenSelect is illustrated in Fig. 4.

The simplest selection function is to determine the criticality of the tokens through the dot product of $\mathbf{Q}$ and $\mathbf{K}_{\text{cache}}$, then select the top-$k$ critical ones as $\mathbf{K}_{\text{select}},\mathbf{V}_{\text{select}}$. The selected indices $\mathcal{I}$ are calculated as follow:

However, as discussed in Sec. 4, this approach is prone to inaccuracies due to disparities in norm of attention logits between heads. To maintain independence between heads, a better approach is to have each head select the top-$k$ most critical tokens, and then determine the final selection through voting among the heads:

2. Optimizing Selection Frequency

Although the aforementioned selection function can reduce the complexity of attention from $O(N^{2})$ to $O(k^{2})$ ($k\ll N$), while maintaining performance, the execution time of the selection function itself still affects the latency of LLMs inference. To further accelerate long-context inference, based on our observations of the similarity of consecutive queries, we design optimization strategies for both the Prefill Stage and the Decode Stage to reduce the selection frequency while ensuring its effectiveness.

3. Efficient Implementation

To ensure that our proposed TokenSelect can be used for real-world Web applications, efficient implementation is crucial. We first analyze the computation time breakdown of representative block-level selective sparse attention method, InfLLM (Xiao et al., 2024). From (1)(2)(3) in Fig. 5, we can observe that although selective sparse attention can significantly reduce the complexity of attention calculations, the actual computation time is still highly dependent on the implementation. The incompatibility with efficient attention implementations such as Flash Attention has resulted in methods requiring historical attention scores (Zhang et al., 2024; Oren et al., 2024; Li et al., 2024; Xiao et al., 2024) being difficult to be applied in real-world Web applications. Through the analysis of InfLLM’s Flash Attention-compatible version, we make several discoveries. The initial motivation for estimating token criticality at the block-level is to reduce the overhead of selection function (mainly considering dot product calculation). However, we find that dot product is not the primary performance bottleneck. Instead, a significant portion of the overhead comes from indexing the KV Cache using selected indices and making them contiguous in GPU memory, which frequently occurs during the updating of KV blocks and the concatenation of selected KV Cache. The extensive I/O required for this operation further exacerbates the memory-bound in LLMs inference. Based on this, we propose that Paged Attention is a more suitable implementation for selective sparse attention. Using Paged KV Cache management (with page size=1 for TokenSelect), we can reduce the I/O volume for selection results from the scale of all selected KV Caches $O(2kd)$ to the scale of their indices $O(k)$. However, by observing (4) in Fig. 5, we find that we encounter another bottleneck under Paged KV Cache management. Since logically contiguous KV Cache is not entirely contiguous in GPU memory, it also needs to be made contiguous before performing computational operations. To address this issue, we draw inspiration from the concept of Paged Attention and implement a Paged Dot Product Kernel using Triton (Tillet et al., 2019), which significantly improves the overall efficiency of TokenSelect.

Experiments

In this section, we first introduce the experimental setup of this paper, and then reveal the performance and efficiency of our TokenSelect in long-context inference through experiments.

To evaluate TokenSelect’s performance on long-context inference, we use the following datasets: (1) InfiniteBench (Zhang et al., 2024): The mainstream long-context benchmark consisting of multi-tasks. The average length of it exceeds 200K tokens. (2) RULER (Hsieh et al., 2024): A challenging long-context benchmark containing 13 different tasks, with subsets of varying lengths up to 128K tokens. (3) LongBench (Bai et al., 2024): Another mainstream long-context benchmark comprising 6 types of tasks. The 95% percentile for its lengths is 31K tokens. For each dataset, we use its recommended metrics, which are presented in the Appendix B.

Baselines.

To demonstrate the state-of-the-art (SOTA) performance of TokenSelect, we include the following methods for comparison: (1) Original models: We select three mainstream open-source LLMs - Qwen2-7B-Instruct (Yang et al., 2024), Llama-3-8B-Instruct (Dubey et al., 2024), and Yi-1.5-6B-Chat (AI et al., 2024) - utilizing their original context lengths without any modifications. (2) NTK-Aware Scaled RoPE: A nonlinear RoPE interpolation method. (3) SelfExtend: A RoPE interpolation method that reuses the position ids across neighboring tokens. (4) StreamingLLM: The SOTA method for long-context inference with predefined sparse patterns. (5) InfLLM: The SOTA method for long-context inference and length extrapolation using a block-level selective sparse attention method. (6) MInference: The SOTA method for long-context prefilling acceleration, utilizing three sparse patterns including block-level sparse attention. It’s worth noting that since MInference doesn’t support length extrapolation, we use an alternative evaluation method, applying it to Llama-3-8B-Instruct-262k (Llama3 after long-text post-training). Additionally, we do not include another SOTA method, QUEST (Tang et al., 2024), as it does not support Grouped Query Attention (GQA).

Implementation details.

In all experiments in this paper, we employ greedy decoding to ensure the reliability of the results. For our TokenSelect, we implement it on SGLang (Zheng et al., 2024), which is a fast serving framework based on Flasherinfer (flashinfer ai, [n. d.]). We implement our method using PyTorch (Paszke et al., 2019) and Triton (Tillet et al., 2019). We follow the baseline approach, including 128 initial tokens and $n_{\text{local}}$ most recent tokens in the attention computation in addition to the $k$ selected tokens. For NTK and SelfExtend, we extend the model’s context length to 128K. For StreamLLM, we set $n_{\text{local}}=\text{4K}$ . For InfLLM, we set $k=\text{4K},n_{\text{local}}=\text{4K}$. For our TokenSelect, we set $k=\text{2K},n_{\text{local}}=\text{512}$ to demonstrate our token-level KV Cache selection allows us to achieve better performance with a smaller token budget. Due to the need to demonstrate the method under different $n_{\text{local}}$ and $k$, we denote the specific token budgets in the form of $k+n_{\text{local}}$ if they differ from the aforementioned settings. We use NVIDIA A100 to conduct all experiments. When inferencing sequences over 1M tokens, we additionally employee tensor parallelism, which is transparent to our TokenSelect.

2. Performance Comparisons

As shown in Table 1, our TokenSelect achieves significantly superior overall performance on InfiniteBench compared to all baseline methods, even though TokenSelect uses the smallest token budget (¡3K). The fact that it significantly outperforms the original models demonstrates TokenSelect’s strong length extrapolation capability. We analyze that this is due to our adoption of a fine-grained KV Cache selection strategy, while considering the equal contribution of each head to selection, which ensures that we can select most critical tokens. Observing the performance of other methods, we find that RoPE interpolation methods (NTK, SelfExtend) generally perform poorly unless used on specially trained models such as Qwen2-7B-Instruct. The better performance of Qwen2-7B-Instruct on the original model can also be attributed to this. The sparse attention method StreamingLLM, based on fixed sparse patterns, can guarantee some of the model’s capabilities, but due to discarding a large amount of long-context information, it performs poorly on retrieval-related tasks (R.PK, R.Num, R.KV). The block-level selection method InfLLM can retain more long-context information compared to StreamingLLM. However, due to its sub-optimal block-level selection, it results in lower performance on most tasks compared to TokenSelect, even though we set a larger token budget for InfLLM. It is worth noting that Yi-1.5-6B does not perform normally on the R.KV task, as it is unable to correctly recite strings like the Universally Unique Identifier.

RULER

To further demonstrate the long-context capability of TokenSelect, we conduct evaluation on the more challenging long-context benchmark RULER. Considering the increased difficulty of RULER and its substantial computational requirements, we include only comparable baseline methods. As shown in Table 2, our TokenSelect maintains significantly superior overall performance compared to other long-context inference methods. For all models, TokenSelect achieves length extrapolation while preserving the model’s original capabilities, benefiting from our efficient utilization of the model’s limited context length. Notably, due to the constraints of model’s context length, TokenSelect experiences performance degradation with larger token budgets (4K+4K) on Llama and Yi. However, its performance with smaller token budgets still significantly surpasses other baseline methods.

LongBench.

Due to space constraints, the results of LongBench are presented in the Appendix C. Although its relatively shorter text length makes it less suitable for evaluating state-of-the-art long-context inference methods, our TokenSelect still demonstrates superior overall performance compared to most baseline methods.

Comparing to methods based-on post-trained model.

In Table 3, we present the performance of the post-trained model and long-context inference method (Jiang et al., 2024) based on it. It shows that even compared to length extrapolation methods requiring additional training, the training-free TokenSelect still exhibits superior performance on most tasks. Although Minference can improve the performance of the original model, it fails to reverse the negative impact of long-text post-training on shorter text tasks (LongBench).

3. Ablation Studies

In ablation studies, we primarily analyze the impact of different Selection Functions $\mathcal{S}$ on performance. To compare the performance of different Selection Functions $\mathcal{S}$ under low token budgets (i.e., token efficiency), we maintain the 2K+512 configuration. From Table 4, we can observe that our proposed head soft vote mechanism performs significantly better across all tasks. This indicates that using the head soft vote mechanism to balance each head’s contribution to token selection results can help us avoid the domination of selection by few heads with large attention logits.

4. Hyper-parameter Analysis

As shown in Table 5, we fixed $n_{\text{local}}$ to a relatively small value ($512$) to compare the performance when selecting different numbers of tokens. First, we observe that even selecting a very small number of tokens (e.g., 128, 256), our TokenSelect still demonstrates very comparable performance. Then, as $k$ increases, the effectiveness of TokenSelect further improves, indicating that more moderately critical tokens also contribute to the retention of long-context information. Finally, we find that when $k$ is set to larger values (e.g., 16K), our TokenSelect shows significant improvements in most tasks, further advancing the performance landscape of long-context inference methods.

Similarity threshold of the Selection Cache θ𝜃\theta.

Fig. 6 shows that the Selection Cache hit rate increases significantly as the similarity threshold $\theta$ decreases, converging around $\theta=0.5$. This suggests potential for further acceleration of TokenSelect’s Decode Stage by reducing $\theta$. Performance sensitivity to $\theta$ varies across tasks. While most tasks exhibit slight performance degradation with decreasing $\theta$, and R.PK in InfiniteBench shows no degradation, more challenging retrieval tasks like R.KV demonstrate significant performance deterioration. This indicates higher dynamicity requirements for token selection in these tasks.

5. Efficiency Comparisons

Fig. 7 demonstrates the significant acceleration of attention computation achieved by TokenSelect during long-context inference. With a KV Cache length of 1M, TokenSelect can provide up to $23.84\times$ speedup compared to FlashInfer, which is the inference kernel library we based on. This substantial improvement is attributed to our efficient kernel design.

End-to-end efficiency.

Fig. 8 compares the end-to-end latency of TokenSelect, InfLLM, and standard attention across various tasks. TokenSelect significantly accelerates long-context inference in real-world scenarios, achieving a maximum speedup of $4.70\times$ over standard attention and $2.28\times$ over the SOTA long-context inference method. Moreover, TokenSelect demonstrates superior performance compared to both of them.

6. Scaling Beyond 1 Million Context Length

To further explore TokenSelect’s performance in extreme long-context scenarios, we design an extended benchmark with different text lengths following InfiniteBench. As illustrated in the Fig. 9, our TokenSelect demonstrates the ability to accurately capture critical information with a small token budget in contexts up to 2M tokens, underscoring its potential in more application scenarios.

Conclusion

In this paper, we introduces TokenSelect, a model-agnostic and training-free approach for efficient long-context inference and length extrapolation. TokenSelect addresses the two major challenges faced by LLMs in processing long texts: the context length limitation from pre-training and the computational complexity of attention. This is achieved through a novel token-level selective sparse attention mechanism. Experimental results demonstrate that TokenSelect can achieve up to $23.84\times$ speedup in attention computation and up to $2.28\times$ acceleration in end-to-end inference latency, while exhibiting superior performance across multiple long-context benchmarks. This approach significantly enhances LLMs’ capability to handle long contexts, paving the way for efficient long-text processing in advancing Web applications.

References

Appendix A Formal Statement and Proof of Lemma 1

Lemma 1 (Invariant Top-$k$ Key Selection under Cosine Similarity Threshold, Formal).

Let $k$ be a positive integer such that $1\leq k\leq N$.

where $\|\cdot\|_{2}$ denotes the Euclidean norm.

Define the top-$k$ selection function based on dot product similarity as: $\mathcal{I}(\mathbf{q})=\arg\max_{S\subseteq\{1,2,\dots,N\},|S|=k}\sum_{i\in S}\mathbf{q}\cdot\mathbf{k}_{i}$. Assume that for any query vectors $\mathbf{q}$, the top-$k$ set $\mathcal{I}(\mathbf{q})$ is uniquely determined.

Let $\epsilon\in(0,1]$ be a predefined threshold.

Lemma Statement: If the cosine similarity between the two query vectors $\mathbf{q}_{1}$ and $\mathbf{q}_{2}$ satisfies

then the indices of the top-$k$ keys selected by $\mathbf{q}_{1}$ and $\mathbf{q}_{2}$ are identical, i.e.,

Proof: We start with the given condition:

To facilitate our analysis, we introduce the following notations:

With these definitions, the original condition becomes:

Next, let $\theta$ denote the angle between $\mathbf{q}_{1}$ and $\mathbf{q}_{2}$, $\cos{\theta}=\hat{\mathbf{q}}_{1}\cdot\hat{\mathbf{q}}_{2}$. We can further define:

then $\sin{\theta}=\hat{\mathbf{p}}_{1}\cdot\hat{\mathbf{q}}_{2}$, and

In order to have Eqn. (11) $>0$, we require

This final inequality establishes a sufficient condition for the original statement to hold, thereby completing the proof.

Appendix B More Information on Dataset and Metrics

For InfiniteBench (Zhang et al., 2024), we use longbook_sum_eng (En.Sum), longbook_qa_eng (En.QA), longbook_choice_eng (En.MC), longdialogue_qa_eng (En.Dia), code_debug (Code.D), math_find (Math.F), passkey (R.PK), number_string (R.Num) and kv_retrieval (R.KV) as evaluation datasets. The corresponding evaluation metrics are shown in Table 6. RULER (Hsieh et al., 2024) consists of various evaluation tasks: Single NIAH (needle in a haystack), Multi-keys NIAH, Multi-values NIAH, Multi-values NIAH, Multi-queries NIAH, Variable Tracking, Common Words Extraction, Frequent Words Extraction and Question Answering. The evaluation metric is match rate. For LongBench, we use all English tasks with evaluation metrics in Table 7.

Appendix C Experimental Results on LongBench

Compared to InfiniteBench and RULER, LongBench has much shorter text lengths. The 95% percentile for its lengths is 31K tokens. Considering that recent LLMs after SFT generally have context lengths of up to 32K tokens (Yang et al., 2024), LongBench is less suitable for evaluating state-of-the-art long-context inference methods. Nevertheless, as shown in Table 8, our TokenSelect still demonstrates superior overall performance compared to most baseline methods. It’s worth noting that Yi-1.5-6B did not yield effective results on the SAMSum task because it failed to correctly follow instructions.