Große Sprachmodelle (LLMs) haben die Verarbeitung natürlicher Sprache (NLP) revolutioniert und in verschiedenen Anwendungsbereichen Spitzenleistungen erzielt. Im Zuge ihrer Weiterentwicklung werden LLMs zunehmend für Aufgaben mit langen Kontexten eingesetzt, wie z. B. Chain-of-Thought-Reasoning, Dokumentenzusammenfassung und Retrieval-Augmented Generation. Die schnelle und effiziente Ausführung von LLMs mit langem Kontext stellt jedoch aufgrund der inhärenten Speicher- und Rechenengpässe in den Transformer-Architekturen eine Herausforderung dar.

LLM-Inferenz und der Bedarf an effizienten Dekodierungsstrategien

Die LLM-Inferenz besteht aus zwei Phasen: Prefilling und Dekodierung. Im Prefilling werden die Aktivierungen für alle Eingabe-Token berechnet und die Schlüssel und Werte für alle Token im Schlüsselwert-Cache (KV-Cache) gespeichert. So kann das LLM diese Schlüssel und Werte wiederverwenden, um die Aufmerksamkeit für zukünftige Token zu berechnen. In jeder Dekodierungsphase dekodiert das LLM ein neues Token unter Verwendung aller Eingabe-Token und zuvor generierten Token. Die Größe des KV-Caches wächst linear mit der Sequenzlänge. Bei einer Kontextlänge von 128.000 Token kann der KV-Cache eines LLama2-7B mit Halbpräzision beispielsweise leicht 64 GB erreichen, was den LLM-Server erheblich belastet. Darüber hinaus ist die LLM-Dekodierungsphase speichergebunden, da für die Dekodierung eines neuen Tokens auf alle vorherigen Token im KV-Cache zugegriffen werden muss, was den KV-Cache-Zugriff zum primären Engpass für die Dekodierung von LLMs mit langem Kontext macht. Diese Speichergebundenheit schränkt die Skalierbarkeit und Effizienz von LLM-Servern erheblich ein. Um dieses Problem zu beheben, wurde in letzter Zeit Sparse Attention eingeführt, das Full Attention durch die Verwendung eines kleinen Teils der Token mit den höchsten Aufmerksamkeitswerten approximiert. Im Vergleich zu Full Attention reduziert Sparse Attention den Rechenaufwand und die Speicherzugriffe, während die generative Leistung des LLMs erhalten bleibt. Bestehende Sparse-Attention-Techniken lassen sich in zwei Kategorien einteilen: Eviction- und Selection-basierte Verfahren.

Eviction-basierte Sparse Attention

Eviction-basierte Sparse Attention reduziert die Speichernutzung für den KV-Cache, indem weniger relevante Token selektiv aus dem KV-Cache entfernt werden, wodurch die Anzahl der in Aufmerksamkeitsmechanismen berechneten Token reduziert wird. Diese Verfahren verringern zwar die Größe des KV-Caches, können aber bei Aufgaben unzureichend sein, bei denen wichtige Informationen von Token getragen werden, die vorzeitig entfernt werden, wie z. B. bei den Needle-in-the-Haystack-Aufgaben.

Selection-basierte Sparse Attention

Selection-basierte Sparse Attention hingegen behält alle Token im KV-Cache, schätzt ihre Aufmerksamkeitswerte und wählt eine kleine Teilmenge von Token aus, die an jedem LLM-Dekodierungsschritt teilnehmen. Dieser Ansatz ist anfällig für Probleme im Zusammenhang mit Distributionsverschiebungen, die durch das Anhängen von spärlich besuchten, verzerrten KV-Darstellungen zurück in den Cache entstehen.

TidalDecode: Ein vielversprechender Ansatz für eine effiziente LLM-Dekodierung

In diesem Zusammenhang stellt sich TidalDecode als ein Algorithmus und System für eine schnelle und präzise LLM-Dekodierung vor, das positionsbeständige Sparse Attention (PPSA) verwendet. Die Kernidee hinter TidalDecode ist die Beobachtung, dass Token, die aufgrund ihrer höchsten Aufmerksamkeitswerte für Sparse Attention ausgewählt wurden, über aufeinanderfolgende Transformer-Schichten innerhalb jeder Dekodierungsphase hinweg eine signifikante Überschneidung aufweisen. Anstatt für jede Schicht unabhängig voneinander Token für Sparse Attention auszuwählen, führt TidalDecode einige Token-Auswahlschichten ein, die Full Attention durchführen, um die Token mit den höchsten Aufmerksamkeitswerten zu identifizieren. Alle übrigen Schichten implementieren positionsbeständige Sparse Attention, bei der nur die von den Token-Auswahlschichten ausgewählten Token für die Aufmerksamkeit aus dem KV-Cache abgerufen werden. Folglich arbeiten alle anderen Schichten zwischen zwei Token-Auswahlschichten auf demselben Satz von Token, wodurch der Overhead für die Token-Auswahl reduziert wird. Dieser Ansatz ermöglicht es TidalDecode, den Overhead der Token-Auswahl für Sparse Attention erheblich zu reduzieren, ohne die Qualität der generierten Ergebnisse zu beeinträchtigen. Um die Verzerrung des KV-Caches zu beheben, führt TidalDecode außerdem einen Cache-Korrekturmechanismus ein, der den KV-Cache regelmäßig mithilfe von Full Attention für alle spärlich dekodierten Token auffüllt, um Verzerrungen in den KV-Darstellungen zu verringern.

Bewertung und Ergebnisse

Umfassende Evaluierungen mit den Modellen LongChat-7b-v1.5-32k, Llama-3-8B, Llama-3-70B und Llama-3.1-8B für die Aufgaben Needle-in-the-Haystack, PG-19 und LongBench zeigen, dass TidalDecode im Vergleich zu den besten existierenden Sparse-Attention-Verfahren durchweg den besten Kompromiss zwischen Leistung und Effizienz erzielen kann. Es wurden benutzerdefinierte GPU-Kernel für PPSA und ein End-to-End-System für TidalDecode implementiert. Im Vergleich zu bestehenden Implementierungen für Full und Sparse Attention reduzierte das System die End-to-End-Inferenzlatenz um das bis zu 2,1-fache bzw. 1,2-fache.

Schlussfolgerung

TidalDecode stellt einen wichtigen Schritt in Richtung einer effizienteren Dekodierung von LLMs dar, insbesondere im Kontext von Aufgaben mit langem Kontext. Durch die Nutzung der räumlichen Kohärenz der Token-Auswahl und die Einführung von positionsbeständiger Sparse Attention bietet TidalDecode eine vielversprechende Lösung für die Bewältigung der Speicher- und Rechenengpässe, die mit großen Sprachmodellen verbunden sind. Die Evaluierungsergebnisse deuten darauf hin, dass TidalDecode das Potenzial hat, die Leistung und Skalierbarkeit von LLM-basierten Anwendungen erheblich zu verbessern und so den Weg für neue Möglichkeiten in der Verarbeitung natürlicher Sprache zu ebnen.

Bibliographie

https://arxiv.org/abs/2410.05076 https://arxiv.org/html/2410.05076 https://www.chatpaper.com/chatpaper/zh-CN/paper/64874 https://www.chatpaper.com/chatpaper/ja/paper/64874 https://github.com/DefTruth/Awesome-LLM-Inference https://github.com/HuangOwen/Awesome-LLM-Compression https://paperreading.club/category?cate=Attention