Ubrzavanje Sonara kroz spekulaciju

Spekulativno dekodiranje ubrzava brzinu generisanja velikih jezičkih modela (LLM) korišćenjem brzog i malog nacrtnog modela za proizvodnju kandidata za završetak koje verifikuje veći ciljni model. Pod ovom šemom, umesto pokretanja skupog ciljnog modela koji proizvodi jedan token, više njih se emituje u jednom koraku. Ovde predstavljamo detalje implementacije različitih vrsta spekulativnog dekodiranja, primenjenih na Perplexity-ju radi smanjenja latencije između tokena na Sonar modelima.

Spekulativno dekodiranje

Spekulativno dekodiranje koristi strukturu prirodnih jezika i autoregresivnu prirodu transformera radi ubrzanja generisanja tokena. Iako veći modeli, kao što je Llama-70B, nose više znanja od manjih, kao što je Llama-1B, na nekim jednostavnijim zadacima oni imaju slične performanse. Ovo preklapanje sugeriše da određene sekvence bolje generišu jeftiniji modeli, prepuštajući složene probleme većima. Izazov leži u utvrđivanju koji završeci su bolji i da li je generisanje manjeg modela istog kvaliteta kao i kod većeg.

Srećom, LLM-ovi su autoregresivni transformeri: kada im se zadat niz tokena, oni daju raspodelu verovatnoće sledećeg tokena. Pored toga, logiti izvedeni iz srednjih karakteristika povezanih sa tokenima u ulaznom nizu takođe ukazuju na to kolika je verovatnoća da model izda upravo te tokene. Ovo svojstvo omogućava spekulaciju: ako niz tokena generiše manji model počevši od ulaznog prefiksa, on se može proslediti kroz veći model kako bi se utvrdilo koliko se dobro poklapa sa ciljnim modelom. Svaki prefiks kandidata se ocenjuje verovatnoćom i bira se najduži iznad praga prihvatanja. Kao bonus, ciljni model takođe besplatno obezbeđuje naknadni token: ako nacrtni model generiše n tokena, do n + 1 se može emitovati u jednom koraku.

Uvećaj

U vreme inferencije, proces spekulativnog uzorkovanja se može grubo podeliti u 4 faze:

Predpopunjavanje (Prefill): i ciljni i nacrtni model moraju se pokrenuti na ulaznom nizu da bi se popunile unosi u KV kešu. Dok neke šeme, kao što je Medusa, koriste jednostavnije guste slojeve za predviđanje, u ovom tekstu se fokusiramo na nacrte bazirane na transformerima kojima su potrebni sopstveni KV keševi.

Generisanje nacrta: nacrtni model iterira da bi proizveo određen broj fiksnih tokena. Nacrtna sekvenca može biti linearna ili model može istražiti strukturu nalik drvetu do date dubine (EAGLE, Medusa). Ovde se fokusiramo na linearne sekvence.

Prihvatanje: ciljni model se pokreće na nacrtnoj sekvenci, gradeći logite koji odgovaraju svakom nacrtnom tokenu. Utvrđuje se dužina najduže prihvatljive sekvence.

Ciljno generisanje: pošto je ciljni model generisao logite, na poziciji neusklađenosti ili na kraju sekvence logiti odgovaraju naknadnom tokenu. Ovi logiti se mogu uzorkovati kako bi se obezbedio robustan token iz ciljnog modela, zaključujući sekvencu.

Postoje različite metode za implementaciju spekulativnog dekodiranja. U ovom tekstu ćemo se fokusirati na šeme koje smo koristili za ubrzavanje Sonar modela pomoću internog 1B modela, kao i na mehanizme predviđanja koje gradimo da bismo ubrzali modele u razmeri DeepSeek-a.

Cilj-Nacrt (Target-Draft)

