LazyLLM: Poda Dinâmica de Tokens para Inferência Eficiente de Contexto Longo em LLM

LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

July 19, 2024

Autores: Qichen Fu, Minsik Cho, Thomas Merth, Sachin Mehta, Mohammad Rastegari, Mahyar Najibi

cs.AI

Resumo

A inferência de modelos de linguagem grandes baseados em transformadores consiste em duas etapas sequenciais: 1) uma etapa de pré-preenchimento para calcular o cache KV de prompts e gerar o primeiro token, e 2) uma etapa de decodificação para gerar tokens subsequentes. Para prompts longos, o cache KV deve ser calculado para todos os tokens durante a etapa de pré-preenchimento, o que pode aumentar significativamente o tempo necessário para gerar o primeiro token. Consequentemente, a etapa de pré-preenchimento pode se tornar um gargalo no processo de geração. Uma questão em aberto permanece se todos os tokens do prompt são essenciais para gerar o primeiro token. Para responder a isso, introduzimos um método inovador, LazyLLM, que calcula seletivamente o KV para tokens importantes para a previsão do próximo token em ambas as etapas de pré-preenchimento e decodificação. Ao contrário das abordagens de poda estática que podam o prompt de uma vez, o LazyLLM permite que os modelos de linguagem selecionem dinamicamente diferentes subconjuntos de tokens do contexto em diferentes etapas de geração, mesmo que tenham sido podados em etapas anteriores. Experimentos extensivos em conjuntos de dados padrão em várias tarefas demonstram que o LazyLLM é um método genérico que pode ser integrado perfeitamente aos modelos de linguagem existentes para acelerar significativamente a geração sem ajuste fino. Por exemplo, na tarefa de questionamento e resposta multi-documentos, o LazyLLM acelera a etapa de pré-preenchimento do modelo LLama 2 7B em 2,34 vezes mantendo a precisão.

English

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.