LazyLLM: DYNAMIC TOKEN PRUNING FOR EFFICIENT LONG CONTEXT LLM INFERENCE

## Abstract

Abstract The inference of transformer-based large language models consists of two sequential stages: 1) a prefilling stage to compute the KV cache of prompts and generate the first token, and 2) a decoding stage to generate subsequent tokens.
For long prompts, the KV cache must be computed for all tokens during the prefilling stage, which can significantly increase the time needed to generate the first token. Consequently, the prefilling stage may become a bottleneck in the generation process. An open question remains whether all prompt tokens are essential for generating the first token. To answer this, we introduce a novel method, LazyLLM , that selectively computes the KV for tokens important for the next token prediction in both the prefilling and decoding stages. Contrary to static pruning approaches that prune the prompt at once, LazyLLM allows language models to dynamically select different subsets of tokens from the context in different generation steps, even though they might be pruned in previous steps. Extensive experiments on standard datasets across various tasks demonstrate that LazyLLM is a generic method that can be seamlessly integrated with existing language models to significantly accelerate the generation without fine-tuning . For instance, in the multi-document question-answering task, LazyLLM accelerates the prefilling stage of the LLama 2 7B model by 2.34 × 2.34\times 2.34 × while maintaining accuracy.

## 1 Introduction

Standard prompt-based LLM inference has two sequential stages: *prefilling* and *decoding*, as shown in Figure 1. During the *prefilling* stage, the model computes and saves the KV cache of each token from the prompt, and predicts the first token. We refer to the time taken during *prefilling* stage as “time-to-first-token” (*TTFT*). Following the *prefilling* stage is the *decoding* stage, where the model reuses cached KVs to decode the next token iteratively until the stop criteria are met.

During the *prefilling* stage, all tokens from the prompt are used by all transformer layers. For long prompts, *TTFT* could be slow because state-of-the-art transformer-based LLMs are both deep and wide (Pope et al., ; Kim et al., ; Aminabadi et al., ), and the cost of computing attention increases quadratically with the number of tokens in the prompts. For instance, Llama 2 (Touvron et al., ), with 7 billion parameters, stacks 32 transformer layers with a model dimension of 4096. In this scenario, *TTFT* requires $21\times$ the walltime of each subsequent decoding step, and accounts for approximately 23% of the total generation time on the LongBench benchmark(^1^11The average LongBench prompt length is $3376$ tokens and the average generation length is $68$ tokens.) (Bai et al., ).
Therefore, optimizing *TTFT* is a critical path toward efficient LLM inference (NVIDIA, ).

While optimizing LLM inference is an active area of research, many methods (Leviathan et al., ; Cai et al., ; Zhang et al., ; Bhendawade et al., ; Li et al., ) have focused on improving inference speed during the *decoding* stage. Yet, there is little attention given to improving *TTFT*. We note that some compression-based works implicitly improve the *TTFT* by reducing the size of LLMs (Frantar et al., ; Sun et al., ; Ma et al., ). However, an orthogonal line of research(Li et al., ; Jiang et al., ; Dao et al., ) investigates how *TTFT* can be improved given a static transformer architecture. Within this line of research, a natural question arises: Are all prompt tokens essential for generating the first token?

Figure: Figure 1: Prompt-based LLM inference can be divided into two sequential stages: *prefilling* and *decoding*. For long prompts, the first token generation during *prefilling* stage could be slow. As an example, for Llama 2 7B model (Touvron et al., ), on average, the time to generate the first token requires $21\times$ the walltime of each subsequent decoding step and accounts for $23\%$ of the total generation time in the LongBench benchmark.
Refer to caption: x1.png

LLM profiling on the LongBench benchmark (Bai et al., ) in Figure 2 reveals that the attention scores of input tokens w.r.t. to the first generated token are very sparse, indicating that many tokens in the input prompt are redundant and can be removed without affecting the next token prediction.
To this end, we propose *LazyLLM*, a novel, simple, yet effective technique tailored for speeding up *prefilling*. As depicted in Figure 3, in each generation step, *LazyLLM* selectively computes the KV for tokens important for the next token prediction and “lazily” defers the computation of remaining tokens to later steps when they become relevant.
We propose using the attention score of the prior transformer layer to measure the importance of tokens and progressively prune tokens along the depth of the transformer.
In contrast to prompt compression works (Li et al., ; Jiang et al., ; Xu et al., ), which permanently reduce the prompt for all the following generation steps, our method allows the model to revive previously pruned tokens, which we found crucial to retain accuracy.
Extending progressive token pruning to all generation steps is non-trivial. Specifically, if a token is pruned at generation step $t$, and is revived at generation step $t^{\prime}>t$, some hidden states would need to be recomputed during step $t^{\prime}$.
To avoid such repetitive computation, we employ an additional caching mechanism, *Aux Cache*, to cache the hidden states of pruned tokens. This enables a computationally efficient pathway to revive pruned tokens, and ensures that the worst runtime of *LazyLLM* is never slower than the baseline.

