Accelerare Sonar tramite la speculazione

La decodifica speculativa accelera la velocità di generazione dei Large Language Models (LLM) utilizzando un modello di bozza rapido e di dimensioni ridotte per produrre candidati di completamento che vengono verificati dal modello target più grande. Con questo schema, invece di una singola esecuzione del costoso modello target per produrre un solo token, ne vengono emessi molteplici in un unico passaggio. Qui presentiamo i dettagli di implementazione di vari tipi di decodifica speculativa, applicati presso Perplexity per ridurre la latenza inter-token sui modelli Sonar.

Decodifica speculativa

La decodifica speculativa sfrutta la struttura delle lingue naturali e la natura autoregressiva dei transformer per accelerare la generazione dei token. Sebbene i modelli più grandi, come Llama-70B, possiedano maggiori conoscenze rispetto a quelli più piccoli, come Llama-1B, su alcuni compiti più semplici le loro prestazioni sono simili. Questa sovrapposizione suggerisce che determinate sequenze siano generate meglio da modelli meno costosi, lasciando i problemi complessi a quelli più grandi. La sfida risiede nel determinare quali completamenti siano migliori e se la generazione del modello più piccolo sia della stessa qualità di quella del modello più grande.

Fortunatamente, gli LLM sono transformer autoregressivi: quando viene fornita una sequenza di token, essi restituiscono la distribuzione di probabilità del token successivo. Inoltre, i logit derivati dalle caratteristiche intermedie associate ai token nella sequenza di input indicano anche quanto sia probabile che il modello emetta esattamente quei token. Questa proprietà abilita la speculazione: se una sequenza di token viene generata da un modello più piccolo partendo da un prefisso di input, può essere eseguita tramite il modello più grande per determinare quanto bene si allinei con il modello target. Ogni prefisso dei candidati viene valutato con una probabilità e viene scelto quello più lungo al di sopra di una soglia di accettazione. Come bonus, il modello target fornisce anche un token successivo gratuitamente: se un modello di bozza genera n token, è possibile emetterne fino a n + 1 in un unico passaggio.

Espandi

In fase di inferenza, il processo di campionamento speculativo può essere suddiviso approssimativamente in 4 stadi:

Prefill: sia il modello target che quello di bozza devono essere eseguiti sulla sequenza di input per popolare le voci della cache KV. Mentre alcuni schemi, come Medusa, utilizzano strati densi più semplici per la previsione, in questo post ci concentriamo su bozze basate su transformer che necessitano delle proprie cache KV.

Generazione della bozza: il modello di bozza itera per produrre un numero di token fissi. La sequenza della bozza può essere lineare o il modello può esplorare una struttura ad albero fino a una data profondità (EAGLE, Medusa). Qui, ci concentriamo sulle sequenze lineari.

Accettazione: il modello target viene eseguito sulla sequenza della bozza, costruendo i logit corrispondenti a ciascun token della bozza. Viene determinata la lunghezza della sequenza accettabile più lunga.

Generazione target: poiché il target ha generato dei logit, nella posizione di mancata corrispondenza o alla fine della sequenza i logit corrispondono a un token successivo. Questi logit possono essere campionati per fornire un token robusto dal target, completando la sequenza.

Esistono vari metodi per implementare la decodifica speculativa. In questo post, ci concentreremo sugli schemi che abbiamo utilizzato per accelerare i modelli Sonar utilizzando un modello 1B interno, nonché sui meccanismi di previsione che stiamo sviluppando per accelerare i modelli su scala DeepSeek.

Target-Bozza

La decodifica speculativa può essere ottenuta accoppiando un piccolo LLM esistente come modello di bozza a un modello target per generare sequenze candidate. In produzione, abbiamo accelerato Sonar utilizzando un modello Llama-1B messo a punto sullo stesso set di dati del target. Sebbene questo approccio non abbia richiesto l'addestramento di una bozza da zero, il modello piccolo utilizza comunque una capacità significativa della cache KV e introduce una leggera sovrapposizione di prefill, aumentando il TTFT.

Con questo schema, il decodificatore specula solo su batch di sola decodifica, generando token tramite campionamento standard durante il prefill o su batch misti di prefill-decodifica. Nella fase di prefill, i logit del target vengono campionati immediatamente per pre-riempire anche il token appena generato nella cache KV della bozza. La bozza non viene ancora campionata, ma i logit che produce vengono riportati alla fase di decodifica.

Espandi

