Ubrzavanje Sonar-a kroz spekulaciju

Speculative decoding ubrzava generisanje velikih jezičkih modela (LLM) korištenjem brzog i malog nacrtnog modela za proizvodnju kandidata za dovršetak koji se provjeravaju od strane većeg ciljnog modela. Pod ovom šemom, umjesto jednog pokretanja skupe ciljne jedinice koja proizvodi jedan token, više njih se emituje u jednom koraku. Ovdje predstavljamo detalje implementacije različitih vrsta spekulativnog dekodiranja, primijenjenih u Perplexity za smanjenje latencije između tokena na Sonar modelima.

Speculative Decoding

Speculative Decoding koristi strukturu prirodnih jezika i autoregresivnu prirodu transformera za ubrzavanje 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 nastupaju slično. Ovo preklapanje sugeriše da određene sekvence bolje generišu jeftiniji modeli, prepuštajući složene probleme većima. Izazov je u utvrđivanju koji dovršetak je bolji i da li je generisanje manjeg modela istog kvaliteta kao i onog većeg.

Srećom, LLM-ovi su autoregresivni transformeri: kada im se da sekvenca tokena, oni daju distribuciju vjerovatnoće sljedećeg tokena. Pored toga, logiti izvedeni iz srednjih karakteristika povezanih sa tokenima u ulaznoj sekvenci takođe ukazuju na to koliko je vjerovatno da će model izdati tačno te pokretačke tokene. Ovo svojstvo omogućava spekulaciju: ako sekvencu tokena generiše manji počevši od ulaznog prefiksa, ona se može provući kroz veći da bi se utvrdilo koliko se dobro slaže sa ciljnim modelom. Svaki prefiks kandidata se boduje vjerovatnoćom i bira se najduži iznad praga prihvatanja. Kao bonus, ciljni model takođe besplatno obezbjeđuje naknadni token: ako nacrtni model generiše n tokena, do n + 1 se može emitovati u jednom koraku.

Увећај

U vrijeme zaključivanja, proces spekulativnog uzorkovanja se grubo može podijeliti u 4 faze:

Unapređenje: i ciljni i nacrtni model moraju biti pokrenuti na ulaznoj sekvenci da bi se popunili unosi KV keša. Dok neke šeme, kao što je Medusa, koriste jednostavnije guste slojeve za predviđanje, u ovom postu se fokusiramo na nacrte zasnovane na transformerima kojima su potrebni sopstveni KV keševi.

Generisanje nacrta: nacrtni model iterira da bi proizveo određeni broj fiksnih tokena. Nacrtna sekvenca može biti linearna ili model može istražiti strukturu nalik stablu do date dubine (EAGLE, Medusa). Ovdje 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 ciljni generiše logite, na poziciji neusklađenosti ili na repu sekvence logiti odgovaraju naknadnom tokenu. Ovi logiti se mogu uzorkovati da obezbjeđuju robustan token iz ciljnog, zatvarajući sekvencu.

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

Target-Draft

Speculativno dekodiranje se može postići uparivanjem postojećeg malog LLM-a kao nacrtnog modela sa ciljnim modelom radi generisanja sekvenci kandidata. U proizvodnji smo ubrzali Sonar koristeći Llama-1B model fino podešen na istom skupu podataka kao i ciljni. Iako ovaj pristup nije zahtijevao obuku nacrta od nule, mali model i dalje koristi značajan kapacitet KV keša i uvodi blagi trošak unapređenja, povećavajući TTFT.

Pod ovom šemom, dekoder spekuliše samo na serijama koje služe samo za dekodiranje, generišući tokene kroz standardno uzorkovanje tokom unapređenja ili na mješovitim serijama unapređenje-dekodiranje. U fazi unapređenja, ciljni logiti se odmah uzorkuju kako bi se takođe unaprijed napunio novogenerisani token u KV kešu nacrta. Nacrt se još ne uzorkuje, ali se logiti koje on proizvodi prenose u fazu dekodiranja.

Увећај

U dekodiranju, nacrtni model napreduje, uzorkujući gornji token u svakoj fazi. Nakon što se dostigne željena dužina nacrta, tokeni se provlače kroz ciljni model da proizvedu logite na osnovu kojih uzorkivač identifikuje prihvaćenu dužinu sekvence. Prihvatanje se utvrđuje upoređivanjem punih distribucija vjerovatnoće iz nacrta i ciljnog. Pošto ciljni uvijek izbacuje jedan skup logita koji prati prihvaćenu sekvencu nacrta, to se uzorkuje da proizvede dodatni izlaz. Pošto nacrtni model još uvijek nije vidio taj prihvaćeni token, on se ponovo pokreće da popuni svoje odgovarajuće unose KV keša u pripremi za sljedeći korak dekodiranja, ponovo prenoseći logite.

EAGLE

EAGLE je šema spekulativnog dekodiranja koja istražuje višestruke nacrtne sekvence, generisane kroz prolazak kroz strukturu nalik stablu vjerovatnih nacrtnih tokena. 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 umjesto praćenja tokena sa najvišom ocjenom u linearnoj sekvenci. Sekvence se zatim boduju i najduža pogodna se bira za nastavak, dodajući takođe dodatni token iz ciljnog.

