LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

This paper was accepted at the Efficient Systems for Foundation Models Workshop at ICML 2024
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…Apple Machine Learning Research