MInference 1.0 : Accélération du pré-remplissage pour les LLM à contexte long via une attention dynamique parcimonieuse

MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention

July 2, 2024

Auteurs: Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, Lili Qiu

cs.AI

Résumé

Les défis computationnels liés à l'inférence des modèles de langage à grande échelle (LLM) restent un obstacle majeur à leur déploiement généralisé, en particulier à mesure que la longueur des prompts continue d'augmenter. En raison de la complexité quadratique du calcul d'attention, il faut 30 minutes à un LLM de 8 milliards de paramètres pour traiter un prompt de 1 million de tokens (c'est-à-dire l'étape de pré-remplissage) sur un seul GPU A100. Les méthodes existantes pour accélérer le pré-remplissage échouent souvent à maintenir une précision ou une efficacité acceptable lorsqu'elles sont appliquées à des LLM à contexte long. Pour combler cette lacune, nous introduisons MInference (Inférence pour Millions de Tokens), une méthode de calcul parcimonieux conçue pour accélérer le pré-remplissage du traitement de séquences longues. Plus précisément, nous identifions trois motifs uniques dans les matrices d'attention à contexte long—la forme en A, la barre verticale et la parcimonie par blocs—qui peuvent être exploités pour un calcul parcimonieux efficace sur les GPU. Nous déterminons le motif optimal pour chaque tête d'attention hors ligne et construisons dynamiquement des indices parcimonieux basés sur le motif attribué lors de l'inférence. Avec le motif et les indices parcimonieux, nous effectuons des calculs d'attention parcimonieux efficaces via nos noyaux GPU optimisés pour réduire significativement la latence lors de l'étape de pré-remplissage des LLM à contexte long. Notre technique proposée peut être directement appliquée aux LLM existants sans aucune modification de la configuration de pré-entraînement ni de fine-tuning supplémentaire. En évaluant sur un large éventail de tâches en aval, notamment InfiniteBench, RULER, PG-19 et Needle In A Haystack, et sur des modèles incluant LLaMA-3-1M, GLM4-1M, Yi-200K, Phi-3-128K et Qwen2-128K, nous démontrons que MInference réduit efficacement la latence d'inférence jusqu'à 10 fois pour le pré-remplissage sur un A100, tout en maintenant la précision. Notre code est disponible à l'adresse https://aka.ms/MInference.

English

The computational challenges of Large Language Model (LLM) inference remain a significant barrier to their widespread deployment, especially as prompt lengths continue to increase. Due to the quadratic complexity of the attention computation, it takes 30 minutes for an 8B LLM to process a prompt of 1M tokens (i.e., the pre-filling stage) on a single A100 GPU. Existing methods for speeding up prefilling often fail to maintain acceptable accuracy or efficiency when applied to long-context LLMs. To address this gap, we introduce MInference (Milliontokens Inference), a sparse calculation method designed to accelerate pre-filling of long-sequence processing. Specifically, we identify three unique patterns in long-context attention matrices-the A-shape, Vertical-Slash, and Block-Sparsethat can be leveraged for efficient sparse computation on GPUs. We determine the optimal pattern for each attention head offline and dynamically build sparse indices based on the assigned pattern during inference. With the pattern and sparse indices, we perform efficient sparse attention calculations via our optimized GPU kernels to significantly reduce the latency in the pre-filling stage of long-context LLMs. Our proposed technique can be directly applied to existing LLMs without any modifications to the pre-training setup or additional fine-tuning. By evaluating on a wide range of downstream tasks, including InfiniteBench, RULER, PG-19, and Needle In A Haystack, and models including LLaMA-3-1M, GLM4-1M, Yi-200K, Phi-3-128K, and Qwen2-128K, we demonstrate that MInference effectively reduces inference latency by up to 10x for pre-filling on an A100, while maintaining accuracy. Our code is available at https://aka.ms/MInference.