Sonar paātrināšana, izmantojot spekulācijas

Spekulatīvā dekodēšana paātrina lielo valodu modeļu (LLM) ģenerēšanas ātrumu, izmantojot ātrāku un mazāku melnraksta modeli, lai izveidotu pabeigšanas kandidātus, kurus pārbauda lielākais mērķa modelis. Saskaņā ar šo shēmu, tā vietā, lai dārgais mērķa modelis izpildes laikā ģenerētu vienu tokenu, vienā solī tiek izvadīti vairāki. Šeit mēs iepazīstinām ar dažādu veidu spekulatīvās dekodēšanas ieviešanas detaļām, kas Perplexity tiek lietotas, lai samazinātu aizturi starp tokeniem Sonar modeļos.

Spekulatīvā dekodēšana

Spekulatīvā dekodēšana izmanto dabisko valodu struktūru un transformatoru autoregresīvo dabu, lai paātrinātu tokenu ģenerēšanu. Lai gan lielākiem modeļiem, piemēram, Llama-70B, ir vairāk zināšanu nekā mazākiem, piemēram, Llama-1B, vienkāršākos uzdevumos tie darbojas līdzīgi. Šī sakritība liecina, ka noteiktas sekvences labāk ģenerē mazāk dārgi modeļi, savukārt sarežģītas problēmas tiek uzticētas lielākiem modeļiem. Izaicinājums slēpjas tajā, lai noteiktu, kuras pabeigšanas versijas ir labākas un vai mazākā modeļa ģenerētais saturs ir tikpat kvalitatīvs kā lielākā modeļa sniegums.

Par laimi, LLM ir autoregresīvi transformatori: ja tiem tiek dota tokenu sekvence, tie izvada nākamā tokena varbūtības sadalījumu. Turklāt logiti, kas iegūti no starpposma pazīmēm, kuras saistītas ar ievades sekvences tokeniem, norāda arī uz to, cik ticams ir tas, ka modelis izdos tieši šos tokenus. Šī īpašība nodrošina spekulāciju: ja tokenu sekvenci ģenerē mazāks modelis, sākot no ievades prefiksa, to var palaist caur lielāko modeli, lai noteiktu, cik labi tas sakrīt ar mērķa modeli. Katrs kandidātu prefikss tiek novērtēts ar varbūtību, un tiek izvēlēts garākais, kas pārsniedz pieņemšanas slieksni. Papildu bonuss ir tas, ka mērķa modelis nodrošina arī nākamo tokenu bez maksas: ja melnraksta modelis ģenerē n tokenus, vienā solī var izvadīt līdz pat n + 1 tokeniem.

Izvērst

Secināšanas laikā spekulatīvās paraugu ņemšanas procesu var aptuveni sadalīt 4 posmos:

Iepriekšēja aizpildīšana (prefill): gan mērķa, gan melnraksta modeļiem jādarbojas ar ievades sekvenci, lai aizpildītu KV kešatmiņas ierakstus. Lai gan dažas shēmas, piemēram, Medusa, prognozēšanai izmanto vienkāršākus blīvos slāņus, šajā rakstā mēs koncentrējamies uz transformatoru balstītiem melnrakstiem, kam nepieciešamas savas KV kešatmiņas.

Melnraksta ģenerēšana: melnraksta modelis iteratīvi ģenerē noteiktu skaitu fiksētu tokenu. Melnraksta sekvence var būt lineāra, vai arī modelis var izpētīt koka veida struktūru līdz noteiktam dziļumam (EAGLE, Medusa). Šeit mēs koncentrējamies uz lineārām sekvencēm.

Pieņemšana: mērķa modelis darbojas ar melnraksta sekvenci, veidojot logitus, kas atbilst katram melnraksta tokenam. Tiek noteikts garākās pieņemamās sekvences garums.

