Beschleunigung von Sonar durch Spekulation

Spekulatives Decoding beschleunigt die Generierungsgeschwindigkeit von Large Language Models (LLMs), indem ein schnelles und kleines Entwurfsmodell verwendet wird, um Vervollständigungskandidaten zu erstellen, die vom größeren Zielmodell verifiziert werden. Unter diesem Schema werden in einem einzigen Schritt mehrere Token emittiert, anstatt dass ein Durchlauf des aufwendigen Zielmodells nur ein einzelnes Token erzeugt. Hier präsentieren wir die Implementierungsdetails verschiedener Arten des spekulativen Decodings, die bei Perplexity angewendet werden, um die Inter-Token-Latenz bei Sonar-Modellen zu reduzieren.

Spekulatives Decoding

Spekulatives Decoding nutzt die Struktur natürlicher Sprachen und die autoregressive Natur von Transformern, um die Token-Generierung zu beschleunigen. Auch wenn größere Modelle wie Llama-70B über mehr Wissen verfügen als kleinere wie Llama-1B, erbringen sie bei einigen einfacheren Aufgaben ähnliche Leistungen. Diese Überschneidung legt nahe, dass bestimmte Sequenzen besser von den kostengünstigeren Modellen generiert werden, während komplexe Probleme den größeren Modellen überlassen werden. Die Herausforderung besteht darin zu bestimmen, welche Vervollständigungen besser sind und ob die Generierung des kleineren Modells von gleicher Qualität ist wie die des größeren.

Glücklicherweise sind LLMs autoregressive Transformer: Wenn sie eine Sequenz von Token erhalten, geben sie die Wahrscheinlichkeitsverteilung des nächsten Tokens aus. Zusätzlich geben die Logits, die aus den mit den Token in der Eingabesequenz assoziierten Zwischenmerkmalen abgeleitet werden, auch an, wie wahrscheinlich es ist, dass das Modell genau diese Token ausgibt. Diese Eigenschaft ermöglicht Spekulation: Wenn eine Sequenz von Token von einem kleineren Modell ausgehend von einem Eingabepräfix generiert wird, kann sie durch das größere Modell laufen, um zu bestimmen, wie gut sie mit dem Zielmodell übereinstimmt. Jedes Präfix der Kandidaten wird mit einer Wahrscheinlichkeit bewertet und das längste, das über einem Akzeptanzschwellenwert liegt, wird ausgewählt. Als Bonus stellt das Zielmodell auch ein nachfolgendes Token kostenlos zur Verfügung: Wenn ein Entwurfsmodell n Token generiert, können bis zu n + 1 in einem Schritt emittiert werden.

Vergrößern

Zum Zeitpunkt der Inferenz kann der spekulative Sampling-Prozess grob in vier Phasen unterteilt werden:

Prefill: Sowohl das Zielmodell als auch das Entwurfsmodell müssen auf der Eingabesequenz ausgeführt werden, um die KV-Cache-Einträge zu füllen. Während einige Schemata, wie Medusa, einfachere dichte Schichten zur Vorhersage verwenden, konzentrieren wir uns in diesem Beitrag auf Transformer-basierte Entwürfe, die ihre eigenen KV-Caches benötigen.

Entwurfsgenerierung: Das Entwurfsmodell iteriert, um eine Anzahl fester Token zu erzeugen. Die Entwurfssequenz kann linear sein oder das Modell kann eine baumartige Struktur bis zu einer gegebenen Tiefe untersuchen (EAGLE, Medusa). Hier konzentrieren wir uns auf lineare Sequenzen.

Akzeptanz: Das Zielmodell läuft auf der Entwurfssequenz und bildet Logits, die jedem Entwurfstoken entsprechen. Die Länge der längsten akzeptablen Sequenz wird bestimmt.

Zielgenerierung: Da das Ziel Logits generierte, entsprechen die Logits an der nicht übereinstimmenden Position oder am Ende der Sequenz einem nachfolgenden Token. Diese Logits können abgetastet werden, um ein robustes Token vom Ziel zu erhalten, das die Sequenz abschließt.

Es existieren verschiedene Methoden zur Implementierung von spekulativem Decoding. In diesem Beitrag konzentrieren wir uns auf die Schemata, die wir zur Beschleunigung von Sonar-Modellen unter Verwendung eines hausinternen 1B-Modells eingesetzt haben, sowie auf die Vorhersagemechanismen, die wir entwickeln, um Modelle in der Größenordnung von DeepSeek zu beschleunigen.

Ziel-Entwurf

