Deagregované Předběžné Vyplnění a Dekódování

Aby bylo možné generovat výstupní tokeny z vstupního promptu, je inference LLM rozdělena do dvou fází: předsázka (prefill) a dekódování. Předsázka se provádí na vstupních tokenech a naplňuje mezipaměti KV, než se přejde do fáze dekódování, která generuje tokeny jeden po druhém.

Zatímco jeden dekódovací krok obvykle trvá desítky milisekund, předsázka trvá výrazně déle. Pokud se obě fáze běží na stejných zařízeních, smíchání předsázky s dekódováním zhoršuje výkon dekódování. V tomto článku zkoumáme osvědčené řešení ve formě oddělené předsázky a dekódování provozovaného na samostatných zařízeních, aby byl maximalizován jak přenos předsázek, tak latence dekódování.

Výkon předsázky vs dekódování

V typickém obslužném engine LLM plánovač dávkování vybírá požadavky k zpracování v každém kroku provádění modelu. Při provozu na jednom zařízení nebo uzlu jsou požadavky na předsázku a dekódování dávkovány dohromady. Cena pozornosti, která se agreguje podle délky sekvence, roste pro obě fáze úměrně délce položek v KV mezipaměti (kv_len). Požadavky na dekódování obvykle předávají jeden token (qo_len=1) s minimálními náklady přes další vrstvy, které fungují nezávisle na tocích sekvencí. Požadavky na předsázku předávají tisíce nebo desítky tisíc tokenů s významnými náklady přes husté vrstvy (velké qo_len).

Latence průchodu dopředu je více ovlivněna počtem nezávislých tokenů procházejících hustými vrstvami (qo_len) než počtem tokenů získaných z KV mezipaměti při pozornosti (kv_len). Pozornost se může paralelizovat jak podle počtu požadavků, tak kv_len úměrného délkám sekvencí, což dosahuje dobrého využití. Předsázka je vázána na výpočetní výkon: qo_len je vysoká, takže lze přidělit dostatečné bloky podél dimenze M k plnému využití výpočetních schopností moderních GPU. Dekódování je vázáno na paměť: vzhledem k obvykle nízkým dávkovým velikostem je počet vstupů podél M obvykle malý, dostatečný jen pro jeden blok. I když jádra Split-K GEMM mohou zlepšit využití SM pro nízké dávkové velikosti tokenů, mezipaměti a jednotky maticové multiplikace zůstávají typicky nedostatečně využité.

Když se smíchají, dávky obsahující požadavky na předsázku způsobují vyšší latenci v průchodu dopředu, což negativně ovlivňuje průtok dekódování v celé instanci. I když smíchání požadavků na předsázku s požadavky na dekódování nebo použití dělené předsázky může mírně zlepšit výkon dekódování, je obtížné udržet dostatečnou kapacitu předsázky ke zpracování dostatečného počtu požadavků na instanci pro maximalizaci průtoku dekódování. V případě velkých modelů s typickou délkou výstupu musí být předsázka provedena dostatečně často, což významně zhoršuje průměrnou latenci a způsobuje koktání ve výstupu.

Tato problematika může být řešena použitím samostatné sady uzlů k provádění předsázky a dekódování. Přidružením předsázkového uzlu s více dekódovacími uzly může být plánováno dostatečné množství požadavků na předsázku k maximalizaci propustnosti a udržení dostatečně velkého počtu souběžných požadavků na dekódovacích uzlech k maximalizaci průtoku dekódování. Předsázkové uzly naplňují KV mezipaměti, které jsou poté přeneseny k dekódovacím uzlům. Protože dekódovací uzly již nemusí přerušit předsázku, latence se stávají mnohem determinističtějšími, jelikož celkový dopad rostoucího kv_len aktivních požadavků je méně výrazný. Cena je zaplacena zvýšením času do prvního tokenu (TTFT), protože přenos KV mezipamětí přes síť může trvat desítky až stovky milisekund.

KV Messenger

V Perplexity je naše implementace oddělené předsázky a dekódování postavena kolem KV messengeru, který interaguje s engine LLM pro orchestraci přenosů KV mezipamětí z předsázkových uzlů k dekódovacím uzlům přes síť. Na straně předsázky messenger přijímá požadavky ze strany dekódovacích uzlů, předává je plánovači dávky a sleduje průběh průchodu k předání KV mezipamětí s co nejmenší latencí. Na dekódovací straně, po přidělení neodstranitelných stránek, messenger zablokuje požadavek od plánování dekódování, dokud nebude oznámeno dokončení přenosu KV mezipamětí a kontextů dekódovacích uzlů.