Mērķa ģenerēšana: tā kā mērķa modelis ģenerēja logitus, neatbilstības pozīcijā vai sekvences beigās logiti atbilst nākamajam tokenam. No šiem logitiem var ņemt paraugus, lai nodrošinātu stabilu tokenu no mērķa modeļa, noslēdzot sekvenci.

Pastāv vairākas metodes spekulatīvās dekodēšanas ieviešanai. Šajā rakstā mēs koncentrēsimies uz shēmām, kuras izmantojām, lai paātrinātu Sonar modeļus, izmantojot iekšējo 1B modeli, kā arī uz prognozēšanas mehānismiem, ko veidojam, lai paātrinātu modeļus DeepSeek mērogā.

Mērķa un melnraksta kombinācija (Target-Draft)

Spekulatīvo dekodēšanu var sasniegt, savienojot esošu mazu LLM kā melnraksta modeli ar mērķa modeli, lai ģenerētu kandidātu sekvences. Produkcijā mēs esam paātrinājuši Sonar, izmantojot Llama-1B modeli, kas ir smalki noregulēts (fine-tuned) ar to pašu datu kopu, ko mērķa modelis. Lai gan šī läeja neprasīja trenēt melnraksta modeli no nulles, mazais modelis joprojām izmanto ievērojamu KV kešatmiņas ietilpību un rada nelielu iepriekšējas aizpildīšanas pieskaitāmo laiku, palielinot TTFT.

Saskaņā ar šo shēmu dekodētājs spekulē tikai uz partijām, kas paredzētas tikai dekodēšanai, ģenerējot tokenus, izmantojot standarta paraugu ņemšanu iepriekšējas aizpildīšanas laikā vai uz jauktām iepriekšējās aizpildīšanas un dekodēšanas partijām. Iepriekšējās aizpildīšanas posmā mērķa logiti tiek nekavējoties ņemti paraugos, lai iepriekš aizpildītu arī jaunizveidoto tokenu melnraksta KV kešatmiņā. No melnraksta vēl netiekņemti paraugi, bet tā ģenerētie logiti tiek pārcelti uz dekodēšanas posmu.

Izvērst

Dekodēšanas laikā melnraksta modelis tiek virzīts uz priekšu, katrā posmā ņemot paraugā populārāko tokenu. Pēc tam, kad ir sasniegts vēlamais melnraksta garums, tokeni tiek palaisti caur mērķa modeli, lai ģenerētu logitus, uz kuru pamata paraugu ņemšanas rīks nosaka pieņemto sekvences garumu. Pieņemšana tiek noteikta, salīdzinot pilnus varbūtības sadalījumus no melnraksta un mērķa modeļa. Tā kā mērķa modelis vienmēr izvada vienu logitu kopu pēc pieņemtās melnraksta sekvences, no tās tiek ņemti paraugi, lai iegūtu papildu izvadi. Tā kā melnraksta modelis vēl nav redzējis šo pieņemto tokenu, tas tiek palaists atkārtoti, lai aizpildītu tā atbilstošos KV kešatmiņas ierakstus, gatavojoties nākamajam dekodēšanas solim un atkal pārvietojot logitus līdzi.

EAGLE

EAGLE ir spekulatīva dekodēšanas shēma, kas pēta vairākas melnraksta sekvences, kuras ģenerētas, izmantojot ticamu melnraksta tokenu koka veida šķērsošanu. Fiksēts (EAGLE) vai dinamiski veidots (EAGLE-2) koks tiek pētnieciski izmantots, secīgi izpildot melnraksta tokenus un ņemot vērā Top-K kandidātus katrā mezglā, nevis sekojot visaugstāk novērtētajam tokenam lineārā sekvencē. Pēc tam sekvences tiek novērtētas, un tiek atlasīta garākā piemērotā sekvence turpināšanai, pievienojot arī papildu tokenu no mērķa modeļa.

Izvērst