Figure: Figure 2: We visualize the attention scores of input tokens in the prompt w.r.t. to the next token for each layer of Llama 2 7BTouvron et al. (). We also plot the distribution of the average attention score across all transformer layers. Result reveals that the attention scores of input tokens w.r.t. to the next token are very sparse, indicating that many tokens in the input prompt are redundant and can be safely removed without affecting the next token prediction.
Refer to caption: x2.png

In summary, the advantages of *LazyLLM* are: (1) Universal: *LazyLLM* can be seamlessly integrated with any existing transformer-based LLM to improve inference speed, (2) Training-free: *LazyLLM* doesn’t require any finetuning and can be directly integrated without any parameter modification, (3) Effective: Empirical results on 16 standard datasets across 6 different language tasks shows *LazyLLM* can improve the inference speed of the LLM during both *prefilling* and *decoding* stages.

## 2 Related Work

The increase in the scale of large language models (LLMs) has greatly enhanced their performance but also introduced challenges with respect to their inference efficiency. The inference of generative LLMs consists of two distinct stages as depicted in Figure 1.
In particular, extensive computation is needed under long context scenarios to calculate the full KV cache during the *prefilling* stage, resulting in a long time-to-first-token (*TTFT*). This delay causes users to wait several seconds after submitting a prompt before receiving any response from the agent, leading to a poor user experience.

Efficient Long Context Inference. Extensive work (Merth et al., ; Chen et al., ; Beltagy et al., ; Kitaev et al., ) has been proposed to improve inference efficiency for long context applications by reducing the memory footprint and total computations. Some works have focused on tailoring the architecture of the transformer for long context input. For instance, (Beltagy et al., ) introduces a drop-in replacement for standard self-attention and combines local windowed attention with task-motivated global attention. In parallel, Reformer (Kitaev et al., ) replaces dot-product attention by one that uses locality-sensitive hashing to reduce its computational complexity. Though the above methods can speed up long context inference, they require significant model architecture change and re-training. This drawback makes them impractical to be applied to existing pre-trained LLMs.
Closer to our work are efficient techniques that optimize the KV cache (Zhang et al., ; Li et al., ; Anagnostidis et al., ; Nawrot et al., ) by minimizing the KV cache size and data transfer. However, these works only focus on accelerating decoding steps, which are not applicable to reducing *TTFT*.

Token Pruning. Previous studies on the sentence classification task (Kim et al., ; Anagnostidis et al., ; He et al., ) has shown that not all tokens (*i.e*. words) in an input sequence are necessary to make a successful prediction. This provides several possibilities for token pruning, which minimizes computational demands by selectively removing less important tokens during inference.
For example, (Kim et al., ) presents Learned Token Pruning which adaptively removes unimportant tokens as an input sequence passes through transformer layers. In parallel, (He et al., ) proposes to reduce width-wise computation via token pruning for transformer-based models such as BERT (Devlin et al., ). These aforementioned approaches were designed for tasks requiring only a single iteration of processing, such as text classification.
In this work, we extend the idea of token pruning to generative LLMs. Specifically, our method allows the model to dynamically choose different sets of tokens at each generation step, which is crucial to retaining the performance.
Furthermore, we also introduce *Aux Cache* to ensure that each token is computed
at most once along the whole generation, and ensure the worst runtime of our method is not slower than the baseline.

## 3 LazyLLM

Figure: Figure 3: Comparison between standard LLM and *LazyLLM*. Instead of computing the KV cache of all input tokens at the *prefilling* stage, *LazyLLM* only selectively computes the tokens that are important to the next token prediction, deferring the computation of remaining tokens to later steps. *LazyLLM* significantly optimizes *TTFT* by reducing the amount of computation during *prefilling*. Moreover, as some tokens in the prompt are never selected by *LazyLLM* during the whole generation process (even though theoretically the model could use all tokens in the prompt), *LazyLLM* also reduces the total amount of computation and accelerates the overall generation.
Refer to caption: x3.png

### 3.1 Background on LLM Inference

Generative LLM inference consists of two stages: *prefilling* and *decoding* (see Figure 1).
In the *prefilling* stage, the model receives the prompt (a sequence of tokens) $\mathcal{T}=\{t_{i}\}_{i=1}^{N}$ of length N, where $t_{i}$ denotes a token and $N$ denotes the length of the prompt, then computes and saves the KV cache of each token, and produces the first token $t_{n+1}$. The transformer architecture commonly used in LLMs is a stack of layers where each layer shares the same architecture with a multiple-head self-attention mechanism followed by a multi-layer perception (MLP).
The time of *prefilling* is referred to as time-to-first-token (*a.k.a*. *TTFT*).
Following the *prefilling* is the *decoding* steps, where the model appends the generated token $t_{n+1}$ to the input, and subsequently decodes the following token. The *decoding* step is repeatedly performed until the stop criteria are met. While the formula of each decoding step is similar to *prefilling*, the amount of its computation is significantly lower thanks to the KV cache. Specifically,
with saved KV cache from *prefilling*, all the previous tokens do not need to pass any linear layers in the model.

### 3.2 Inference with LazyLLM

