Artikel
Beschleunigung von Sonar durch Spekulation
Spekulative Dekodierung beschleunigt die Generierungsgeschwindigkeit von Großes Sprachmodellen (LLMs), indem ein schnelles und kleines Entwurfsmodell verwendet wird, um Abschluskandidaten zu produzieren, die vom größeren Zielmodell verifiziert werden.
Gemäß diesem Schema werden anstelle eines Laufs des teuren Ziels, der ein einziges Token produziert, mehrere in einem einzigen Schritt erzeugt. Hier präsentieren wir die Implementierungsdetails verschiedener Arten der spekulativen Dekodierung, die bei Perplexity angewendet wurden, um die Latenz zwischen den Tokens bei Sonar-Modellen zu reduzieren.
Spekulative Dekodierung
Spekulative Dekodierung nutzt die Struktur von natürlichen Sprachen und die autoregressive Natur von Transformatoren, um die Token-Generierung zu beschleunigen. Auch wenn größere Modelle, wie Llama-70B, mehr Wissen tragen als kleinere, wie Llama-1B, erbringen sie bei einigen einfacheren Aufgaben ähnliche Leistungen. Diese Überschneidung deutet darauf hin, dass bestimmte Sequenzen besser von den weniger teuren Modellen generiert werden, während komplexe Probleme größeren Modellen überlassen werden. Die Herausforderung besteht darin, zu bestimmen, welche Abschlüsse besser sind und ob die Generierung des kleineren Modells von der gleichen Qualität ist wie die des größeren.
Glücklicherweise sind LLMs autoregressive Transformatoren: Wenn sie eine Sequenz von Tokens erhalten, geben sie die Wahrscheinlichkeitsverteilung des nächsten Tokens aus. Darüber hinaus zeigen die Logits, die aus den Zwischenmerkmalen abgeleitet werden, die mit den Tokens in der Eingabesequenz verbunden sind, auch an, wie wahrscheinlich es ist, dass das Modell diese genauen Tokens ausgibt. Diese Eigenschaft ermöglicht Spekulation: Wenn eine Sequenz von Tokens von einem kleineren Modell ausgehend von einem Eingabepräfix generiert wird, kann sie durch das größere Modell geführt werden, um festzustellen, wie gut sie mit dem Zielmodell übereinstimmt. Jedes Präfix der Kandidaten wird mit einer Wahrscheinlichkeit bewertet und das längste über einem Akzeptanzschwellenwert wird ausgewählt. Als Bonus stellt das Zielmodell auch ein nachfolgendes Token kostenlos zur Verfügung: Wenn ein Entwurfsmodell n Tokens generiert, können in einem Schritt bis zu n + 1 ausgegeben werden.
Zur Inferenzzeit kann der spekulative Stichprobenprozess grob in 4 Phasen unterteilt werden:
Vorauffüllen: Sowohl das Ziel- als auch das Entwurfsmodell müssen auf der Eingabesequenz ausgeführt werden, um die KV-Cacheeinträge zu füllen. Während einige Schemata, wie Medusa, einfachere dichte Schichten zur Vorhersage verwenden, konzentrieren wir uns in diesem Beitrag auf Entwürfe, die transformatorbasiert sind und eigene KV-Caches benötigen.
Entwurfsproduktion: Das Entwurfsmodell iteriert, um eine Anzahl fester Tokens zu produzieren. Die Entwurfssequenz kann linear sein oder das Modell kann eine baumartige Struktur bis zu einer gegebenen Tiefe erkunden (EAGLE, Medusa). Hier konzentrieren wir uns auf lineare Sequenzen.
Akzeptanz: Das Zielmodell läuft über die Entwurfssequenz und erstellt Logits, die jedem Entwurfstoken entsprechen. Die Länge der längsten akzeptablen Sequenz wird bestimmt.
Zielerstellung: Da das Ziel Logits erzeugt hat, stimmen die Logits an der unpassenden Position oder am Ende der Sequenz mit einem nachfolgenden Token überein. Diese Logits können abgetastet werden, um ein robustes Token vom Ziel zu liefern, das die Sequenz abrundet.
Verschiedene Methoden existieren zur Implementierung der spekulativen Dekodierung. In diesem Beitrag konzentrieren wir uns auf die Schemata, die wir zur Beschleunigung von Sonar-Modellen mit einem internen 1B-Modell verwendet haben, sowie auf die Vorhersagemechanismen, die wir entwickeln, um Modelle im Umfang von DeepSeek zu beschleunigen.
Ziel-Entwurf
Spekulative Dekodierung kann erreicht werden, indem ein bestehendes kleines LLM als Entwurfsmodell mit einem Zielmodell gekoppelt wird, um Kandidatensequenzen zu generieren. In der Produktion haben wir Sonar mit einem Llama-1B-Modell beschleunigt, das auf demselben Dataset wie das Ziel feingetunt wurde. Obwohl bei diesem Ansatz kein Entwurf von Grund auf neu trainiert werden musste, verwendet das kleine Modell dennoch signifikante KV-Cache-Kapazität und führt zu einem leichten Vorfüll-Overhead, der TTFT erhöht.
In diesem Schema spekuliert der Decoder nur über Decode-Only-Chargen, wobei Tokens durch Standardstichproben während des Vorfüllens oder auf gemischten Vorfüll-Decode-Chargen generiert werden. Im Vorfüllstadium werden die Ziel-Logits sofort abgetastet, um auch das neu generierte Token im KV-Cache des Entwurfs vorzubefüllen. Der Entwurf wird noch nicht abgetastet, aber die Logits, die er produziert, werden in die Decode-Phase übertragen.
Im Decode wird das Entwurfsmodell vorangetrieben und das oberste Token in jeder Phase abgetastet. Nachdem die gewünschte Entwurfslänge erreicht ist, werden die Tokens durch das Zielmodell geführt, um die Logits zu produzieren, anhand derer der Abtaster die akzeptierte Sequenzlänge identifiziert. Die Akzeptanz wird bestimmt, indem die vollständigen Wahrscheinlichkeitsverteilungen des Entwurfs und des Ziels verglichen werden. Da das Ziel immer einen Satz von Logits nach der akzeptierten Entwurfssequenz ausgibt, wird dieser abgetastet, um eine zusätzliche Ausgabe zu erzeugen. Da das Entwurfsmodell dieses akzeptierte Token noch nicht gesehen hat, wird es neu gestartet, um seine entsprechenden KV-Cacheeinträge für den nächsten Decode-Schritt zu füllen und die Logits erneut zu übertragen.
EAGLE
EAGLE ist ein spekulatives Dekodierungsschema, das mehrere Entwurfssequenzen erkundet, die durch eine baumartige Traversierung wahrscheinlicher Entwurfstokens generiert werden. Ein fester (EAGLE) oder dynamisch geformter (EAGLE-2) Baum wird durch aufeinanderfolgende Ausführungen der Entwurfstokens erkundet, wobei die Top-K-Kandidaten an jedem Knoten in Betracht gezogen werden, anstatt dem höchstbewerteten Token in einer linearen Sequenz zu folgen. Die Sequenzen werden dann bewertet und die längste geeignete wird ausgewählt, um fortzufahren, wobei auch ein zusätzliches Token vom Ziel hinzugefügt wird.
Um genauere Vorhersagen zu erreichen, sagt ein EAGLE-Entwurfsmodell nicht nur auf Basis von Tokens, sondern auch unter Verwendung der Zielmerkmale (versteckte Zustände der letzten Schicht) des Zielmodells vorher. Der Nachteil von EAGLE ist die Notwendigkeit, benutzerdefinierte, kleine Entwurfsmodelle zu trainieren, die genau genug sind, um geeignete Kandidaten innerhalb eines niedrigen Latenzbudgets zu generieren. Typischerweise ist ein Entwurfsmodell eine einzelne Transformator-Schicht, die identisch mit einer Decoder-Schicht des Originalmodells ist, die eng mit dem Ziel durch Kopplung mit dessen Einbettungen und lm_head-Projektionen verbunden ist. Da dies weniger KV-Cache-Kapazität erfordert, hat EAGLE einen geringeren Speicherbedarf.
Um baumartige Sequenzen im Zielmodell zu überprüfen, müssen benutzerdefinierte Aufmerksamkeitsmasken verwendet werden. Leider verlangsamt die Verwendung einer benutzerdefinierten Aufmerksamkeitsmaske für eine ganze Sequenz die Aufmerksamkeit bei realistischen Eingabelängen erheblich (um bis zu 50 %), was einen Teil der durch Spekulation erreichbaren Beschleunigung zunichtemacht. Aus diesem Grund haben wir die vollständige Baumexploration noch nicht in der Produktion eingesetzt, sondern uns auf den Spezialfall der Einzel-Token-Vorhersage über MTP-ähnliche Schemata konzentriert, die im DeepSeek-V3 Technical Report vorgestellt wurden.
MTP
Dieses Schema ähnelt der Entwurf-Ziel-Dekodierung, mit der Ausnahme, dass versteckte Zustände neben Tokens zur Vorhersage verwendet werden. Etwas mehr Arbeit muss sowohl in den Vorfüll- als auch in den Dekodiestufen im Vergleich zur regulären Entwurf-Ziel-Spekulation geleistet werden. Das Entwurfsmodell verwendet sowohl Tokens als auch versteckte Zustände: Token t_{i+1} wird aus den Logits L_i abgetastet, die dem Token t_i entsprechen, welches wiederum aus den versteckten Zuständen H_i abgeleitet wird. Folglich müssen die Eingabetokenpuffer relativ zu den von dem Ziel ausgegebenen versteckten Zustandvektoren um einen Schritt nach links verschoben werden. Die Abbildung unten markiert die für das Training verwendeten Korrespondenzen sowie die Verschiebung während der Inferenz.
Der Dekodierungsfluss ist der der Entwurf-Ziel-Dekodierung sehr ähnlich, mit der Ausnahme, dass sowohl versteckte Zustände als auch Logits übertragen werden. Unsere Implementierung teilt alle damit verbundenen Stichproben- und Logit-Verarbeitungslogits und spezialisiert nur die Vorwärtsaufrufe des Modells. Wenn mehrere Tokens vorhergesagt werden, verwendet das Entwurfsmodell die versteckten Entwurfszustände zur Vorhersage, füllt auch die KV-Cacheeinträge basierend auf seinen eigenen Merkmalen. Auf lange Sicht kann dies die Genauigkeit verschlechtern. Anschließend, wenn das Entwurfsmodell ausgeführt wird, um die KV-Cacheeinträge für die Zielvorhersage zu füllen, führen wir es auf der gesamten Sequenz aus, wobei die genaueren versteckten Zustände des Ziels als Eingaben verwendet werden. Da diese Entwurfsmodelle klein sind, sind die zusätzlichen Kosten für die Verarbeitung der zusätzlichen Tokens vernachlässigbar.
MTP Kopf Training
Um von MTP zu profitieren, bauten wir die Infrastruktur auf, die erforderlich ist, um MTP-Köpfe an unsere feingetunten Modelle auf den Perplexity-Datasets zu trainieren, die auf einem Knoten mit 8xH100-Geräten laufen. In etwa einem Tag können wir Köpfe für Modelle von Llama-1B bis Llama-70B und DeepSeek V2-Lite erstellen. Bei größeren Modellen verlassen wir uns auf MTP-Köpfe, die während des Feintunings aufgebaut wurden.
Das Ziel des MTP-Trainings ist es, die versteckten Entwurfszustände und Logits, die von den versteckten Zielzuständen extrapoliert werden, mit den nächsten Token-Logits und versteckten Zuständen des Ziels abzugleichen. Da die Inferenz versteckter Zustände teuer ist, berechnen wir sie mit unserer Inferenz-optimierten Implementierung des Zielmodells vor, die beim Training verwendet wird. Um jedoch die Inferenz-MTP-Implementierung zu validieren und sicherzustellen, dass numerische Unterschiede aufgrund von Quantisierung oder Optimierungen die Ergebnisse nicht behindern, verwenden wir für Validierungsverlust und Genauigkeitsschätzung vollständig die Inferenzimplementierung sowohl des Ziel- als auch des Entwurfsmodells erneut.
Beim Hochskalieren vom ShareGPT-Datensatz, der im Originalpapier verwendet wurde, zu größeren Proben, bemerkten wir, dass die MTP-Kopfarchitektur, die im EAGLE-Papier skizziert und implementiert wurde, für Modelle in der Größe von 70B nicht trainierbar war. Anders als ShareGPT, das eine größere Anzahl kürzerer Sequenzen enthielt, trainieren wir auf einer etwas kleineren Anzahl von erheblich längeren Eingaben. Da die ursprünglichen EAGLE-Köpfe in der Struktur leicht von einem typischen Transformator abwichen, führten wir einige RMS-Normalisierungsschichten, die entfernt worden waren, wieder ein. Wir stellten fest, dass dies nicht nur das Training zum Konvergieren brachte, sondern auch die Genauigkeit der Köpfe um einige Prozentpunkte steigerte.
Nicht nur erleichtern Schichtnormen das Training, sie wieder einzuführen ist auch mathematisch intuitiv. MTP-Köpfe wiederverwenden die Einbettungen und die Logit-Projektionen des Zielmodells, da sie erheblich groß sein können (etwa 2 GB für Llama 70B). Während des Trainings sind diese eingefroren und die Erwartung ist, dass die MTP-Schicht lernt, Vorhersagen in den gleichen Vektorraum zu betten, den die Projektionsschicht des ursprünglichen Modells während des Trainings gelernt hat. Durch das Weglassen der Normen wird erwartet, dass ein einzelnes MLP die gleiche Funktion lernt wie ein MLP, gefolgt von einer Norm, was das Abgleichen der versteckten Zustände der Entwurfs- und Zielmodelle behindert.
Inferenz mit Spekulativer Dekodierung
Im Inferenz-Engine, um Tokens für Eingabesequenzen zu generieren, müssen sie zunächst in einigermaßen große Chargen gruppiert werden, dann müssen Seiten im KV-Cache für die nächsten Tokens zugewiesen werden. Die Eingabetokens und die KV-Seiteninformation werden dann in einen Puffer gepackt, der an alle parallelen Ränge übertragen wird, die das Modell ausführen. Schließlich werden die Metadaten in den GPU-Speicher kopiert und das Modell wird ausgeführt, um die Logits zu erzeugen, aus denen das nächste Token abgetastet wird.
Anders als bestimmte Implementierungen, die einen Entwurfs- und Ziel-Inferenz-Server locker durch einen Wrapper koppeln, der Anfragen zwischen ihnen orchestriert, sind unsere Entwurf-Ziel-Paare eng gekoppelt und schreiten synchron durch die Generierung. Batch-Scheduling und KV-Seitenzuweisung werden zwischen den Modellen für alle Formen der spekulativen Dekodierung geteilt: Dies vereinheitlicht die Logik, die ein Modell mit dem übergreifenden Inferenz-Server verbindet, da sie alle dieselbe Schnittstelle bieten.
Die Inferenz-Laufzeit bei Perplexity ist um FlashInfer herum gestaltet, das die Metadaten bestimmt, die gebaut werden müssen, um den Aufmerksamkeitskernel zu konfigurieren und zu planen. Angesichts einiger Eingabesequenzen, die eine Charge bilden, müssen für das Vorauffüllen, Dekodieren oder die Überprüfung CPU-seitig Arbeiten durchgeführt werden, um Zwischenpuffer zuzuweisen und bestimmte konstante Puffer zu füllen, die in der Aufmerksamkeit verwendet werden. Diese Arbeiten kommen zu den Kosten des Batch-Schedulings und der KV-Seitenzuweisung hinzu, die ebenfalls Latenzen verursachen, die versteckt werden müssen, um die GPU-Auslastung zu maximieren.
Während wir die CPU-seitigen und GPU-seitigen Arbeiten für die Inferenz ohne Spekulation vollständig parallelisiert hatten, stellten wir fest, dass das CPU-GPU-Verhältnis für spekulative Dekodierung komplexer ist. Die Hauptherausforderung ergibt sich aus der Tatsache, dass die Anzahl der akzeptierten Tokens die Sequenzlänge für einen nachfolgenden Lauf bestimmt und somit einen unvermeidbaren GPU-zu-CPU-Synchronisationspunkt einführt. Wir experimentierten mit verschiedenen Planungs-Schemata, um die Latenz der CPU-Arbeit bestmöglich zu verbergen.
Entwurf-Ziel-Zeitplan
Obwohl es kleiner als ein Zielmodell ist, führt ein Ganzes LLM als Entwurf immer noch zu erheblicher Latenz auf der GPU, was etwas Spielraum bietet, um teure CPU-Operationen zu verbergen. Da kleinere Modelle nicht von der Tensor-Parallelisierung profitieren, gibt es ein Missverhältnis zwischen der Anzahl der Ränge, die ein Ziel und ein Entwurf überlappen. In unserer Implementierung wird das Entwurfsmodell nur im führenden Rang einer TP-Gruppe ausgeführt.
Wie bereits erwähnt, trägt ein Dekodierungsschritt Logits in den nächsten Lauf über. Dies ermöglicht es uns, eine Ausführung des Entwurfsmodells mit der CPU-seitigen Batch-Scheduling-Arbeit zu überlappen. Nach dem Zusammenstellen der Batch rufen wiederholte Aufrufe zum Abtaster und Entwurf die Entwurfstokens ab. Parallel wird die Charge zur Verifizierung für das Zielmodell zusammengestellt und mit den parallelle Arbeitern synchronisiert. Die Ziel-Logits werden verifiziert und abgetastet, um die akzeptierten Sequenzlängen zu bestimmen. An diesem Punkt ist die GPU-zu-CPU-Synchronisation notwendig, um die nachfolgenden Sequenzlängen zu bestimmen. Da das Entwurfsmodell nur auf dem führenden Knoten ausgeführt wird, wird seine Batch sequenziert und seine Ausführung gestartet, um KV-Cacheeinträge mit dem zusätzlichen Token des Ziels zu füllen. Die Logits, die in diesem Entwurfs-Durchlauf erzeugt werden, werden im nächsten Lauf abgeprobt, um das erste Entwurfs-Token zu erzeugen. Am wichtigsten ist, dass während der Ausführung des Entwurfsmodells die nächste Charge geplant werden kann.
MTP-Plan für ein Einzelnes Token
Während die Laufzeit noch keine Eagle-ähnliche Entwurfsbaumerforschung bietet, haben wir einen Sonderfall dieses Schemas implementiert, der eine lineare Sequenz von Entwurfstokens in Betracht zieht, die von einem Modell in der Größe einer einzelnen Transformator-Decoder-Schicht erzeugt wird. Dieses Schema kann für die Entwurfsvorhersage unter Verwendung der quelloffenen Gewichte von DeepSeek R1 verwendet werden. Der Unterfall der Vorhersage eines einzelnen Tokens ist interessant, da große MTP-Schichten ausreichend hohe Akzeptanzraten erreichen, um ihren Overhead zu rechtfertigen.
MTP-Planung ist etwas komplexer, da das Entwurfsmodell wesentlich schneller ist und weniger CPU-seitige Latenz verbirgt. Darüber hinaus wird der Entwurf parallel zum Zielmodell aufgeteilt, was gemeinsame Speicherübertragungen für Batchinformationen erfordert. Ein Lauf beginnt damit, Batch Infos zu übertragen und das erste Token aus übertragenen Logits abzutasten, ähnlich dem vorherigen Schema. Als Nächstes wird das Ziel ausgeführt, um Tokens zu validieren, wobei 2 * D Tokens verarbeitet werden, wobei D die Dekodiestammlänge ist. Dies ist ideal für Mikrobatching in Mixture-of-Experts (MoE)-Modellen über langsamere Interkonnekte, wie InfiniBand, da die Batch gleichmäßig in zwei Hälften geteilt wird. Die versteckten Zustände des Ziels werden zum nächsten Entwurfslauf hinübergetragen, während die Logits in den Abtaster zur Verifizierung übergeben werden.
Indem eine begrenzte Menge zusätzlicher Arbeit auf der GPU durchgeführt wird, vermeiden wir CPU-zur-GPU-Synchronisation nach der Akzeptanz der Entwurfssequenz. Nachdem die Eingabetokens der Ziele verschoben wurden, setzt ein Kernel die nächsten Ziel-Tokens an ihre entsprechenden Positionen ein. Der Entwurf wird dann mit den gleichen Batchinformationen wie das Ziel erneut ausgeführt, KV-Cacheeinträge füllend und die Logits und versteckten Zustände für den nächsten Lauf aufbauend, einige redundante Arbeiten an Tokens ausführend, die nicht akzeptiert wurden. In diesen Situationen ist die Latenz der ungenutzten Arbeit aufgrund der kleinen Größe des Entwurfsmodells kaum messbar. Parallel zum Entwurfslauf werden die Sequenzlängen auf der CPU bestimmt und die Planung der nächsten Batch eingeleitet, ohne auf das Ende der GPU-Arbeit zu warten.
Der Overhead zusätzlicher Arbeit in der Entwurfsschicht ist bei der Aufmerksamkeit nicht spürbar, jedoch sind MLP-Schichten problematischer. Da Matrix-Multiplikationsanweisungen zu einer Grenze von 64 an der Dimension der Anzahl der Tokens auffüllen, wenn das Verdoppeln nicht erheblich mehr Blöcke erfordert, bleibt der Overhead verborgen. Bei längeren Entwurfssequenzen ist der Overhead teurer und das Schema, das für reguläre Entwurf-Ziel-Modelle verwendet wird, funktioniert besser.
Referenzen
Schnelle Inferenz von Transformatoren mittels Spekulativer Dekodierung
EAGLE: Spekulative Stichprobe erfordert ein Umdenken der Merkmalsunsicherheit
EAGLE-2: Schnellere Inferenz von Sprachmodellen mit dynamischen Entwurfsbäumen
Medusa: Einfaches LLM-Inferenzbeschleunigungs-Framework mit mehreren Dekodierköpfen
FlashInfer: Effiziente und anpassbare Aufmerksamkeitsengine für LLM-Inferenzdienste