Urýchlenie modelu Sonar prostredníctvom špekulácií

Špekulatívne dekódovanie urýchľuje rýchlosť generovania veľkých jazykových modelov (LLM) použitím rýchleho a malého návrhového modelu na vytvorenie kandidátov na dokončenie, ktoré overí väčší cieľový model. V rámci tohto prístupu sa namiesto jedného spustenia nákladného cieľového modelu, ktoré vytvorí jediný token, v jednom kroku vygeneruje viacero tokenov. Tu uvádzame podrobnosti implementácie rôznych druhov špekulatívneho dekódovania, ktoré sú v Perplexity použité na zníženie latencie medzi tokenmi pri modeloch Sonar.

Špekulatívne dekódovanie

Špekulatívne dekódovanie využíva štruktúru prirodzených jazykov a autoregresívnu povahu transformátorov na urýchlenie generovania tokenov. Aj keď väčšie modely, ako napríklad Llama-70B, obsahujú viac znalostí než menšie modely, ako napríklad Llama-1B, pri niektorých jednoduchších úlohách dosahujú podobné výsledky. Tento prekrývajúci sa priestor naznačuje, že určité sekvencie lepšie generujú lacnejšie modely, pričom zložité problémy sa ponechávajú na väčšie modely. Výzva spočíva v určení toho, ktoré dokončenia sú lepšie a či je generovanie menšieho modelu rovnako kvalitné ako generovanie väčšieho modelu.

Našťastie sú LLM autoregresívne transformátory: keď dostanú sekvenciu tokenov, vygenerujú pravdepodobnostné rozdelenie nasledujúceho tokenu. Logity odvodené zo stredných príznakov spojených s tokenmi v vstupnej sekvencii navyše naznačujú, aké pravdepodobné je pre model vydať presne tieto tokeny. Táto vlastnosť umožňuje špekuláciu: ak menší model vygeneruje sekvenciu tokenov začínajúcu od vstupnej predpony, možno ju prehnať väčším modelom, aby sa zistilo, do akej miery sa zhoduje s cieľovým modelom. Každá predpona kandidátov je ohodnotená pravdepodobnosťou a vyberie sa tá najdlhšia, ktorá prekročí prah akceptácie. Ako bonus cieľový model poskytuje aj dodatočný token zadarmo: ak model návrhu vygeneruje n tokenov, v jednom kroku možno vyemitovať až n + 1 tokenov.

Zväčšiť

V čase inferencie možno proces špekulatívneho vzorkovania rozdeliť približne na 4 fázy:

Predbežné naplnenie (prefill): cieľový aj návrhový model musia byť spustené na vstupnej sekvencii, aby sa naplnili položky vyrovnávacej pamäte KV. Zatiaľ čo niektoré schémy, ako napríklad Medusa, používajú na predikciu jednoduchšie husté vrstvy, v tomto príspevku sa zameriavame na návrhové modely založené na transformátoroch, ktoré potrebujú vlastné vyrovnávacie pamäte KV.

Generovanie návrhu: návrhový model iteruje s cieľom vytvoriť určitý počet pevných tokenov. Návrhová sekvencia môže byť lineárna alebo model môže preskúmať štruktúru podobnú stromu do určitej hĺbky (EAGLE, Medusa). Tu sa zameriavame na lineárne sekvencie.

Akceptácia: cieľový model beží na návrhovej sekvencii a vytvára logity zodpovedajúce každému návrhovému tokenu. Určí sa dĺžka najdlhšej akceptovateľnej sekvencie.

Cieľové generovanie: keďže cieľový model vygeneroval logity, na mieste nezhody alebo na konci sekvencie logity zodpovedajú následnému tokenu. Tieto logity možno vzorkovať, čím sa poskytne robustný token z cieľového modelu, ktorý uzavrie sekvenciu.

Existujú rôzne metódy implementácie špekulatívneho dekódovania. V tomto príspevku sa zameriavame na schémy, ktoré sme použili na urýchlenie modelov Sonar pomocou interného 1B modelu, ako aj na predikčné mechanizmy, ktoré budujeme na urýchlenie modelov v rozsahu DeepSeek.

Cieľ-Návrh (Target-Draft)