Spekulativno dekodiranje se može postići uparivanjem postojećeg malog LLM-a kao nacrtnog modela sa ciljnim modelom radi generisanja kandidata sekvenci. U proizvodnji, ubrzali smo Sonar pomoću Llama-1B modela fino podešenog na istom skupu podataka kao i ciljni model. Iako ovaj pristup nije zahtevao obuku nacrta od nule, mali model i dalje koristi značajan kapacitet KV keša i uvodi blago opterećenje predpopunjavanja, povećavajući TTFT.

Pod ovom šemom, dekoder spekuliše samo na serijama koje služe samo za dekodiranje (decode-only), generišući tokene kroz standardno uzorkovanje tokom predpopunjavanja ili na mešovitim serijama predpopunjavanje-dekodiranje. U fazi predpopunjavanja, ciljni logiti se odmah uzorkuju kako bi se takođe predpopunio novogenerisani token u KV kešu nacrta. Nacrt se još uvek ne uzorkuje, ali se logiti koje proizvodi prenose u fazu dekodiranja.

Uvećaj

U dekodiranju, nacrtni model napreduje, uzorkujući vodeći token u svakoj fazi. Nakon što se dostigne željena dužina nacrta, tokeni se prosleđuju kroz ciljni model da bi se proizveli logiti na osnovu kojih uzorkovač identifikuje prihvaćenu dužinu sekvence. Prihvatanje se utvrđuje upoređivanjem potpunih raspodela verovatnoće iz nacrta i ciljnog modela. Pošto ciljni model uvek daje jedan skup logita koji prati prihvaćenu nacrtnu sekvencu, on se uzorkuje da bi se proizveo dodatni izlaz. Pošto nacrtni model još uvek nije video taj prihvaćeni token, on se ponovo pokreće da bi popunio odgovarajuće unose u KV kešu pripremajući se za sledeći korak dekodiranja, ponovo prenoseći logite.

EAGLE

EAGLE je šema spekulativnog dekodiranja koja istražuje više nacrtnih sekvenci, generisanih kroz obilazak verovatnih nacrtnih tokena nalik stablu. Fiksno (EAGLE) ili dinamički oblikovano (EAGLE-2) stablo se istražuje uzastopnim izvršavanjima nacrtnih tokena, razmatrajući Top-K kandidate na svakom čvoru umesto praćenja tokena sa najvišom ocenom u linearnoj sekvenci. Sekvence se zatim ocenjuju i bira se najduža pogodna sekvenca za nastavak, uz dodavanje dodatnog tokena iz ciljnog modela.

Uvećaj

U cilju postizanja preciznijeg predviđanja, EAGLE nacrtni model predviđa ne samo na osnovu tokena, već i korišćenjem ciljnih karakteristika (skrivenih stanja poslednjeg sloja) ciljnog modela. Nedostatak EAGLE-a je potreba za obukom prilagođenih, malih nacrtnih modela koji su dovoljno precizni da generišu pogodne kandidate unutar budžeta niske latencije. Obično je nacrtni model jedinstveni transformerski sloj identičan sloju dekodera originalnog modela, koji je čvrsto povezan sa ciljnim modelom vezivanjem za njegove ugradnje (embeddings) i lm_head projekcije. Pošto ovo zahteva manji kapacitet KV keša, EAGLE ima manji memorijski otisak.

Da bi se proverile sekvence nalik stablu u ciljnom modelu, moraju se koristiti prilagođene maske pažnje. Nažalost, korišćenje prilagođene maske pažnje za čitavu sekvencu značajno usporava pažnju za realistične dužine unosa (za do 50%), poništavajući deo ubrzanja koje se može postići spekulacijom. Iz tog razloga još uvek nismo primenili puno istraživanje stabla u proizvodnji, već se fokusiramo na specijalan slučaj predviđanja jednog tokena preko MTP-sličnih šema predstavljenih u tehničkom izveštaju DeepSeek-V3.

MTP