The overview of the proposed *LazyLLM* framework is illustrated in Figure 4. *LazyLLM* starts with the full context and progressively prunes tokens to gradually reduce the number of computations towards the end of the model. Note, *LazyLLM* allows the model to select different subsets of tokens from the context in different generation steps,
even though some of them may be pruned in previous steps. Compared to static pruning which prunes all the tokens at once, dynamic pruning optimizes the next token prediction in each generation step, which is crucial to retaining the performance.

Progressive Token Pruning.
Prior to this work, token pruning has been successfully applied to optimize LLM inference (Zhang et al., ; Li et al., ; Adnan et al., ; Nawrot et al., ). However, these approaches require accumulating the full attention maps of predicting the first few tokens to profile the importance of prompt tokens before starting pruning. Consequently, they are not applicable to reduce *TTFT* as they still require computing all the KV cache at the *prefilling* stage.

In contrast, *LazyLLM* only “lazily” computes the tokens that are important to predict the next token by starting from the first iteration of the inference (the *prefilling* step). A key challenge to pruning tokens in the first iteration is determining their importance.
Inspired by the early exiting work (Elhoushi et al., ) which shows the token hidden states gradually evolve through the transformer layers, we apply layer-wise token pruning in each generation step.
Specifically, we use the attention map of the layer $A^{l}\in\mathcal{R}^{H\times N\times N}$ to determine the importance of input token $t_{i}$ w.r.t. the next token to be predicted as

$$ $s_{i}^{l}=\frac{1}{H}\sum_{h=1}^{H}A^{l}_{h,i,N}$ (1) $$

where $H$ denotes number of attention heads, $N$ is the sequence length, and $A_{h,i,j}$ is the attention probability of the token $t_{j}$ attending to token $t_{i}$ at $h^{th}$ head.

After computing the confidence scores of tokens, it is challenging to determine the threshold value to prune the token. Concretely, the threshold can change as the distribution of the attention scores varies between different layers and different tasks. We address this challenge by using the top-$k$ percentile selection strategy to prune tokens. Specifically, token $t_{i}$ is pruned at layer $l+1$ if its confidence score $s^{l}_{i}$ is smaller than $k^{l}$th percentile among the input tokens. Once the token is pruned, it is excluded from the computation of all successive layers. In other words, the tokens used in the later layers will be a subset of previous layers.

Our study in Section 5.4 shows the performance changes with different locations of pruning layers and the number of tokens pruned. In particular, when pruning at the same transformer layer, the model’s performance gradually decreases as fewer
tokens are kept. We also found pruning at later transformer layers consistently has better performance than pruning at earlier layers, suggesting that later layers are less sensitive to token pruning. To achieve a better balance of speedup and accuracy, as shown in Figure 4, we apply progressive pruning that keeps more tokens at earlier transformer layers and gradually reduces the number of tokens towards the end of the transformer.

Figure: Figure 4: Overview of the *LazyLLM* framework. *LazyLLM* starts with the full context and progressively prunes tokens to gradually reduce the number of computations towards the end of the model. *LazyLLM* allows the model to select different subsets of tokens from the context in different generation steps, which is crucial to retaining the performance.
Refer to caption: x4.png

Aux Cache. In the prefilling stage, there is no KV cache and every token is represented by hidden states. Thus, progressive token pruning can be implemented by removing pruned tokens’ hidden states. However, extending the progressive token pruning to the following *decoding* steps is non-trivial. This is because each *decoding* step leverages the KV cache computed in the *prefilling* to compute attention. As the *LazyLLM* performs progressive token pruning at the *prefilling* stage, the KV of tokens pruned at layer $l$ (*e.g*. $T4$ in Figure 4) will not exist in the KV cache of layer $l+1$. As a reminder, the *LazyLLM* framework allows each generation step to pick a different subset set of tokens from the full input token sequences in every step, regardless of whether they are pruned in previous generation steps or not. For example, during the following *decoding* steps, those pruned tokens (*e.g*. $T4$) that do not exist in the KV cache of layer $l+1$ may be re-selected to compute attention. In such cases, the model can not retrieve the KV cache of these tokens. An intuitive solution is to pass those tokens again from the beginning of the transformer. However,
that would cause repetitive computation for the same token, and eventually slow down the whole
generation.

To tackle this challenge, we introduce *Aux Cache* in addition to the original KV cache, which stores the hidden states of those pruned tokens (*e.g*. $T4$ and $T7$ in Figure 4) if their KV is not present in the following layer’s KV cache, which could be potentially retrieved for the following iterations. As shown in Figure 4, in each *decoding* step, each transformer layer (*e.g*. layer $l+1$) first retrieves the KV cache of past tokens if they exist (*e.g*. $T1$ and $T8$). For those tokens that do not exist in the KV cache (*e.g*. $T3$), we could retrieve their hidden states from the *Aux Cache* of its previous layer directly instead of passing through previous layers again. The introduction of *Aux Cache* ensures that each token is computed at most once in every transformer layer, and ensures the worst runtime of *LazyLLM* is not slower than the baseline.

