NVIDIA radí modelům sdílet KV hlavy. Testy ale neříkají, na které GPU běžely

Dlouhé kontextové okno není jedna úloha, která se jen zvětšuje. Když model nejdřív načte celé zadání, počítá jinak než ve chvíli, kdy po jednom přidává další tokeny. Technický rozbor NVIDIA zveřejněný 31. července z tohoto rozdílu odvozuje čtyři pravidla pro návrh modelů. Nejsou to ale univerzální rychlostní tabulky: část závěrů je analytická a část stojí na měření, u kterého chybí přesná testovací karta.

Pozornost, anglicky attention, je mechanismus, kterým model při zpracování textu zjišťuje, ke kterým dřívějším tokenům se má vrátit. Pro každý token vytváří dotaz a dvojici klíč a hodnota. Klíče a hodnoty z minulých tokenů si při generování ukládá do KV cache, tedy pracovní paměti pozornosti. Čím delší rozhovor nebo dokument, tím větší objem dat musí tahle paměť držet a opakovaně číst.

Načtení zadání a další token nejsou stejná práce

První fázi NVIDIA označuje jako prefill. Model v ní dostane všechny tokeny zadání najednou. U husté pozornosti porovnává každý token s ostatními, takže množství práce roste přibližně s druhou mocninou délky vstupu. Výpočetní jednotky GPU mají dost dat najednou a obvykle jsou hlavním omezením právě jejich počet a rychlost.

Ve druhé fázi, decode, přibývá vždy jediný nový token. Samotného počítání je méně, karta však pro každý krok znovu načítá klíče a hodnoty uložené v KV cache. Úzkým místem se proto stává propustnost paměti, tedy kolik dat karta dokáže za sekundu přečíst. Výjimkou může být načtení krátkého nezapamatovaného konce zadání po použití prefixové cache; takový běh se podle NVIDIA chová spíš jako generování.

Rozdíl je vidět i v číslech výrobce. Na grafu pro DeepSeek-R1 připadá na pozornost při délce 4 000 tokenů 18 procent času načtení zadání, při 128 000 tokenech už 85 procent. Jde o měření NVIDIA, ne o nezávislý test. Směr růstu ale odpovídá vlastnosti husté pozornosti: počet porovnání ve fázi prefill roste s druhou mocninou délky.

První praktické pravidlo proto zní: používat jádro typu FlashAttention, které nesestavuje celou matici pozornosti v pomalé paměti karty. Původní práce FlashAttention popisuje dlaždicové zpracování, při němž se malé části dotazů, klíčů a hodnot přesouvají do rychlé paměti přímo uvnitř GPU a mezivýsledky se průběžně slučují. Výsledek je matematicky přesný, k hlavní paměti se ale chodí méně často.

Jedna sada klíčů a hodnot pro více hlav

Další pravidlo míří na počet hlav pozornosti. Klasická vícehlavá pozornost má pro každou hlavu dotazů také vlastní hlavu klíčů a hodnot. Multi-query attention nechává všem hlavám dotazů jednu společnou dvojici. Grouped-query attention, zkráceně GQA, stojí mezi nimi: několik hlav dotazů sdílí jednu KV hlavu.

NVIDIA poměr označuje písmenem G a počítá ho jako počet hlav dotazů dělený počtem KV hlav. Vyšší G znamená méně KV hlav, menší cache a méně dat načítaných při každém novém tokenu. V měření výrobce se čas generování zhruba půlí s každým zdvojnásobením G, dokud nezačnou převládat jiné režijní náklady. Doba načtení dlouhého zadání se mění málo, protože v této fázi má GPU dost souběžné práce už díky počtu tokenů.

Původní práce o GQA nezávisle potvrzuje, že jde o prostřední cestu mezi vícehlavou a multi-query pozorností. Sdílení KV hlav ale není přepínač, který lze bez následků zapnout v hotovém modelu. Mění se architektura i trénink a návrhář musí hlídat také kvalitu odpovědí. Z grafu rychlosti proto nejde vyvozovat, že nejvyšší možné G je nejlepší pro každý model.

Dvě čísla rozhodují o škálování

Třetí doporučení se týká velikosti jedné hlavy. NVIDIA pro svá optimalizovaná jádra doporučuje 128 nebo 256 prvků. Taková velikost se dobře skládá do dlaždic výpočtu a do přenosů pamětí, které hardware provádí po blocích. Nejde o matematickou nutnost modelu, ale o sladění architektury s konkrétními jádry a GPU. Jiná implementace může mít jiné optimum.

Čtvrté pravidlo říká zmenšit stav pozornosti tam, kde model nepotřebuje plný globální pohled na všechny předchozí tokeny. Možností je posuvné okno, řídká pozornost nebo kombinace vrstev s globální a omezenou pozorností. Ušetřená KV cache pomáhá hlavně při generování, protože právě tam se čte při každém kroku.

Počet KV hlav má ještě jeden méně nápadný důsledek. Když se model rozděluje mezi více GPU pomocí tensorového paralelismu, je nejčistší přidělit každé kartě celé KV hlavy. NVIDIA proto doporučuje, aby počet paralelních částí nepřekročil počet KV hlav. Při větším rozdělení se musí některá KV data duplikovat, takže část úspory mizí. Model s malou cache tedy může být rychlý na jedné kartě, ale hůř se dělit na mnoho karet.

Měření bez názvu karty

Rozbor uvádí použité formáty FP8, rozměry hlav, počty hlav i délky vstupu. U křivek naměřeného času však neříká, na kterém modelu GPU běžely ani jaká verze knihoven a ovladače byla použita. Bez těchto údajů nejde test zopakovat a zjistit, zda se stejné poměry objeví třeba na starší kartě nebo v jiné implementaci.

To nevyvrací analytickou část textu. Prefill skutečně nabízí mnoho souběžných operací, decode skutečně opakovaně čte KV cache a méně KV hlav skutečně znamená méně uložených dat. Konkrétní násobky z grafů je ale přesnější číst jako měření dodavatele na nepopsané sestavě, ne jako slib výkonu pro každou GPU.

Pro posouzení modelu s dlouhým kontextem tak nestačí údaj o maximálním počtu tokenů. Důležitý je počet hlav dotazů a KV hlav, velikost jedné hlavy, podíl vrstev s globální pozorností a způsob, jakým se model rozdělí mezi karty. Právě tato čtyři čísla napoví, zda dlouhé okno zpomalí hlavně načtení dokumentu, nebo každé další slovo odpovědi.