Accelerarea Sonar prin speculație

Decodarea speculativă accelerează viteza de generare a modelelor lingvistice de mari dimensiuni (LLM) utilizând un model preliminar rapid și mic pentru a produce candidați de completare care sunt verificați de modelul țintă mai mare. În cadrul acestei scheme, în loc ca o rulare a modelului țintă costisitor să producă un singur token, mai multe tokenuri sunt emise într-un singur pas. Aici prezentăm detaliile de implementare ale diverselor tipuri de decodare speculativă, aplicate la Perplexity pentru a reduce latența dintre tokenuri pentru modelele Sonar.

Decodare speculativă

Decodarea speculativă valorifică structura limbilor naturale și natura autoregresivă a transformatoarelor pentru a accelera generarea de tokenuri. Chiar dacă modelele mai mari, cum ar fi Llama-70B, dețin mai multe cunoștințe decât cele mai mici, cum ar fi Llama-1B, pe unele sarcini mai simple ele au o performanță similară. Această suprapunere sugerează că anumite secvențe sunt generate mai bine de modelele mai puțin costisitoare, lăsând problemele complexe în seama celor mai mari. Provocarea constă în determinarea completărilor care sunt mai bune și stabilirea dacă generarea modelului mai mic este de aceeași calitate ca cea a modelului mai mare.

Din fericire, modelele LLM sunt transformatoare autoregresive: atunci când li se oferă o secvență de tokenuri, ele scot în evidență distribuția de probabilitate a următorului token. În plus, loțiunile derivate din caracteristicile intermediare asociate cu tokenurile din secvența de intrare indică, de asemenea, cât de probabil este ca modelul să emită acele tokenuri exacte. Această proprietate permite speculația: dacă o secvență de tokenuri este generată de un model mai mic pornind de la un prefix de intrare, aceasta poate fi rulată prin modelul mai mare pentru a determina cât de bine se aliniază cu modelul țintă. Fiecare prefix al candidaților este punctat cu o probabilitate, iar cel mai lung prefix situat deasupra unui prag de acceptare este selectat. Ca un bonus, modelul țintă furnizează și un token subsecvent în mod gratuit: dacă un model preliminar generează n tokeni, până la n + 1 pot fi emiși într-un singur pas.

Extinde

În faza de inferență, procesul de eșantionare speculativă poate fi împărțit în aproximativ 4 etape:

Precompletare: atât modelul țintă, cât și modelul preliminar trebuie rulate pe secvența de intrare pentru a popula intrările din memoria cache KV. În timp ce unele scheme, cum ar fi Medusa, utilizează straturi dense mai simple pentru predicție, în această postare ne concentrăm pe modelele preliminare bazate pe transformatoare care își necesită propriile memorii cache KV.

Generarea preliminară: modelul preliminar efectuează iterații pentru a produce un număr de tokenuri fixe. Secvența preliminară poate fi liniară sau modelul poate explora o structură asemănătoare unui arbore până la o adâncime dată (EAGLE, Medusa). Aici ne concentrăm pe secvențele liniare.

Acceptare: modelul țintă rulează pe secvența preliminară, construind logite corespunzătoare fiecărui token preliminar. Se determină lungimea celei mai lungi secvențe acceptabile.

Generarea țintă: deoarece modelul țintă a generat logite, în poziția necorespunzătoare sau la sfârșitul secvenței, logitele corespund unui token subsecvent. Aceste logite pot fi eșantionate pentru a furniza un token robust din partea țintă, încheind secvența.

Există diverse metode de implementare a decodării speculative. În această postare, ne vom concentra pe schemele pe care le-am folosit pentru a accelera modelele Sonar folosind un model intern de 1B, precum și pe mecanismele de predicție pe care le construim pentru a accelera modelele la scara DeepSeek.

Țintă-Preliminar