## 4 Implementations Details

We implement *LazyLLM* on Llama 2 (Touvron et al., ) and XGen (Nijkamp et al., ) and evaluate it on the LongBench (Bai et al., ) using HuggingFace(^2^22[https://github.com/huggingface/transformers/](https://github.com/huggingface/transformers/)). We follow the official GitHub repository(^3^33[https://github.com/THUDM/LongBench](https://github.com/THUDM/LongBench)) of LongBench for data preprocessing and prompting in all experiments. The LongBench benchmark consists of multiple datasets in different tasks, where each task may have different metrics, including ROUGE-L, F1, Accuracy, and Edit Sim. Following the official evaluation pipeline, we categorize all results over major task categories by computing the macro-average score.

As previously noted, the proposed *LazyLLM* doesn’t require any training. Thus, *LazyLLM* uses the exact same existing checkpoints as the baseline, for all models. For inference, we conduct all experiments on NVIDIA A100 GPUs. We measure and report the speedup based on the empirical walltime improvement. Specifically, for *TTFT Speedup*, we measure the empirical walltime between when the prompt is fed to the model, and when the model generates the first token. For *Generation Speedup*, we measure the empirical walltime between when the prompt is fed to the model, and when the model finished generating all output tokens. We add 5 warmup runs for each experiment before starting the time measurement to remove the noise such as loading model parameters.

## 5 Experiments

We examine our method using two large language models: Llama 2 7B and XGen 7B. We compare our method with baselines using the same publicly released pretrained checkpoints, without employing any additional training. We perform experiments using LongBench, a multi-task benchmark for long content understanding. The LongBench comprises 16 datasets and covers 6 tasks including single-doc QA, multi-doc QA, summarization, few-shot learning, synthetic tasks, and code completion.

For the metrics, we primarily evaluate the effectiveness and efficiency of each method in the *TTFT* speedup *vs*. accuracy trade-off. Following LongBench, the accuracy (*score*) denotes the macro-averaged scores across datasets in each task. The *TTFT* speedup measures the wall time improvement w.r.t. to the baseline for generating the first token. In analysis, we also assess the impact of our method on *$\%$ of Prompt Token Computed* and *Generation* speedup. The *$\%$ of Prompt Token Computed* measures the accumulated percent of prompt tokens computed at the end of the generation, which indicates the save of total computation. The *Generation* speedup measures the walltime change w.r.t. to the baseline for completing the entire generation process.

**Table 1: Comparisons of *TTFT* speedup *vs*. accuracy on various tasks. Without requiring any training/finetuning, *LazyLLM* consistently achieves better *TTFT* speedup with negligible accuracy drop. Note that the prompt compression approach fails at improving *TTFT* because the overhead of running LLMs to compress the prompt is very computationally expensive.**
| Tasks | Method | Llama 2 |  | XGen |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Score | TTFT Speedup ($\times$) |  | Score | TTFT Speedup ($\times$) |  |  |
| Single-Document QA | Baseline | $\mathbf{25.79}$ | $1.00$ |  | $\mathbf{25.19}$ | $1.00$ |
| Random Token Drop | $20.05$ | $1.20$ |  | $18.32$ | $1.58$ |  |
| Static Token Pruning | $21.89$ | $1.18$ |  | $19.30$ | $1.61$ |  |
| Prompt Compression | $22.88$ | $0.12$ |  | $15.31$ | $0.20$ |  |
| *LazyLLM (Ours)* | $25.59$ | $\mathbf{1.36}$ |  | $25.00$ | $\mathbf{1.96}$ |  |
| Multi-Document QA | Baseline | $\mathbf{22.43}$ | $1.00$ |  | $\mathbf{20.71}$ | $1.00$ |
| Random Token Drop | $16.77$ | $1.19$ |  | $14.86$ | $1.37$ |  |
| Static Token Pruning | $19.93$ | $2.16$ |  | $17.23$ | $2.11$ |  |
| Prompt Compression | $8.42$ | $0.13$ |  | $11.56$ | $0.19$ |  |
| *LazyLLM (Ours)* | $22.31$ | $\mathbf{2.34}$ |  | $20.68$ | $\mathbf{2.65}$ |  |
| Summarization | Baseline | $24.65$ | $1.00$ |  | $\mathbf{24.85}$ | $1.00$ |
| Random Token Drop | $24.39$ | $1.39$ |  | $24.47$ | $1.70$ |  |
| Static Token Pruning | $24.59$ | $1.33$ |  | $24.46$ | $1.65$ |  |
| Prompt Compression | $25.16$ | $0.12$ |  | $24.57$ | $0.17$ |  |
| *LazyLLM (Ours)* | $\mathbf{24.75}$ | $\mathbf{1.46}$ |  | $24.74$ | $\mathbf{1.91}$ |  |
| Few-shot Learning | Baseline | $\mathbf{62.90}$ | $1.00$ |  | $\mathbf{56.40}$ | $1.00$ |
| Random Token Drop | $53.93$ | $1.19$ |  | $46.35$ | $1.62$ |  |
| Static Token Pruning | $56.54$ | $2.16$ |  | $51.93$ | $3.17$ |  |
| Prompt Compression | $24.18$ | $0.10$ |  | $23.72$ | $0.15$ |  |
| *LazyLLM (Ours)* | $62.81$ | $\mathbf{2.19}$ |  | $56.12$ | $\mathbf{3.42}$ |  |
| Synthetic | Baseline | $4.97$ | $1.00$ |  | $5.40$ | $1.00$ |
| Random Token Drop | $3.57$ | $1.18$ |  | $2.53$ | $1.13$ |  |
| Static Token Pruning | $2.81$ | $2.15$ |  | $3.00$ | $4.14$ |  |
| Prompt Compression | $3.20$ | $0.12$ |  | $1.42$ | $0.17$ |  |
| *LazyLLM (Ours)* | $\mathbf{4.98}$ | $\mathbf{2.89}$ |  | $\mathbf{5.66}$ | $\mathbf{4.77}$ |  |
| Code Completion | Baseline | $\mathbf{55.18}$ | $1.00$ |  | $\mathbf{36.49}$ | $1.00$ |
| Random Token Drop | $44.92$ | $1.23$ |  | $32.34$ | $1.57$ |  |
| Static Token Pruning | $37.51$ | $1.84$ |  | $32.27$ | $2.97$ |  |
| Prompt Compression | $17.45$ | $0.49$ |  | $11.38$ | $0.69$ |  |
| *LazyLLM (Ours)* | $53.30$ | $\mathbf{1.94}$ |  | $36.47$ | $\mathbf{3.47}$ |  |

### 5.1 Results

Table 1 presents the *TTFT* speedup *vs*. accuracy comparisons between *LazyLLM*, standard LLM, and other baselines. In the table, the “baseline” refers to the standard LLM inference. The “random token drop” baseline is based on (Yao et al., ) that randomly prunes the prompt tokens before feeding them to the LLMs. We report the average metrics across 5 runs for the “random token drop” baseline. Our “static token pruning” baseline prunes input tokens at once based on their attention score of the first few transformer layers during the *prefilling* stage. We also compare with the prompt compression method (Li et al., ) which pruning redundancy in the input context using LLMs.
Table 1 shows *LazyLLM* consistently achieves better *TTFT* speedup with negligible accuracy drop across multiple tasks. It is worth noting that the overhead of running LLMs to compress the prompt is very computationally expensive. Even though the inference on the reduced prompt is faster, the actual *TTFT* of the “prompt compression” baseline is longer than the baseline.

### 5.2 TTFT Speedup vs . Accuracy

The inference efficiency of *LazyLLM* is controlled using three parameters: 1) the number of pruning layers, 2) the locations of these pruning layers, and 3) the number of tokens pruned within these layers. Increasing the number of pruning layers and pruning more tokens optimize computation by processing fewer tokens, and pruning tokens at earlier layers can save the computations for the successive layers. Prompting these factors will give more overall computation reduction, and offer better *TTFT* speedup. As a side effect, excessively pruning tokens may cause information loss and eventually lead to performance degradation. Similarly, the *TTFT* speedup and accuracy of baselines can vary with different hyperparameters.