Oddělení předsázky vyžaduje vysokorychlostní, nízkolatenční připojení, proto je naše implementace přizpůsobena pro RDMA, podporující jak EFA, tak síťové řadiče ConnectX (NICs). KV Messenger je postaven na libfabric, pomocí našich obalů fabric-lib, aby poskytovala vyšší úroveň nízko-latenčních abstrakce nad primitivy Remote Direct Memory Access (RDMA), implementuje efektivní přenos stránek a metadat, spolu s nízkou latencí signalizací. Na pozadí fabric-lib koordinuje GPU a své přímo připojené NICs, aby kopíroval data z předsázkového uzlu na dekódovací uzel.

Po doručení předsázkový uzel přiděluje odpovídající sadu zdrojových KV stránek a plánuje požadavek na předsázku pomocí svého místního engine. Aby byla minimalizována latence, přenosy nečekají na průchod dopředu: místo toho se kopie stránek KV zahájí, jakmile model dokončí připojení položek KV mezipamětí k KV mezipaměti pro jednotlivé vrstvy. Jelikož požadavky na předsázku mohou být rozšířeny na části, plánovač dávky upozorňuje KV messenger o aktuálně plánovaných částech před provedením. Na podporu CUDA grafů a schopnost sledovat vrstvy si messenger udržuje vyhrazené vlákno pro sledování pultu, který se zvyšuje po projekci pozornosti výstupu. Pult je udržován pouze na hlavním uzlu v prostředí se sdílením: i když jsou položky KV mezipaměti platné po doplnění a před pozorností, výstupní projekce je zredukována napříč hodnostmi, což implicitně synchronizuje je. Jakmile je pozorována změna v pultu, messenger je upozorněn a volá fabric-lib k zahájení přenosu vrstvy.

Po dokončení přenosu poslední části jsou také zkopírována všechna další metadata: spekulativní dekódování nebo MTP vyžadují, aby byly přeneseny logity a skryté stavy k dekodéru. Tyto kopie jsou také realizovány přes RDMA, k a z předalokovaných vyrovnávacích pamětí.

Po dokončení všech čekajících přenosů poslední části předsázkový uzel likviduje KV stránky a uzavírá požadavek. Dekódovací uzel není výslovně upozorněn: místo toho používá okamžité počítadla k sledování počtu dokončených operací. Počet RDMA operací na straně předsázky je úměrný počtu převedených stránek. Po dokončení známého počtu kopii stránek a kontextů fabric-lib zavolá KV messenger, aby signalizoval, že požadavek je připraven kdekódování. Messenger likviduje jakýkoli kontext a předává požadavek na engine LLM.

Sharded KV cache Transfers

Pokud se předsázkový a dekódovací systém spoléhají na tensor pararelismus (TP) a shardují nebo replikují KV mezipaměti identicky, jediný přenosový engine koordinuje více zařízení k odesílání a přijímání stránek všech replik. Aby bylo možné použít jediného messengera a přenosový engine navzdory tomu, že vykonavatel modelu je replikován napříč několika zařízeními a procesy, jsou použity cuMem a cuMemImportFromShareableHandle k přidělení paměti zařízení podporující KV mezipaměti a mapování je do hlavního procesu. Přenosový engine zkoumá topologii uzlu k nalezení NICs a CPU v nejbližších NUMA uzlech pro použití pro přenosy každého z KV mezipaměťových kousků.

Pokud zdroj a cíl shardují identicky, jsou přenosy triviální, jelikož existuje jedno-na-jeden mapovaní ze zařízení a stránek zdroje a cíle. V této situaci shardování implicitně pomáhá přenosovým latencím: použitím více GPU, mohou být zaměstnány více asociované NICs, čímž se přiblíží k plnému využití šířky pásma. Avšak pokud je nesoulad, přenosový engine musí stránky sdělit nebo rekonstruovat podle poměru mezi zdrojovými a cílovými částmi.

Pokud předsázkový systém rozdělí KV mezipaměť na více zařízení, celé stránky jsou rekonstruovány na dekodéru tím, že se odpovídající části vycházející vstříc z předsázkových zařízení. Pokud má dekodér více shardů, přijímá stránky z několika zdrojů. Dekodér musí znát shardový plán předsázkového systému, aby mohl spočítat počet RDMA zápisů, které se čeká, že obdrží. Pokud se jedná o replikaci, předsázkový systém sdružuje zařízení do sestav replik, které replikují celou KV mezipaměť uvnitř nich. Cílové repliky jsou náhodně přiřazené jedna z výchozích sestav, aby použily všechna dostupná zařízení k iniciování RDMA zápisů.

Rozdělené přenosy vyžadují mírné vylepšení KV mezipamětí. Ve výchozím nastavení FlashInfer spoléhá na rozložení NHD, které objednává tokeny uvnitř stránky v rámci hlav. Jelikož se mezipaměti nejpravděpodobněji rozdělují podle počtu pozorovacích hlav, vytváří to přerušení uvnitř hlavy. RDMA přenosy nepodporují implicitně záznamová strided zápisy, což vyžaduje jednu operaci na hlavu, aby se provedl přenos. Místo toho, ve snaze snížit počet interakcí s libfabric, organizujeme KV mezipaměti pomocí rozložení HND, které umístí dimenzi hlavy před počet tokenů. To zajišťuje kontinuitu, což umožňuje kopírování stránky s jedním zápisem.