Špekulatívne dekódovanie možno dosiahnuť spojením existujúceho malého modelu LLM ako návrhového modelu s cieľovým modelom na generovanie kandidátskych sekvencií. V produkcii sme urýchlili Sonar pomocou modelu Llama-1B jemne doladeného na rovnakej databáze ako cieľový model. Hoci tento prístup nevyžadoval trénovanie návrhu od nuly, malý model stále využíva významnú kapacitu vyrovnávacej pamäte KV a zavádza miernu réžiu predbežného naplnenia, čím sa zvyšuje TTFT.

V rámci tohto prístupu dekodér špekuluje iba pri dávkach určených len na dekódovanie, pričom generuje tokeny prostredníctvom štandardného vzorkovania počas predbežného naplnenia alebo pri zmiešaných dávkach predbežného naplnenia a dekódovania. Vo fáze predbežného naplnenia sa cieľové logity okamžite vzorkujú, aby sa predbežne naplnil aj novo vygenerovaný token vo vyrovnávacej pamäti KV návrhového modelu. Návrhový model sa zatiaľ nevzorkuje, ale logity, ktoré produkuje, sa prenášajú do fázy dekódovania.

Zväčšiť

Pri dekódovaní sa návrhový model posúva vpred a v každej fáze vzorkuje špičkový token. Po dosiahnutí požadovanej dĺžky návrhu sa tokeny preženú cieľovým modelom, aby sa vytvorili logity, na základe ktorých vzorkovač identifikuje dĺžku akceptovanej sekvencie. Akceptácia sa určí porovnaním úplných pravdepodobnostných rozdelení z návrhového a cieľového modelu. Keďže cieľový model vždy vydá jednu sadu logitov po akceptovanej návrhovej sekvencii, táto sada sa vzorkuje, aby sa vytvoril dodatočný výstup. Keďže návrhový model tento akceptovaný token ešte nevidel, spustí sa znova, aby sa naplnili jeho zodpovedajúce položky vyrovnávacej pamäte KV na prípravu na ďalší krok dekódovania, pričom sa logity opäť prenesú.

EAGLE

EAGLE je schéma špekulatívneho dekódovania, ktorá skúma viacero návrhových sekvencií generovaných prostredníctvom stromového prechodu pravdepodobných návrhových tokenov. Pevný (EAGLE) alebo dynamicky formovaný (EAGLE-2) strom sa skúma pomocou po sebe idúcich spustení návrhových tokenov, pričom v každom uzle sa berú do úvahy kandidáti Top-K namiesto nasledovania tokenu s najvyšším skóre v lineárnej sekvencii. Sekvencie sa potom ohodnotia a vyberie sa najdlhšia vhodná sekvencia na pokračovanie, pričom sa pridá aj dodatočný token z cieľového modelu.

Zväčšiť

Na dosiahnutie presnejšej predikcie predikuje návrhový model EAGLE nielen na základe tokenov, ale aj pomocou cieľových príznakov (skryté stavy poslednej vrstvy) cieľového modelu. Nevýhodou systému EAGLE je potreba trénovať vlastné malé návrhové modely, ktoré sú dostatočne presné na generovanie vhodných kandidátov v rámci nízkeho rozpočtu na latenciu. Návrhový model je zvyčajne jedna transformátorová vrstva identická s dekodérovou vrstvou pôvodného modelu, ktorá je pevne spojená s cieľovým modelom prostredníctvom väzby na jeho vnorenia a projekcie lm_head. Keďže to vyžaduje menšiu kapacitu vyrovnávacej pamäte KV, EAGLE má menšiu pamäťovú stopu.

Na overenie stromových sekvencií v cieľovom modeli sa musia použiť vlastné masky pozornosti. Použitie vlastnej masky pozornosti pre celú sekvenciu však nanešťastie výrazne spomaľuje pozornosť pri reálnych dĺžkach vstupov (až o 50 %), čím sa ruší časť zrýchlenia dosiahnuteľného prostredníctvom špekulácií. Z tohto dôvodu sme zatiaľ nenasadili úplné prieskumy stromov do produkcie, namiesto toho sa zameriavame na špeciálny prípad predikcie jedného tokenu prostredníctvom schém podobných MTP, ktoré sú prezentované v technickej správe DeepSeek-V3.

MTP

