Apple chce zlacniť prevádzku LLM cez zdieľanie KV cache medzi vrstvami

Apple Machine Learning Research zverejnilo v máji novú prácu Stochastic KV Routing, ktorá sa sústreďuje na jednu z najdrahších častí prevádzky veľkých jazykových modelov: na pamäť potrebnú pre KV cache pri autoregresívnom generovaní. Pre bežného používateľa je KV cache neviditeľná, no pre infra tím rozhoduje o tom, koľko dopytov sa zmestí na GPU, akú latenciu bude mať odpoveď a koľko bude stáť serving pri dlhších kontextoch. Doterajšia debata sa často točila okolo kvantizácie, kompresie alebo agresívnejšej práce s časovou osou kontextu. Apple teraz tvrdí, že veľký priestor na úspory sa skrýva aj v inom smere: nie v dĺžke histórie, ale v tom, koľko vrstiev si musí držať vlastnú kópiu cache.

Základná intuícia výskumu je pomerne priamočiara. V transformerových modeloch sa pri každej vrstve ukladajú kľúče a hodnoty, aby model pri ďalšom generovaní nemusel znovu prepočítavať už spracované tokeny. Táto cache však rastie s počtom vrstiev aj s dĺžkou kontextu a pri produkčnom nasadení predstavuje významnú časť spotreby pamäte. Apple argumentuje, že plná samostatná cache pre každú vrstvu nemusí byť vždy nevyhnutná. Namiesto toho skúma zdieľanie cache naprieč vrstvami, teda depth-wise cache sharing, pri ktorom časť vrstiev môže pri inferencii čítať aj z cache susednej či predchádzajúcej vrstvy bez toho, aby sa systém rozpadol na neakceptovateľnej strate kvality.

Podstatné je, že autori sa nesnažia len dodatočne „odrezať“ cache po tréningu a dúfať, že model to ustojí. Navrhujú tréningový postup nazvaný random cross-layer attention. Počas tréningu si vrstvy náhodne vyberajú, či budú pozerať do vlastných KV stavov alebo do stavov predchádzajúcej vrstvy. Tento stochastický režim má model pripraviť na to, že v ostrej prevádzke nebude mať vždy k dispozícii plnú ideálnu cache pre každú vrstvu. Inými slovami, nejde len o kompresný trik po dokončení modelu, ale o adaptáciu správania už počas tréningu alebo fine-tuningu tak, aby bol model robustný voči rôznym servisným rozpočtom a hardvérovým limitom.

Podľa Apple je práve táto robustnosť prakticky dôležitá. Veľa starších návrhov na zdieľanie cache síce sľubovalo úsporu pamäte, no v praxi narážali buď na horšiu priepustnosť, alebo na zhoršenie času do prvej odpovede. To je kľúčový problém, pretože prevádzkovatelia nechcú vymeniť jeden bottleneck za druhý. Zverejnený text tvrdí, že pri ich prístupe možno cache niektorej vrstvy vynechať bez straty informácie, pričom model si zachováva flexibilitu pri voľbe stratégie nasadenia. To je atraktívne najmä v prostrediach, kde sa deployment musí dynamicky prispôsobovať rozdielnym GPU konfiguráciám alebo kolísajúcemu zaťaženiu.

Výsledky podľa abstraktu naznačujú, že schéma funguje naprieč viacerými rodinami modelov a že ju možno použiť buď pri predtréningu, alebo pri následnom doladení. Ešte zaujímavejšie je tvrdenie, že pri väčších modeloch a dátovo obmedzených nastaveniach sa tento postup môže správať aj ako forma regularizácie. To znamená, že model nemusí len menej pamätať, ale v niektorých prípadoch si môže udržať alebo dokonca mierne zlepšiť kvalitu. Ak sa podobné zistenie potvrdí aj v širšej praxi, zmena by bola dôležitá: úspora pamäte by nebola iba bolestivý kompromis, ale potenciálne aj spôsob, ako tréning stabilizovať a deployment zlacniť bez priamej obete vo výstupe.

Načasovanie tejto práce nie je náhodné. Celé odvetvie dnes tlačí na dlhšie kontexty, vyššiu paralelizáciu a lepšie využitie GPU pamäte, pričom práve KV cache sa čoraz častejšie stáva jedným z hlavných limitov. Vidno to v open-source serving stackoch, v komerčných inference platformách aj pri návrhoch nových architektúr. Apple síce v tejto práci nepredstavuje hotový produkčný produkt, ale ukazuje výskumný smer, ktorý zapadá do širšieho trendu: namiesto nekonečného zvyšovania hardvérových nárokov sa firmy snažia hľadať jemnejšie architektonické a tréningové zásahy, ktoré zlepšia ekonomiku inferencie bez rozbitia používateľskej skúsenosti.

Pre infra a MLOps tímy je praktický význam zrejmý. Ak sa depth-wise cache sharing ukáže ako spoľahlivý aj mimo výskumných benchmarkov, môže znížiť počet GPU potrebných na rovnaký objem prevádzky alebo umožniť obslúžiť dlhší kontext na tej istej pamäťovej stope. To je zvlášť dôležité pri modeloch, kde náklady na serving nerastú lineárne len s počtom používateľov, ale aj s tým, ako dlhé a stavové sú konverzácie. Takéto úspory by mohli mať väčší reálny dosah než kozmetické zmeny v benchmarkoch, pretože priamo menia jednotkovú cenu nasadenia.

Pre Apple je zároveň typické, že výskum smeruje k efektívnosti, nie iba k hrubej škále. Firma tým pripomína, že boj o ďalšiu generáciu AI nebude len o tom, kto natrénuje väčší model, ale aj o tom, kto ho dokáže udržať v produkcii lacnejšie, stabilnejšie a s menšími kompromismi. Stochastic KV Routing preto nie je len technická poznámka o internom detaile transformerov. Je to signál, že ďalšia vlna konkurencie sa presúva hlbšie do samotnej ekonomiky inferencie, kde sa rozhoduje o tom, ktoré modely budú vedieť fungovať pri masovom nasadení bez explózie nákladov.