Decodarea speculativă poate fi obținută prin cuplarea unui LLM mic existent ca model preliminar la un model țintă pentru a genera secvențe candidate. În producție, am accelerat Sonar folosind un model Llama-1B finisat pe același set de date ca și ținta. Deși această abordare nu a necesitat antrenarea unui model preliminar de la zero, modelul mic utilizează totuși o capacitate semnificativă a memoriei cache KV și introduce o ușoară suprasarcină de precompletare, mărind TTFT.

În cadrul acestei scheme, decodorul speculează numai pe loturi destinate exclusiv decodării, generând tokenuri prin eșantionare standard în timpul precompletării sau pe loturi mixte de precompletare și decodare. În etapa de precompletare, logitele țintă sunt eșantionate imediat pentru a precompleta, de asemenea, tokenul nou generat în memoria cache KV a modelului preliminar. Modelul preliminar nu este eșantionat încă, dar logitele pe care le produce sunt transferate în etapa de decodare.

Extinde

În decodare, modelul preliminar este avansat, eșantionând cel mai bun token în fiecare etapă. După ce se atinge lungimea preliminară dorită, tokenurile sunt trecute prin modelul țintă pentru a produce logitele pe baza cărora eșantionatorul identifică lungimea acceptată a secvenței. Acceptarea este determinată prin compararea distribuțiilor de probabilitate complete din partea modelului preliminar și a celui țintă. Deoarece modelul țintă emite întotdeauna un set de logite care urmează secvenței preliminare acceptate, acesta este eșantionat pentru a produce o ieșire suplimentară. Deoarece modelul preliminar nu a văzut încă acel token acceptat, el este rulat din nou pentru a popula intrările corespunzătoare din memoria cache KV în pregătirea următorului pas de decodare, transferând din nou logitele.

EAGLE

EAGLE este o schemă de decodare speculativă care explorează mai multe secvențe preliminare, generate printr-o parcurgere de tip arbore a tokenurilor preliminare probabile. Un arbore fix (EAGLE) sau cu formă dinamică (EAGLE-2) este explorat utilizând execuții consecutive ale tokenurilor preliminare, luând în considerare candidații Top-K la fiecare nod, în loc să urmeze tokenul cu cel mai mare scor dintr-o secvență liniară. Secvențele sunt apoi punctate și cea mai lungă secvență adecvată este selectată pentru continuare, adăugând de asemenea un token suplimentar din partea modelului țintă.

Extinde

Pentru a obține o predicție mai precisă, un model preliminar EAGLE prezice nu numai pe baza tokenurilor, ci și utilizând caracteristicile țintă (stările ascunse ale ultimului strat) ale modelului țintă. Dezavantajul EAGLE este necesitatea de a antrena modele preliminare personalizate, mici, care să fie suficient de precise pentru a genera candidați adecvați în limitele unui buget redus de latență. De obicei, un model preliminar este un singur strat de transformator identic cu un strat decodor al modelului original, care este cuplat strâns de țintă prin conectarea la încorporările sale și la proiecțiile lm_head. Deoarece acest lucru necesită mai puțină capacitate a memoriei cache KV, EAGLE are o amprentă de memorie mai mică.

Pentru a verifica secvențele asemănătoare unui arbore în modelul țintă, trebuie utilizate măști de atenție personalizate. Din păcate, utilizarea unei măști de atenție personalizate pentru o secvență întreagă încetinește semnificativ atenția pentru lungimi de intrare realiste (cu până la 50%), anulând o parte din accelerarea care poate fi obținută prin speculație. Din acest motiv, nu am implementat încă în producție explorarea completă a arborelui, concentrându-ne în schimb pe cazul special al predicției unui singur token prin scheme similare MTP prezentate în raportul tehnic DeepSeek-V3.

MTP

