SageMaker smeruje podobné prompty na rovnaký server a zrýchľuje prvú odpoveď LLM

Amazon Web Services pridal do real-time endpointov služby SageMaker Inference stratégiu PREFIX_AWARE. Namiesto náhodného rozdeľovania požiadaviek medzi dostupné inštancie skúma začiatok každého vstupu a požiadavky s rovnakým prefixom sa snaží posielať na rovnaký server. Cieľom je zachovať už vypočítané hodnoty v lokálnej KV cache a neopakovať spracovanie rozsiahlej časti promptu pri každom volaní. Zmena sa týka infraštruktúrnej vrstvy pred modelovým serverom, takže aplikácia môže naďalej používať existujúce rozhrania InvokeEndpoint, streamované odpovede aj rozhranie kompatibilné s OpenAI API.

Prečo samotná KV cache vo väčšom klastri nestačí

Pri generovaní odpovede model najskôr spracuje vstupný kontext vo fáze označovanej ako prefill. Z neho vytvorí dvojice kľúčov a hodnôt, ktoré si môže uložiť do KV cache. Ak ďalšia požiadavka začína rovnakými tokenmi, kompatibilný inferenčný server nemusí túto časť počítať znova. Dokumentácia vLLM uvádza ako typické príklady opakované otázky nad tým istým dlhým dokumentom a viacnásobné kolá jednej konverzácie. Úspora sa prejavuje najmä v čase do prvého vygenerovaného tokenu a v priepustnosti prefillingu, nie v rýchlosti následného generovania nových tokenov.

Problém vzniká po rozdelení prevádzky medzi viacero GPU inštancií. Bežný load balancer môže jednotlivé požiadavky s rovnakým systémovým promptom alebo dokumentom posielať zakaždým inam. Každý server potom vidí daný prefix iba sporadicky, vytvorí si vlastnú oddelenú cache a často musí výpočet zopakovať. SageMaker preto presúva informáciu o podobnosti vstupov do smerovacej vrstvy. Rovnaký začiatok požiadavky má stabilne smerovať na rovnakú inštanciu, kde už príslušné bloky KV cache pravdepodobne existujú.

AWS zároveň pridáva ochranu pred preťažením. Správca nastaví maximálny počet súbežne spracúvaných požiadaviek pre cieľovú inštanciu. Keď je limit dosiahnutý, ďalšie volanie môže byť presmerované na menej vyťažený server aj za cenu straty konkrétneho zásahu do cache. Smerovanie má zostať relatívne stabilné aj pri pridávaní alebo odoberaní inštancií, aby horizontálne škálovanie zbytočne nerozbilo väčšinu zahriatych cache. Ide teda o kompromis medzi lokalitou dát, rovnomerným využitím GPU a ochranou latencie počas špičky.

Výsledky AWS a hranice ich platnosti

AWS testovalo funkciu na modeli Llama 3.1 70B Instruct so siedmimi inštanciami ml.p5.48xlarge a serverom vLLM so zapnutým prefixovým cachovaním. Pri hodinovom zaťažení so spoločným prefixom dlhým 8 000 tokenov firma namerala skrátenie mediánu času do prvého tokenu o 71 až 77 percent. Hodnota P90 klesla o 33 až 37 percent, priepustnosť stúpla o 15 až 16 percent a podiel zásahov do KV cache sa podľa AWS posunul približne z 25 na 82 percent.

Pri kratších konverzačných vstupoch podobných dátam ShareGPT bol efekt menší. Medián času do prvého tokenu klesol o 13 až 16 percent, P90 o 24 až 37 percent a priepustnosť sa zvýšila o 1,7 až 2 percentá. Samotné rozhodovanie smerovača pridalo podľa meraní AWS približne 1,3 až 1,9 milisekundy na požiadavku. Tieto čísla sú výsledkami dodávateľa v konkrétnej konfigurácii, nie nezávislým benchmarkom ani zárukou pre každý model, typ GPU a profil prevádzky.

Rozdiel medzi dlhým a krátkym kontextom je podstatný. Prefixová cache odstraňuje opakovaný výpočet vstupu, ale neskracuje dekódovanie dlhej odpovede. Ak aplikácia posiela prevažne unikátne prompty alebo model trávi väčšinu času generovaním množstva nových tokenov, zisk môže byť nízky. Rovnako nepomôže samotné smerovanie, ak použitý modelový server prefixové cachovanie nepodporuje alebo ho nemá aktivované. vLLM túto optimalizáciu opisuje ako opätovné použitie úplných blokov cache so zhodným tokenovým prefixom.