Ova šema je slična dekodiranju nacrt-cilj, sa izuzetkom što se skrivena stanja koriste uz tokene za predviđanje. Mora se obaviti nešto više posla i u fazi predpopunjavanja i u fazi dekodiranja u poređenju sa običnom nacrt-cilj spekulacijom. Nacrtni model koristi i tokene i skrivena stanja: token t_{i+1} se uzorkuje iz logita L_i koji odgovaraju tokenu t_i, koji se zauzvrat izvode iz skrivenih stanja H_i. S tim u vezi, baferi ulaznih tokena moraju biti pomaknuti za jedan korak ulevo u odnosu na vektore skrivenih stanja koje daje ciljni model. Slika ispod označava korelacije korišćene za obuku, kao i pomeranje tokom inferencije.

Uvećaj

Tok dekodiranja je prilično sličan dekodiranju nacrt-cilj, sa izuzetkom što se prenose i skrivena stanja i logiti. Naša implementacija deli sve povezane logite uzorkovanja i obrade logita, specijalizujući samo pozive prosleđivanja modela (forward). Kada se predviđa više tokena, nacrtni model koristi skrivena stanja nacrta za predviđanje, takođe popunjavajući unose KV keša na osnovu sopstvenih karakteristika. Dugoročno gledano, ovo može narušiti tačnost. Naknadno, pri pokretanju nacrtnog modela radi popunjavanja unosa KV keša za ciljno predviđanje, pokrećemo ga na čitavoj sekvenci uzimajući preciznija skrivena stanja ciljnog modela kao ulaze. Pošto su ovi nacrtni modeli mali, dodatni trošak obrade dodatnih tokena je zanemarljiv.

Obuka MTP glava

U cilju ostvarivanja koristi od MTP-a, izgradili smo infrastrukturu potrebnu za obuku MTP glava povezanih sa našim fino podešenim modelima na Perplexity-jevim skupovima podataka, koja se izvršava na jednom čvoru sa 8xH100 uređaja. Za otprilike jedan dan možemo izgraditi glave za modele u rasponu od Llama-1B do Llama-70B i DeepSeek V2-Lite. Za veće modele oslanjamo se na MTP glave izgrađene tokom procesa finog podešavanja.

Cilj MTP obuke je da se usklade skrivena stanja nacrta i logiti ekstrapolirani iz skrivenih stanja ciljnog modela sa logitima sledećeg tokena i skrivenim stanjima ciljnog modela. Pošto je inferencija za skrivena stanja skupa, mi ih unapred izračunavamo koristeći našu implementaciju ciljnog modela optimizovanu za inferenciju, koja će se koristiti tokom obuke. Međutim, da bismo potvrdili MTP implementaciju inferencije i osigurali da numeričke razlike usled kvantizacije ili optimizacija ne ometaju rezultate, za gubitak validacije i procenu tačnosti u potpunosti ponovo koristimo implementaciju inferencije i ciljnog i nacrtnog modela.

Prilikom skaliranja sa ShareGPT skupa podataka korišćenog u originalnom radu na veće uzorke, primetili smo da arhitektura MTP glave navedena i implementirana u EAGLE radu nije uspela da se obučava za modele veličine 70B. Za razliku od ShareGPT-a koji je sadržao veći broj kraćih sekvenci, mi vršimo obuku na nešto manjem broju znatno dušćih prompova. Pošto su se originalne EAGLE glave neznatno razlikovale po strukturi od tipičnog transformera, ponovo smo uveli neke slojeve RMS normalizacije koji su bili uklonjeni. Utvrdili smo da ovo nije samo omogućilo konvergenciju obuke, već je i povećalo tačnost glava za nekoliko procentnih poena.

Uvećaj

Ne samo da normalizacije slojeva olakšavaju obuku, već je ponovno uvođenje normalizacija i matematički intuitivno. MTP glave ponovo koriste ugradnje i projekcije logita ciljnog modela, pošto One mogu biti znatne veličine (oko 2 GB za Llama 70B). Tokom obuke, one su zamrznute i očekivanje je da MTP sloj uči da ugradi predviđanja u isti vektorski prostor koji je sloj projekcije originalnog modela naučio tokom obuke. Izbacivanjem normalizacija, očekuje se da jedan MLP nauči istu funkciju kao MLP praćen normalizacijom, što otežava usklađivanje između skrivenih stanja nacrtnog i ciljnog modela.