We compare *TTFT* speedup *vs*. accuracy in Figure 5 with different hyperparameters. The visualization shows that, without any training, the proposed *LazyLLM* retains the accuracy better than baselines under the same *TTFT* speedup. For example, our method can offer $2.34\times$ *TTFT* speedup in the multi-document question-answering task with negligible ($\leq 1\%$) performance loss. By controlling the pruning parameters, *LazyLLM* provides a good trade-off between accuracy and inference speed as compared to baseline methods. For instance, *LazyLLM* can achieve $3.0\times$ *TTFT* speedup in the multi-document question-answering task with $\leq 10\%$ degradation in accuracy. On the other hand, baseline methods accuracy degrades significantly for similar *TTFT* speed-up. Note that the prompt compression approaches fail at improving *TTFT* because of the compression overhead.

### 5.3 Impact on Overall Generation Speed

To evaluate the impact of the proposed method on the overall generation process, we also profile the *$\%$ of Prompt Token Computed* and *Generation* speedup in Table 2. We can find the *$\%$ of Token Computed* of *LazyLLM* is less than 100$\%$, indicating that not all tokens in the prompt are selected by *LazyLLM* at the end of the generation, even though theoretically the model could use all tokens. Computations in the FFN layers increase linearly, while those in the attention layers grow quadratically with the *$\%$ of Token Computed*. A lower *$\%$ of Token Computed* indicates *LazyLLM* reduces the total computation, consequently offering additional speedup to the overall generation process across diverse tasks.

Figure: Figure 5: *TTFT* speedup *vs*. accuracy comparison for Llama 2 7B across different tasks.
Refer to caption: x5.png