Увећај

Kako bi se postiglo preciznije predviđanje, EAGLE nacrtni model predviđa ne samo na osnovu tokena, već i koristeći ciljne karakteristike (skrivena stanja posljednjeg 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 jedan sloj transformera identičan sloju dekodera originalnog modela, koji je čvrsto povezan sa ciljnim povezivanjem sa njegovim ugrađivanjima i lm_head projekcijama. Pošto ovo zahtijeva manje kapaciteta KV keša, EAGLE ima manji otisak memorije.

Da bi se provjerile sekvence nalik stablu u ciljnom modelu, moraju se koristiti prilagođene maske pažnje. Nažalost, korištenje prilagođene maske pažnje za cijelu sekvencu značajno usporava pažnju za realistične dužine unosa (za do 50%), poništavajući dio ubrzanja koje se može postići spekulacijom. Iz tog razloga još nismo primijenili puno istraživanje stabla u proizvodnji, fokusirajući se umjesto toga na poseban slučaj predviđanja s jednim tokenom putem MTP šema predstavljenih u tehničkom izvještaju za DeepSeek-V3.

MTP

Ova šema je slična draft-target dekodiranju, s izuzetkom toga što se skrivena stanja koriste zajedno sa tokenima za predikciju. Nešto više posla mora da se obavi i u fazi prefill i u fazi dekodiranja u poređenju sa običnom draft-target spekulacijom. Draft model koristi i tokene i skrivena stanja: token t_{i+1} se uzorkuje iz logita L_i koji odgovaraju tokenu t_i, a koji su pak izvedeni iz skrivenih stanja H_i. Shodno tome, ulazni baferi tokena moraju da budu pomjereni za jedan korak ulijevo u odnosu na vektore skrivenih stanja koje daje target. Donja ilustracija označava odgovarajuće parove koji se koriste za obuku, kao i pomjeranje tokom inferencije.

Увећај

Tok dekodiranja je prilično sličan dekodiranju nacrt-cilj, sa izuzetkom prenošenja i skrivenih stanja i logita. Naša implementacija dijeli sve povezane logite uzorkovanja i obrade logita, specijalizujući samo pozive prosljeđivanja modela. 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 da popuni unos KV keša za ciljno predviđanje, pokrećemo ga na cijeloj sekvenci uzimajući tačnija skrivena stanja ciljnog kao ulaze. Pošto su ovi nacrtni modeli mali, dodatni trošak obrade dodatnih tokena je zanemarljiv.

Obuka MTP glava

Kako bismo imali koristi od MTP-a, izgradili smo infrastrukturu potrebnu za obuku MTP glava pridruženih našim fino podešenim modelima na skupovima podataka Perplexity, koji rade 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 sa logitima sljedećeg tokena i skrivenim stanjima ciljnog. Pošto je zaključivanje za skrivena stanja skupo, mi ih unaprijed izračunavamo koristeći našu implementaciju ciljnog modela optimizovanu za zaključivanje, koja će se koristiti tokom obuke. Međutim, da bismo potvrdili MTP implementaciju zaključivanja i osigurali da numeričke razlike uslijed kvantizacije ili optimizacija ne ometaju rezultate, za gubitak validacije i procjenu tačnosti u potpunosti ponovo koristimo implementaciju zaključivanja i ciljnog i nacrtnih modela.

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

Увећај

Ne samo da normativi slojeva olakšavaju obuku, ponovno uvođenje normativa je takođe matematički intuitivno. MTP glave ponovo koriste ugrađivanja i projekcije logita ciljnog modela, pošto mogu biti značajne veličine (oko 2 GB za Llama 70B). Tokom obuke, oni su zamrznuti 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. Odbacivanjem normativa, očekuje se da jedan MLP uči istu funkciju kao MLP praćen normativom, što otežava usklađivanje između skrivenih stanja nacrta i ciljnih modela.

Zaključivanje sa spekulativnim dekodiranjem

U mašini za zaključivanje, da bi se generisali tokeni za ulazne sekvence, oni se prvo moraju grupirati u serije razumne veličine, zatim se stranice moraju dodijeliti u KV kešu za sljedeće tokene. Ulazni tokeni i informacije o KV stranici se zatim pakuju u međuspremnik koji se emituje svim paralelnim rangovima koji pokreću model. Na kraju, metapodaci se kopiraju u GPU memoriju i model se izvršava da proizvede logite iz kojih se uzorkuje sljedeći token.

Za razliku od određenih implementacija koje labavo povezuju nacrtni i ciljni server za zaključivanje preko omotača koji orkestrira zahtjeve između njih, naši parovi nacrt-cilj su čvrsto povezani i prolaze kroz generisanje u unisonu. Zakazivanje serija i dodjela KV stranica se dijele između modela za sve oblike spekulativnog dekodiranja: ovo unificira logiku koja povezuje model sa sveobuhvatnim serverom za zaključivanje, pošto svi oni izlažu isti interfejs.

Vrijeme izvršavanja zaključivanja u Perplexity je oblikovano oko FlashInfer-a, koji određuje metapodatke koje je potrebno izgraditi kako bi se konfigurisalo i zakazalo jezgro pažnje. S obzirom na neke ulazne sekvence koje čine seriju, za unapređenje, dekodiranje ili provjeru, rad na strani CPU-a se mora obaviti kako bi se dodijelili međuprostori i popunili određeni konstantni međuprostori koji se koriste u pažnji. Ovaj rad je dodat trošku zakazivanja serija i dodjele KV stranica, koji takođe izazivaju latencije koje se moraju sakriti kako bi se maksimiziralo iskorištenje GPU-a.

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

Raspored nacrt-cilj

Uprkos tome što je manji od ciljnog modela, kada se cijeli LLM koristi kao nacrt, on i dalje uvodi znatnu latenciju na GPU-u, obezbjeđujući određeni prostor za skrivanje skupih operacija CPU-a. Pošto manji modeli nemaju koristi od tenzorskog paralelizma, postoji neusklađenost između broja rangova po kojima su ciljni i nacrt podijeljeni. U našoj implementaciji, nacrtni model radi samo na vodećem rang-mjestu TP grupe.

Увећај

Kao što je ranije naznačeno, korak dekodiranja prenosi logite u sljedeće pokretanje. Ovo nam omogućava da preklopimo jedno izvršavanje nacrtnog modela sa radom zakazivanja serije na strani CPU-a. Nakon što se serija sastavi, ponovljeni pozivi uzorkivaču i nacrtu proizvode tokene nacrta. Paralelno, serija za provjeru se sastavlja za ciljni model i sinhronizuje sa paralelnim radnicima. Ciljni logiti se provjeravaju i uzorkuju da se utvrde prihvaćene dužine sekvenci. U ovom trenutku, sinhronizacija GPU-CPU je neophodna kako 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 popunili unosi KV keša dodatnim tokenom koji je ciljni proizveo. Logiti koje proizvodi ovo pokretanje nacrta u trenutnom pokretanju koristit će se za uzorkovanje prvog nacrtnog tokena u naknadnom pokretanju. Što je najvažnije, dok nacrt radi, sljedeća serija se može zakazati.

MTP raspored za jedan token

Iako vrijeme izvršavanja još uvijek ne nudi istraživanje nacrtnog stabla u stilu Eagle-a, implementirali smo poseban slučaj ove šeme, razmatrajući linearnu sekvencu nacrtnih tokena proizvedenih od strane modela veličine jednog sloja dekodera transformera. Ova šema se može koristiti za predviđanje nacrta korištenjem otvoreno dostupnih težina za DeepSeek R1. Podslučaj predviđanja jednog tokena je interesantan, pošto veliki MTP slojevi postižu dovoljno visoke stope prihvatanja da opravdaju svoje dodatne troškove.

MTP zakazivanje je donekle složenije, pošto je nacrtni model mnogo brži, sakrivajući manje latencije na strani CPU-a. Pored toga, nacrt je podijeljen zajedno sa ciljnim modelom, zahtijevajući transfere dijeljene memorije za informacije o seriji. Pokretanje počinje prenosom informacija o seriji i uzorkovanjem prvog tokena iz prenesenih logita, slično prethodnoj šemi. Zatim se ciljni pokreće da provjeri tokene, obrađujući 2 * D tokena, gdje je D veličina serije dekodiranja. Ovo je idealno za mikro-serijsku obradu u modelima mješavine stručnjaka (MoE) preko sporijih međuveza kao što je InfiniBand, pošto se serija ravnomjerno dijeli na dva dijela. Skrivena stanja ciljnog se prenose na sljedeće pokretanje nacrta, dok se logiti proslijeđuju u uzorkivač radi provjere.

Увећај

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

Dodatni trošak dodatnog rada u sloju nacrta nije primjetan u pažnji, međutim MLP slojevi su problematičniji. Pošto instrukcije množenja matrica popunjavaju do granice od 64 duž dimenzije broja tokena, ako udvostručavanje ne zahtijeva značajno više blokova, dodatni trošak je sakriven. Za duže nacrtne sekvence dodatni trošak je skuplji i šema koja se koristi za regularne nacrt-cilj modele radi bolje.

Reference

Brzo zaključivanje iz transformera putem spekulativnog dekodiranja

EAGLE: Spekulativno uzorkovanje zahtijeva ponovno razmatranje neizvjesnosti karakteristika

EAGLE-2: Brže zaključivanje jezičkih modela sa dinamičkim nacrtnim stablima

EAGLE-3: Skaliranje ubrzanja zaključivanja velikih jezičkih modela putem testiranja u vrijeme obuke

Medusa: Jednostavan okvir za ubrzavanje LLM zaključivanja sa višestrukim glavama dekodiranja

Normalizacija sloja

FlashInfer: Efikasan i prilagodljiv mehanizam pažnje za pružanje LLM zaključivanja

Tehnički izvještaj za DeepSeek-V3

Цитат

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