LazyLLM: Dynamisches Token-Pruning zur effizienten Inferenz von Long-Context-LLMs

LazyLLM: Dynamic Token Pruning for Efficient Long Context LLM Inference

July 19, 2024

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

cs.AI

Zusammenfassung

Die Inferenz von auf Transformer basierenden großen Sprachmodellen besteht aus zwei aufeinanderfolgenden Phasen: 1) einer Vorausfüllphase zur Berechnung des KV-Caches von Eingaben und zur Generierung des ersten Tokens und 2) einer Decodierungsphase zur Generierung nachfolgender Tokens. Bei langen Eingaben müssen während der Vorausfüllphase der KV-Cache für alle Tokens berechnet werden, was die Zeit zur Generierung des ersten Tokens signifikant erhöhen kann. Folglich kann die Vorausfüllphase zum Engpass im Generierungsprozess werden. Es bleibt eine offene Frage, ob alle Eingabetokens für die Generierung des ersten Tokens unerlässlich sind. Um dies zu beantworten, stellen wir eine neue Methode namens LazyLLM vor, die selektiv den KV für Tokens berechnet, die für die Vorhersage des nächsten Tokens in beiden Phasen, der Vorausfüllung und der Decodierung, wichtig sind. Im Gegensatz zu statischen Pruning-Ansätzen, die die Eingabe auf einmal kürzen, ermöglicht LazyLLM Sprachmodellen, in verschiedenen Generierungsschritten dynamisch unterschiedliche Teilmengen von Tokens aus dem Kontext auszuwählen, auch wenn sie in vorherigen Schritten gekürzt wurden. Umfangreiche Experimente mit Standarddatensätzen in verschiedenen Aufgaben zeigen, dass LazyLLM eine generische Methode ist, die nahtlos in bestehende Sprachmodelle integriert werden kann, um die Generierung signifikant zu beschleunigen, ohne Feinabstimmung. Beispielsweise beschleunigt LazyLLM in der Aufgabe des Multi-Dokumenten-Frage-Antwortens die Vorausfüllphase des LLama 2 7B-Modells um das 2,34-fache, während die Genauigkeit beibehalten wird.

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.