Această schemă este similară decodării preliminar-țintă, cu excepția faptului că stările ascunse sunt utilizate alături de tokenuri pentru predicție. Trebuie efectuat puțin mai mult efort atât în etapele de precompletare, cât și în cele de decodare, în comparație cu speculația regulată preliminar-țintă. Modelul preliminar utilizează atât tokenuri, cât și stări ascunse: tokenul t_{i+1} este eșantionat din logitele L_i corespunzătoare tokenului t_i, care la rândul lor sunt derivate din stările ascunse H_i. În consecință, tamponele de tokenuri de intrare trebuie deplasate cu un pas spre stânga în raport cu vectorii de stare ascunși emisi de modelul țintă. Figura de mai jos marchează corespondențele utilizate pentru antrenare, precum și deplasarea în timpul inferenței.

Extinde

Fluxul de decodare este destul de similar cu cel al decodării preliminar-țintă, cu excepția faptului că atât stările ascunse, cât și logitele sunt transferate. Implementarea noastră partajează toate logitele de eșantionare și de procesare a logitelor asociate, specializând doar apelurile înainte ale modelului. Atunci când sunt prezise mai multe tokenuri, modelul preliminar utilizează stări ascunse preliminare pentru predicție, populând de asemenea intrările din memoria cache KV pe baza propriilor caracteristici. Pe termen lung, acest lucru poate reduce precizia. Ulterior, când rulăm modelul preliminar pentru a popula intrarea din memoria cache KV pentru predicția țintă, îl rulăm pe întreaga secvență luând ca intrări stările ascunse țintă mai precise. Deoarece aceste modele preliminare sunt mici, costul suplimentar al procesării tokenurilor suplimentare este neglijabil.

Antrenarea capetelor MTP

Pentru a beneficia de MTP, am construit infrastructura necesară pentru a antrena capete MTP atașate la modelele noastre finisate pe seturile de date Perplexity, rulând pe un singur nod cu 8 dispozitive H100. În aproximativ o zi, putem construi capete pentru modele variind de la Llama-1B la Llama-70B și DeepSeek V2-Lite. Pentru modelele mai mari, ne bazăm pe capetele MTP construite în timpul procesului de finisare.

Scopul antrenamentului MTP este de a alinia stările ascunse preliminare și logitele extrapolate din stările ascunse țintă cu logitele și stările ascunse ale următorului token din modelul țintă. Deoarece inferența pentru stările ascunse este costisitoare, le precalculăm utilizând implementarea optimizată pentru inferență a modelului țintă, pentru a fi utilizate în timpul antrenamentului. Cu toate acestea, pentru a valida implementarea MTP pentru inferență și a ne asigura că diferențele numerice datorate cuantificării sau optimizărilor nu împiedică rezultatele, pentru pierderea de validare și estimarea preciziei reutilizăm pe deplin implementarea pentru inferență a ambelor modele, atât cel țintă, cât și cel preliminar.

La trecerea de la setul de date ShareGPT utilizat în lucrarea originală la eșantioane mai mari, am observat că arhitectura capului MTP conturată și implementată în lucrarea EAGLE nu a reușit să se antreneze pentru modele de dimensiunea 70B. Spre deosebire de ShareGPT, care conținea un număr mai mare de secvențe mai scurte, noi ne antrenăm pe un număr ușor mai mic de prompturi substanțial mai lungi. Deoarece capetele EAGLE originale au deviat ușor ca structură de la un transformator tipic, am reintrodus unele straturi de normalizare RMS care fuseseră eliminate. Am constatat că acest lucru nu numai că a permis convergența antrenamentului, dar a și crescut precizia capetelor cu câteva puncte procentuale.

Extinde

Normele de strat nu numai că facilitează antrenamentul, dar reintroducerea lor este și intuitivă din punct de vedere matematic. Capetele MTP reutilizează încorporările și proiecțiile de logite ale modelului țintă, deoarece acestea pot fi de dimensiuni considerabile (aproximativ 2 GB pentru Llama 70B). În timpul antrenamentului, acestea sunt înghețate, iar așteptarea este ca stratul MTP să învețe să încorporeze predicțiile în același spațiu vectorial pe care stratul de proiecție al modelului original l-a învățat în timpul antrenamentului. Prin renunțarea la norme, se așteaptă ca un singur MLP să învețe aceeași funcție ca un MLP urmat de o normă, ceea ce împiedică potrivirea dintre stările ascunse ale modelului preliminar și cele ale modelului țintă.

