A Sonar felgyorsítása spekuláció révén

A spekulatív dekódolás felgyorsítja a nagy nyelvi modellek (LLM-ek) generálási sebességét egy gyors és kis vázlatmodell használatával, amely olyan kitöltési jelölteket hoz létre, amelyeket a nagyobb célmodell ellenőriz. Ezen a sémán keresztül a költséges célmodell egyetlen token előállítása helyett egyszerre többet bocsát ki egyetlen lépésben. Itt bemutatjuk a különböző típusú spekulatív dekódolások implementációs részleteit, amelyeket a Perplexitynél alkalmazunk a Sonar modellek tokenek közötti késleltetésének csökkentése érdekében.

Spekulatív dekódolás

A spekulatív dekódolás a természetes nyelvek szerkezetét és a transzformátorok autoregresszív jellegét használja fel a tokenek generálásának felgyorsítására. Annak ellenére, hogy a nagyobb modellek, mint például a Llama-70B, több tudással rendelkeznek, mint a kisebbek, mint például a Llama-1B, bizonyos egyszerűbb feladatokon hasonlóan teljesítenek. Ez az átfedés azt sugallja, hogy bizonyos szekvenciákat jobban generálnak a kevésbé költséges modellek, míg a összetett problémákat a nagyobbakra bízzák. A kihívás abban rejlik, hogy megállapítsák, melyik kitöltés a jobb, és hogy a kisebb modell generálása ugyanolyan minőségű-e, mint a nagyobbé.

Szerencsére az LLM-ek autoregresszív transzformátorok: ha egy token-sorozatot kapnak, kiadják a következő token valószínűségi eloszlását. Emellett a bemeneti sorozatban lévő tokenekhez kapcsolódó köztes jellemzőkből származtatott logitok azt is jelzik, hogy mekkora valószínűséggel adja ki a modell pontosan azokat a tokeneket. Ez a tulajdonság teszi lehetővé a spekulációt: ha egy token-sorozatot egy kisebb modell generál egy bemeneti előtagból kiindulva, az futtatható a nagyobb modellen annak megállapítására, hogy mennyire vág egybe a célmodellel. A jelöltek minden előtagját valószínűséggel pontozzák, és a leghosszabbat választják ki, amely az elfogadási küszöb felett van. Bónuszként a célmodell egy további tokent is biztosít ingyen: ha egy vázlatmodell n tokent generál, legfeljebb n + 1 bocsátható ki egyetlen lépésben.

Nagyítás

A következtetési időben a spekulatív mintavételezési folyamat nagyjából 4 szakaszra osztható:

Előtöltés (prefill): a célmodellt és a vázlatmodellt is le kell futtatni a bemeneti sorozaton a KV-gyorsítótár bejegyzéseinek feltöltéséhez. Míg egyes sémák, mint például a Medusa, egyszerűbb sűrű rétegeket használnak a előrejelzéshez, ebben a bejegyzésben a transzformátor alapú vázlatokra összpontosítunk, amelyeknek saját KV-gyorsítótárra van szükségük.

Vázlatgenerálás: a vázlatmodell iterál, hogy számos rögzített tokent hozzon létre. A vázlat-sorozat lehet lineáris, vagy a modell felfedezhet egy faszerű struktúrát egy adott mélységig (EAGLE, Medusa). Itt a lineáris sorozatokra összpontosítunk.

Elfogadás: a célmodell fut a vázlat-sorozaton, logitokat hozva létre, amelyek megfelelnek az egyes vázlat-tokeneknek. Meghatározzák a leghosszabb elfogadható sorozat hosszát.

Célgenerálás: mivel a célmodell logitokat generált, az eltérő pozícióban vagy a sorozat végén a logitok egy utólagos tokennek felelnek meg. Ezekből a logitokból mintát lehet venni, hogy egy robusztus tokent biztosítsunk a céltól, lezárva a sorozatot.

Különböző módszerek léteznek a spekulatív dekódolás megvalósítására. Ebben a bejegyzésben arra a sémákra fókuszálunk, amelyeket a Sonar modellek felgyorsítására használtunk egy saját fejlesztésű 1B-s modellel, valamint azokra az előrejelzési mechanizmusokra, amelyeket a DeepSeek léptékű modellek felgyorsítására építünk.

Cél-vázlat (Target-Draft)