**Table 2: The *$\%$ of Prompt Token Computed* and *Generation* speedup of *LazyLLM* on various tasks. Reported values are based on the same setting as Table 1. A lower *$\%$ of Token Computed* indicates *LazyLLM* reduces the total computation, consequently offering additional speedup to the overall generation process across diverse tasks.**
| Tasks | $\%$ of Prompt Token Computed |  | Overall Generation Speedup |  |  |
| --- | --- | --- | --- | --- | --- |
| Llama 2 | XGen |  | Llama 2 | XGen |  |
| Single-Document QA | $87.31$ | $89.16$ |  | $1.34$ | $1.33$ |
| Multi-Document QA | $63.94$ | $69.60$ |  | $1.56$ | $1.70$ |
| Summarization | $99.59$ | $96.11$ |  | $1.02$ | $1.09$ |
| Few-shot Learning | $69.98$ | $65.30$ |  | $1.28$ | $1.59$ |
| Synthetic | $63.73$ | $40.54$ |  | $1.79$ | $3.16$ |
| Code Completion | $68.57$ | $72.61$ |  | $1.01$ | $1.16$ |

### 5.4 Drop Rate in Different Layers

In this section, we analyze the effect of the locations of pruning layers, and the number of tokens pruned. In particular, we report a series of experiments using a simplified version of *LazyLLM* that prunes tokens just once within the transformer. For each trial, we position the pruning layer at various levels of the transformer stack and apply different pruning ratios. We perform the experiments for both Llama 2 and XGen, and visualize the results in Figure 6.

The results show both models share a similar trend. As expected, when pruning at the same transformer layer, the model’s performance gradually decreases as fewer tokens are kept. Furthermore, pruning at later transformer layers consistently yields better performance compared to pruning at earlier layers, suggesting that later layers are less sensitive to token pruning.
Based on these observations, we propose progressive token pruning in Section 3.2, which strategically prunes more tokens in later layers while preserving more in the earlier layers, optimizing the balance between efficiency and performance retention.

Figure: Figure 6: Effect of the locations of pruning layers, and the number of tokens pruned. The results of both Llama 2 7B Touvron et al. () and XGen 7B Nijkamp et al. () share a similar trend: 1) when pruning at the same transformer layer, the model’s performance gradually decreases as fewer tokens are kept, and 2) Pruning at later transformer layers consistently has better performance than pruning at earlier layers, suggesting that later layers are less sensitive to token pruning.
Refer to caption: x6.png

Figure: Figure 7: Statistics on number of tokens processed during generation using our LazyLLM technique with Llama 2 7B (Touvron et al., ). We visualize the statistics of 1000 samples randomly sampled from LongBench. The $x$-axis represents the (absolute) generation time step, and the $y$-axis represents the number of prompt tokens processed at that time step (normalized by the prompt size). We visualize these statistics for various stages within the network. Note that cumulative token usage is upper-bounded by the baseline (evident with early layers).
Refer to caption: x7.png

### 5.5 Progressive KV Growth

In this section, we characterize the internals of the model with the token pruning logic.
Specifically, we seek to understand what fractions of prompt tokens are cumulatively used and, inversely, not used.
This “cumulative token usage” can be equivalently defined as the KV cache size at each given step.
Figure 7 presents these cumulative prompt token usage numbers for each of the stages of the LazyLLM.

Our analysis supports the hypothesis that many tokens are never selected by the model (even though theoretically the model could use all tokens in the prompt).
Since this model retains accuracy on the task(s), we can conclude that the model effectively drops the tokens which do not affect the output quality.

## 6 Conclusion

In this work, we proposed a novel *LazyLLM* technique for efficient LLM inference, in particular under long context scenarios. *LazyLLM* selectively computes the KV for tokens important for the next token prediction and “lazily” defers the computation of remaining tokens to later steps, when they become relevant. We carefully examine *LazyLLM* on various tasks, where we observed the proposed method effectively reduces *TTFT* with negligible performance loss.
It is worth noting that our method can be seamlessly integrated with existing transformer-based LLMs to improve their inference speed without requiring any fine-tuning.

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: * I feel like caching hidden states of pruned tokens in AuxCache is similar to prior work LESS [1]. how is AuxCache different from the prior work?
* In tab 2, code completion computes more tokens in llama2 (68.57%) compared to single-document QA (87.31%) but speedup is marginal (x1.01) compared to single-doc QA(x1.34). Why is it so? is the saving not linear?

[1] https://arxiv.org/pdf/2402.09398v2
soundness: 2
strengths: * Evaluated the proposed method in diverse task.
summary: In this work, authors propose to selectively calculate KV cache instead of computing KV cache of all tokens. Unlike static pruning one in prior work, this work provide dynamic pruning of tokens in different generation steps..
weaknesses: * Problem seems not very general and universal to all context. Authors should be clear about when TTFT becomes x21 compared to decoding. In a large scale system, decoding and prefilling is happening in a different server so it is not a big problem. Also prefilling usually computes more token than decoding so if we normalize the latency by number of tokens, we can’t say it is completely doing wrong although optimizing it helps anyway. 
* Figures are confusing especially fig 4. 
* Methods are compared to Token drop, static token prune, prompt compression but standard technique to reduce TTFT is using parallel computation. Paper lacks comparing the method with those method. [2] These method does not lose any accuracy and effectively accelerate TTFT.

