Sonari kiirendamine spekulatsiooni kaudu

Spekulatiivne dekodeerimine kiirendab suurte keelemudelite (LLM-ide) genereerimiskiirust, kasutades kiiret ja väikest mustandmudelit, et luua lõpetamiskandidaate, mida suurem sihtmudel kontrollib. Selle skeemi kohaselt ei genereeri kalli sihtmudeli käitamine ühte märki, vaid ühes sammus väljastatakse mitu märki. Siin esitame erinevat tüüpi spekulatiivse dekodeerimise rakendamise üksikasjad, mida rakendatakse Perplexity's Sonari mudelite märkidevahelise latentsuse vähendamiseks.

Spekulatiivne dekodeerimine

Spekulatiivne dekodeerimine kasutab loomulike keelte struktuuri ja transformerite autoregressiivset olemust märkide genereerimise kiirendamiseks. Kuigi suurematel mudelitel, nagu Llama-70B, on rohkem teadmisi kui väiksematel, nagu Llama-1B, toimivad nad lihtsamate ülesannete puhul sarnaselt. See kattuvus viitab sellele, et teatud järjestusi on odavamatel mudelitel parem genereerida, jättes keerulised probleemid suurematele mudelitele. Väljakutse seisneb kindlakstegemises, millised lõpetamised on paremad ja kas väiksema mudeli genereering on sama kvaliteediga kui suuremal mudelil.

Õnneks on LLM-id autoregressiivsed transformerid: märkide järjestuse esitamisel väljastavad nad järgmise märgi tõenäosusjaotuse. Lisaks näitavad sisendjärjestuse märkidega seotud vahepealsetest tunnustest tuletatud logitid, kui tõenäoline on, et mudel väljastab täpselt need märgid. See omadus võimaldab spekuleerimist: kui märkide järjestus genereeritakse väiksema mudeli abil sisendi eesliitest alates, saab selle käivitada suurema mudeli kaudu, et teha kindlaks, kui hästi see sihtmudeliga ühtib. Kandidaatide igale eesliitele antakse tõenäosus ja valitakse pikim aktsepteerimiskünnise ületav eesliit. Boonusena pakub sihtmudel ka järgmise märgi tasuta: kui mustandmudel genereerib n märki, saab ühes sammus väljastada kuni n + 1 märki.

Suurenda

Järelduse tegemise ajal saab spekulatiivse valimi võtmise protsessi jagada umbes 4 etappi:

Eeltäitmine (prefill): KV vahemälu kirjete täitmiseks tuleb sisendjärjestusel käivitada nii sihtmudel kui ka mustandmudel. Kuigi mõned skeemid, näiteks Medusa, kasutavad ennustamiseks lihtsamaid tihedaid kihte, keskendume selles postituses transformeripõhistele mustanditele, mis vajavad oma KV vahemälu.

Mustandi genereerimine: mustandmudel teeb korduvaid samme teatud arvu fikseeritud märkide tootmiseks. Mustandijärjestus võib olla lineaarne või mudel võib uurida puutaolist struktuuri kuni antud sügavuseni (EAGLE, Medusa). Siin keskendume lineaarsetele järjestustele.

Aktsepteerimine: sihtmudel käivitub mustandijärjestusel, ehitades üles iga mustandimärgiga seotud logitid. Määratakse kindlaks pikima aktsepteeritava järjestuse pikkus.

Sihtmärgi genereerimine: kuna sihtmudel genereeris logitid, vastavad logitid mittevastavas asukohas või järjestuse lõpus järgmisele märgile. Neist logititest saab võtta näidiseid, et saada sihtmärgilt tugev märk, mis lõpetab järjestuse.

Spekulatiivse dekodeerimise rakendamiseks on mitmesuguseid meetodeid. Selles postituses keskendume skeemidele, mida kasutasime Sonari mudelite kiirendamiseks sisemise 1B mudeli abil, samuti ennustusmehhanismidele, mida loome mudelite kiirendamiseks DeepSeeki skaalal.

Sihtmärk-mustand (Target-Draft)