A spekulatív dekódolás úgy érhető el, hogy egy meglévő kis LLM-et vázlatmodellként egy célmodellhez kapcsolunk jelölt sorozatok generálásához. A gyártásban a Sonar modellt egy Llama-1B modellel gyorsítottuk fel, amelyet ugyanazon az adatkészleten finomhangoltunk, mint a célt. Bár ez a megközelítés nem igényelte vázlat betanítását a semmiből, a kis modell még mindig jelentős KV-gyorsítótár kapacitást használ, és enyhe előtöltési többletterhelést okoz, növelve a TTFT-t.

Ezen a sémán keresztül a dekódoló csak a csak dekódolású kötegeken spekulál, standard mintavételezéssel generálva tokeneket az előtöltés során vagy a vegyes előtöltés-dekódolás kötegeken. Az előtöltési szakaszban a cél logitjaiból azonnal mintát vesznek, hogy előtöltsék az újonnan generált tokent a vázlat KV-gyorsítótárában is. A vázlatból még nem vesznek mintát, de az általa előállított logitok átkerülnek a dekódolási szakaszba.

Nagyítás

A dekódolás során a vázlatmodell előrelép, és minden szakaszban mintát vesz a legfelső tokenből. A kívánt vázlat hosszának elérése után a tokenek átfutnak a célmodellen, hogy létrehozzák azokat a logitokat, amelyek alapján a mintavevő azonosítja az elfogadott sorozathosszat. Az elfogadást a vázlat és a cél teljes valószínűségi eloszlásának összehasonlításával határozzák meg. Mivel a célmodell mindig egy sor logitot ad ki az elfogadott vázlat sorozat után, abból mintát vesznek egy további kimenet előállításához. Mivel a vázlatmodell még nem látta ezt az elfogadott tokent, újra lefut, hogy feltöltse a megfelelő KV-gyorsítótár bejegyzéseit a következő dekódolási lépésre felkészülve, ismét átvíve a logitokat.

EAGLE

Az EAGLE egy spekulatív dekódolási séma, amely több vázlat-sorozatot vizsgál meg, amelyeket a valószínű vázlat-tokenek faszerű bejárásával generálnak. Egy rögzített (EAGLE) vagy dinamikusan alakított (EAGLE-2) fát vizsgálnak a vázlat-tokenek egymást követő végrehajtásával, a Top-K jelölteket figyelembe véve az egyes csomópontokon a legmagasabb pontszámú token követése helyett egy lineáris sorozatban. A sorozatokat ezután pontozzák, és a leghosszabb megfelelőt választják ki a folytatáshoz, és egy további tokent is hozzáfűznek a céltól.

Nagyítás

A pontosabb előrejelzés elérése érdekében egy EAGLE vázlatmodell nemcsak tokenek alapján jelez előre, hanem a célmodell céljellemzőit (az utolsó réteg rejtett állapotait) is felhasználva. Az EAGLE hátránya, hogy egyedi, kis vázlatmodelleket kell betanítani, amelyek elég pontosak ahhoz, hogy megfelelő jelölteket generáljanak alacsony késleltetési kereteken belül. Jellemzően a vázlatmodell egyetlen transzformátor réteg, amely megegyezik az eredeti modell dekódoló rétegével, és szorosan kapcsolódik a célhoz azáltal, hogy hozzá van kötve annak beágyazásaihoz és lm_head vetületeihez. Mivel ez kevesebb KV-gyorsítótár kapacitást igényel, az EAGLE alacsonyabb memóriaterheléssel rendelkezik.

A célmodell faszerű sorozatainak ellenőrzéséhez egyedi figyelemmaszkokat kell használni. Sajnos egy egész sorozathoz tartozó egyedi figyelemmaszk használata jelentősen lelassítja a figyelmet a reális bemeneti hosszúságoknál (akár 50%-kal), semmissé téve a spekuláció révén elérhető sebességnövekedés egy részét. Emiatt még nem telepítettük a teljes fafeltárást a gyártásba, hanem a DeepSeek-V3 technikai jelentésben bemutatott MTP-szerű sémákon keresztül történő egytokenes előrejelzés speciális esetére összpontosítunk.

MTP

