Prečo sú vstupné tokeny drahé?

Nedávno som narazil na článok, ktorý tvrdí, že vstupné tokeny sú extrémne predražené. Argument znie, že spracovanie vstupu v jazykových modeloch (LLM) je limitované výpočtami a výpočty sú lacné, takže vstup by mal stáť okolo jedného centu za milión tokenov aj pri špičkových modeloch. Nie som expert na prevádzkovanie jazykových modelov, ale aj mne je jasné, že je to zložitejšie.

Stručne k tomu, ako prebieha inferencia

Prácu, ktorú musí GPU vykonať pri spracovaní dotazu pre jazykový model, si môžete predstaviť ako obdĺžnik. Na vodorovnej osi je pozícia tokenu v sekvencii. Na zvislej osi sú vrstvy modelu zoradené odspodu nahor. Tokenov je oveľa viac (~100 tisíc) než vrstiev (~100), takže obdĺžnik je široký a nízky. Delí sa na dve časti: vstup a výstup.

GPU spracúva vstupnú časť tohto obdĺžnika po riadkoch (vrstvách modelu). Začína naspodku a rýchlo sa dostane po vrch. Tento proces je vysoko paralelizovateľný a limitovaný výpočtovým výkonom, takže GPU v ňom vyniká.

Výstup sa spracúva inak. GPU spracúva stĺpce (tokeny) jeden po druhom. V rámci každého tokenu postupuje bunku po bunke (vrstvu po vrstve) odspodu nahor. Až keď dokončí aktuálny token, môže začať spracovávať ďalší.

Náklady na spracovanie parametrov modelu

Každá vrstva v modeli pozostáva z niekoľkých veľkých matíc, ktoré treba načítať z pamäte. Pri vstupných tokenoch sa to robí raz pre celý riadok (vrstvu). Pri výstupných tokenoch sa to robí pre každú bunku (kombináciu tokenu a vrstvy).

Ak spracúvame I vstupných tokenov a O výstupných tokenov a model má celkovo N parametrov a A aktívnych parametrov, potom spracovanie vstupných tokenov vyžaduje prenesenie O(N) dát z pamäte, zatiaľ čo výstupné tokeny potrebujú preniesť O(OA) dát. Výpočtová zložitosť je O(IA) pre vstup a O(OA) pre výstup. Zatiaľ to vyzerá takto:

Keďže GPU vykonávajú výpočty rádovo rýchlejšie než prístupy do pamäte, výstupné tokeny limituje rýchlosť pamäte, kým vstupné tokeny limituje výpočtový výkon. Batching pomáha, ale v moderných MoE modeloch len v obmedzenej miere.

Náklady na spracovanie KV cache

Vyššie uvedený naivný výpočet, ktorý zdá sa používa aj autor spomínaného článku, ignoruje KV cache. Tá uchováva informácie o už spracovaných tokenoch. Pri krátkych dotazoch (otázka a odpoveď) to až tak nebolí, ale pri dlhom kontexte (napr. agentické nástroje pre programátorov) môže KV cache ľahko prerásť aktívne parametre modelu.

Z pohľadu GPU sa KV cache správa ako sada dynamicky generovaných matíc, ktoré pri každom tokene narastú o jeden riadok/stĺpec. KV cache tak spotrebúva výpočty aj rýchlosť pamäte podobne ako parametre modelu, len táto spotreba rastie s dĺžkou sekvencie.

Aby som bol presný, mechanizmus pozornosti tiež používa parametre modelu na výpočet vektorov pozornosti (queries, keys, values) a na projekciu výsledku späť do embeddingu, ale náklady na tieto operácie už sú započítané vo vyššie opísanom spracovaní parametrov. Čo doteraz nie je zahrnuté, je samotná operácia pozornosti, ktorá porovnáva dotazové vektory s kľúčovými vektormi a kombinuje vektory hodnôt.

Ak sú I a O definované ako vyššie a C je počet skalárnych hodnôt, ktoré každý token pridá do KV cache, potom veľkosť KV cache v pamäti je O((I+O)C). Ak je vstup oveľa dlhší než výstup, môžeme to zjednodušiť na O(IC). Spracovanie KV cache si vyžiada prenesenie O(IC) dát pre vstupné tokeny a O(OIC) dát pre výstupné tokeny. Vspomeňme si, že je to tak preto, že vstup sa spracúva po vrstvách, zatiaľ čo výstup po tokenoch, takže každý výstupný token musí znova a znova načítať celú KV cache z pamäte. Výpočtová zložitosť je pre vstup O(IIC) (dobre známa kvadratická pozornosť) a na výstupe O(OIC). Pre samotnú KV cache teda platí:

Úzkym hrdlom zostáva prenos dát z pamäte pri spracovaní výstupných tokenov. Spojme to s nákladmi na spracovanie parametrov:

Oproti naivnej analýze, ktorá rátala len s parametrami modelu, sú dôležité dva rozdiely. Po prvé, každá sekvencia má vlastnú KV cache, takže batching tu nepomôže vôbec. Po druhé, všetky štyri limity zložitosti teraz rastú s dĺžkou sekvencie.