Nella decodifica, il modello di bozza viene avanzato, campionando il token principale in ogni stadio. Dopo aver raggiunto la lunghezza della bozza desiderata, i token vengono eseguiti attraverso il modello target per produrre i logit in base ai quali il campionatore identifica la lunghezza della sequenza accettata. L'accettazione è determinata dal confronto tra le distribuzioni di probabilità complete della bozza e del target. Poiché il target emette sempre una serie di logit dopo la sequenza di bozza accettata, questa viene campionata per produrre un output aggiuntivo. Poiché il modello di bozza non ha ancora visto quel token accettato, viene rieseguito per popolare le sue voci della cache KV corrispondenti in preparazione per il successivo passaggio di decodifica, riportando nuovamente i logit.

EAGLE

EAGLE è uno schema di decodifica speculativa che esplora molteplici sequenze di bozza, generate tramite una traversata ad albero dei token di bozza probabili. Un albero fisso (EAGLE) o a forma dinamica (EAGLE-2) viene esplorato utilizzando esecuzioni consecutive dei token di bozza, considerando i candidati Top-K in ogni nodo invece di seguire il token con il punteggio più alto in una sequenza lineare. Le sequenze vengono quindi valutate e viene selezionata la più lunga adatta per continuare, aggiungendo anche un token aggiuntivo dal target.

Espandi