Lai sasniegtu precīzāku prognozēšanu, EAGLE melnraksta modelis prognozē ne tikai pamatojoties uz tokeniem, bet izmanto arī mērķa modeļa mērķa pazīmes (pēdējā slāņa slēptos stāvokļus). EAGLE trūkums ir nepieciešamība apmācīt pielāgotus, mazus melnraksta modeļus, kas ir pietiekami precīzi, lai ģenerētu piemērotus kandidātus zemā aiztures budžetā. Parasti melnraksta modelis ir viens transformatora slānis, kas ir identisks oriģinālā modeļa dekodētāja slānim un ir cieši savienots ar mērķa modeli, piesaistot to tā iegulumiem un lm_head projekcijām. Tā kā tam ir nepieciešams mazāks KV kešatmiņas apjoms, EAGLE ir mazāks atmiņas patēriņš.

Lai pārbaudītu koka veida sekvences mērķa modelī, ir jāizmanto pielāgotas uzmanības maskas (attention masks). Diemžēl pielāgotas uzmanības maskas izmantošana visai sekvencei būtiski palēnina uzmanības mehānismu reālistiskiem ievades garumiem (līdz pat 50%), anulējot daļu no paātrinājuma, ko var sasniegt, izmantojot spekulācijas. Šī iemesla dēļ mēs vēl neesam ieviesuši pilnu koka izpēti produkcijā, tā vietā koncentrējoties uz īpašu gadījumu — viena tokena prognozēšanu, izmantojot MTP līdzīgas shēmas, kas aprakstītas DeepSeek-V3 tehniskajā ziņojumā.

MTP

Šī shēma ir līdzīga melnraksta un mērķa dekodēšanai, ar izņēmumu, ka prognozēšanai līdzās tokeniem tiek izmantoti arī slēptie stāvokļi. Salīdzinot ar parasto melnraksta un mērķa spekulāciju, gan iepriekšējās aizpildīšanas, gan dekodēšanas posmā ir jāveic nedaudz vairāk darba. Melnraksta modelis izmanto gan tokenus, gan slēptos stāvokļus: tokens t_{i+1} tiek ņemts paraugā no logitiem L_i, kas atbilst tokenam t_i, kuri savukārt ir iegūti no slēptajiem stāvokļiem H_i. Līdz ar to ievades tokenu buferi ir jābīda par vienu soli pa kreisi attiecībā pret slēpto stāvokļu vektoriem, ko izvada mērķa modelis. Zemāk redzamajā attēlā ir atzīmētas apmācībai izmantotās atbilstības, kā arī nobīde secināšanas laikā.

Izvērst

Dekodēšanas plūsma ir diezgan līdzīga melnraksta un mērķa dekodēšanai, ar izņēmumu, ka tiek pārvietoti gan slēptie stāvokļi, gan logiti. Mūsu ieviešanā tiek koplietoti visi saistītie paraugu ņemšanas un logitu apstrādes soļi, specializējot tikai modeļa forward izsaukumus. Kad tiek prognozēti vairāki tokeni, melnraksta modelis prognozēšanai izmanto melnraksta slēptos stāvokļus, kā arī aizpilda KV kešatmiņas ierakstus, pamatojoties uz savām pazīmēm. Ilgtermiņā tas var pasliktināt precizitāti. Pēc tam, palaižot melnraksta modeli, lai aizpildītu KV kešatmiņas ierakstu mērķa prognozei, mēs to palaižam visai sekvencei, par ievadi ņemot precīzākus mērķa slēptos stāvokļus. Tā kā šie melnraksta modeļi ir mazi, papildu izmaksas par papildu tokenu apstrādi ir niecīgas.

MTP galvu (Heads) apmācība

Lai gūtu labumu no MTP, mēs izveidojām infrastruktūru, kas nepieciešama, lai apmācītu MTP galvas, kuras pievienotas mūsu smalki noregulētajiem modeļiem uz Perplexity datu kopām, darbinot tās uz viena mezkla ar 8xH100 ierīcēm. Apmēram vienas dienas laikā mēs varam izveidot galvas modeļiem, sākot no Llama-1B līdz Llama-70B un DeepSeek V2-Lite. Lielākiem modeļiem mēs paļaujamies uz MTP galvām, kas izveidotas smalkas regulēšanas (fine-tuning) procesā.