Ako započítať náklady na KV cache

Najdrahšia časť spracovania dotazu je prenos O(OA+OIC) dát pri generovaní výstupu. Zložka O(OA) sú pevné náklady za každý výstupný token, ktoré sú zahrnuté v cene výstupných tokenov.

Zaujímavá je zložka O(OIC). Všimnite si, že rastie s dĺžkou vstupu aj výstupu. Majú ju dodávatelia účtovať do ceny vstupu, výstupu, alebo do oboch? A čo ak klienti začnú posielať dotazy s veľmi dlhým vstupom aj veľmi dlhým výstupom? Náklady potom môžu veľmi rýchlo vystreliť.

V praxi býva vstup oveľa dlhší než výstup a je aj variabilnejší, preto dáva zmysel zahrnúť túto časť nákladov do ceny vstupných tokenov. Ak sa nad tým zamyslíte, je to celkom logické: každý vstupný token pridá dáta do KV cache, ktoré sa potom musia znova a znova načítať pre každý výstupný token, teda dokopy približne O(O) krát. To je kumulatívne veľmi veľa prenesených dát. Zdanlivo výpočtovo limitované vstupné tokeny tak v skutočnosti veľmi zaťažujú zbernicu do pamäte.

Prečo je dĺžká výstupu obmedzená?

Ak ste si prezerali ponuky poskytovateľov jazykových modelov, možno vás prekvapilo, že obmedzujú dĺžku výstupu. Technický dôvod na to nie je. Je tu však ekonomický dôvod: aj keď sa dĺžka výstupu mení menej než dĺžka vstupu, výstup môže byť stále veľmi dlhý, čo by kvôli O(OIC) kvadratickej zložitosti dátových prenosov vyšlo poskytovateľov veľmi draho. Aby klienti nemohli dotazmi s dlhým vstupom a dlhým výstupom vytlačiť náklady nad cenu API, poskytovatelia si stanovili limit na dĺžku výstupu. Tento limit je len vecou obchodnej politiky.

Využitie pamäte

Pamäťová stopa je O(IC) pre jeden dotaz a O(BIC) pre batch B súbežne spracúvaných dotazov. Táto pamäť musí zostať alokovaná až do vygenerovania posledného výstupného tokenu. Koľko je O(IC) v praxi? Špičkové modely sú veľké a veľké modely mávajú aj väčšiu KV cache. Na druhej strane, vývojári modelov agresívne optimalizujú KV cache. Odhadujem, že pri 100-tisíc vstupných tokenoch môže KV cache zaberať niekde medzi 10 GB a 100 GB. To určite obmedzuje veľkosť batchu aj na priestranných GPU od AMD.

Batching je v MoE modeloch menej účinný, ale stále pomáha. Parametre súvisiace s pozornosťou a router expertov sú spoločné. Niektoré modely majú aj spoločných expertov, ktorí bežia vždy. Každý dotaz si volí vlastných expertov, ale pri väčšom batchi tam býva aspoň čiastočný prekryv.

Ak však jeden dotaz zožerie všetku pamäť GPU, batching nemôžete použiť. Alebo sa vám na GPU zmestia len 2–3 dotazy, čo prínos batchingu výrazne zníži. Všimnite si, že dĺžka vstupu v O(BIC) je pod kontrolou užívateľa. Tým, že klienti posielajú dlhšie vstupy, znižujú efektivitu batchingu a zvyšujú náklady. Vstupné tokeny tak vyťažujú nielen rýchlosť pamäte, ale aj kapacitu pamäte. Aj to treba premietnuť do ceny vstupných tokenov.

Čo sa týka optimalizovanej pozornosti

V celom texte som predpokladal jednoduchý mechanizmus pozornosti s pevnými nákladmi na každý token. Vývojári modelov ale pozornosť niekedy optimalizujú. V niektorých vrstvách čítajú z KV cache len posledných ~1 000 tokenov. Niekedy dokonca staršie časti KV cache dealokujú alebo presunú do prompt cache mimo GPU. Môžu použiť aj nejakú formu subkvadratickej pozornosti. Tieto optimalizácie znižujú dopad dlhého kontextu, ale často zároveň zhoršujú výkon modelu, takže rozsah optimalizácií je limitovaný. Aj napriek týmto optimalizáciám majú vstupné tokeny vplyv na náklady a nikdy nebudú zadarmo.

Reálna nákladová štruktúra inferencie

Ako som ukázal, superlacné vstupné tokeny zo spomínaného článku sú ilúzia, ktorá vznikne, keď nezapočítate náklady na KV cache. Reálne náklady inferencie sú úmerné O(IO). Je to obdĺžnik, ktorého stranami sú dĺžka vstupu a dĺžka výstupu. Poskytovatelia spoplatňujú vstupné tokeny a obmedzujú dĺžku výstupu, aby nad týmito O(IO) nákladmi nestratili kontrolu.