Kde môže nové smerovanie priniesť najväčší úžitok

Najzreteľnejším prípadom sú RAG aplikácie, v ktorých veľa používateľov kladie rôzne otázky nad rovnakou príručkou, zmluvou alebo interným dokumentom. Dokument tvorí dlhú spoločnú časť promptu a mení sa až otázka na jeho konci. Podobný vzor majú firemní asistenti s rozsiahlymi systémovými pravidlami, zákaznícke chatboty s opakovanou sadou politík, agenti používajúci stabilný opis dostupných nástrojov a programátorské asistenty, ktoré opakovane posielajú obsah rovnakého súboru.

Pre slovenské tímy môže byť funkcia zaujímavá najmä tam, kde prevádzkujú vlastný open-weight model na viacerých GPU a pracujú s dlhými podnikovými podkladmi. Potenciálnym výsledkom nie je len svižnejšia prvá odpoveď. Vyšší počet zásahov do cache môže uvoľniť čas akcelerátorov na nové požiadavky, čo môže znížiť potrebnú kapacitu alebo oddialiť ďalšie škálovanie. Reálnu ekonomiku však treba merať na vlastných dátach: rozhodujú dĺžky prefixov, miera ich opakovania, frekvencia škálovania, veľkosť cache aj pomer prefillingu k dekódovaniu.

Nasadenie vyžaduje konzistentné vstupy

Stratégia sa nastavuje pre produkčný variant endpointu. Parameter PrefixLength určuje, akú veľkú počiatočnú časť požiadavky použije SageMaker pri rozhodovaní. API referencia povoľuje rozsah od 1 024 do 65 536. Pri natívnom SageMaker API ide o bajty od začiatku tela požiadavky, zatiaľ čo pri rozhraní kompatibilnom s OpenAI sa počítajú znaky textového obsahu správ. Druhý parameter, ConcurrencyThreshold, určuje hranicu súbežných požiadaviek, po ktorej sa uplatní ochrana pred preťažením.

Voľba dĺžky prefixu nie je kozmetické nastavenie. Príliš krátka hodnota môže spojiť nesúvisiacu prevádzku a vytvoriť horúce miesto na jednej inštancii. Príliš dlhý prefix môže zahrnúť meniace sa polia a rozptýliť požiadavky, ktoré mali využívať tú istú cache. Pri natívnom API záleží aj na serializácii JSON: rozdielne poradie kľúčov, medzery alebo dynamické metadáta na začiatku tela môžu zmeniť smerovací vstup, hoci text promptu je významovo rovnaký. Producent požiadaviek by preto mal používať deterministický formát.

SageMaker umožňuje oddeliť rovnaké prompty rôznych zákazníkov pomocou identifikátora. Natívne volania môžu použiť hlavičku X-Amzn-SageMaker-Prefix-Aware-Id a OpenAI-kompatibilné požiadavky pole prompt_cache_key. AWS ich kombinuje s prefixom pri výbere inštancie. Je to dôležité v multitenantnom prostredí, no nemožno to automaticky zamieňať za úplnú bezpečnostnú izoláciu: prevádzkovateľ musí ďalej preveriť správanie modelového servera, životný cyklus cache, oprávnenia endpointu a požiadavky na oddelenie citlivých dát.

Čo zatiaľ nevieme

Zverejnené testy nepokrývajú všetky modely, menšie GPU zostavy, prudko kolísajúcu prevádzku ani dlhodobé automatické škálovanie. Nie je tiež jasné, ako sa prínos zmení pri veľmi veľkom počte takmer rovnakých tenantov, pri častých aktualizáciách systémových promptov alebo pri cache, ktorá sa pod tlakom pamäte rýchlo prepisuje. AWS uvádza vyvážené rozdelenie testovacej prevádzky, ale produkčný výsledok bude závisieť od popularity jednotlivých prefixov a správne nastavenej hranice súbehu.

Novinka preto nie je univerzálnym prepínačom na zrýchlenie každého LLM endpointu. Je to cielenejšia forma plánovania prevádzky, ktorá prepája load balancing so stavom uloženým v lokálnej KV cache. Pre opakujúce sa dlhé kontexty môže odstrániť významnú časť redundantného výpočtu; pri unikátnych promptoch alebo dlhých odpovediach bude efekt obmedzený. Rozumné nasadenie by malo začať meraním času do prvého tokenu, podielu zásahov do cache, priepustnosti a rozdelenia požiadaviek medzi inštancie a až potom porovnať PREFIX_AWARE s náhodným smerovaním alebo stratégiou najmenšieho počtu rozpracovaných požiadaviek.