AWS ukázalo zdieľanú KV cache pre vLLM cez HyperPod a distribuované NVMe

AWS zverejnilo referenčnú architektúru viacúrovňovej KV cache pre inferenciu veľkých jazykových modelov v Amazon SageMaker HyperPod. Návrh kombinuje natívnu prefixovú cache vLLM v pamäti GPU, odkladanie blokov cez LMCache do operačnej pamäte a zdieľanú vrstvu na lokálnych NVMe diskoch, ktoré spája distribuovaný súborový systém Curvine. Cieľom je umožniť viacerým inferenčným replikám opakovane používať už vypočítané kľúče a hodnoty pozornosti namiesto toho, aby každá replika spracovala rovnaký úvod promptu od začiatku.

KV cache uchováva medzivýsledky mechanizmu pozornosti pre tokeny, ktoré model už spracoval. Pri generovaní zabraňuje opakovanému výpočtu celej predchádzajúcej sekvencie a prefixová cache môže rovnaké bloky využiť aj medzi požiadavkami so spoločným začiatkom. Typickým príkladom je dlhý systémový prompt, rovnaký kontext z RAG pipeline alebo história viacotáčkového rozhovoru. Problém nastáva pri horizontálnom škálovaní: každá replika vLLM má vlastnú izolovanú cache a presmerovanie požiadavky na iný pod sa správa ako studený štart.

AWS rozdelilo hierarchiu na tri úrovne. L0 tvorí HBM pamäť GPU a obsahuje najčastejšie používané bloky s najnižšou latenciou. L1 používa hostiteľskú operačnú pamäť, kam LMCache presúva bloky vytlačené z GPU; táto vrstva je rýchla, ale zostáva lokálna pre konkrétny pod. L2 tvorí zdieľaný distribuovaný pool NVMe diskov. Curvine ho pripája cez FUSE ako spoločný adresár s režimom ReadWriteMany, takže blok, ktorý vytvorí jedna replika, môže načítať ktorákoľvek iná replika.

Takáto cache je dôležitá najmä pri väčších modeloch a lacnejších GPU inštanciách. AWS uvádza príklad inštancie ml.g6e.4xlarge so 48 GB pamäte na GPU. Sedemmiliardový model v bf16 zaberie približne 14 GB a ponechá relatívne veľa priestoru pre KV bloky. Tridsaťdvojmiliardový model však potrebuje približne 64 GB iba na váhy, musí sa rozdeliť medzi GPU a pre cache zostáva omnoho menej miesta. Pri vysokej súbežnosti sa bloky rýchlejšie vytláčajú a zdieľané úložisko môže znížiť počet drahých opakovaných prefill výpočtov.

Architektúra zároveň používa inteligentné smerovanie v HyperPod Inference Operator. Prefix-aware režim vyberá repliku podľa zhody úvodu promptu a hodí sa na konverzácie či spoločné systémové inštrukcie. KV-aware režim zohľadňuje stav cache jednotlivých workerov a cieli na dlhé dokumenty alebo relácie. Round-robin zostáva možnosťou pre bezstavové dávkové spracovanie. Router sa snaží najprv dosiahnuť zásah v L0 alebo L1; ak to nie je možné, zdieľaná L2 stále umožní vyhnúť sa úplnému prefillu.

V testovom nasadení AWS uvádza až 100-percentnú úspešnosť zásahov pri zdieľaní cache medzi podmi, približne 56-milisekundovú latenciu čítania L2 medzi uzlami pre prompt s asi 1 900 tokenmi a najviac 2,7-násobné zlepšenie času do prvého tokenu. GPU zásah má podľa opisu latenciu pod jednu milisekundu, operačná pamäť rádovo jednotky milisekúnd a vzdialená L2 desiatky milisekúnd. Aj najpomalšia úroveň však môže byť výhodnejšia než opätovný výpočet dlhého prefixu.

Zverejnené čísla nemožno prenášať na každú prevádzku. AWS testovalo sériové jednotlivé požiadavky, aby izolovalo cestu cache, a výsledok závisí od generácie GPU, priepustnosti siete, dĺžky kontextu a podielu spoločných úvodných tokenov. Firma odporúča uvažovať o L2 najmä pri prekrývaní promptov nad približne 40 percent. Pri vstupoch kratších než približne tisíc tokenov môže byť sieťové čítanie menej výhodné a prevádzka by sa mala spoliehať skôr na L0 a L1. Uvádzané úspory nákladov preto nie sú garantované.

Curvine v návrhu používa lokálne NVMe disky GPU uzlov, takže kapacita zdieľanej cache rastie s počtom workerov bez samostatnej ceny za blokové úložisko. Metadátový uzol a žurnál však zostávajú na trvalom EBS zväzku. Strata obsahu worker cache nie je kritickou stratou dát, pretože KV bloky možno znovu vypočítať. Produkčné nasadenie napriek tomu pridáva ďalšiu distribuovanú komponentu, FUSE pripojenia, CSI ovládač a požiadavky na monitorovanie siete, metadát a obsadenosti diskov.

Implementácia zatiaľ nie je úplne deklaratívna. Pole pre L2 backend v InferenceEndpointConfig natívne akceptuje iba Redis alebo štandardné tiered storage, preto AWS používa dočasnú hodnotu a následne mení premennú LMCACHE_REMOTE_URL na súborovú cestu Curvine. Priama úprava vygenerovaného Kubernetes Deploymentu môže byť prepísaná reconcile slučkou operátora, a postup preto počas patchovania operátor pozastavuje. Dokument tiež upozorňuje na potrebu pevne nastaviť PYTHONHASHSEED, aby rôzne pody vytvárali pre rovnaký prompt totožné kľúče cache, a odporúča stabilnejší serializér pre súborový konektor.

Praktický význam návrhu spočíva v možnosti zmenšiť GPU inštancie bez úplnej straty výhod prefixovej cache. To môže byť zaujímavé pri flotile viacerých modelov, RAG aplikáciách, agentoch s rozsiahlymi systémovými pokynmi a konverzáciách s opakovanou históriou. Najväčší prínos však nevzniká automaticky: tím musí zmerať distribúciu dĺžok promptov, mieru spoločných prefixov, latenciu siete, počet replík a pomer zásahov na každej úrovni. Bez dostatočného opakovania môže L2 iba zvýšiť prevádzkovú zložitosť.

Referenčná architektúra zároveň ukazuje širší trend v inferenčnej infraštruktúre. Optimalizácia sa už netýka iba kvantizácie modelových váh alebo rýchlosti samotného GPU. Čoraz väčšiu úlohu zohráva presun stavových dát medzi pamäťovými úrovňami, smerovanie podľa obsahu cache a zdieľanie výpočtu medzi replikami. Pre prevádzkovateľov veľkých modelov môže byť práve práca s KV cache jednou z hlavných ciest k nižšiemu času do prvého tokenu a lepšiemu využitiu drahého hardvéru.