LazyLLM: Динамическое обрезание токенов для эффективного вывода LLM с длинным контекстом

LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

July 19, 2024

Авторы: Qichen Fu, Minsik Cho, Thomas Merth, Sachin Mehta, Mohammad Rastegari, Mahyar Najibi

cs.AI

Аннотация

Выводы моделей больших языков на основе трансформеров состоят из двух последовательных этапов: 1) этап предварительного заполнения для вычисления кэша KV подсказок и генерации первого токена, и 2) этап декодирования для генерации последующих токенов. Для длинных подсказок кэш KV должен быть вычислен для всех токенов во время этапа предварительного заполнения, что может значительно увеличить время, необходимое для генерации первого токена. Следовательно, этап предварительного заполнения может стать узким местом в процессе генерации. Остается открытым вопрос, являются ли все токены подсказки необходимы для генерации первого токена. Для ответа на этот вопрос мы представляем новый метод, LazyLLM, который выборочно вычисляет KV для токенов, важных для предсказания следующего токена как на этапе предварительного заполнения, так и на этапе декодирования. В отличие от статических методов обрезки, которые сразу обрезают подсказку, LazyLLM позволяет моделям языка динамически выбирать различные подмножества токенов из контекста на разных этапах генерации, даже если они были обрезаны на предыдущих этапах. Обширные эксперименты на стандартных наборах данных по различным задачам показывают, что LazyLLM является универсальным методом, который может быть легко интегрирован с существующими языковыми моделями для значительного ускорения генерации без дополнительной настройки. Например, в задаче мультидокументного вопросно-ответного тестирования LazyLLM ускоряет этап предварительного заполнения модели LLama 2 7B в 2,34 раза, сохраняя точность.

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.