Spekulatives Decoding kann erreicht werden, indem ein existierendes kleines LLM als Entwurfsmodell an ein Zielmodell gekoppelt wird, um Kandidatensequenzen zu generieren. In der Produktion haben wir Sonar unter Verwendung eines Llama-1B-Modells beschleunigt, das auf demselben Datensatz wie das Zielmodell feinabgestimmt wurde. Obwohl dieser Ansatz kein Training eines Entwurfs von Grund auf erforderte, verbraucht das kleine Modell dennoch erhebliche KV-Cache-Kapazität und führt einen leichten Prefill-Overhead ein, was die TTFT erhöht.

Unter diesem Schema spekuliert der Decoder nur auf Decodier-Batch-Vorgängen und generiert Token durch Standard-Sampling während des Prefills oder bei gemischten Prefill-Decode-Batches. In der Prefill-Phase werden die Ziel-Logits sofort abgetastet, um auch das neu generierte Token im KV-Cache des Entwurfs vorab zu füllen. Der Entwurf wird noch nicht abgetastet, aber die von ihm produzierten Logits werden in die Decodier-Phase übernommen.

Vergrößern

Beim Decodieren wird das Entwurfsmodell vorangetrieben, wobei in jeder Stufe das Top-Token abgetastet wird. Nachdem die gewünschte Entwurfslänge erreicht ist, werden die Token durch das Zielmodell geleitet, um die Logits zu erzeugen, auf deren Basis der Sampler die akzeptierte Sequenzlänge identifiziert. Die Akzeptanz wird durch den Vergleich der vollständigen Wahrscheinlichkeitsverteilungen des Entwurfs und des Ziels bestimmt. 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 erneut ausgeführt, um seine entsprechenden KV-Cache-Einträge in Vorbereitung auf den nächsten Decodierschritt zu füllen, wobei die Logits erneut übernommen werden.

EAGLE

EAGLE ist ein spekulatives Decoding-Schema, das mehrere Entwurfssequenzen untersucht, die durch eine baumartige Traversierung wahrscheinlicher Entwurfstoken generiert werden. Ein fester (EAGLE) oder dynamisch geformter (EAGLE-2) Baum wird durch aufeinanderfolgende Ausführungen der Entwurfstoken untersucht, wobei die Top-K-Kandidaten an jedem Knoten berücksichtigt werden, anstatt dem Token mit der höchsten Punktzahl in einer linearen Sequenz zu folgen. Die Sequenzen werden dann bewertet und die längste geeignete wird ausgewählt, um fortzufahren, wobei zusätzlich ein weiteres Token vom Ziel angehängt wird.

Vergrößern

Um eine genauere Vorhersage zu erreichen, sagt ein EAGLE-Entwurfsmodell nicht nur basierend auf Token vorher, sondern nutzt auch die Zielmerkmale (verborgene Zustände der letzten Schicht) des Zielmodells. Der Nachteil von EAGLE ist die Notwendigkeit, benutzerdefinierte, kleine Entwurfsmodelle zu trainieren, die genau genug sind, um geeignete Kandidaten innerhalb eines geringen Latenzbudgets zu generieren. Typischerweise ist ein Entwurfsmodell eine einzelne Transformer-Schicht, die identisch mit einer Decodier-Schicht des ursprünglichen Modells ist und eng an das Ziel gekoppelt ist, indem sie an dessen Einbettungen und lm_head-Projektionen gebunden ist. Da dies weniger KV-Cache-Kapazität erfordert, hat EAGLE einen geringeren Speicherbedarf.

Um baumartige Sequenzen im Zielmodell zu verifizieren, müssen benutzerdefinierte Aufmerksamkeitsmasken verwendet werden. Leider verlangsamt die Verwendung einer benutzerdefinierten Aufmerksamkeitsmaske für eine gesamte Sequenz die Aufmerksamkeit bei realistischen Eingabelängen erheblich (um bis zu 50 %), wodurch ein Teil der durch Spekulation erreichbaren Beschleunigung zunichtegemacht wird. Wir haben aus diesem Grund noch keine vollständige Baumexploration in die Produktion gebracht und konzentrieren uns stattdessen auf den Spezialfall der Einzel-Token-Vorhersage mittels MTP-ähnlicher Schemata, die im DeepSeek-V3 Technical Report vorgestellt werden.

MTP