[2] https://arxiv.org/abs/2409.17264

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: * Regarding the progressive token pruning by attention-based top-k percentile selection, could the authors consider measures to prevent the unintentional removal of tokens that might prove essential in later layers? Accumulating attention scores across multiple layers could potentially address inconsistencies in token selection. Would such an approach be integrated in LazyLLM and mitigate issues associated with layer-wise pruning?

* How many previously evicted tokens can be revived in later decoding steps, and does this impact the semantic structure of the input sequence? Understanding the limitations on reviving tokens would help clarify whether token revival affects the semantic coherence of the sequence.

* Does the Aux Cache and subsequent KV computation for revived tokens introduce memory or latency overhead during decoding?
Would it be possible to apply additional pruning or compression techniques at decoding time to further optimize inference speed with LazyLLM?
soundness: 3
strengths: * **Sharp focus on dynamic token pruning for the prefilling stage**.
This paper proposes an innovative approach to tackle the TTFT problem by shifting part of the prompt token computation to the decoding stage. The dynamic token pruning at different decoding steps allows for the selective retention of previously pruned but relevant tokens.

* **Effective and flexible layer-wise pruning strategy**.
The progressive token pruning from earlier to later layers is well-justified, offering a flexible approach to managing the trade-offs between efficiency and performance.

* **Comprehensive analysis and convincing results**.
The experiments span multiple datasets and models, providing a convincing case for LazyLLM’s efficiency. The ablation study on token drop rates and locations offers insights into the performance-speedup trade-offs and memory and computation needs.
summary: This paper addresses the inference latency at the prefilling stage of long-context LLM inference. 

The authors propose LazyLLM, a training-free method that enables dynamic token pruning across transformer layers and at different decoding steps. 

Through experiments on various long-context inference datasets, they demonstrate that LazyLLM significantly reduces inference time while maintaining performance levels comparable to baseline models.
weaknesses: * **Limited detail on hyperparameter settings and implementation strategies**.
The approach introduces numerous hyperparameters, particularly with progressive token pruning and token revival, which could impact implementation. Providing additional details on the decision-making process for these hyperparameters would enhance transparency and offer insights into LazyLLM’s effectiveness and generalizability.
    - Top-$k$ percentile selection strategy: Unless I missed something, it appears that different values of $k^l$ are set at the corresponding layer $l$. Clarifying how these values were determined and their generalizability across different tasks and decoding steps would be beneficial.
    - For reviving tokens, the authors skip the KV updating of tokens before and after the revived tokens, instead appending the revived tokens to simplify KV computation. This strategy, however, breaks the sequential dependency of tokens, potentially affecting performance due to misalignment with training data. An extended discussion on the effect and trade-offs behind this implementation could be beneficial and provide insights into the role of token orders in inference speed and performance.

* **Lack of discussion of the potential bias enhanced by attention-based pruning**.
Selecting tokens to prune based on attention scores could inadvertently amplify inherent biases in LLMs, particularly when uncertainty is high.
    - Pruning a fixed number of tokens according to attention scores can lead to unintentional bias, as high-entropy distributions may cause equally significant tokens to be pruned due to marginal differences in scores. This issue could be exacerbated in challenging tasks, where hyperparameter sensitivity may impact LazyLLM’s reliability.
    - Previous works (Li et al. 2023) suggest that the attention mechanism in transformers may focus on different tokens across layers. LazyLLM’s early pruning might inadvertently exclude tokens relevant at later stages. Exploring this potential limitation would strengthen the analysis and illuminate LazyLLM’s effectiveness across diverse text generation tasks.

* **Possible memory overhead from the Aux Cache during decoding**.
The authors propose Aux Cache to store the hidden states of pruned tokens for efficient future KV computation. However, this can pose challenges to further reducing the memory footprint in the subsequent decoding stage. Given the utilization of the attention mechanism of LazyLLM in common with existing works on optimizing KV cache (Xiao et al. 2024; Liu et al. 2023, Zhang et al. 2023), perhaps the authors could further discuss how we can optimize the Aux cache for long context inference in general. 

---
Kenneth Li, Oam Patel, Fernanda B. Viégas, Hanspeter Pfister, Martin Wattenberg: Inference-Time Intervention: Eliciting Truthful Answers from a Language Model. NeurIPS 2023

Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, Mike Lewis: Efficient Streaming Language Models with Attention Sinks. ICLR 2024