Inferencija sa spekulativnim dekodiranjem

U mehanizmu inferencije, da bi se generisali tokeni za ulazne sekvence, one se prvo moraju grupisati u serije razumne veličine, a zatim se moraju alocirati stranice u KV kešu za sledeće tokene. Ulazni tokeni i informacije o KV stranicama se zatim pakuju u bafer koji se emituje na sve paralelne rangove koji izvršavaju model. Na kraju, metapodaci se kopiraju u GPU memoriju i model se izvršava radi proizvodnje logita iz kojih se uzorkuje sledeći token.

Za razliku od određenih implementacija koje labavo povezuju server inferencije nacrta i ciljnog modela preko omotača (wrapper) koji upravlja zahtevima između njih, naši parovi nacrt-cilj su čvrsto povezani i prolaze kroz generisanje u unisonu. Raspoređivanje serija i alokacija KV stranica se dele između modela za sve oblike spekulativnog dekodiranja: ovo unifikuje logiku koja povezuje model sa sveobuhvatnim serverom inferencije, pošto svi oni izlažu isti interfejs.

Vreme izvršavanja inferencije na Perplexity-ju je oblikovano oko FlashInfer-a, koji određuje metapodatke koje je potrebno izgraditi radi konfigurisanja i raspoređivanja jezgra pažnje (attention kernel). S obzirom na neke ulazne sekvence koje čine seriju, za predpopunjavanje, dekodiranje ili verifikaciju, rad na strani CPU-a se mora obaviti kako bi se alocirali srednji baferi i popunili određeni konstantni baferi koji se koriste u pažnji. Ovaj rad je dodat trošku raspoređivanja serija i alokacije KV stranica, koji takođe izazivaju latencije koje se moraju sakriti da bi se maksimiziralo iskorišćenje GPU-a.

Iako smo u potpunosti paralelizovali rad na strani CPU-a i GPU-a za inferenciju bez spekulacije, otkrili smo da je balans između CPU-a i GPU-a za spekulativno dekodiranje složeniji. Glavni izazov proizilazi iz činjenice da broj prihvaćenih tokena određuje dužinu sekvence za naknadno pokretanje, uvodeći tačku sinhronizacije između GPU-a i CPU-a koju je teško izbeći. Eksperimentisali smo sa različitim šemama raspoređivanja kako bismo najbolje sakrili latenciju CPU rada.

Raspored nacrt-cilj (Draft-Target Schedule)

Uprkos tome što je manji od ciljnog modela, kada se čitav LLM koristi kao nacrt, on i dalje uvodi znatnu latenciju na GPU-u, pružajući izvesni prostor za skrivanje skupih operacija na CPU-u. Pošto manji modeli nemaju koristi od tenzorskog paralelizma, postoji neusklađenost između broja rangova na koje su podeljeni ciljni i nacrtni model. U našoj implementaciji, nacrtni model se pokreće samo na vodećem rangu TP grupe.

Uvećaj

Kao što je ranije nagovešteno, korak dekodiranja prenosi logite u sledeće pokretanje. Ovo nam omogućava da preklopimo jedno izvršavanje nacrtnog modela sa radom raspoređivanja serija na strani CPU-a. Nakon što je serija sastavljena, ponovljeni pozivi uzorkovaču i nacrtu proizvode nacrtne tokene. Paralelno, serija za verifikaciju se sastavlja za ciljni model i sinhronizuje sa paralelnim radnicima. Ciljni logiti se verifikuju i uzorkuju radi utvrđivanja prihvaćenih dužina sekvenci. U ovom trenutku, sinhronizacija između GPU-a i CPU-a je neophodna da bi se utvrdile naknadne dužine sekvenci. Pošto se nacrtni model pokreće samo na vodećem čvoru, njegova serija se postavlja sekvencijalno i njegovo izvršavanje se pokreće kako bi se unosi u njegovom KV kešu popunili dodatnim tokenom koji je ciljni model proizveo. Logiti proizvedeni ovim pokretanjem nacrta u trenutnom pokretanju koristiće se za uzorkovanje prvog nacrtnog tokena u naknadnom pokretanju. Što je najvažnije, dok je nacrt u toku, sledeća serija se može rasporediti.