Táto schéma je podobná dekódovaniu typu návrh-cieľ, s výnimkou toho, že na predikciu sa okrem tokenov používajú aj skryté stavy. Vo fázach predbežného naplnenia aj dekódovania je potrebné vykonať o niečo viac práce v porovnaní s bežnou špekuláciou návrh-cieľ. Návrhový model používa tokeny aj skryté stavy: token t_{i+1} sa vzorkuje z logitov L_i zodpovedajúcich tokenu t_i, ktoré sú zase odvodené zo skrytých stavov H_i. V dôsledku toho sa vstupné vyrovnávacie pamäte tokenov musia posunúť o jeden krok doľava vzhľadom na vektory skrytých stavov vydané cieľovým modelom. Nasledujúci obrázok označuje zhody použité na trénovanie, ako aj posun počas inferencie.

Zväčšiť

Tok dekódovania je dosť podobný dekódovaniu typu návrh-cieľ, s výnimkou prenášania skrytých stavov aj logitov. Naša implementácia zdieľa všetky súvisiace vzorkovania a spracovania logitov, pričom špecializuje iba vyvolania dopredného prechodu modelu. Keď sa predikuje viacero tokenov, návrhový model používa na predikciu skryté stavy návrhu a taktiež zapĺňa položky vyrovnávacej pamäte KV na základe vlastných príznakov. Z dlhodobého hľadiska to môže zhoršiť presnosť. Následne pri spúšťaní návrhového modelu na naplnenie položky vyrovnávacej pamäte KV pre cieľovú predikciu ho spustíme na celej sekvencii, pričom ako vstupy použijeme presnejšie skryté stavy cieľového modelu. Keďže tieto návrhové modely sú malé, dodatočné náklady na spracovanie ďalších tokenov sú zanedbateľné.

Trénovanie MTP hláv

Aby sme mohli ťažiť z výhod MTP, vybudovali sme infraštruktúru potrebnú na trénovanie MTP hláv pripojených k našim jemne doladeným modelom na databázach Perplexity, bežiacich na jednom uzle so zariadeniami 8xH100. Za približne jeden deň dokážeme vytvoriť hlavy pre modely od Llama-1B až po Llama-70B a DeepSeek V2-Lite. V prípade väčších modelov spoliehame na MTP hlavy vytvorené počas procesu jemného doladenia.

Cieľom trénovania MTP je priradiť skryté stavy návrhu a logity extrapolované zo skrytých stavov cieľového modelu k logitom nasledujúceho tokenu a skrytým stavom cieľového modelu. Keďže inferencia pre skryté stavy je nákladná, vopred ich vypočítame pomocou našej na inferenciu optimalizovanej implementácie cieľového modelu, ktorá sa použije počas trénovania. Na overenie straty a presnosti odhadu však plne opätovne používame inferenčnú implementáciu cieľového aj návrhového modelu, aby sme overili inferenčnú implementáciu MTP a zaistili, že numerické rozdiely spôsobené kvantizáciou alebo optimalizáciami nebudú brániť výsledkom.

Pri prechode z databázy ShareGPT použitej v pôvodnom článku na väčšie vzorky sme si všimli, že štruktúra MTP hláv popísaná a implementovaná v článku EAGLE sa nepodarilo natrénovať pre modely veľkosti 70B. Na rozdiel od ShareGPT, ktorá obsahovala väčší počet kratších sekvencií, trénujeme na o niečo menšom počte podstatne dlhších výziev. Keďže pôvodné hlavy EAGLE sa štruktúrou mierne odlišovali od typického transformátora, opätovne sme zaviedli niektoré vrstvy normalizácie RMS, ktoré boli odstránené. Zistili sme, že to nielen umožnilo konvergenciu trénovania, ale taktiež to zvýšilo presnosť hláv o niekoľko percentuálnych bodov.

Zväčšiť