Inferența cu decodare speculativă

În motorul de inferență, pentru a genera tokenuri pentru secvențele de intrare, acestea trebuie mai întâi grupate în loturi de dimensiuni rezonabile, apoi trebuie alocate pagini în memoria cache KV pentru următoarele tokenuri. Tokenurile de intrare și informațiile despre paginile KV sunt apoi împachetate într-un tampon transmis prin difuzare (broadcast) către toate rândurile paralele care rulează modelul. În cele din urmă, metadatele sunt copiate în memoria GPU și modelul este executat pentru a produce logitele din care este eșantionat următorul token.

Spre deosebire de anumite implementări care cuplează lax un server de inferență preliminar și țintă printr-un intermediar care orchestrează cererile între ele, perechile noastre preliminar-țintă sunt cuplate strâns și parcurg generarea în unison. Planificarea loturilor și alocarea paginilor KV sunt partajate între modele pentru toate formele de decodare speculativă: acest lucru unifică logica care face legătura între un model și serverul de inferență general, deoarece toate expun aceeași interfață.

Timpul de execuție a inferenței la Perplexity este creat în jurul FlashInfer, care determină metadatele ce trebuie construite pentru a configura și planifica nucleul de atenție. Având în vedere unele secvențe de intrare care formează un lot, pentru precompletare, decodare sau verificare, trebuie efectuat un efort la nivel de CPU pentru a aloca tampoane intermediare și a popula anumite tampoane constante utilizate în atenție. Această muncă se adaugă costului de planificare a loturilor și de alocare a paginilor KV, care generează, de asemenea, latențe ce trebuie ascunse pentru a maximiza utilizarea GPU-ului.

În timp ce am paralelizat pe deplin lucrul la nivel de CPU și GPU pentru inferență fără speculație, am constatat că echilibrul CPU-GPU pentru decodarea speculativă este mai complex. Principala provocare decurge din faptul că numărul de tokenuri acceptate determină lungimea secvenței pentru o rulare ulterioară, introducând un punct de sincronizare GPU-CPU greu de evitat. Am experimentat cu diferite scheme de planificare pentru a ascunde cel mai bine latența muncii CPU.

Programul Preliminar-Țintă

În ciuda faptului că este mai mic decât un model țintă, atunci când un LLM întreg este utilizat ca model preliminar, acesta generează totuși o latență considerabilă pe GPU, oferind o anumită marjă pentru a ascunde operațiuni CPU costisitoare. Deoarece modelele mai mici nu beneficiază de paralelismul tensorilor, există o neconcordanță între numărul de rânduri pe care sunt divizate un model țintă și un model preliminar. În implementarea noastră, modelul preliminar rulează numai pe rândul principal al unui grup TP.

Extinde

Așa cum s-a menționat anterior, un pas de decodare transferă logitele în rularea următoare. Acest lucru ne permite să suprapunem o execuție a modelului preliminar cu activitatea de planificare a loturilor la nivel de CPU. După ce lotul este pus cap la cap, apelurile repetate către eșantionator și modelul preliminar produc tokenurile preliminare. În paralel, lotul pentru verificare este pus cap la cap pentru modelul țintă și sincronizat cu lucrătorii paraleli. Logitele țintă sunt verificate și eșantionate pentru a determina lungimile acceptate ale secvenței. În acest moment, sincronizarea GPU-CPU este necesară pentru a determina lungimile ulterioare ale secvenței. Deoarece modelul preliminar rulează numai pe nodul principal, lotul său este configurat secvențial, iar execuția sa este lansată pentru a popula intrările din memoria cache KV cu tokenul suplimentar pe care l-a produs modelul țintă. Logitele produse de această rulare preliminară în rularea curentă vor fi utilizate pentru a eșantiona primul token preliminar în rularea ulterioară. Cel mai important, în timp ce modelul preliminar rulează, următorul lot poate fi planificat.

