LazyLLM: Dynamische Tokenpruning voor Efficiënte Inferentie van LLM's met Lange Contexten

LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

July 19, 2024

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

cs.AI

Samenvatting

De inferentie van transformer-gebaseerde grote taalmodellen bestaat uit twee opeenvolgende fasen: 1) een prefilling-fase om de KV-cache van prompts te berekenen en het eerste token te genereren, en 2) een decodeerfase om de daaropvolgende tokens te genereren. Voor lange prompts moet de KV-cache voor alle tokens tijdens de prefilling-fase worden berekend, wat de tijd die nodig is om het eerste token te genereren aanzienlijk kan verlengen. Hierdoor kan de prefilling-fase een knelpunt worden in het generatieproces. Een open vraag blijft of alle prompt-tokens essentieel zijn voor het genereren van het eerste token. Om dit te beantwoorden introduceren we een nieuwe methode, LazyLLM, die selectief de KV berekent voor tokens die belangrijk zijn voor de voorspelling van het volgende token, zowel in de prefilling- als in de decodeerfase. In tegenstelling tot statische pruning-benaderingen die de prompt in één keer snoeien, stelt LazyLLM taalmodellen in staat om dynamisch verschillende subsets van tokens uit de context te selecteren in verschillende generatiestappen, zelfs als ze in vorige stappen zijn gesnoeid. Uitgebreide experimenten op standaarddatasets voor diverse taken tonen aan dat LazyLLM een generieke methode is die naadloos kan worden geïntegreerd met bestaande taalmodellen om de generatie aanzienlijk te versnellen zonder fine-tuning. Zo versnelt LazyLLM in de taak van multi-document vraagbeantwoording de prefilling-fase van het LLama 2 7B-model met 2,34x terwijl de nauwkeurigheid behouden blijft.

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.