MTP apmācības mērķis ir saskaņot melnraksta slēptos stāvokļus un logitus, kas ekstrapolēti no mērķa slēptajiem stāvokļiem, ar mērķa nākamā tokena logitiem un slēptajiem stāvokļiem. Tā kā slēpto stāvokļu secināšana ir dārga, mēs tos iepriekš aprēķinām, izmantojot mūsu secināšanai optimizēto mērķa modeļa ieviešanu, lai tos izmantotu apmācības laikā. Tomēr, lai apstiprinātu MTP secināšanas ieviešanu un nodrošinātu, ka skaitliskās atšķirības kvantēšanas vai optimizāciju dēļ netraucē rezultātiem, validācijas zuduma un precizitātes novērtēšanai mēs pilnībā atkārtoti izmantojam gan mērķa, gan melnraksta modeļu secināšanas ieviešanu.

Mērogojot no ShareGPT datu kopas, kas tika izmantota oriģinālajā rakstā, uz lielākiem paraugiem, mēs pamanījām, ka MTP galvu arhitektūra, kas aprakstīta un ieviesta EAGLE rakstā, neizdevās apmācīt 70B izmēra modeļiem. Atšķirībā no ShareGPT, kas saturēja lielāku skaitu īsāku seku, mēs veicam apmācību ar nedaudz mazāku skaitu ievērojami garāku uzvedņu (prompts). Tā kā oriģinālās EAGLE galvas nedaudz atšķīrās pēc struktūras no tipiska transformatora, mēs atgriezām dažus RMS normalizācijas slāņus, kas bija noņemti. Mēs atklājām, ka tas ļāva ne tikai panākt apmācības konverģenci, bet arī uzlaboja galvu precizitāti par dažiem procentpunktiem.

Izvērst

Slāņu normalizācija ne tikai veicina apmācību, bet arī šo normalizāciju atkārtota ieviešana ir matemātiski intuitīva. MTP galvas atkārtoti izmanto mērķa modeļa iegulumus un logitu projekcijas, jo tās var būt ievērojamas pēc izmēra (apmēram 2 GB Llama 70B gadījumā). Apmācības laikā tās tiek iesaldētas, un tiek sagaidīts, ka MTP slānis iemācās iegult prognozes tajā pašā vektora telpā, ko sākotnējā modeļa projekcijas slānis apguva apmācības laikā. Noņemot normalizācijas, tiek sagaidīts, ka viens MLP iemācīsies to pašu funkciju, ko MLP, kuram seko normalizācija, kas apgrūtina melnraksta un mērķa modeļu slēpto stāvokļu savstarpējo atbilstību.

Secināšana ar spekulatīvo dekodēšanu

Secināšanas dzinējā, lai ģenerētu tokenus ievades sekvencēm, tās vispirms ir jāgrupē saprātīga izmēra partijās, pēc tam KV kešatmiņā ir jāpiešķir lapas nākamajiem tokeniem. Ievades tokeni un KV lapas informācija pēc tam tiek sapakota buferī, kas tiek pārraidīts visiem paralēlajiem ratiem (ranks), kuri darbina modeli. Visbeidzot, metadati tiek kopēti GPU atmiņā un modelis tiek izpildīts, lai iegūtu logitus, no kuriem tiek ņemts nākamais tokens.

Atšķirībā no noteiktām ieviešanām, kas vāji savieno melnraksta un mērķa secināšanas serveri, izmantojot starpprogrammatūru (wrapper), kura orķestrē pieprasījumus starp tiem, mūsu melnraksta un mērķa pāri ir cieši savienoti un veic ģenerēšanas soļus saskaņoti. Partiju plānošana un KV lapu piešķiršana tiek koplietota starp modeļiem visiem spekulatīvās dekodēšanas veidiem: tas apvieno loģiku, kas savieno modeli ar vispārējo secināšanas serveri, jo tie visi nodrošina vienu un to pašu saskarni.