Spekulatiivset dekodeerimist saab saavutada olemasoleva väikese LLM-i sidumisega mustandmudelina sihtmudeliga, et genereerida kandidaatide järjestusi. Tootmises oleme Sonarit kiirendanud Llama-1B mudeli abil, mis on peenhäälestatud samal andmestikul kui sihtmudel. Kuigi see lähenemine ei nõudnud mustandi nullist treenimist, kasutab väike mudel siiski märkimisväärset KV vahemälu mahtu ja toob kaasa väikese eeltäitmise lisakoormuse, mis suurendab TTFT-d.

Selle skeemi kohaselt spekuleerib dekooder ainult dekodeerimisele mõeldud pakettidega, genereerides märke eeltäitmise ajal või segatud eeltäitmise-dekodeerimise pakettidel standardse valimi võtmise kaudu. Eeltäitmise etapis võetakse sihtmärgi logititest viivitamatult näidiseid, et eeltäita ka äsja genereeritud märk mustandi KV vahemälus. Mustandist veel näidiseid ei võeta, kuid selle toodetud logitid kantakse üle dekodeerimise etappi.

Suurenda

Dekodeerimisel mustandmudel liigub edasi, võttes igas etapis näidise tipu märgist. Pärast soovitud mustandi pikkuse saavutamist lastakse märgid läbi sihtmudeli, et toota logitid, mille alusel valimivõtja tuvastab aktsepteeritud järjestuse pikkuse. Aktsepteerimine määratakse mustandi ja sihtmudeli täielike tõenäosusjaotuste võrdlemise teel. Kuna sihtmudel väljastab alati ühe logitite komplekti pärast aktsepteeritud mustandijärjestust, võetakse sellest lisaväljundi tootmiseks näidis. Kuna mustandmudel ei ole seda aktsepteeritud märki veel näinud, käivitatakse see uuesti, et täita selle vastavad KV vahemälu kirjed järgmise dekodeerimissammu jaoks, kandes logitid uuesti üle.

EAGLE

EAGLE on spekulatiivse dekodeerimise skeem, mis uurib mitut mustandijärjestust, mis on genereeritud tõenäoliste mustandimärkide puutaolise läbimise teel. Fikseeritud (EAGLE) või dünaamilise kujuga (EAGLE-2) puud uuritakse mustandimärkide järjestikuste täitmiste abil, võttes arvesse Top-K kandidaate igas sõlmes, selle ainsa kõrgeima skooriga märgi järel liikumise asemel lineaarses järjestuses. Järjestused hinnatakse seejärel punktidega ja jätkamiseks valitakse pikim sobiv järjestus, lisades lisaks sihtmärgist täiendava märgi.

Suurenda

Täpsema ennustuse saavutamiseks ei ennusta EAGLE mustandmudel mitte ainult märkide põhjal, vaid kasutab ka sihtmudeli sihttunnuseid (viimase kihi peidetud olekuid). EAGLE puuduseks on vajadus treenida kohandatud väikesi mustandmudeleid, mis on piisavalt täpsed sobivate kandidaatide genereerimiseks madala latentsuse eelarve piires. Tavaliselt on mustandmudel üksainus transformerikiht, mis on identne originaalmudeli dekoodrikihiga ja mis on sihtmudeliga tihedalt seotud, olles seotud selle manuste ja lm_head’i projektsioonidega. Kuna see nõuab vähem KV vahemälu mahtu, on EAGLE-l väiksem mälujalajälg.

Puutaoliste järjestuste kontrollimiseks sihtmudelis tuleb kasutada kohandatud tähelepanumaske (attention masks). Kahjuks aeglustab kohandatud tähelepanumaski kasutamine terve järjestuse jaoks reaalsete sisendpikkuste puhul tähelepanu märkimisväärselt (kuni 50%), tühistades osa spekuleerimise teel saavutatavast kiiruse kasvust. Me ei ole sel põhjusel täielikku puu uurimist veel tootmisse juurutanud, keskendudes selle asemel ühe märgi ennustamise erijuhule MTP-laadsete skeemide kaudu, mis on esitatud DeepSeek-V3 tehnilises aruandes.

MTP

