Spekulativní dekódování urychluje generování Velkých jazykových modelů (LLM) pomocí rychlého a malého návrhového modelu, který vytváří kandidáty na dokončení, které ověřuje větší cílový model.

Podle tohoto schématu místo jedné drahé akce cíle, která produkuje jediný token, se v jednom kroku vysílá více tokenů. Zde představujeme detaily implementace různých druhů spekulativního dekódování, aplikovaných v Perplexity ke snížení latence mezi tokeny na modelech Sonar.

Spekulativní dekódování

Spekulativní dekódování využívá struktury přirozených jazyků a autoregresivní povahy transformátorů pro urychlení generování tokenů. Přestože větší modely, jako Llama-70B, obsahují více znalostí než menší modely, jako Llama-1B, u některých jednodušších úloh fungují podobně. Tato podobnost naznačuje, že určité sekvence jsou lépe generovány méně nákladnými modely, přičemž složité problémy nechávají většímu modelu. Výzvou je určit, které dokončení jsou lepší a zda generování menšího modelu je stejně kvalitní jako u většího.

Naštěstí jsou LLM autoregresivní transformátory: při zadání sekvence tokenů vytvoří rozložení pravděpodobnosti pro následující token. Navíc logity odvozené z mezivýsledků spojených s tokeny v vstupní sekvenci také indikují, jak pravděpodobné je, že model vydá právě tyto tokeny. Tato vlastnost umožňuje spekulaci: pokud je sekvence tokenů generována menším modelem začínajícím z prefixu, může být spuštěna přes větší model pro určení, jak dobře se vyrovnává s cílovým modelem. Každý prefix kandidátů je hodnocen pravděpodobností a nejdelší nad akceptačním prahem je vybraná. Jako bonus cílový model také poskytuje následný token zdarma: pokud návrhový model generuje tokeny n, může být v jednom kroku vysíláno až n + 1.

Při času analýzy může být proces spekulativního vzorkování rozdělen zhruba do 4 etap:

1. Předvolba: jak cílový, tak i návrhový model musí být spuštěny na vstupní sekvenci, aby se naplnila KV mezipaměť. Zatímco některé schémata, jako Medúza, používají jednodušší husté vrstvy pro predikci, v tomto příspěvku se zaměřujeme na návrhové koncepty založené na transformátorech, které potřebují své vlastní cache KV.

2. Generování návrhu: návrhový model iteruje, aby vytvořil počet pevných tokenů. Návrhová sekvence může být lineární nebo model může prozkoumat strukturu podobnou stromu až do určité hloubky (EAGLE, Medúza). Zde se zaměřujeme na lineární sekvence.

3. Přijetí: cílový model běží na návrhové sekvenci, vytváří logity odpovídající jednotlivým návrhovým tokenům. Stanovena je délka nejdelší přijatelné sekvence.

4. Generování cíle: jelikož cílový model generoval logity, na nesrovnalém místě nebo na konci sekvence logity odpovídají následujícímu tokenu. Tyto logity mohou být vzorkovány, aby poskytly robustní token z cíle, zakončující sekvenci.

Existuje několik metod pro implementaci spekulativního dekódování. V tomto příspěvku se zaměříme na schémata, která jsme použili k zrychlení modelů Sonar pomocí našeho interního modelu 1B, stejně jako na predikční mechanismy, které budujeme k urychlení modelů měřítka DeepSeek.

Cílový-Návrh

Spekulativního dekódování lze dosáhnout spojením stávajícího malého LLM jako návrhového modelu s cílovým modelem pro generování kandidátských sekvencí. V produkci jsme zrychlili Sonar pomocí modelu Llama-1B doladěného na stejnou datovou sadu jako cíl. I když tento přístup nevyžadoval náročné trénování návrhu od nuly, malý model stále vyžaduje značnou kapacitu paměti KV a zvyšuje mírný předvolací overhead, což zvyšuje TTFT.

Podle tohoto schématu dekodér spekuluje pouze o šaržích dekódování, generováním tokenů pomocí standardního vzorkování během předvolby nebo na smíšených šaržích předvolby-dekódování. Ve fázi předvolby jsou cílové logity okamžitě vzorkovány, aby také předvolily nově vygenerovaný token v cache KV návrhu. Návrh ještě není vzorkován, ale logity, které produkuje, jsou přeneseny do fáze dekódování.