Secināšanas izpildlaiks Perplexity ir balstīts uz FlashInfer, kas nosaka metadatus, kuri jāizveido, lai konfigurētu un ieplānotu uzmanības kodolu (attention kernel). Ņemot vērā dažas ievades sekvences, kas veido partiju iepriekšējai aizpildīšanai, dekodēšanai vai verificēšanai, ir jāveic darbs CPU pusē, lai piešķirtu starpbuferus un aizpildītu noteiktus konstantes buferus, ko izmanto uzmanības mehānismā. Šis darbs papildina partiju plānošanas un KV lapu piešķiršanas izmaksas, kas rada arī aizturi, kura ir jānoslēpj, lai maksimāli palielinātu GPU izmantošanu.

Lai gan mēs pilnībā paralelizējām CPU puses un GPU puses darbu secināšanai bez spekulācijām, mēs atklājām, ka CPU un GPU līdzsvars spekulatīvajai dekodēšanai ir sarežģītāks. Galvenais izaicinājums izriet no tā, ka pieņemto tokenu skaits nosaka sekvences garumu nākamajai izpildei, ieviešot grūti novēršamu GPU un CPU sinhronizācijas punktu. Mēs eksperimentējām ar dažādām plānošanas shēmām, lai vislabāk noslēptu CPU darba aizturi.

Melnraksta un mērķa grafiks (Draft-Target Schedule)

Neskatoties uz to, ka tas ir mazāks par mērķa modeli, ja kā melnraksts tiek izmantots viss LLM, tas joprojām rada ievērojamu aizturi GPU, nodrošinot zināmu rezervi dārgu CPU operāciju noslēpšanai. Tā kā mazāki modeļi negūst labumu no tenzoru paralēlisma, pastāv neatbilstība starp rangu skaitu, starp kuriem ir sadalīts mērķa un melnraksta modelis. Mūsu ieviešanā melnraksta modelis darbojas tikai TP grupas līdera rangā.

Izvērst

Kā minēts iepriekš, dekodēšanas solis pārvieto logitus uz nākamo izpildi. Tas ļauj mums pārklāt vienu melnraksta modeļa izpildi ar CPU puses partijas plānošanas darbu. Pēc partijas izveidošanas atkārtoti izsaukumi uz paraugu ņemšanas rīku un melnraksta modeli ģenerē melnraksta tokenus. Paralēli tam mērķa modelim tiek izveidota partija pārbaudei un sinhronizēta ar paralēli strādājošajiem procesiem. Mērķa logiti tiek verificēti un ņemti paraugos, lai noteiktu pieņemtos sekvences garumus. Šajā brīdī ir nepieciešama GPU un CPU sinhronizācija, lai noteiktu nākamos sekvences garumus. Tā kā melnraksta modelis darbojas tikai līdera mezglā, tā partija tiek iestatīta secīgi, un tā izpilde tiek sākta, lai aizpildītu tā KV kešatmiņas ierakstus ar papildu tokenu, ko mērķa modelis ir ģenerējis. Logiti, ko ģenerējis šis melnraksta modelis pašreizējā izpildē, tiks izmantoti, lai ņemtu paraugos pirmo melnraksta tokenu nākamajā izpildē. Pats svarīgākais ir tas, ka, kamēr melnraksts darbojas, var ieplānot nākamo partiju.

MTP grafiks vienam tokenam