Normalizácie vrstiev nielenže uľahčujú trénovanie, ale ich opätovné zavedenie je tiež matematicky intuitívne. MTP hlavy opätovne používajú vnorenia a projekcie logitov cieľového modelu, pretože môžu byť značne veľké (približne 2 GB pre Llama 70B). Počas trénovania sú tieto časti zmrazené a očakáva sa, že MTP vrstva sa naučí vkladať predikcie do rovnakého vektorového priestoru, aký sa projekčná vrstva pôvodného modelu naučila počas trénovania. Vynechaním normalizácií sa od jedinej vrstvy MLP očakáva, že sa naučí rovnakú funkciu ako MLP nasledovaná normalizáciou, čo bráni zhodovaniu medzi skrytými stavmi návrhu a cieľovými modelmi.

Inferencia so špekulatívnym dekódovaním

V inferenčnom engine je potrebné vstupné sekvencie na generovanie tokenov najprv zoskupiť do dávok primeranej veľkosti a potom v pamäti KV prideliť stránky pre nasledujúce tokeny. Vstupné tokeny a informácie o stránkach KV sa potom zbalia do vyrovnávacej pamäte, ktorá sa vysiela do všetkých paralelných úrovní spúšťajúcich model. Nakoniec sa metadáta skopírujú do pamäte GPU a model sa spustí, aby sa vytvorili logity, z ktorých sa vzorkuje nasledujúci token.

Na rozdiel od určitých implementácií, ktoré voľne spájajú návrhový a cieľový inferenčný server cez obal, ktorý medzi nimi orchestruje požiadavky, sú naše páry návrh-cieľ pevne spojené a prechádzajú generovaním spoločne. Plánovanie dávok a prideľovanie stránok KV sa zdieľa medzi modelmi pre všetky formy špekulatívneho dekódovania: to zjednocuje logiku, ktorá premostí model so zastrešujúcim inferenčným serverom, pretože všetky vystavujú rovnaké rozhranie.

Inferenčný beh v Perplexity je vytvorený okolo FlashInfer, ktorý určuje metadáta potrebné na konfiguráciu a plánovanie jadra pozornosti. Vzhľadom na niektoré vstupné sekvencie tvoriace dávku na predbežné naplnenie, dekódovanie alebo overenie sa musí vykonať práca na strane CPU na pridelenie medzipamätí a naplnenie určitých konštantných vyrovnávacích pamätí používaných v pozornosti. Táto práca je navyše k nákladom na plánovanie dávok a prideľovanie stránok KV, ktoré taktiež spôsobujú latencie, ktoré je nutné skryť, aby sa maximalizovalo využitie GPU.

Zatiaľ čo prácu na strane CPU a GPU sme pre inferenciu bez špekulácií plne paralelizovali, zistili sme, že rovnováha medzi CPU a GPU pri špekulatívnom dekódovaní je zložitejšia. Hlavná výzva vyplýva zo skutočnosti, že počet akceptovaných tokenov určuje dĺžku sekvencie pre následný beh, čo zavádza ťažko odstrániteľný bod synchronizácie medzi GPU a CPU. Experimentovali sme s rôznymi schémami plánovania, aby sme čo najlepšie skryli latenciu práce CPU.

Plán pre návrh a cieľ (Draft-Target Schedule)

Napriek tomu, že je menší než cieľový model, keď sa celý LLM použije ako návrh, stále spôsobuje na GPU značnú latenciu, čím poskytuje určitý priestor na skrytie náročných operácií CPU. Keďže menšie modely neťažia z paralelizmu tenzorov, existuje nezhoda medzi počtom úrovní, medzi ktorými je cieľový a návrhový model rozdelený. V našej implementácii beží návrhový model iba na hlavnej úrovni skupiny TP.

Zväčšiť

Ako bolo uvedené predtým, krok dekódovania prenáša logity do nasledujúceho behu. To nám umožňuje prekrývať jedno spustenie návrhového modelu s prácou plánovania dávok na strane CPU. Po zostavení dávky vygenerujú opakované volania vzorkovača a návrhového modelu návrhové tokeny. Súbežne sa pre cieľový model zostaví dávka na overenie a zosynchronizuje sa s paralelnými pracovnými procesmi. Cieľové logity sa overia a vzorkujú na určenie dĺžok akceptovaných sekvencií. V tomto bode je potrebná synchronizácia medzi GPU a CPU na určenie následných dĺžok sekvencií. Keďže návrhový model beží iba na hlavnom uzle, jeho dávka sa nastavuje postupne a jeho spustenie sa inicializuje, aby sa naplnili položky vyrovnávacej pamäte KV dodatočným tokenom, ktorý cieľový model vytvoril. Logity vytvorené týmto spustením návrhu v aktuálnom behu sa použijú na vzorkovanie prvého návrhového tokenu v nasledujúcom behu. Najdôležitejšie je, že kým návrh beží, možno naplánovať ďalšiu dávku.