Spekulativní dekódování

Spekulativní dekódování vyžaduje mírné úpravy oddělené předsázky-dekódování. V naší implementaci nejsou předsázkovým uzlům dovoleno vzorkovat tokeny. Jelikož Sonar modely Perplexity podporují strukturovaný výstup, nechceme si přidělávat komplikace synchronizace implementace zpracovatelů schémat napříč předsázkami a dekodéry. V mechanismech MTP a spekulativního dekódování zahrnuje předsázkování návrhového modelu vzorkování tokenů z cílového modelu.

Abychom obešli tyto problémy, předsázka nezahrnuje poslední token vstupní sekvence. Skryté stavy nebo logity z předsázky předcházející poslednímu tokenu se přenáší a jsou považovány za dekódovací token v dalším kroku na dekodéru. I když to mírně zvyšuje latenci, protože je třeba provést plný dekódovací krok po předsázce k vydání prvního tokenu, složitost implementace se výrazně zjednodušuje.

Oddělená nasazení

Nasaďovali jsme nebo experimentovali s několika oddělenými konfiguracemi s různými modely, abychom podpořili buď produkční provoz, nebo doma prováděné výkonnostní úlohy. Na základě velikosti a mechaniky pozornosti modelů jsme zvolili vhodné shardovací plány pro předsázkové a dekódovací uzly k nejlepšímu využití GPU.

DeepSeek-R1

U DeepSeek jsme zvažovali jak nasazení s tensorovým paralelismem (TP), tak datovým paralelismem (DP). Jak bylo diskutováno v předchozích blogových příspěvcích, TP nasazení poskytují lepší latenci, avšak za cenu nižší průchodnosti, což vyžaduje více GPU k obsloužení vysokého provozu. DP nasazení se lépe škáluje na zatížení, avšak jejich špičková průchodnost je nižší kvůli nákladům na komunikaci mezi zařízeními nebo uzly.

DeepSeek spoléhá na Multi-Head Latent Attention, komprimující KV mezipaměti. Protože všechny KV hlavy jsou komprimovány do jednoho latentního vektoru, TP nemůže shardovat KV mezipaměti, protože místo toho musí replikovat latentní vektory na všech hodnostech. Shardování probíhá až po dekompresi, protože každá hodnost může extrahovat různé hlavy z toho samého latentního vyobrazení. V důsledku toho jsou všechny shardované KV mezipaměti identické jak v předsázkových, tak dekódovacích shardech.

S intranodovým uspořádáním TP jsou jak předsázky, tak dekódovací zařízení shardovány identicky. Přenosy jsou vysílány z všech hodností k plnému využití všech dostupných NICs. Avšak s nasazením DP, kde je velikost TP hodnosti nižší nebo každá DP hodnost je přidělena ke jednomu GPU, může vkládat jakékoli předsázkové zařízení, které obsahuje kopii replikované KV mezipaměti. Aby se vyrovnaly požadavky mezi všechna dostupná NICs, automaticky vybíráme GPU a NIC k odeslání KV mezipaměti z předsázky do dekoderu.

S mixovanou předsázka-dekodérem měla naše nasazení R1 potíže s pravidelným přesahováním 50 TPS kvůli častým přerušením předsázky v řádu stovek milisekund. Naopak, oddělením předsázky jsme měli cca 100ms zpoždění k TTFT pro každý požadavek, ale jeden předsázkový uzel mohl udržovat konstantní dávkové velikosti na 3 dekódovacích uzlech, poskytujíc průtok přes 90 TPS při obsluze zátěže asi 1 QPS na dekódovací uzel. S nasazeními založenými na datech bylo TPS mírně nižší přibližně na 50, avšak instance mohla obsloužit zátěž 1 QPS na hodnost, s 8 hodnostmi na jeden uzel.

Qwen3-Coder

Tento model 480B používá Grouped-Query Attention (GQA), takže lze pozornost snadno shardovat a může těžit z tensor pararelismu, aniž by obětovala paměť pro KV mezipaměti. V důsledku toho jsme mohli model rozšířit na 8 GPU jak pro předsázku, tak dekódování, a spojit asi 3 dekódovací uzly s jedním předsázkovým uzlem. Jelikož je pozornost shardovaná, spoléháme se na HND rozložení KV mezipamětí k shardování předsázkových a dekódovacích KV mezipamětí, spojujíc předsázkové hodnosti s dekodérovými hodnostmi a plně využívající všechny NICs k přenosu skladeb paralelně.