See skeem on sarnane mustandi ja sihtmärgi vahelisele dekodeerimisele, välja arvatud see, et ennustamiseks kasutatakse märkide kõrval ka peidetud olekuid. Nii eeltäitmise kui ka dekodeerimise etapis tuleb teha veidi rohkem tööd võrreldes tavalise mustandi ja sihtmärgi vahelise spekuleerimisega. Mustandmudel kasutab nii märke kui ka peidetud olekuid: märk t_{i+1} võetakse näidisena logititest L_i, mis vastavad märgile t_i, mis omakorda on tuletatud peidetud olekutest H_i. Jõudmaks sihtmudeli väljastatud peidetud oleku vektoreiteni, tuleb sisendimärgi puhvreid vastavalt ühe sammu võrra vasakule nihutada. Allolev joonis tähistab treenimiseks kasutatavaid vastavusi ja ka nihet järelduse tegemise ajal.

Suurenda

Dekodeerimisvoog on üsna sarnane mustandi ja sihtmärgi vahelisele dekodeerimisele, välja arvatud see, et kaasa kantakse nii peidetud olekuid kui ka logitit. Meie rakenduses jagatakse kõiki seotud valimi võtmise ja logitite töötlemise funktsioone, spetsialiseerides ainult mudeli edasisuunamise (forward) väljakutseid. Kui ennustatakse mitu märki, kasutab mustandmudel ennustamiseks mustandi peidetud olekuid, täites KV vahemälu kirjeid ka oma tunnuste põhjal. Pikemas perspektiivis võib see täpsust halvendada. Seejärel, kui käivitame mustandmudeli sihtmärgi ennustuse KV vahemälu kirje täitmiseks, käivitame selle kogu järjestusel, võttes sisenditeks täpsemad sihtmudeli peidetud olekud. Kuna need mustandmudelid on väikesed, on täiendavate märkide töötlemise lisakulu tühine.

MTP peade treenimine

MTP-st kasu saamiseks ehitasime vajaliku infrastruktuuri Perplexity andmestikel meie peenhäälestatud mudelitega seotud MTP peade treenimiseks, mis töötab ühel 8xH100 seadmega sõlmel. Umbes ühe päevaga saame luua päid mudelitele vahemikus Llama-1B kuni Llama-70B ja DeepSeek V2-Lite. Suuremate mudelite puhul tugineme peenhäälestusprotsessi käigus ehitatud MTP peadele.

MTP treenimise eesmärk on sobitada sihtmudeli peidetud olekutest ekstrapoleeritud mustandi peidetud olekud ja logitid sihtmudeli järgmise märgi logitite ja peidetud olekutega. Kuna peidetud olekute järelduse tegemine on kallis, eelarvutame need sihtmudeli järelduse tegemiseks optimeeritud rakendust kasutades, et neid treeningu ajal kasutada. Siiski, et valideerida MTP järelduse tegemise rakendust ja tagada, et kvantimisest või optimeerimistest tulenevad numbrilised erinevused tulemusi ei takistaks, taaskasutame valideerimiskao ja täpsuse hindamise jaoks täielikult nii sihtmudeli kui ka mustandmudeli järelduse tegemise rakendust.

Skaleerimisel algses artiklis kasutatud ShareGPT andmestikult suurematele valimitele märkasime, et EAGLE artiklis kirjeldatud ja rakendatud MTP pea arhitektuur ei suutnud 70B suuruste mudelite jaoks treenida. Erinevalt ShareGPT-st, mis sisaldas suuremat arvu lühemaid järjestusi, treenime me veidi väiksema arvu oluliselt pikemate prompte peal. Kuna algsete EAGLE peade struktuur erines tüüpilisest transformerist veidi, toimus eemaldatud RMS-normaliseerimise kihtide uuesti sisse toomine. Leidsime, et see mitte ainult ei võimaldanud treenimisel koonduda, vaid suurendas ka peade täpsust mõne protsendipunkti võrra.

Suurenda

Kihi normid ei hõlbusta mitte ainult treenimist, vaid normide uuesti sisse toomine on ka matemaatiliselt intuitiivne. MTP pead taaskasutavad sihtmudeli manuseid ja logitite projektsioone, kuna need võivad olla märkimisväärse suurusega (Llama 70B puhul umbes 2 GB). Treeningu ajal need külmutatakse ja ootus on, et MTP kiht õpib manustama ennustusi samasse vektoriavarusse, mida originaalmudeli projektsioonikiht treeningu ajal õppis. Normide eemaldamisega eeldatakse, et üks MLP õpib sama funktsiooni mis normile järgnev MLP, mis takistab mustand- ja sihtmudelite peidetud olekute vahelist sobitamist.

