NVIDIA: öt irányelv gyorsíthatja az LLM-ek spekulatív dekódolását

Mit csinál a spekulatív dekódolás?

Az NVIDIA fejlesztői blogján megjelent bejegyzés a vállalat AI model co-design sorozatának harmadik része. A cikk szerzői Tiyasa Mitra, Maximilian Golub, Maor Ashkenazi, Samkit Jain, Nidhi Bhatia és Bita Darvish Rouhani. A téma az, hogyan lehet gyorsítani a nagy nyelvi modellek, vagyis LLM-ek inferenciájának autoregresszív dekódolási szakaszát úgy, hogy a kimenet pontossága megmaradjon.

A spekulatív dekódolás lényege, hogy egy kisebb draft modell egyszerre több lehetséges következő tokent javasol. Ezeket a nagyobb célmodell egyetlen párhuzamos ellenőrzési lépésben vizsgálja meg. A célmodell sorban elfogadja a javasolt tokeneket addig, amíg el nem éri az első eltérést, majd a következő ciklus onnan folytatódik.

Az NVIDIA leírása szerint mivel csak a célmodell által elfogadott tokenek maradnak meg, a módszer ugyanazt a kimeneti szekvenciát állítja elő, mint a hagyományos dekódolás, kivéve ha az elfogadási feltételeket szándékosan lazítják. A megközelítés csökkenti a dekódolási iterációk számát, miközben növeli a célmodell aritmetikai intenzitását, nagyobb konkurencia igénye nélkül.

A bejegyzés két fontos fogalmat különít el. A draft length, jelölése D, azt mutatja meg, hány tokent javasol a kisebb modell egy célmodell-iterációban. Az acceptance length, jelölése AL, azt jelenti, hány token készül el, vagyis hány tokent fogad el a rendszer egy célmodell-iterációban. Az AL értéke 1 és 1 + D között lehet, mert a célmodell mindig képes legalább egy új helyes tokent előállítani az elfogadott draft tokenek mellett.

Öt irányelv a megfelelő draft length kiválasztásához

Az NVIDIA szerint a gyorsulás maximalizálásához a D, az AL és a draftolási késleltetés megfelelő kombinációját kell megtalálni. A bejegyzés első irányelve az, hogy a spekulatív dekódolás draft length értékét addig érdemes növelni, amíg a GEMM-műveletek a compute-bound tartományba kerülnek, anélkül hogy nőne a KV cache kapacitásnyomása.

A cikk példája szerint spekulációval a célmodell lineáris rétegeiben a GEMM-M érték M-ről M × (1 + D)-re nő. Egy 6144 x 6144 méretű reprezentatív expert GEMM esetén a bejegyzés azt mutatja, hogy nagyobb draft length mellett a GEMM-ek kisebb effektív batch méretnél is elérhetik a csúcsteljesítményt. Az NVIDIA kiemeli, hogy D=7 mellett a D=0 állapothoz képest a batch méret egynyolcada is elég ahhoz, hogy a működés compute-bound legyen.

A második irányelv az attention által dominált dekódolási időre vonatkozik. Az NVIDIA szerint ilyen esetben a választandó draft length: D = 128/G - 1, ahol G azt jelöli, hány query head osztozik egy KV headen. A bejegyzés szerint a decode attention aritmetikai intenzitása körülbelül 2 × G, spekulációval pedig 2 × G × (1 + D), mivel a spekulált tokenek ugyanazt a KV cache-t használják.

A harmadik irányelv az attention kernel tile-határaihoz igazítja a választást. Ha D nagyobb, mint 128/G - 1, akkor az NVIDIA olyan értékeket javasol, ahol G × (1 + D) a 128 többszöröse. A cikk szerint a benchmarkolt attention kernel szoftveres tile mérete 128, és ha az érték két tile-határ közé esik, az utolsó tile csak részben hasznosul, de költsége nagyjából egy teljes tile-éval egyezik meg.

A bejegyzés további irányelvei az alacsony késleltetésű helyzetekre és a draft mechanizmus kiválasztására vonatkoznak. Nagyon alacsony késleltetésnél az NVIDIA szerint a draft length csak addig növelendő, amíg az elfogadási nyereség indokolja a többlet draft költséget. A draft mechanizmusnál pedig az acceptance length, a draft overhead, valamint a tanítási és telepítési költség közötti egyensúlyt kell figyelembe venni.

Több draft mechanizmust hasonlít össze az NVIDIA

A bejegyzés több spekulatív dekódolási megoldást is felsorol. Ezek között szerepelnek külső draft modellek, az EAGLE-3, az MTP, a DFlash, a DSpark, valamint suffix vagy n-gram alapú módszerek. Az NVIDIA szerint ezek eltérő kompromisszumokat kínálnak a tanítási költség, a kiszolgáláskori memóriaigény és a draft tokenekre jutó spekulációs költség terén.

A vállalat a SPEED-Bench nevű benchmarkot is megemlíti, amelyet az NVIDIA fejlesztett a spekulatív dekódolás mérésére. A benchmark a bejegyzés szerint valószerű munkaterheléseket ad az acceptance length méréséhez, több feladatterületen, például kódolásban és összefoglalásban.

A gyakorlati kipróbáláshoz az NVIDIA/Model-Optimizer repóban elérhető, futtatásra kész tanítási példákat említ az EAGLE-3, a DFlash és a DSpark esetében. A forrás szerint ezek finomhangolási és kvantálási munkafolyamatokat is tartalmaznak, amelyeket Nemotron 3.5 Lightning modellen demonstrálnak. A bejegyzés következő lépésként a Model-Optimizer repó felfedezését, a SPEED-Bench használatát, valamint a draft overhead NVIDIA TensorRT LLM-mel történő számszerűsítését javasolja.

Mit jelent ez a fejlesztőknek és a piacnak?

Az NVIDIA bejegyzése alapján a spekulatív dekódolás nem egyszerűen bekapcsolható gyorsítótrükk, hanem olyan tervezési döntés, amelyet a modell, a hardver és a kiszolgálási cél együtt határoz meg. A megfelelő D érték más lehet ott, ahol a lineáris rétegek, és más ott, ahol az attention adja a futási idő nagy részét.

A gyakorlati üzenet az, hogy az LLM-inferencia gyorsítása méréshez kötött optimalizáció. A fejlesztőknek nemcsak azt kell nézniük, hány tokent tud javasolni egy draft mechanizmus, hanem azt is, mennyit fogad el a célmodell, mekkora késleltetést ad hozzá a draftolás, és milyen költsége van a tanításnak, a memóriának és a telepítésnek.

A SPEED-Bench és az NVIDIA/Model-Optimizer példái ezért a forrás szerint abban segíthetnek, hogy a spekulatív dekódolás hatását realisztikus promptokon, eltérő feladatterületeken és konkrét munkafolyamatokon mérjék. Ez különösen azoknál a rendszereknél lehet fontos, ahol az interaktivitás, a hosszú kontextus és a hardver kihasználtsága egyszerre számít.