Dieses Schema ähnelt dem Entwurf-Ziel-Decoding, mit der Ausnahme, dass verborgene Zustände zusammen mit Token für die Vorhersage verwendet werden. Im Vergleich zum regulären Entwurf-Ziel-Spekulieren muss sowohl in der Prefill- als auch in der Decodier-Phase etwas mehr Arbeit geleistet werden. Das Entwurfsmodell verwendet sowohl Token als auch verborgene Zustände: Token t_{i+1} wird aus den Logits L_i abgetastet, die dem Token t_i entsprechen, welche wiederum aus den verborgenen Zuständen H_i abgeleitet werden. Folglich müssen die Eingabe-Token-Puffer relativ zu den vom Ziel ausgegebenen Vektoren der verborgenen Zustände um einen Schritt nach links verschoben werden. Die Abbildung unten markiert die für das Training verwendeten Entsprechungen sowie die Verschiebung während der Inferenz.

Vergrößern

Der Decodier-Fluss ist dem Entwurf-Ziel-Decoding recht ähnlich, mit der Ausnahme, dass sowohl verborgene Zustände als auch Logits übernommen werden. Unsere Implementierung teilt alle zugehörigen Sampling- und Logit-Verarbeitungsschritte und spezialisiert nur die Modell-Forward-Aufrufe. Wenn mehrere Token vorhergesagt werden, verwendet das Entwurfsmodell verborgene Zustände des Entwurfs für die Vorhersage und füllt auch KV-Cache-Einträge basierend auf seinen eigenen Merkmalen. Auf lange Sicht kann dies die Genauigkeit verschlechtern. Wenn wir anschließend das Entwurfsmodell ausführen, um den KV-Cache-Eintrag für die Zielvorhersage zu füllen, führen wir es auf der gesamten Sequenz aus und verwenden die genaueren verborgenen Zustände des Ziels als Eingaben. Da diese Entwurfsmodelle klein sind, sind die zusätzlichen Kosten für die Verarbeitung der zusätzlichen Token vernachlässigbar.

Training von MTP-Köpfen

Um von MTP zu profitieren, haben wir die Infrastruktur aufgebaut, die erforderlich ist, um MTP-Köpfe zu trainieren, die an unsere feinabgestimmten Modelle auf Perplexity-Datensätzen angehängt sind, und führen dies auf einem Knoten mit 8xH100-Geräten aus. In etwa einem Tag können wir Köpfe für Modelle von Llama-1B bis Llama-70B und DeepSeek V2-Lite bauen. Für größere Modelle verlassen wir uns auf MTP-Köpfe, die während des Feinabstimmungsprozesses erstellt wurden.

Das Ziel des MTP-Trainings ist es, die verborgenen Zustände des Entwurfs und die aus den verborgenen Zuständen des Ziels extrapolierten Logits an die Logits und verborgenen Zustände des nächsten Tokens des Ziels anzupassen. Da die Inferenz für verborgene Zustände teuer ist, berechnen wir sie vorab unter Verwendung unserer für die Inferenz optimierten Implementierung des Zielmodells, um sie während des Trainings zu verwenden. Um jedoch die MTP-Implementierung der Inferenz zu validieren und sicherzustellen, dass numerische Unterschiede aufgrund von Quantisierung oder Optimierungen die Ergebnisse nicht behindern, verwenden wir für die Validierung von Verlust und Genauigkeitsschätzung die Inferenzimplementierung sowohl des Ziel- als auch des Entwurfsmodells vollständig wieder.

Bei der Skalierung vom ShareGPT-Datensatz, der im ursprünglichen Dokument verwendet wurde, auf größere Stichproben stellten wir fest, dass die im EAGLE-Dokument skizzierte und implementierte MTP-Kopfarchitektur bei Modellen der 70B-Größe nicht trainiert werden konnte. Im Gegensatz zu ShareGPT, das eine größere Anzahl kürzerer Sequenzen enthielt, trainieren wir auf einer etwas geringeren Anzahl wesentlich längerer Prompts. Da die ursprünglichen EAGLE-Köpfe in ihrer Struktur leicht von einem typischen Transformer abwichen, haben wir einige RMS-Normalisierungsschichten wieder eingeführt, die zuvor entfernt wurden. Wir stellten fest, dass dies nicht nur die Konvergenz des Trainings ermöglichte, sondern auch die Genauigkeit der Köpfe um einige Prozentpunkte steigerte.

Vergrößern