Järelduse tegemine spekulatiivse dekodeerimisega

Järelduse mootoris tuleb sisendjärjestuste jaoks märkide genereerimiseks need kõigepealt rühmitada mõistliku suurusega pakettideks ning seejärel eraldada KV vahemälust leheküljed järgmiste märkide jaoks. Seejärel pakitakse sisendimärgid ja KV lehekülje teave puhvrisse, mis edastatakse ringhäälinguna kõigile mudelit käitavatele paralleelsetele järkudele (ranks). Lõpuks kopeeritakse metaandmed GPU mällu ja mudelit käivitatakse logitite tootmiseks, millest järgmine märk valitakse.

Erinevalt teatud rakendustest, mis ühendavad mustandi ja sihtmudeli järelduse serveri lõdvalt ümbrise abil, mis korraldab nendevahelisi päringuid, on meie mustand-sihtmärk paarid tihedalt seotud ja liiguvad genereerimisel ühes rütmis. Pakettide ajastamine ja KV lehekülgede eraldamine jagatakse mudelite vahel kõigi spekulatiivse dekodeerimise vormide puhul: see ühtlustab loogikat, mis ühendab mudelit üldise järelduse serveriga, kuna need kõik paljastavad sama liidese.

Perplexity järelduse käitusaeg on üles ehitatud FlashInferi ümber, mis määrab kindlaks metaandmed, mis tuleb tähelepanutuumiku (attention kernel) konfigureerimiseks ja ajastamiseks luua. Arvestades mõningaid paketti moodustavaid sisendjärjestusi eeltäitmise, dekodeerimise või kinnitamise jaoks, tuleb CPU poolel teha tööd vahepealsete puhvrite eraldamiseks ja teatud tähelepanus kasutatavate konstantsete puhvrite täitmiseks. See töö lisandub pakettide ajastamise ja KV lehekülgede eraldamise kulule, mis põhjustavad samuti latentsusi, mis tuleb GPU kasutusaste maksimeerimiseks varjata.

Kuigi me paralleelsime CPU- ja GPU-poolse töö järelduste tegemiseks ilma spekulatsioonita täielikult, leidsime, et spekulatiivse dekodeerimise CPU ja GPU tasakaal on keerulisem. Peamine väljakutse tuleneb tõsiasjast, et aktsepteeritud märkide arv määrab järgmise käivitamise järjestuse pikkuse, mis toob kaasa raskesti välditava GPU-CPU sünkroonimispunkti. Katsetasime erinevaid ajastamisskeeme, et CPU töö latentsust kõige paremini varjata.

Mustandi ja sihtmärgi ajastamine (Draft-Target Schedule)

Vaatamata sellele, et mustandmudel on sihtmudelist väiksem, tekitab terve LLM-i kasutamisel mustandina see GPU-s siiski märkimisväärset latentsust, pakkudes veidi varu kulukate CPU operatsioonide varjamiseks. Kuna väiksemad mudelid ei saa tensorparalleelsusest kasu, esineb ebakõla järkude arvus, mille vahel sihtmärk ja mustand on šarditud (sharded). Meie rakenduses töötab mustandmudel ainult TP-grupi liidri järgul (leader rank).

Suurenda

Nagu varem märgitud, kannab dekodeerimissamm logitit järgmisesse käivitusse. See võimaldab meil mustandmudeli ühe täitmise kattuda CPU-poolse pakettide ajastamise tööga. Pärast paketi kokkupanemist luuakse valimivõtja ja mustandi korduvate päringute abil mustandimärgid. Paralleelselt pannakse sihtmudeli jaoks kokku kontrollimise pakett ja sünkroonitakse see paralleelsete töötajatega. Sihtmudeli logitid kontrollitakse ja neist võetakse aktsepteeritud järjestuse pikkuste määratlemiseks näidiseid. Selles punktis on järgmiste järjestuse pikkuste määratlemiseks vajalik GPU-CPU sünkroonimine. Kuna mustandmudel töötab ainult liidersõlmes, seatakse selle pakett üles järjestikku ja selle täitmine käivitatakse, et täita selle KV vahemälu kirjed täiendava märgiga, mille sihtmudel tootis. Logitid, mida see mustandikäivitus praeguses käivituses toodab, kasutatakse esimese mustandimärgi võtmiseks järgmises käivituses. Mis kõige tähtsam, mustandi töötamise ajal saab järgmise paketi ajastada.