Ez a séma hasonló a vázlat-cél dekódoláshoz, azzal a kivétellel, hogy az előrejelzéshez a rejtett állapotokat is használják a tokenek mellett. Kicsit több munkát kell végezni mind az előtöltési, mind a dekódolási szakaszban a szokásos vázlat-cél spekulációhoz képest. A vázlatmodell tokeneket és rejtett állapotokat is használ: a t_{i+1} tokent a t_i tokennek megfelelő L_i logitokból mintázzák, amelyek viszont a H_i rejtett állapotokból származnak. Következésképpen a bemeneti token pufferöket egy lépéssel balra kell eltolni a cél által kimenetként adott rejtett állapotvektorokhoz képest. Az alábbi ábra jelöli a betanításhoz használt megfeleltetéseket, valamint az eltolást a következtetés során.

Nagyítás

A dekódolási folyamat nagyon hasonló a vázlat-cél dekódoláshoz, azzal a kivétellel, hogy mind a rejtett állapotok, mind a logitok átkerülnek. Megvalósításunk megosztja az összes kapcsolódó mintavételezési és logit-feldolgozási logitot, és csak a modell előre hívásait teszi egyedivé. Ha több tokent jeleznek előre, a vázlatmodell a vázlat rejtett állapotait használja az előrejelzéshez, és a KV-gyorsítótár bejegyzéseit is a saját jellemzői alapján tölti fel. Hosszú távon ez csökkentheti a pontosságot. Ezt követően, amikor a vázlatmodellt futtatják a cél előrejelzéséhez tartozó KV-gyorsítótár bejegyzés feltöltésére, azt a teljes sorozaton futtatják, bemenetként a pontosabb cél rejtett állapotokat használva. Mivel ezek a vázlatmodellek kicsik, a további tokenek feldolgozásának többletköltsége elhanyagolható.

MTP fejek betanítása

Az MTP előnyeinek kihasználása érdekében kiépítettük azt az infrastruktúrát, amely a Perplexity adatkészletein finomhangolt modelleinkhez csatolt MTP fejek betanításához szükséges, és 8xH100 eszközt tartalmazó egyetlen csomóponton fut. Körülbelül egy nap alatt képesek vagyunk fejeket építeni a Llama-1B-től a Llama-70B-ig és a DeepSeek V2-Lite-ig terjedő modellekhez. Nagyobb modellek esetén a finomhangolási folyamat során felépített MTP fejekre támaszkodunk.

Az MTP betanítás célja, hogy a cél rejtett állapotaiból extrapolált vázlat rejtett állapotokat és logitokat összehangolja a cél következő tokenjének logitjaival és rejtett állapotaival. Mivel a rejtett állapotok következtetése költséges, azt a célmodell következtetésre optimalizált megvalósításával előre kiszámítjuk a betanítás során történő felhasználáshoz. Az MTP következtetési megvalósítás érvényesítéséhez és annak biztosításához, hogy a kvantálás vagy optimalizálások miatti numerikus különbségek ne akadályozzák az eredményeket, a validációs veszteség és a pontosság becslése érdekében azonban teljesen újrahasznosítjuk mind a cél-, mind a vázlatmodell következtetési megvalósítását.

Amikor az eredeti cikkben használt ShareGPT adatkészkeletről nagyobb mintákra tértünk át, észrevettük, hogy az EAGLE cikkben vázolt és megvalósított MTP fej architektúra nem tudott betanulni a 70B méretű modelleken. Ellentétben a ShareGPT-vel, amely nagyobb számú rövidebb sorozatot tartalmazott, mi kissé kisebb számú, lényegesen hosszabb prompton tanítunk. Mivel az eredeti EAGLE fejek szerkezetükben kissé eltértek a tipikus transzformátortól, visszavezettünk néhány eltávolított RMS Normalizálási réteget. Azt tapasztaltuk, hogy ez nemcsak lehetővé tette a betanítás konvergálását, hanem néhány százalékponttal meg is növelte a fejek pontosságát.

Nagyítás

A réteg-normalizálások nemcsak megkönnyítik a betanítást, de a normalizálások visszavezetése matematikailag is intuitív. Az MTP fejek újrahasznosítják a célmodell beágyazásait és logit vetületeit, mivel méretük jelentős lehet (körülbelül 2 GB a Llama 70B esetén). A betanítás során ezeket fagyasztják, és az elvárás az, hogy az MTP réteg megtanulja beágyazni az előrejelzéseket ugyanabba a vektortérbe, mint amit az eredeti modell vetítési rétege megtanult a betanítás során. A normalizálások elhagyásával egyetlen MLP-nek kellene megtanulnia ugyanazt a függvényt, mint egy MLP-nek, amelyet egy normalizálás követ, ami akadályozza a vázlat és a célmodellek rejtett állapotai közötti megfeleltetést.