Zichang Liu, Aditya Desai, Fangshuo Liao, Weitao Wang, Victor Xie, Zhaozhuo Xu, Anastasios Kyrillidis, Anshumali Shrivastava: Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time. NeurIPS 2023

Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark W. Barrett, Zhangyang Wang, Beidi Chen: H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models. NeurIPS 2023

Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, Jianfeng Gao: Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs. ICLR 2024

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 3
questions: Please check the weaknesses section.
soundness: 3
strengths: A key strength of LazyLLM lies in its dynamic, training-free approach to token pruning, which allows it to be easily integrated into existing transformer-based LLMs without requiring model fine-tuning or architectural changes. By selectively computing only the most important tokens for each generation step, LazyLLM not only optimizes the time-to-first-token (TTFT) but also reduces the overall computation during inference. This results in significant speedups across various tasks and model configurations while maintaining accuracy. Furthermore, the method’s adaptive token selection and use of an auxiliary cache (Aux Cache) enable the model to “revive” previously pruned tokens when needed, ensuring that efficiency gains do not come at the cost of degraded model performance.
summary: The paper introduces LazyLLM, a dynamic token pruning technique designed to improve the efficiency of large language model (LLM) inference, particularly in long-context scenarios. LazyLLM selectively computes key-value (KV) pairs for tokens that are crucial for the prediction of the next token, thereby deferring the computation of less important tokens to later stages. Unlike static pruning, LazyLLM dynamically selects subsets of tokens at each generation step, allowing it to maintain accuracy while achieving speed improvements.

This approach requires no additional training and can be integrated with existing transformer-based models seamlessly. Experimental results demonstrate that LazyLLM significantly accelerates the inference process, especially in the initial token generation phase (prefilling). For instance, in multi-document question-answering tasks on the LLaMA 2 7B model, LazyLLM achieves a 2.34x speedup during prefilling with minimal accuracy loss
weaknesses: 1.	**Additional GPU memory usage for Aux Cache**: If the Aux Cache is retained on the GPU, it will increase GPU memory consumption. As a result, the actual GPU memory footprint of LazyLLM should account for both the retained KV cache and the Aux Cache. This design might limit LazyLLM’s applicability in scenarios with high memory demands.

2.	**Alignment of GPU memory costs in experiments**: It is crucial to clarify whether the GPU memory usage for each method was fairly aligned in the experiments, especially if LazyLLM’s actual GPU memory footprint surpasses that of the baseline methods. Any advantages in speed and efficiency for LazyLLM need to be re-evaluated under fair memory cost comparisons to ensure experimental fairness.

3.	**Lack of comparisons with more advanced baselines**: The paper does not compare LazyLLM with recent efficient cache management methods, such as H2O[1] and SnapKV[2]. These methods offer innovations in areas like KV cache reduction strategies and generation speed optimization. Comparing LazyLLM with them would more comprehensively illustrate its performance strengths and limitations.

> Reference:  
> [1] Zhang, Zhenyu, et al. "H2o: Heavy-hitter oracle for efficient generative inference of large language models." Advances in Neural Information Processing Systems 36 (2024)  
> [2] Li, Yuhong, et al. "Snapkv: Llm knows what you are looking for before generation." arXiv preprint arXiv:2404.14469 (2024).

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: Please check the weaknesses.
soundness: 2
strengths: The following are the strengths.
* The paper introduces an efficient way to speed up LLM inference.
* The paper involves an aux cache to be sure that each token is calculated at once at most.
* The author conducts experiments to support their claims.
summary: The paper proposes a new method to speed up the LLM inference via selectively computing the KV for tokens important for the next token prediction in both the prefilling and decoding stages. The method is training-free. The author also conducts experiments on Llama 2 and XGen.
weaknesses: My concerns mainly come from two parts

First, it seems that the paper template has some problems. To be more specific, the "Anonymous authors Paper under double-blind review" at the Top of Page 1 is in the middle, where it actually should be on the left side. I wonder whether such template change is allowed or not.

Second, there are several parts that the author needs to explain further.
* If possible, could the author provide Python code or pseudo code to explain the idea? This will make it easier for the reader to understand the core idea.
* Question for the KV cache part. 
   * If for the previous generation step, Layer i has tokens [T1, T4, T5, T7] and Layer i+1 has tokens [ T4, T5]. 
   * And for the current generation step,  Layer i has tokens [T1, T2, T4, T5], and Layer i+1 has tokens [ T4, T5]. 
   * **Question**: for the current generation step, do we directly use the [T4, T5] from the KV cache?
   * The paper claims that "each token is computed at most once along
the whole generation" so it seems that the answer is Yes. Then it seems that the previous generation step  Layer i+1 tokens representation [ T4, T5] is different from the current generation step  Layer i+1 tokens representation [ T4, T5]. Why the work could directly use it?
* For Table 1, LazyLLM could even achieve better performance than the baseline which is standard LLM inference. Could you explain such an observation because the baseline should be the performance upper bound of LazyLLM?

Meta-review: This paper proposes LazyLLM, a method to accelerate LLM inference by selectively computing KV cache for important tokens during both prefilling and decoding stages. The key strength is its training-free, dynamic approach that allows different tokens to be selected at different generation steps, achieving significant speedup (e.g., 2.34x for prefilling in multi-document QA) while maintaining accuracy. The paper's main weakness is the lack of comparison with parallel computation techniques, which are standard approaches for reducing TTFT. Additionally, the method introduces memory overhead from the Aux Cache during decoding, and the paper lacks detailed discussion on hyperparameter settings and implementation strategies. While the proposed method shows promise in accelerating inference, there are concerns about potential bias from attention-based pruning and the need for more comprehensive experimental validation.

Decision: Reject