MTP ajastamine ühe märgi jaoks

Kuigi käitusaeg ei paku veel Eagle'i stiilis mustandi puu uurimist, rakendasime selle skeemi erijuhtu, arvestades ühe transformeripõhise dekoodri kihi suuruse mudeli poolt toodetud mustandimärkide lineaarset järjestust. Seda skeemi saab kasutada mustandi ennustamiseks DeepSeek R1 avatud lähtekoodiga kaalude abil. Ühe märgi ennustamise alamjuhtum on huvitav, kuna suured MTP kihid saavutavad piisavalt kõrged aktsepteerimismäärad, et õigustada nende lisakoormust.

MTP ajastamine on mõnevõrra keerulisem, kuna mustandmudel on palju kiirem, varjates vähem CPU-poolset latentsust. Lisaks on mustand šarditud koos sihtmudeliga, mis nõuab paketi teabe jaoks jagatud mälu edastusi. Käivitus algab paketi teabe edastamise ja esimese märgi võtmisega ülekantavatest logititest, sarnaselt eelmisele skeemile. Järgmisena käivitatakse sihtmärk märkide valideerimiseks, töödeldes 2 * D märki, kus D on dekodeerimise paketi suurus. See on ideaalne mikropakettide loomiseks Ekspertide Segu (Mixture-of-Experts, MoE) mudelites aeglasemate ühenduste, näiteks InfiniBandi kaudu, kuna pakett jaguneb ühtlaselt kaheks pooleks. Sihtmudeli peidetud olekud kantakse üle järgmisse mustandikäivitusse, samas kui logitid edastatakse valideerimiseks valimivõtjale.

Suurenda

Teostades GPU-s piiratud hulgal täiendavat tööd, väldime GPU-CPU sünkroonimist pärast mustandijärjestuse aktsepteerimist. Pärast sihtmärkide sisendimärkide nihutamist lisab tuumik (kernel) järgmised sihtmärgi märgid nende vastavatesse asukohtadesse. Seejärel käivitatakse mustand uuesti sama paketiteabega nagu sihtmudel, täites KV vahemälu kirjed ning luues logitid ja peidetud olekud järgmise käivitamise jaoks, tehes ümbermõtestatud tööd märkidega, mida ei aktsepteeritud. Nendes olukordades on kasutamata töö latentsus mustandmudeli väikese suuruse tõttu vaevalt mõõdetav. Mustandikäivitusega paralleelselt määratakse CPU-l järjestuse pikkused ja käivitatakse järgmise paketi ajastamine, ilma et peaks ootama GPU töö lõppemist.

Mustandikihi lisatöö lisakoormus ei ole tähelepanus märgatav, kuid MLP kihid on problemaatilisemad. Kuna maatrikskorrutamise juhised polsterdatakse (pad) märkide arvu dimensiooni ulatuses 64 piirini, on lisakoormus varjatud, kui kahekordistamine ei nõua oluliselt rohkem plokke. Pikemate mustandijärjestuste puhul on lisakoormus kulukam ja tavaliste mustand-sihtmärk mudelite jaoks kasutatav skeem toimib paremini.

Viited

Kiire järeldus transformeritest spekulatiivse dekodeerimise kaudu (Fast Inference from Transformers via Speculative Decoding)

EAGLE: Spekulatiivne valimi võtmine nõuab tunnuste ebakindluse ümbermõtestamist (EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty)

EAGLE-2: Keelte mudelite kiirem järeldus dünaamiliste mustandipuudega (EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees)

EAGLE-3: Suurte keelemudelite järelduse kiirenduse skaleerimine treeninguaegse testi kaudu (EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test)

Medusa: Lihtne LLM-i järelduse kiirendamise raamistik mitme dekodeerimispeaga (Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads)

Kihi normaliseerimine (Layer Normalization)

FlashInfer: Tõhus ja kohandatav tähelepanumootor LLM-i järelduste teenindamiseks (FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving)

DeepSeek-V3 tehniline aruanne (DeepSeek-V3 Technical Report)

Tsitaat

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