Schichtnormen erleichtern nicht nur das Training, die Wiedereinführung der Normen ist auch mathematisch intuitiv. MTP-Köpfe verwenden die Einbettungen und die Logit-Projektionen des Zielmodells wieder, da sie in ihrer Größe beträchtlich sein können (etwa 2 GB für Llama 70B). Während des Trainings werden diese eingefroren und es wird erwartet, dass die MTP-Schicht lernt, Vorhersagen in denselben Vektorraum einzubetten, den die Projektionsschicht des ursprünglichen Modells während des Trainings gelernt hat. Durch das Weglassen der Normen wird von einem einzelnen MLP erwartet, dass es dieselbe Funktion lernt wie ein MLP gefolgt von einer Norm, was die Übereinstimmung zwischen den verborgenen Zuständen des Entwurfs- und des Zielmodells behindert.

Inferenz mit spekulativem Decoding

In der Inferenz-Engine müssen Token für Eingabesequenzen zunächst in Batches angemessener Größe gruppiert werden, dann müssen Seiten im KV-Cache für die nächsten Token zugewiesen werden. Die Eingabetoken und die KV-Seiteninformationen 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 ausgeführt, um die Logits zu erzeugen, aus denen das nächste Token abgetastet wird.

Im Gegensatz zu bestimmten Implementierungen, die einen Entwurfs- und Ziel-Inferenzserver lose über einen Wrapper koppeln, der Anfragen zwischen ihnen orchestriert, sind unsere Entwurf-Ziel-Paare eng gekoppelt und durchlaufen die Generierung gemeinsam. Batch-Scheduling und KV-Seitenzuweisung werden zwischen den Modellen für alle Formen des spekulativen Decodings geteilt: Dies vereinheitlicht die Logik, die eine Brücke zwischen einem Modell und dem übergeordneten Inferenzserver schlägt, da alle dieselbe Schnittstelle exponieren.

Die Inferenz-Laufzeit bei Perplexity basiert auf FlashInfer, das die Metadaten bestimmt, die erstellt werden müssen, um den Aufmerksamkeits-Kernel zu konfigurieren und zu planen. Angesichts einiger Eingabesequenzen, die einen Batch bilden, muss für Prefill, Decodieren oder Verifizierung CPU-seitige Arbeit geleistet werden, um Zwischenpuffer zuzuweisen und bestimmte konstante Puffer zu füllen, die bei der Aufmerksamkeit verwendet werden. Diese Arbeit kommt zu den Kosten für Batch-Scheduling und KV-Seitenzuweisung hinzu, die ebenfalls Latenzen verursachen, die verborgen werden müssen, um die GPU-Auslastung zu maximieren.

Während wir CPU-seitige und GPU-seitige Arbeit für die Inferenz ohne Spekulation vollständig parallelisiert haben, stellten wir fest, dass das CPU-GPU-Gleichgewicht für spekulatives Decoding komplexer ist. Die größte Herausforderung ergibt sich aus der Tatsache, dass die Anzahl der akzeptierten Token die Sequenzlänge für einen nachfolgenden Durchlauf bestimmt, was einen schwer zu vermeidenden GPU-zu-CPU-Synchronisationspunkt einführt. Wir haben mit verschiedenen Scheduling-Schemata experimentiert, um die Latenz der CPU-Arbeit bestmöglich zu verbergen.

Zeitplan für Entwurf-Ziel

Obwohl ein ganzes LLM als Entwurf kleiner als ein Zielmodell ist, führt es dennoch zu einer beträchtlichen Latenz auf der GPU und bietet einen gewissen Spielraum, um teure CPU-Operationen zu verbergen. Da kleinere Modelle nicht von Tensor-Parallelität profitieren, gibt es eine Diskrepanz zwischen der Anzahl der Ränge, auf die ein Ziel und ein Entwurf verteilt sind. In unserer Implementierung läuft das Entwurfsmodell nur auf dem führenden Rang einer TP-Gruppe.

Vergrößern

Wie zuvor angegeben, übernimmt ein Decodierschritt Logits in den nächsten Durchlauf. Dies ermöglicht es uns, eine Ausführung des Entwurfsmodells mit der CPU-seitigen Batch-Scheduling-Arbeit zu überlappen. Nachdem der Batch zusammengestellt wurde, erzeugen wiederholte Aufrufe an den Sampler und den Entwurf die Entwurfstoken. Parallel dazu wird der Batch für die Verifizierung für das Zielmodell zusammengestellt und mit den parallelen Workern synchronisiert. Die Ziel-Logits werden verifiziert und abgetastet, um die akzeptierten Sequenzlängen zu bestimmen. An diesem Punkt ist eine GPU-zu-CPU-Synchronisation erforderlich, um nachfolgende Sequenzlängen zu bestimmen. Da das Entwurfsmodell nur auf dem führenden Knoten läuft, wird sein Batch sequenziell eingerichtet und seine Ausführung gestartet, um seine KV-Cache-Einträge mit dem zusätzlichen Token zu füllen, das das Ziel produziert hat. Die von diesem Entwurfsdurchlauf im aktuellen Durchlauf produzierten Logits werden verwendet, um das erste Entwurfstoken im nachfolgenden Durchlauf abzutasten. Am wichtigsten ist, dass während der Entwurf läuft, der nächste Batch geplant werden kann.