Při dekódování je návrhový model pokročilý, vzorkuje top token v každé fázi. Poté, co je dosažena požadovaná délka návrhu, jsou tokeny spuštěny skrz cílový model, aby vytvořily logity, na základě kterých vzorkovač identifikuje akceptovanou délku sekvence. Přijetí je určeno porovnáním plných rozložení pravděpodobnosti mezi návrhem a cílem. Jelikož cíl vždy poskytuje jednu sadu logitů následujících po akceptované návrhové sekvenci, je to vzorkováno k vytvoření dalšího výstupu. Jelikož návrhový model ještě neviděl tento přijatý token, je znovu spuštěn, aby naplnil odpovídající záznamy cache KV pro přípravu na další dekódovací krok, znovu přenášející logity.

EAGLE

EAGLE je schéma spekulativního dekódování, které prozkoumává více návrhových sekvencí vytvořených prostřednictvím stromovitého průchodu pravděpodobnými návrhovými tokeny. Pevný (EAGLE) nebo dynamicky tvarovaný (EAGLE-2) strom je prozkoumán pomocí po sobě jdoucích exekucí návrhových tokenů, přičemž se na každém uzlu berou v úvahu Top-K kandidáti namísto sledování nejvyššího skórujícího tokenu v lineární sekvenci. Sekvence jsou pak hodnoceny a nejdelší vhodná je vybrána k pokračování, také se přidává další token z cíle.

Aby bylo dosaženo přesnější predikce, EAGLE návrhový model předpovídá nejen na základě tokenů, ale také s využitím cílových charakteristik (skryté stavy poslední vrstvy) cílového modelu. Nevýhodou EAGLE je potřeba vytrénovat vlastní malé návrhové modely, které jsou dostatečně přesné, aby generovaly vhodné kandidáty ve stanoveném latencí rozpočtu. Typicky je návrhový model jednou transformátorovou vrstvou identickou k dekódovací vrstvě původního modelu, která je úzce spojena s cílem navázáním na jeho embedování a lm_head projekce. Jelikož to vyžaduje méně capacity cache KV, EAGLE má nižší paměťovou stopu.

Aby se ověřily stromové sekvence v cílovém modelu, musí být používány vlastní masky pozornosti. Bohužel používání vlastní masky pozornosti pro celou sekvenci významně zpomaluje pozornost pro realistické délky vstupu (až o 50%), což některé zrychlení dosažené skrze spekulaci ruší. Z tohoto důvodu jsme dosud nenasadili celou stromovou průzkum do produkce, místo toho jsme se soustředili na speciální případ predikce jednotlivých tokenů prostřednictvím MTP-like schémat představených v Technické zprávě DeepSeek-V3.

MTP

Toto schéma je podobné dekódování návrh-cíl, s výjimkou použití skrytých stavů společně s tokeny pro predikci. Mírně více práce musí být vykonáno jak v etapách předvolby, tak u dekódování ve srovnání s běžnou spekulací návrh-cíl. Návrhový model používá jak tokeny, tak i skryté stavy: token t_{i+1} je vzat z logitů L_i, které odpovídají tokenu t_i, které jsou zase odvozeny ze skrytých stavů H_i. Z tohoto důvodu se musí posunout levé bufry vstupních tokenů ve vztahu k výstupním vektorům skrytého stavu cíle. Následující graf znázorňuje použití pro trénink, stejně jako posun během inferenčního zpracování.

Tok dekódování je velmi podobný dekódování návrh-cíl, s výjimkou přenášení jak skrytých stavů, tak logitů. Naše implementace sdílí veškeré spojené vzorkování a zpracování logitů, specifické jsou pouze přímé volání modelu. Když je předpovězeno více tokenů, návrhový model používá návrhové skryté stavy pro predikci, také naplňuje záznamy mezipaměti KV na základě vlastních charakteristik. Z dlouhodobého hlediska to může degradovat přesnost. Následně, když se spustí návrhový model pro naplnění záznamů cache KV pro predikci cíle, spustíme jej na celé sekvenci využívajíc přesnější cílové skryté stavy jako vstupy. Jelikož jsou tyto návrhové modely malé, dodatečné zpracování tokenů je zanedbatelné.

Trénování MTP hlav

Abychom využili MTP, vybudovali jsme infrastrukturu potřebnou k trénování MTP hlav napojených na naše doladěné modely na Perplexity’s datech, spuštěné na jednom uzlu s 8xH100 zařízeními. Během jednoho dne můžeme vytvořit hlavy pro modely v rozsahu od Llama-1B po Llama-70B a DeepSeek V2-Lite. Pro větší modely se spoléháme na MTP hlavy vybudované během doladěného procesu.

