LazyLLM : Élagage dynamique de tokens pour une inférence efficace des LLM à contexte long

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

Résumé

L'inférence des grands modèles de langage basés sur des transformateurs se compose de deux étapes séquentielles : 1) une étape de préremplissage pour calculer le cache KV des prompts et générer le premier token, et 2) une étape de décodage pour générer les tokens suivants. Pour les prompts longs, le cache KV doit être calculé pour tous les tokens lors de l'étape de préremplissage, ce qui peut considérablement augmenter le temps nécessaire pour générer le premier token. Par conséquent, l'étape de préremplissage peut devenir un goulot d'étranglement dans le processus de génération. Une question ouverte demeure : tous les tokens du prompt sont-ils essentiels pour générer le premier token ? Pour y répondre, nous introduisons une nouvelle méthode, LazyLLM, qui calcule sélectivement le KV pour les tokens importants pour la prédiction du token suivant, à la fois lors des étapes de préremplissage et de décodage. Contrairement aux approches d'élagage statique qui élaguent le prompt en une seule fois, LazyLLM permet aux modèles de langage de sélectionner dynamiquement différents sous-ensembles de tokens du contexte à différentes étapes de génération, même s'ils ont pu être élagués lors des étapes précédentes. Des expériences approfondies sur des ensembles de données standard pour diverses tâches démontrent que LazyLLM est une méthode générique qui peut être intégrée de manière transparente avec les modèles de langage existants pour accélérer significativement la génération sans nécessiter de réglage fin. Par exemple, dans la tâche de réponse à des questions sur plusieurs documents, LazyLLM accélère l'étape de préremplissage du modèle LLama 2 7B par un facteur de 2,34 tout en maintenant la précision.

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.