Következtetés spekulatív dekódolással

A következtetési motorban a bemeneti sorozatokhoz tartozó tokenek generálásához először ésszerű méretű kötegekbe kell őket csoportosítani, majd oldalakat kell lefoglalni a KV-gyorsítótárban a következő tokenek számára. A bemeneti tokenek és a KV-oldal információi ezután egy olyan pufferbe csomagolódnak, amely szétkerül a modellt futtató összes párhuzamos ranghoz. Végül a metaadatok átmásolódnak a GPU memóriájába, és a modell végrehajtódik, hogy előállítsa a logitokat, amelyekből a következő tokent mintázzák.

Ellentétben bizonyos megvalósításokkal, amelyek lazán kapcsolnak össze egy vázlat- és cél-következtetési szervert egy burkolón keresztül, amely vezényli a kéréseket közöttük, a mi vázlat-cél párjaink szorosan össze vannak kapcsolva, és egységben lépkednek át a generáláson. A köteg-ütemezés és a KV-oldal lefoglalás megoszlik a modellek között a spekulatív dekódolás minden formájához: ez egységesíti azt a logikát, amely áthidalja a modellt az átfogó következtetési szerverrel, mivel mindannyian ugyanazt az interfészt teszik közzé.

A Perplexitynél a következtetési futtatókörnyezet a FlashInfer köré épül, amely meghatározza azokat a metaadatokat, amelyeket fel kell építeni a figyelem-kernel konfigurálásához és ütemezéséhez. Adott bemeneti sorozatok esetén, amelyek egy köteget alkotnak az előtöltéshez, dekódoláshoz vagy ellenőrzéshez, a CPU-oldali munkát kell végezni a köztes pufferek lefoglalásához és a figyelemben használt bizonyos állandó pufferek feltöltéséhez. Ez a munka kiegészíti a köteg-ütemezés és a KV-oldal lefoglalás költségét, amelyek szintén késéseket okoznak, amelyeket el kell rejteni a GPU kihasználtságának maximalizálása érdekében.

Míg a spekuláció nélküli következtetéshez teljes mértékben párhuzamosítottuk a CPU-oldali és GPU-oldali munkát, azt találtuk, hogy a spekulatív dekódolás CPU-GPU egyensúlya összetettebb. A fő kihívás abból adódik, hogy az elfogadott tokenek száma határozza meg a következő futtatás sorozathosszát, ami egy nehezen elkerülhető GPU-CPU szinkronizációs pontot vezet be. Különböző ütemezési sémákkal kísérleteztünk a CPU-munka késleltetésének legjobb elrejtése érdekében.

Vázlat-cél ütemezés

Annak ellenére, hogy kisebb, mint a célmodell, ha egy egész LLM-et használnak vázlatként, az még mindig jelentős késleltetést okoz a GPU-n, némi mozgásteret biztosítva a költséges CPU-műveletek elrejtésére. Mivel a kisebb modellek nem részesülnek előnyben a tenzorpárhuzamosságból, eltérés van azon rangok száma között, amelyekre egy cél és egy vázlat fel van darabolva. Megvalósításunkban a vázlatmodell csak egy TP-csoport vezető rangján fut.

Nagyítás

Ahogy korábban jeleztük, a dekódolási lépés átviszi a logitokat a következő futtatásba. Ez lehetővé teszi számunkra, hogy átfedjük a vázlatmodell egyetlen végrehajtását a CPU-oldali köteg-ütemezési munkával. Miután a köteg összeállt, a mintavevő és a vázlat ismételt hívásai létrehozzák a vázlat-tokeneket. Ezzel párhuzamosan az ellenőrzési köteg összeáll a célmodell számára, és szinkronizálódik a párhuzamos feldolgozókkal. A cél logitjait ellenőrzik és mintázzák az elfogadott sorozathosszak meghatározásához. Ezen a ponton GPU-CPU szinkronizáció szükséges a későbbi sorozathosszak meghatározásához. Mivel a vázlatmodell csak a vezető csomóponton fut, a kötege sorozatosan van beállítva, és a végrehajtása elindul, hogy feltöltse a KV-gyorsítótár bejegyzéseit azzal a további tokennel, amelyet a cél előállított. Az aktuális futtatásban ezen vázlatfuttatás által előállított logitokat fogják felhasználni az első vázlat-token mintavételezésére a következő futtatásban. A legfontosabb, hogy amíg a vázlat fut, a következő köteg ütemezhető.