MTP-Zeitplan für ein einzelnes Token

Während die Laufzeit noch keine Entwurfsbaumexploration im Eagle-Stil bietet, haben wir einen Spezialfall dieses Schemas implementiert, der eine lineare Sequenz von Entwurfstoken betrachtet, die von einem Modell der Größe einer einzelnen Transformer-Decodierschicht produziert werden. Dieses Schema kann für die Entwurfsvorhersage unter Verwendung der Open-Source-Gewichte von DeepSeek R1 verwendet werden. Der Unterfall der Vorhersage eines einzelnen Tokens ist interessant, da große MTP-Schichten ausreichend hohe Akzeptanzraten erzielen, um ihren Overhead zu rechtfertigen.

Das MTP-Scheduling ist etwas komplexer, da das Entwurfsmodell viel schneller ist und weniger CPU-seitige Latenz verbirgt. Zusätzlich wird der Entwurf zusammen mit dem Zielmodell gesharded, was Transfers im gemeinsam genutzten Speicher für Batch-Informationen erfordert. Ein Durchlauf beginnt mit der Übertragung von Batch-Informationen und dem Abtasten des ersten Tokens aus übernommenen Logits, ähnlich dem vorherigen Schema. Als Nächstes wird das Ziel ausgeführt, um Token zu validieren, wobei 2 * D Token verarbeitet werden, wobei D die Decodier-Batch-Größe ist. Dies ist ideal für Micro-Batching in Mixture-of-Experts (MoE)-Modellen über langsamere Verbindungen wie InfiniBand, da sich der Batch gleichmäßig in zwei Hälften teilt. Die verborgenen Zustände des Ziels werden in den nächsten Entwurfsdurchlauf übernommen, während die Logits zur Verifizierung an den Sampler weitergegeben werden.

Vergrößern

Durch die Durchführung einer begrenzten Menge zusätzlicher Arbeit auf der GPU vermeiden wir eine CPU-zu-GPU-Synchronisation nach der Akzeptanz der Entwurfssequenz. Nachdem die Eingabetoken der Ziele verschoben wurden, fügt ein Kernel die nächsten Zieltoken an ihre entsprechenden Stellen ein. Der Entwurf wird dann erneut mit denselben Batch-Informationen wie das Ziel ausgeführt, wobei KV-Cache-Einträge gefüllt und die Logits und verborgenen Zustände für den nächsten Durchlauf erstellt werden, wobei redundante Arbeit an Token geleistet wird, die nicht akzeptiert wurden. In diesen Situationen ist die Latenz der ungenutzten Arbeit aufgrund der geringen Größe des Entwurfsmodells kaum messbar. Parallel zum Entwurfsdurchlauf werden die Sequenzlängen auf der CPU bestimmt und das Scheduling des nächsten Batch gestartet, ohne auf den Abschluss der GPU-Arbeit warten zu müssen.

Der Overhead der zusätzlichen Arbeit in der Entwurfsschicht ist bei der Aufmerksamkeit nicht spürbar, jedoch sind MLP-Schichten problematischer. Da Anweisungen zur Matrixmultiplikation auf eine Grenze von 64 entlang der Dimension der Token-Anzahl aufgefüllt werden, ist der Overhead verborgen, wenn eine Verdoppelung nicht wesentlich mehr Blöcke erfordert. Für längere Entwurfssequenzen ist der Overhead teurer und das für reguläre Entwurf-Ziel-Modelle verwendete Schema funktioniert besser.

Referenzen

Fast Inference from Transformers via Speculative Decoding

EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty

EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees

EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test

Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads

Layer Normalization

FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving

DeepSeek-V3 Technical Report

Zitieren

@online{perplexity2025AcceleratingSonarThroughSpeculation,
  author = {{Perplexity} and Nandor Licker and Erin McKnight},
  title = {{Accelerating Sonar through speculation}},
  date = {2025-06-10},
  url = {https://perplexity.coolkeypoint.com/hub/blog/accelerating-sonar-through-speculation}
}