Lai gan izpildlaiks vēl nenodrošina Eagle stila melnraksta koka izpēti, mēs ieviesām īpašu šīs shēmas gadījumu, ņemot vērā melnraksta tokenu lineāru sekvenci, ko ģenerējis modelis, kura izmērs ir vienāds ar vienu transformatora dekodētāja slāni. Šo shēmu var izmantot melnraksta prognozēšanai, izmantojot DeepSeek R1 atvērtā koda svarus. Viena tokena prognozēšanas apakšgadījums ir interesants, jo lielie MTP slāņi sasniedz pietiekami augstus pieņemšanas rādītājus, lai attaisnotu to pieskaitāmās izmaksas.

MTP plānošana ir nedaudz sarežģītāka, jo melnraksta modelis ir daudz ātrāks, tādējādi noslēpjot mazāku CPU puses aizturi. Turklāt melnraksts tiek sadalīts kopā ar mērķa modeli, kas prasa koplietotās atmiņas pārsūtīšanu partijas informācijai. Izpilde sākas, pārsūtot partijas informāciju un ņemot paraugā pirmo tokenu no pārvietotajiem logitiem, līdzīgi kā iepriekšējā shēmā. Nākamais solis ir mērķa modeļa palaišana tokenu validēšanai, apstrādājot 2 * D tokenus, kur D ir dekodēšanas partijas izmērs. Tas ir ideāli piemērots mikropartiju veidošanai Mixture-of-Experts (MoE) modeļos, izmantojot lēnākus starpsavienojumus, piemēram, InfiniBand, jo partija vienmērīgi sadalās uz pusēm. Mērķa slēptie stāvokļi tiek pārcelti uz nākamo melnraksta izpildi, savukārt logiti tiek padoti paraugu ņemšanas rīkam verificēšanai.

Izvērst

Veicot ierobežotu papildu darba apjomu GPU, mēs izvairāmies no GPU un CPU sinhronizācijas pēc melnraksta sekvences pieņemšanas. Pēc tam, kad mērķa modeļa ievades tokeni ir pārbīdīti, kodols ievieto nākamos mērķa tokenus to atbilstošajās vietās. Pēc tam melnraksts tiek palaists atkārtoti ar to pašu partijas informāciju kā mērķa modelim, aizpildot KV kešatmiņas ierakstus un veidojot logitus un slēptos stāvokļus nākamajai izpildei, veicot zināmu lieku darbu pie tokeniem, kuri netika pieņemti. Šādās situācijās neizmantotā darba aizture ir tik tikko izmērāma, pateicoties melnraksta modeļa mazajam izmēram. Paralēli melnraksta izpildei sekvences garumi tiek noteikti CPU, un tiek sākta nākamās partijas plānošana, negaidot GPU darba beigšanos.

Papildu darba pieskaitāmās izmaksas melnraksta slānī nav pamanāmas uzmanības mehānismā, taču MLP slāņi ir problemātiskāki. Tā kā matricas reizināšanas instrukcijas aizpilda līdz 64 robežai tokenu skaita dimensijā, ja dubultošanai nav nepieciešams ievērojami vairāk bloku, pieskaitāmās izmaksas ir noslēptas. Garākām melnraksta sekvencēm pieskaitāmās izmaksas ir dārgākas, un parastajiem melnraksta un mērķa modeļiem izmantotā shēma darbojas labāk.

Atsauces

Ātra secināšana no transformatoriem, izmantojot spekulatīvo dekodēšanu

EAGLE: Spekulatīvā paraugu ņemšana prasa pārskatīt pazīmju nenoteiktību

EAGLE-2: Ātrāka valodu modeļu secināšana ar dinamiskiem melnraksta kokiem

EAGLE-3: Lielo valodu modeļu secināšanas paātrinājuma mērogošana, izmantojot testa laika apmācību

Medusa: Vienkārša LLM secināšanas paātrinājuma sistēma ar vairākām dekodēšanas galvām

Slāņu normalizācija

FlashInfer: Efektīvs un pielāgojams uzmanības dzinējs LLM secināšanas apkalpošanai

DeepSeek-V3 tehniskais ziņojums

Citāts

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