MTP raspored za jedan token

Iako vreme izvršavanja još uvek ne obezbeđuje istraživanje stabla nacrta u stilu Eagle-a, implementirali smo specijalan slučaj ove šeme, razmatrajući linearnu sekvencu nacrtnih tokena proizvedenih od strane modela veličine jednog sloja transformerskog dekodera. Ova šema se može koristiti za predviđanje nacrta korišćenjem otvorenih težina (open-source weights) modela DeepSeek R1. Podslučaj predviđanja jednog tokena je interesantan, pošto velike MTP glave postižu dovoljno visoke stope prihvatanja da opravdaju svoje opterećenje.

MTP raspored je nešto složeniji, pošto je nacrtni model mnogo brži, sakrivajući manju latenciju na strani CPU-a. Pored toga, nacrt je podeljen zajedno sa ciljnim modelom, što zahteva transfere deljene memorije za informacije o seriji. Pokretanje počinje prenosom informacija o seriji i uzorkovanjem prvog tokena iz prenetih logita, slično prethodnoj šemi. Zatim se ciljni model pokreće radi validacije tokena, obrađujući 2 * D tokena, gde je D veličina serije dekodiranja. Ovo je idealno za mikro-serije u modelima sa mešavinom stručnjaka (MoE) preko sporijih međusobnih veza kao što je InfiniBand, pošto se serija ravnomerno deli na dve polovine. Skrivena stanja ciljnog modela se prenose na sledeće pokretanje nacrta, dok se logiti prosleđuju u uzorkovač radi verifikacije.

Uvećaj

Obavljanjem ograničene količine dodatnog rada na GPU-u, izbegavamo sinhronizaciju između CPU-a i GPU-a nakon prihvatanja nacrtne sekvence. Nakon što se ulazni tokeni ciljnih modela pomere, jezgro ubacuje sledeće ciljne tokene na njihove odgovarajuće lokacije. Nacrt se zatim ponovo pokreće sa istim informacijama o seriji kao i ciljni model, popunjavajući unose KV keša i gradeći logite i skrivena stanja za sledeće pokretanje, obavljajući izvestan suvišan rad na tokenima koji nisu prihvaćeni. U ovim situacijama, latencija neiskorišćenog rada je jedva merljiva zbog male veličine nacrtnog modela. Paralelno sa pokretanjem nacrta, dužine sekvenci se utvrđuju na CPU-u i pokreće se raspoređivanje sledeće serije, bez potrebe za čekanjem da se rad GPU-a završi.

Opterećenje dodatnog rada u nacrtnom sloju nije primetno u pažnji, međutim MLP slojevi su problematičniji. Pošto instrukcije mnoštva matrica vrše dopunu (pad) do granice od 64 duž dimenzije broja tokena, ako udvostručavanje ne zahteva znatno više blokova, opterećenje je skriveno. Za duže nacrtne sekvence opterećenje je skuplje i šema koja se koristi za obične nacrt-cilj modele radi bolje.

Reference

Brza inferencija iz transformera pomoću spekulativnog dekodiranja

EAGLE: Spekulativno uzorkovanje zahteva preispitivanje neizvesnosti karakteristika

EAGLE-2: Brža inferencija jezičkih modela sa dinamičkim stablima nacrta

EAGLE-3: Skaliranje ubrzanja inferencije velikih jezičkih modela putem testiranja tokom obuke

Medusa: Jednostavan okvir za ubrzanje LLM inferencije sa višestrukim glavama dekodiranja

Normalizacija slojeva

FlashInfer: Efikasan i prilagodljiv mehanizam pažnje za pružanje LLM inferencije

Tehnički izveštaj DeepSeek-V3

Citiranje

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