Az NVIDIA szerint a modelltervezés is kulcs a hosszú kontextushoz

A hosszú kontextusnál egyre fontosabb a figyelem

Az NVIDIA Developer július 31-én közzétett elemzése azt vizsgálja, hogyan befolyásolja a sűrű figyelem teljesítményét a csoportméret, a fejdimenzió, a sorozathossz és a tenzorpárhuzamosítás stratégiája. A sűrű figyelem minden lekérdezéshez a teljes sorozat kulcsait és értékeit használja.

A vállalat kiindulópontja a modell és a hardver közös tervezése. Ennek lényege, hogy a modell architektúráját ahhoz igazítják, ahogyan a GPU végrehajtja a számításokat. Az NVIDIA szerint a DeepSeek-R1 esetében a figyelem részaránya a prefillel töltött időn belül 18 százalékról 85 százalékra nőtt, amikor a kontextus hossza 4K-ról 128K-ra emelkedett.

Más szűk keresztmetszet határozza meg a prefilt és a dekódolást

A prefilt a teljes bemenetet párhuzamosan dolgozza fel, ezért a nagy mátrixszorzások miatt elsősorban számításigényes. A dekódolás ezzel szemben alapértelmezés szerint tokenenként halad, és a kulcsérték-gyorsítótár, vagyis a KV cache beolvasása miatt főként memóriaigényes. Spekulatív dekódolás esetén nagyobbá válhat az egyszerre feldolgozott mátrixrész, így a dekódolás közelebb kerülhet a számításigényes működéshez.

Prefixgyorsítótár használatakor egy új kör rövid bemenetet is jelenthet úgy, hogy közben hosszú, korábban eltárolt előzményt kell figyelembe venni. Az NVIDIA megjegyzése szerint ilyen helyzetben a prefilt viselkedése a dekódoláséhoz hasonlíthat.

A FlashAttention a teljes figyelemmátrix létrehozása nélkül dolgozik. A lekérdezések, kulcsok és értékek csempéit a nagy sávszélességű memóriából a lapkán belüli memóriába továbbítja, majd egy menetben végzi el a pontszámítást, az online softmaxot és az értékek összesítését.

A nagyobb csoportméret főként a dekódolást gyorsítja

A csoportméret azt jelenti, hány lekérdezési fej osztozik egyetlen KV-fejen. A többfejes figyelemnél ez az érték 1, a csoportos lekérdezési figyelemnél, vagyis a GQA-nál például 4, 8 vagy 16, a több lekérdezési fejes figyelemnél, az MQA-nál pedig a lekérdezési fejek számával egyezik meg.

Az NVIDIA elemzése szerint a csoportméret növelése a prefilt futási idejét alig befolyásolja. 32K-s bemenetnél a 8-ról 16-ra emelés 6 százaléknál kisebb javulást hoz a számítási intenzitásban, a mérésekben pedig az 1 és 64 közötti értékek változtatása 1 százaléknál kisebb futásidő-különbséget okozott.

A dekódolásnál ezzel szemben a nagyobb csoportméret javítja a számítási intenzitást, mivel több lekérdezési fej osztozik ugyanazon betöltött KV-adaton. Az NVIDIA szerint 1-ről 8-ra emelve ez nyolcszoros elméleti növekedést jelenthet, a mért dekódolási futásidő pedig nagyjából felére csökken minden duplázáskor. A vállalat példaként a Nemotron 3 modellt említi, amely két KV-fejes GQA-t használ a hatékonyabb dekódolásért.

A fejméret és a párhuzamosítás is meghatározó

Az NVIDIA szerint a 128 és 256 elemű fejdimenziók illeszkednek a GPU csempeméreteihez és a 128 bájtos memóriaátvitelekhez, miközben beleférnek a tenzormemória kapacitásába. A 64 elemű fejméret továbbra is 128 széles csempét használ, az 512 elemű pedig már megközelíti a hardver korlátait.

A tenzorpárhuzamosításnál a vállalat irányelve szerint a párhuzamos egységek száma ne haladja meg a KV-fejek számát. A kevés KV-fejet használó modellek ezért gyorsan elérhetik a tenzorpárhuzamosítás határát. Az NVIDIA ilyen esetekre a figyelem adatpárhuzamosítását vagy a KV-párhuzamosítást javasolja szakértői párhuzamosítással kombinálva. A TensorRT-LLM-ben ezekhez a megközelítésekhez a Wide EP és a Helix Parallelism kapcsolódik.

Az irányelvek a modellfejlesztőknek szólnak, akik NVIDIA GPU-kon szeretnék növelni a következtetés áteresztőképességét és az interaktivitást. Az NVIDIA következő, erre vonatkozó bejegyzésében a ritka figyelem vizsgálatát ígéri.