MTP ütemezés egyetlen tokenhez

Míg a futtatókörnyezet még nem biztosít Eagle-stílusú vázlatszektor-felfedezést, megvalósítottuk ezen séma egy speciális esetét, figyelembe véve egyetlen transzformátor dekódoló réteg méretű modell által előállított vázlat-tokenek lineáris sorozatát. Ez a séma használható vázlat-előrejelzéshez a DeepSeek R1 nyílt forráskódú súlyainak felhasználásával. Az egyetlen token előrejelzésének al-esete érdekes, mivel a nagy MTP rétegek elég magas elfogadási arányt érnek el ahhoz, hogy igazolják a többletterhelésüket.

Az MTP ütemezés valamivel összetettebb, mivel a vázlatmodell sokkal gyorsabb, kevesebb CPU-oldali késleltetést rejtve el. Ezenkívül a vázlat a célmodell mellett van darabolva, ami megosztott memóriás átviteleket igényel a köteg-információkhoz. A futtatás a köteg-információk átvitelével és az első token mintavételezésével kezdődik a továbbvitt logitokból, hasonlóan az előző sémához. Ezután a cél lefut a tokenek érvényesítésére, 2 * D tokent feldolgozva, ahol a D a dekódolási köteg mérete. Ez ideális a MoE (Mixture-of-Experts) modellekben lévő mikrokötegeléshez lassabb összeköttetéseken, például az InfiniBanden keresztül, mivel a köteg egyenletesen oszlik meg két félre. A cél rejtett állapota átkerül a következő vázlatfuttatásba, míg a logitok bekerülnek a mintavevőbe ellenőrzésre.

Nagyítás

A GPU-n végzett korlátozott mennyiségű további munka elvégzésével elkerüljük a GPU-CPU szinkronizációt a vázlat-sorozat elfogadása után. A célok bemeneti tokenjeinek eltolása után egy kernel beilleszti a következő cél-tokeneket a megfelelő helyükre. A vázlatot ezután újra lefuttatják a céggel megegyező köteg-információkkal, feltöltve a KV-gyorsítótár bejegyzéseit, és felépítve a következő futtatás logitjait és rejtett állapotait, miközben némi redundáns munkát végeznek olyan tokeneken, amelyek nem lettek elfogadva. Ilyen helyzetekben a fel nem használt munka késleltetése alig mérhető a vázlatmodell kis mérete miatt. A vázlatfuttatással párhuzamosan a sorozathosszak meghatározásra kerülnek a CPU-n, és elindul a következő köteg ütemezése anélkül, hogy meg kellene várni a GPU-munka befejeződését.

A vázlatrétegben végzett további munka többletterhelése nem észreveszhető a figyelemben, azonban az MLP rétegek problémásabbak. Mivel a mátrixszorzási utasítások 64-es határig igazodnak a tokenek számának dimenziója mentén, ha a duplázás nem igényel jelentősen több blokkot, a többletterhelés rejtve marad. Hosszabb vázlat-sorozatok esetén a többletterhelés költségesebb, és a szokásos vázlat-cél modellekhez használt séma jobban működik.

Hivatkozások

Gyors következtetés transzformátorokból spekulatív dekódoláson keresztül

EAGLE: A spekulatív mintavételezés a jellembizonytalanság újragondolását igényli

EAGLE-2: Nyelvi modellek gyorsabb következtetése dinamikus vázlatfákkal

EAGLE-3: Nagy nyelvi modellek következtetési gyorsításának skálázása betanítási idejű teszttel

Medusa: Egyszerű LLM következtetés-gyorsítási keretrendszer több dekódoló fejjel

Réteg-normalizálás

FlashInfer: Hatékony és testreszabható figyelem-motor LLM következtetési kiszolgáláshoz

DeepSeek-V3 technikai jelentés

Idézet

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