LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

TL;DR

The paper tackles the TTFT bottleneck in long-context LLM inference by introducing LazyLLM, a dynamic token-pruning method that selectively computes KV caches for tokens important to predicting the next token and lazily defers others. It uses layer-wise attention-based importance scoring and a top-k threshold, combined with an Aux Cache to allow revival of pruned tokens without recomputation, ensuring worst-case runtime does not exceed the baseline. LazyLLM is training-free and demonstrated across Llama 2 and XGen on LongBench, delivering significant TTFT and overall generation speedups with negligible accuracy loss. This approach provides a practical, model-agnostic enhancement for efficient long-context generation without architectural changes or finetuning.

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.34x while maintaining accuracy.

Figures (7)

• 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 touvron2023llama, 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.
• 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 7Btouvron2023llama. 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.
• 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.
• 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.
• Figure 5: TTFT speedup vs. accuracy comparison for Llama 2 7B across different tasks.