Program MTP pentru un singur token

În timp ce timpul de execuție nu oferă încă o explorare a arborelui preliminar în stil Eagle, am implementat un caz special al acestei scheme, luând în considerare o secvență liniară de tokenuri preliminare produsă de un model de dimensiunea unui singur strat decodor de transformator. Această schemă poate fi utilizată pentru predicția preliminară folosind ponderile open-source ale DeepSeek R1. Subcazul prezicerii unui singur token este interesant, deoarece straturile MTP mari ating rate de acceptare suficient de ridicate pentru a justifica suprasarcinile lor.

Planificarea MTP este oarecum mai complexă, deoarece modelul preliminar este mult mai rapid, ascunzând mai puțină latență la nivel de CPU. În plus, modelul preliminar este fragmentat alături de modelul țintă, necesitând transferuri de memorie partajată pentru informațiile despre lot. O rulare începe prin transferul informațiilor despre lot și eșantionarea primului token din logitele transferate, în mod similar schemei anterioare. În continuare, modelul țintă este rulat pentru a valida tokenurile, procesând 2 * D tokeni, unde D este dimensiunea lotului de decodare. Acest lucru este ideal pentru micro-batched în modelele Mixture-of-Experts (MoE) prin interconexiuni mai lente, cum ar fi InfiniBand, deoarece lotul se împarte în mod egal în două jumătăți. Stările ascunse ale modelului țintă sunt transferate la următoarea rulare preliminară, în timp ce logitele sunt transmise în eșantionator pentru verificare.

Extinde

Prin efectuarea unei cantități limitate de lucru suplimentar pe GPU, evităm sincronizarea CPU-GPU după acceptarea secvenței preliminare. După ce tokenurile de intrare ale modelelor țintă sunt deplasate, un nucleu introduce următoarele tokenuri țintă în locațiile lor corespunzătoare. Modelul preliminar este apoi rulat din nou cu aceleași informații despre lot ca și modelul țintă, populând intrările din memoria cache KV și construind logitele și stările ascunse pentru următoarea rulare, efectuând o muncă redundantă asupra tokenurilor care nu au fost acceptate. În aceste situații, latența lucrului neutilizat este abia măsurabilă datorită dimensiunii mici a modelului preliminar. În paralel cu rularea modelului preliminar, lungimile secvenței sunt determinate pe CPU, iar planificarea următorului lot este lansată, fără a fi nevoie să se aștepte terminarea lucrului pe GPU.

Suprasarcinile de lucru suplimentar din stratul preliminar nu sunt vizibile în atenție, însă straturile MLP sunt mai problematice. Deoarece instrucțiunile de înmulțire a matrioștelor se aliniază la o limită de 64 de-a lungul dimensiunii numărului de tokenuri, dacă dublarea nu necesită semnificativ mai multe blocuri, suprasarcina este ascunsă. Pentru secvențe preliminare mai lungi, suprasarcina este mai costisitoare, iar schema utilizată pentru modelele obișnuite preliminar-țintă funcționează mai bine.

Referințe

Inferență rapidă din transformatoare prin decodare speculativă

EAGLE: Eșantionarea speculativă necesită regândirea incertitudinii caracteristicilor

EAGLE-2: Inferență mai rapidă a modelelor lingvistice cu arbori preliminari dinamici

EAGLE-3: Scalarea accelerării inferenței modelelor lingvistice de mari dimensiuni prin testare în timpul antrenamentului

Medusa: Cadru simplu de accelerare a inferenței LLM cu mai multe capete de decodare

Normalizarea stratului

FlashInfer: Motor de atenție eficient și personalizabil pentru servirea inferenței LLM

Raport tehnic DeepSeek-V3

Citare

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