Cílem trénování MTP je sladit návrhové skryté stavy a logity extrapolované z cílových skrytých stavů k dalším logitům tokenů a skrytým stavům cíle. Jelikož je inferenční provádění skrytých stavů nákladné, předpočítáváme je pomocí naší optimalizované implementace cílového modelu pro trénink. Nicméně pro ověření inferenční implementace MTP a zajištění, že numerické rozdíly v důsledku kvantizace nebo optimalizací nebrání výsledkům, při odhadu ztráty ověření a přesnosti plně znovu použijeme implementaci inference cílových i návrhových modelů.

Při škálování z datek ShareGPT použitého v původní práci na větší vzorky jsme si všimli, že architektura MTP hlav načrtnutá a implementovaná v dokumentu EAGLE se pro modely o velikosti 70B nedokázala trénovat. Na rozdíl od ShareGPT, který obsahoval větší počet kratších sekvencí, trénujeme na o něco menším počtu podstatně delších výzev. Vzhledem k tomu, že původní EAGLE hlavy se mírně lišily od struktury typického transformátoru, znovu jsme zavedli některé vrstvy normalizace RMS, které byly odstraněny. Zjistili jsme, že to nejen umožnilo trénovala, ale také zvýšilo přesnost hlav o několik procentních bodů.

Nejen, že normalizace vrstev usnadňuje tréninku, znovuzavedení normalizace je také matematicky intuitivní. MTP hlavy znovu používají vkládání a logitové projekce cílového modelu, jelikož mohou být podstatné velikostí (asi 2 GB pro Llama 70B). Během tréninku jsou tyto zmraženy a očekává se, že MTP vrstva se naučí vsazovat do stejného vektorového prostoru, jaký vrstvy projekce původního modelu naučily během tréninku. Opomenutím normalizace se očekává, že jedna MLP vrstva naučí stejnou funkci jako MLP následovaná normalizací, což narušuje vyrovnání mezi návrhovými a cílovými skrytými stavy modelu.

Inferenční zpracování se spekulativním dekódováním

V inferenčním enginu, aby byly generovány tokeny pro vstupní sekvence, je třeba je nejprve seskupit do rozumně velkých šarží, poté musí být přiděleny stránky v mezipaměti KV pro další tokeny. Vstupní tokeny a informace o stránkách KV jsou pak zabaleny do bufferu vysílaného na všechny paralelní hodnosti běžící model. Nakonec je metadata zkopírována do paměti GPU a model je spuštěn, aby vytvořil logity, z kterých je vzorkován další token.

Na rozdíl od některých implementací, které volně spojují návrhový a cílový inferenční server jako obalu, který organizuje požadavky mezi nimi, jsou naše páry návrh-cíl úzce spojeny a procházejí generování jednotně. Plánování šarží a přidělení stránek KV je sdíleno mezi modely pro všechny formy spekulativního dekódování: to unifikuje logiku, která propojuje model s centrálním inferenčním serverem, jak všechny vystavují stejné rozhraní.

Čas běhu inference v Perplexity je strukturován kolem FlashInfer, který určuje metadata, která je třeba sestavit, aby byla konfigurována a plánována jádro pozornosti. Vzhledem k některým vstupním sekvencím tvořícím šarži, pro prefill, dekódování nebo ověření musí být provedená práce na straně CPU pro přidělení mezibufrů a naplnění určitých konstantních bufrových prvků použitých v pozornosti. Tato práce je navíc ke nákladům na plánování šarží a přidělení stránek KV, což také způsobuje latenci, kterou je třeba skrýt pro maximalizaci využití GPU.

Zatímco jsme zcela paralelizovali práci na straně CPU a GPU pro inference bez spekulace, zjistili jsme, že rovnováha CPU-GPU pro spekulativní dekódování je složitější. Hlavní výzva vychází ze skutečnosti, že počet přijatých tokenů určuje délku sekvence pro následující běh, zavádějící obtížně se vyhnout bod synchronizace GPU-CPU. Zkoumali jsme různé plány, abychom co nejlépe skryli latenci práce CPU.

Harmonogram návrh-cíl

Přestože je menší než cílový model, když je celé LLM použito jako návrh, stále zavádí značnou latenci na GPU, což poskytuje určitou rezervu pro skrytí drahých operací na CPU. Protože menší modely nevyužívají paralelismus tensoru, dochází k neshodě mezi počtem hodností, které jsou cílové a návrhové modely sdíleny. V naší implementaci běží návrhový model pouze na vedoucí hodnosti TP skupiny.