Plán MTP pre jeden token

Zatiaľ čo bežiaci systém zatiaľ neposkytuje prieskum stromov návrhov v štýle Eagle, implementovali sme špeciálny prípad tejto schémy, pričom berieme do úvahy lineárnu sekvenciu návrhových tokenov produkovaných modelom veľkosti jednej vrstvy transformátorového dekodéra. Túto schému možno použiť na predikciu návrhu pomocou open-source váh DeepSeek R1. Podprípad predikcie jedného tokenu je zaujímavý, pretože veľké MTP vrstvy dosahujú dostatočne vysokú mieru akceptácie na ospravedlnenie svojej réžie.

Plánovanie MTP je o niečo zložitejšie, pretože návrhový model je oveľa rýchlejší, čím skrýva menšiu latenciu na strane CPU. Návrhový model je navyše rozdelený popri cieľovom modeli, čo si vyžaduje prenosy v zdieľanej pamäti pre informácie o dávkach. Beh sa začína prenosom informácií o dávke a vzorkovaním prvého tokenu z prenesených logitov, podobne ako v predchádzajúcej schéme. Ďalej sa spustí cieľový model na overenie tokenov, pričom sa spracujú 2 * D tokeny, kde D je veľkosť dávky dekódovania. To je ideálne pre dávkovanie v modeloch typu Mixture-of-Experts (MoE) cez pomalšie prepojenia, ako je InfiniBand, pretože dávka sa rovnomerne rozdelí na dve polovice. Skryté stavy cieľového modelu sa prenesú do nasledujúceho behu návrhu, zatiaľ čo logity sa odovzdajú do vzorkovača na overenie.

Zväčšiť

Vykonaním obmedzeného množstva dodatočnej práce na GPU sa vyhneme synchronizácii medzi CPU a GPU po akceptovaní návrhovej sekvencie. Po posunutí vstupných tokenov cieľových modelov vloží jadro ďalšie cieľové tokeny na ich zodpovedajúce miesta. Návrhový model sa potom spustí znova s rovnakými informáciami o dávke ako cieľový model, pričom sa naplnia položky vyrovnávacej pamäte KV a vytvoria sa logity a skryté stavy pre nasledujúci beh, pričom sa vykoná určitá nadbytočná práca na tokenoch, ktoré neboli akceptované. V týchto situáciách je latencia nepoužitej práce takmer merateľná kvôli malej veľkosti návrhového modelu. Súbežne so spustením návrhového modelu sa na CPU určia dĺžky sekvencií a spustí sa plánovanie ďalšej dávky bez toho, aby bolo nutné čakať na ukončenie práce GPU.

Réžia dodatočnej práce v návrhovej vrstve nie je v pozornosti badateľná, avšak vrstvy MLP sú problematickejšie. Keďže inštrukcie na násobenie matíc zarovnávajú na hranicu 64 pozdĺž dimenzie počtu tokenov, ak zdvojnásobenie nevyžaduje výrazne viac blokov, táto réžia zostáva skrytá. Pre dlhšie návrhové sekvencie je réžia nákladnejšia a schéma použitá pre bežné modely návrh-cieľ funguje lepšie.

Referencie

Rýchla inferencia z transformátorov pomocou špekulatívneho dekódovania

EAGLE: Špekulatívne vzorkovanie si vyžaduje prehodnotenie neistoty príznakov

EAGLE-2: Rýchlejšia inferencia jazykových modelov s dynamickými návrhovými stromami

EAGLE-3: Zvyšovanie rozsahu zrýchlenia inferencie veľkých jazykových modelov prostredníctvom testu v čase trénovania

Medusa: Jednoduchý rámec na akceleráciu inferencie LLM s viacerými dekódovacími hlavami

Normalizácia vrstiev

FlashInfer: Efektívny a prispôsobiteľný engine pozornosti pre inferenčné poskytovanie LLM

Technická správa DeepSeek-V3

Citácia

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