Per ottenere una previsione più accurata, un modello di bozza EAGLE non effettua previsioni solo basandosi sui token, ma utilizza anche le caratteristiche target (stati nascosti dell'ultimo livello) del modello target. Lo svantaggio di EAGLE è la necessità di addestrare modelli di bozza piccoli e personalizzati che siano abbastanza accurati da generare candidati adatti entro un basso budget di latenza. Tipicamente, un modello di bozza è un singolo strato transformer identico a uno strato di decodifica del modello originale, strettamente accoppiato al target legandolo ai suoi embedding e alle proiezioni lm_head. Poiché ciò richiede meno capacità della cache KV, EAGLE ha un ingombro di memoria inferiore.

Per verificare sequenze simili ad alberi nel modello target, è necessario utilizzare maschere di attenzione personalizzate. Sfortunatamente, l'utilizzo di una maschera di attenzione personalizzata per un'intera sequenza rallenta significativamente l'attenzione per lunghezze di input realistiche (fino al 50%), annullando parte dell'accelerazione ottenibile tramite speculazione. Non abbiamo ancora implementato l'esplorazione ad albero completa in produzione per questo motivo, concentrandoci invece sul caso speciale della previsione a singolo token tramite schemi simili a MTP presentati nel rapporto tecnico di DeepSeek-V3.

MTP

Questo schema è simile alla decodifica bozza-target, con l'eccezione che gli stati nascosti vengono utilizzati insieme ai token per la previsione. È necessario svolgere un po' più di lavoro sia nelle fasi di prefill che di decodifica rispetto alla normale speculazione bozza-target. Il modello di bozza utilizza sia token che stati nascosti: il token t_{i+1} viene campionato dai logit L_i corrispondenti al token t_i, che a loro volta sono derivati dagli stati nascosti H_i. Di conseguenza, i buffer dei token di input devono essere spostati di un passo a sinistra rispetto ai vettori di stato nascosti emessi dal target. La figura sottostante segna le corrispondenze utilizzate per l'addestramento, così come lo spostamento durante l'inferenza.

Espandi

Il flusso di decodifica è abbastanza simile alla decodifica bozza-target, con l'eccezione che sia gli stati nascosti che i logit vengono trasportati. La nostra implementazione condivide tutti i campionamenti associati e i logit di elaborazione dei logit, specializzando solo le invocazioni in avanti del modello. Quando vengono previsti più token, il modello di bozza utilizza stati nascosti di bozza per la previsione, popolando anche le voci della cache KV in base alle proprie caratteristiche. A lungo termine, ciò può degradare l'accuratezza. Successivamente, quando si esegue il modello di bozza per popolare la voce della cache KV per la previsione del target, lo eseguiamo sull'intera sequenza prendendo gli stati nascosti target più accurati come input. Poiché questi modelli di bozza sono piccoli, il costo aggiunto dell'elaborazione dei token aggiuntivi è trascurabile.

Addestramento delle teste MTP

Per beneficiare di MTP, abbiamo costruito l'infrastruttura necessaria per addestrare le teste MTP collegate ai nostri modelli messi a punto sui set di dati di Perplexity, in esecuzione su un nodo con dispositivi 8xH100. In circa un giorno, possiamo costruire teste per modelli che vanno da Llama-1B a Llama-70B e DeepSeek V2-Lite. Per i modelli più grandi, ci affidiamo alle teste MTP costruite durante il processo di messa a punto.

L'obiettivo dell'addestramento MTP è far corrispondere gli stati nascosti della bozza e i logit estrapolati dagli stati nascosti del target ai logit del token successivo e agli stati nascosti del target. Poiché l'inferenza per gli stati nascosti è costosa, li pre-calcoliamo utilizzando la nostra implementazione del modello target ottimizzata per l'inferenza, da utilizzare durante l'addestramento. Tuttavia, per convalidare l'implementazione MTP di inferenza e garantire che le differenze numeriche dovute alla quantizzazione o alle ottimizzazioni non ostacolino i risultati, per la perdita di convalida e la stima dell'accuratezza riutilizziamo completamente l'implementazione di inferenza sia del modello target che di quello di bozza.

Quando si è passati dal set di dati ShareGPT utilizzato nel documento originale a campioni più grandi, abbiamo notato che l'architettura della testa MTP delineata e implementata nel documento EAGLE non è riuscita ad addestrare per modelli di dimensioni 70B. A differenza di ShareGPT, che conteneva un numero maggiore di sequenze più brevi, ci addestriamo su un numero leggermente inferiore di prompt sostanzialmente più lunghi. Poiché le teste EAGLE originali divergevano leggermente nella struttura da un tipico transformer, abbiamo reintrodotto alcuni strati di Normalizzazione RMS che erano stati rimossi. Abbiamo scoperto che ciò non solo ha permesso all'addestramento di convergere, ma ha anche aumentato l'accuratezza delle teste di alcuni punti percentuali.

Espandi

Non solo le norme di strato facilitano l'addestramento, ma la reintroduzione delle norme è anche matematicamente intuitiva. Le teste MTP riutilizzano gli embedding e le proiezioni dei logit del modello target, poiché possono essere di dimensioni sostanziali (circa 2 GB per Llama 70B). Durante l'addestramento, questi vengono congelati e l'aspettativa è che lo strato MTP impari a incorporare le previsioni nello stesso Spazio vettoriale di ciò che lo strato di proiezione del modello originale ha appreso durante l'addestramento. Eliminando le norme, si prevede che un singolo MLP impari la stessa funzione di un MLP seguito da una norma, il che ostacola il confronto tra gli stati nascosti della bozza e dei modelli target.

Inferenza con decodifica speculativa

Nel motore di inferenza, per generare token per sequenze di input, devono essere prima raggruppati in batch di dimensioni ragionevoli, quindi le pagine devono essere allocate nella cache KV per i token successivi. I token di input e le informazioni sulla pagina KV vengono quindi compressi in un buffer trasmesso a tutti i ranghi paralleli che eseguono il modello. Infine, i metadati vengono copiati nella memoria della GPU e il modello viene eseguito per produrre i logit da cui viene campionato il token successivo.

A differenza di alcune implementazioni che accoppiano liberamente un server di inferenza di bozza e target tramite un wrapper che orchestra le richieste tra loro, le nostre coppie bozza-target sono strettamente accoppiate e procedono nella generazione all'unisono. La pianificazione dei batch e l'allocazione delle pagine KV sono condivise tra i modelli per tutte le forme di decodifica speculativa: questo unifica la logica che collega un modello al server di inferenza generale, poiché tutti espongono la stessa interfaccia.

Il runtime di inferenza presso Perplexity è modellato attorno a FlashInfer, che determina i metadati che devono essere costruiti per configurare e pianificare il kernel di attenzione. Date alcune sequenze di input che formano un batch, per il prefill, la decodifica o la verifica, deve essere svolto lavoro lato CPU per allocare buffer intermedi e popolare alcuni buffer costanti utilizzati nell'attenzione. Questo lavoro si aggiunge al costo della pianificazione dei batch e dell'allocazione delle pagine KV, che comportano anche latenze che devono essere nascoste per massimizzare l'utilizzo della GPU.

Mentre abbiamo completamente parallelizzato il lavoro lato CPU e lato GPU per l'inferenza senza speculazione, abbiamo riscontrato che l'equilibrio CPU-GPU per la decodifica speculativa è più complesso. La sfida principale nasce dal fatto che il numero di token accettati determina la lunghezza della sequenza per un'esecuzione successiva, introducendo un punto di sincronizzazione GPU-CPU difficile da evitare. Abbiamo sperimentato diversi schemi di pianificazione per nascondere al meglio la latenza del lavoro della CPU.

Programma bozza-target

Nonostante sia più piccolo di un modello target, quando un intero LLM viene utilizzato come bozza, introduce comunque una latenza considerevole sulla GPU, fornendo un certo margine per nascondere costose operazioni della CPU. Poiché i modelli più piccoli non beneficiano del parallelismo tensoriale, c'è una discrepanza tra il numero di ranghi su cui sono suddivisi un target e una bozza. Nella nostra implementazione, il modello di bozza viene eseguito solo sul rango leader di un gruppo TP.

Espandi

Come indicato in precedenza, un passaggio di decodifica trasporta i logit nell'esecuzione successiva. Ciò ci consente di sovrapporre un'esecuzione del modello di bozza al lavoro di pianificazione dei batch lato CPU. Dopo aver assemblato il batch, chiamate ripetute al campionatore e alla bozza producono i token della bozza. In parallelo, il batch per la verifica viene assemblato per il modello target e sincronizzato con i lavoratori paralleli. I logit del target vengono verificati e campionati per determinare le lunghezze delle sequenze accettate. A questo punto, è necessaria la sincronizzazione GPU-CPU per determinare le lunghezze delle sequenze successive. Poiché il modello di bozza viene eseguito solo sul nodo leader, il suo batch viene impostato sequenzialmente e la sua esecuzione viene avviata per popolare le sue voci della cache KV con il token aggiuntivo che il target ha prodotto. I logit prodotti da questa esecuzione di bozza nell'esecuzione corrente verranno utilizzati per campionare il primo token di bozza nell'esecuzione successiva. Soprattutto, mentre la bozza è in esecuzione, il batch successivo può essere pianificato.

Programma MTP per un singolo token

Mentre il runtime non fornisce ancora l'esplorazione dell'albero delle bozze in stile Eagle, abbiamo implementato un caso speciale di questo schema, considerando una sequenza lineare di token di bozza prodotti da un modello delle dimensioni di un singolo strato di decodifica transformer. Questo schema può essere utilizzato per la previsione della bozza utilizzando i pesi open-source di DeepSeek R1. Il sottocaso della previsione di un singolo token è interessante, poiché i grandi strati MTP raggiungono tassi di accettazione sufficientemente elevati da giustificare il loro sovraccarico.

La pianificazione MTP è alquanto più complessa, poiché il modello di bozza è molto più veloce, nascondendo meno latenza lato CPU. Inoltre, la bozza è suddivisa insieme al modello target, richiedendo trasferimenti di memoria condivisa per le informazioni sui batch. Un'esecuzione inizia trasferendo le informazioni sul batch e campionando il primo token dai logit riportati, in modo simile allo schema precedente. Successivamente, il target viene eseguito per convalidare i token, elaborando 2 * D token, dove D è la dimensione del batch di decodifica. Questo è ideale per il micro-batching nei modelli Mixture-of-Experts (MoE) su interconnessioni più lente come InfiniBand, poiché il batch si divide equamente in due metà. Gli stati nascosti del target vengono riportati alla successiva esecuzione della bozza, mentre i logit vengono passati al campionatore per la verifica.

Espandi

Eseguendo una quantità limitata di lavoro aggiuntivo sulla GPU, evitiamo la sincronizzazione CPU-GPU dopo l'accettazione della sequenza di bozza. Dopo che i token di input dei target vengono spostati, un kernel collega i token target successivi nelle loro posizioni corrispondenti. La bozza viene quindi rieseguita con le stesse informazioni di batch del target, popolando le voci della cache KV e costruendo i logit e gli stati nascosti per l'esecuzione successiva, svolgendo un po' di lavoro ridondante sui token che non sono stati accettati. In queste situazioni, la latenza del lavoro inutilizzato è a malapena misurabile a causa delle piccole dimensioni del modello di bozza. In parallelo con l'esecuzione della bozza, le lunghezze delle sequenze vengono determinate sulla CPU e viene avviata la pianificazione del batch successivo, senza dover attendere il termine del lavoro della GPU.

Il sovraccarico di lavoro aggiuntivo nello strato di bozza non è evidente nell'attenzione, tuttavia gli strati MLP sono più problematici. Poiché le istruzioni di moltiplicazione di matrici si riempiono fino a un confine di 64 lungo la dimensione del numero di token, se il raddoppio non richiede significativamente più blocchi, il sovraccarico viene nascosto. Per sequenze di bozza più lunghe il sovraccarico è più costoso e lo schema utilizzato per i normali modelli bozza-target funziona meglio.

Riferimenti

Inferenza rapida da Transformer tramite decodifica speculativa

EAGLE: Il campionamento speculativo richiede di ripensare l'incertezza delle caratteristiche

EAGLE-2: Inferenza più rapida dei modelli linguistici con alberi di bozza dinamici

EAGLE-3: Scalare l'accelerazione dell'inferenza di Large Language Models tramite test in fase di addestramento

Medusa: Semplice framework di accelerazione dell'inferenza LLM con teste di decodifica multiple

Normalizzazione dello strato

FlashInfer: Motore di attenzione efficiente e personalizzabile per l'inferenza LLM

Rapporto tecnico DeepSeek-V3

Citazione

@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}
}