Jak již bylo uvedeno dříve, dekódovací krok přenáší šarže logitů do dalšího běhu. To nám umožňuje překrývat jednu exekuci návrhového modelu s prací plánování šarže na straně CPU. Po sestavení šarže opakovanými voláními vzorkovače a návrhu produkovat návrhové tokeny. Paralelně se šarže pro ověření sestaví pro cílový model a synchronizují s paralelními pracovníky. Logity cíle jsou ověřeny a vzorkovány k určení délek přijatých sekvencí. V tomto bodě je nutná synchronizace GPU-CPU, abychom určili délky následujících sekvencí. Protože návrhový model běží pouze na vedoucím uzlu, jeho šarže je nastavena po sobě a jeho exekuce je spuštěna k naplnění jeho odpovídajících záznamů cache KV se dalším tokenem, který cíl produkoval. Logity produkované tímto návrhem v aktuálním běhu budou použity pro vzorkování prvního návrhového tokenu v následujícím běhu. Nejvýznamnější je, zatímco návrh běží, může být plánována další šarže.

MTP harmonogram pro jediný token

Přestože runtime dosud neposkytuje exploraci stromu návrhu ve stylu Eagle, implementovali jsme speciální případ tohoto schématu, zvažující lineární sekvenci návrhových tokenů produkovaných modelem velikosti jediné vrstvy dekodéru. Toto schéma může být použito pro predikci návrhu pomocí open-source vah DeepSeek R1. Podřípad predikce jediného tokenu je zajímavý, jelikož velké MTP vrstvy dosahují dostatečně vysoké míry přijetí, aby ospravedlnily jejich režii.

MTP plánování je poněkud složitější, protože návrhový model je mnohem rychlejší a skryje menší latenci na straně CPU. Kromě toho je návrh shodován spolu s cílovým modelem, což vyžaduje přenosy sdílené paměti pro informace o šarži. Běh začíná přenosem informací o šarži a vzorkováním prvního tokenu z logitů přenesených, podobně jako v předchozím schématu. Dále je cílový model spuštěn k ověření tokenů, zpracování 2 * D tokenů, kde D je velikost dekódovací šarže. To je ideální pro micro-batching v modelech MoE přes pomalejší propojení jako InfiniBand, protože šarže se rovnoměrně rozdělí na dvě poloviny. Skryté stavy cíle se přenesou do dalšího běhu návrhu, zatímco logity se přenášejí do vzorkovače pro ověření.

Prováděním omezeného množství další práce na GPU se vyhýbáme synchronizaci CPU-GPU po přijetí sekvence návrhu. Po posunutí vstupních tokenů cílů zasažen, algodní jádro umístí následující cílové tokeny na jejich odpovídající místa. Návrh se pak znovu spustí se stejnými informacemi o šarži jako cíl, naplnění záznamů cache KV a budování logitů a skrytých stavů pro další běh, provádějící nějakou přebytečnou práci na tokeny, které nebyly přijaty. V těchto situacích je latence nevyužité práce stěží měřitelná kvůli malé velikosti návrhového modelu. Paralelně s během návrhu se na CPU určují délky sekvencí a spustí se plánování další šarže, aniž by bylo třeba čekat na ukončení práce na GPU.

Režie dodatečné práce ve vrstvě návrhu není pozorovatelná v pozornosti, avšak vrstvy MLP jsou problematičtější. Vzhledem k tomu, že pokyny pro násobení matic vyplňují na hranici 64 podél dimenze počtu tokenů, pokud zdvojnásobení nevyžaduje výrazně více skupin, je režie skrytá. Pro delší návrhové sekvence je režie dražší a schéma použité pro běžné modely návrh-cíl funguje lépe.

Reference

• Rychlá inference z transformátorů skrze spekulativní dekódování

• EAGLE: Spekulativní vzorkování vyžaduje přepracování nejistoty funkcionality

• EAGLE-2: Rychlejší inference jazykových modelů s dynamickými návrhovými stromy

• EAGLE-3: Škálování inference zrychlení velkých jazykových modelů skrze tréninkový test

• Medusa: Jednoduchý rámec pro zrychlení inference LLM s více dekódovacími hlavami

• Normalizace vrstvy

• FlashInfer: Účinný a přizpůsobitelný stroj pozornosti pro LLM inference servicing

• Technická zpráva DeepSeek-V3