MInference 1.0: Beschleunigung der Vorabfüllung für Long-Context LLMs durch dynamische Sparse Attention

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

July 2, 2024

Autoren: 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

Zusammenfassung

Die Rechenherausforderungen der Inferenz von Large Language Models (LLMs) bleiben eine bedeutende Hürde für ihren weitreichenden Einsatz, insbesondere da die Promptlängen weiter zunehmen. Aufgrund der quadratischen Komplexität der Aufmerksamkeitsberechnung benötigt ein 8B LLM 30 Minuten, um einen Prompt mit 1M Tokens (d. h. die Vorausfüllungsphase) auf einer einzelnen A100 GPU zu verarbeiten. Bestehende Methoden zur Beschleunigung der Vorausfüllung scheitern oft daran, eine akzeptable Genauigkeit oder Effizienz bei der Anwendung auf LLMs mit langem Kontext aufrechtzuerhalten. Um diese Lücke zu schließen, führen wir MInference (Milliontokens Inferenz) ein, eine sparse Berechnungsmethode, die darauf abzielt, die Vorausfüllung bei der Verarbeitung von langen Sequenzen zu beschleunigen. Konkret identifizieren wir drei einzigartige Muster in den Aufmerksamkeitsmatrizen mit langem Kontext - die A-Form, den Vertikalen Strich und den Block-Sparse -, die für eine effiziente sparse Berechnung auf GPUs genutzt werden können. Wir bestimmen das optimale Muster für jeden Aufmerksamkeitskopf offline und erstellen während der Inferenz dynamisch sparse Indizes basierend auf dem zugewiesenen Muster. Mit dem Muster und den sparse Indizes führen wir effiziente sparse Aufmerksamkeitsberechnungen über unsere optimierten GPU-Kerne durch, um die Latenz in der Vorausfüllungsphase von LLMs mit langem Kontext signifikant zu reduzieren. Unsere vorgeschlagene Technik kann direkt auf bestehende LLMs angewendet werden, ohne dass Änderungen am Pre-Training-Setup oder zusätzliches Fine-Tuning erforderlich sind. Durch Evaluation an einer Vielzahl von Aufgaben, einschließlich InfiniteBench, RULER, PG-19 und Needle In A Haystack, und Modellen wie LLaMA-3-1M, GLM4-1M, Yi-200K, Phi-3-128K und Qwen2-128K, zeigen wir, dass MInference die Inferenzlatenz bei der Vorausfüllung auf einer A100 um bis zu 10x reduziert, während die Genauigkeit beibehalten wird. Unser Code ist unter https://aka.ms/MInference verfügbar.

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.