Gedisaggregeerde prefill en decode

Om uitvoertokens van een invoerprompt te genereren, wordt LLM-inferentie opgesplitst in twee fasen: prefill (vooraf vullen) en decode (decoderen). Prefill draait op de invoertokens en vult de KV-caches, voordat de decode-fase begint die tokens één voor één genereert. Waar een enkele decode-stap doorgaans tientallen milliseconden in beslag neemt, duurt prefill aanzienlijk langer. Indien uitgevoerd op dezelfde apparaten, vermindert het mengen van prefill met decode de prestaties van het decoderen. In dit artikel onderzoeken we een beproefde oplossing in de vorm van gedisaggregeerde prefill en decode, waarbij deze op afzonderlijke apparaten worden uitgevoerd om zowel de prefill-doorvoer als de decode-latentie te maximaliseren.

Prefill versus Decode prestaties

In een typische LLM-serving-engine selecteert de batch-scheduler verzoeken die moeten worden verwerkt in elke uitvoeringsstap van een model. Bij uitvoering op een enkel apparaat of node worden zowel prefill- als decode-verzoeken samen gebatched. De kosten van attention, die accumuleren langs de sequentielengte, nemen voor zowel prefill als decode toe, evenredig aan de lengte van de vermeldingen in de KV-cache (kv_len). Decode-verzoeken sturen meestal één token door (qo_len=1), tegen minimale kosten via andere lagen die onafhankelijk op de tokens van een sequentie werken. Prefill-verzoeken sturen duizenden of tienduizenden tokens door tegen aanzienlijke kosten via dense lagen (grote qo_len).

De latentie van een forward pass wordt sterker beïnvloed door het aantal onafhankelijke tokens dat door dense lagen wordt gestuurd (qo_len) dan door het aantal tokens dat tijdens attention uit de KV-cache wordt opgehaald (kv_len). Attention kan zowel parallel worden uitgevoerd over het aantal verzoeken als over de kv_len evenredig aan de sequentielengten, waardoor een goede benutting wordt bereikt. Prefill is rekenintensief: aangezien qo_len hoog is, kunnen GEMM-kernels voldoende blokken toewijzen langs de M-dimensie om de rekenkracht van moderne GPU's volledig te benutten. Decode is geheugenintensief: vanwege de doorgaans lage batchgrootten is het aantal invoeren langs M meestal klein, wat slechts volstaat voor één blok. Hoewel Split-K GEMM-kernels de SM-benutting voor kleine token-batchgroottes kunnen verbeteren, blijven de caches en de matrixvermenigvuldigingseenheden doorgaans onderbenut.

Vergroten

Wanneer deze worden gemengd, leiden batches met prefill-verzoeken tot hogere latenties tijdens de forward pass, wat de decode-doorvoer van de gehele instantie negatief beïnvloedt. Hoewel het mengen van prefill-verzoeken met decode-verzoeken of het gebruik van chunked prefill de decode-prestaties enigszins kan verbeteren, is het moeilijk om voldoende prefill-doorvoer te behouden om genoeg verzoeken op een instantie te verwerken en zo de decode-doorvoer te maximaliseren. Bij grote modellen, met typische uitvoerlengten, moet prefill vaak genoeg worden uitgevoerd om een grote batchgrootte voor decode te behouden, wat de gemiddelde latentie aanzienlijk verslechtert en tot stotteren in de uitvoer leidt.

Vergroten

Vergroten

Deze problemen kunnen worden aangepakt door een afzonderlijke set nodes te gebruiken voor prefill en decode. Door een prefiller-node te koppelen aan meerdere decoder-nodes, kunnen voldoende verzoeken worden ingepland voor prefill om de doorvoer te maximaliseren en een voldoende groot aantal gelijktijdige verzoeken op de decoder-nodes te behouden om ook de decode-doorvoer te maximaliseren. De prefiller-nodes vullen de KV-caches, die vervolgens naar de decoder-nodes worden overgedragen. Omdat de decoders niet langer hoeven te stoppen voor prefill, worden latenties veel deterministischer, aangezien de algehele impact van de groeiende kv_len van actieve verzoeken veel minder uitgesproken is. De prijs hiervan is een toename in Time to First Token (TTFT), aangezien de overdracht van KV-caches via het netwerk tientallen tot honderden milliseconden in beslag kan nemen.

KV Messenger

Bij Perplexity is onze implementatie voor gedisaggregeerde prefill en decode gebouwd rond een KV-messenger die interactie heeft met de LLM-engine om de overdracht van KV-caches van prefiller-nodes naar de decoder-nodes via een netwerk te orkestreren. Aan de kant van de prefiller accepteert de messenger verzoeken van decoder-nodes, draagt deze over aan de batch-scheduler en houdt de uitvoering van de forward pass bij om KV-caches met zo min mogelijk latentie te versturen. Aan de kant van de decoder, nadat pagina's die niet kunnen worden verwijderd zijn toegewezen, blokkeert de messenger het verzoek om ingepland te worden voor decode totdat deze op de hoogte wordt gesteld van de voltooiing van de overdracht van de KV-cache en decoder-context.

Vergroten

Het disaggregeren van prefill vereist verbindingen met een hoge doorvoer en lage latentie, daarom is onze implementatie afgestemd op RDMA, waarbij zowel EFA- als ConnectX Network Interface Controllers (NIC's) worden ondersteund. De KV Messenger is gebouwd op libfabric, waarbij gebruik wordt gemaakt van onze fabric-lib-wrappers om op een hoger niveau abstracties met lage latentie te bieden boven de Remote Direct Memory Access (RDMA)-primitieven, en waarbij efficiënte pagina- en metadata-overdrachten worden geïmplementeerd, samen met signalering met lage latentie. Op de achtergrond coördineert fabric-lib een GPU en de direct verbonden NIC's om gegevens van de prefiller-node naar de decoder-node te kopiëren.

Na ontvangst wijst de prefiller-node een overeenkomstige set bron-KV-pagina's toe en plant het verzoek voor prefill in met behulp van de lokale engine. Om de latentie te minimaliseren, wachten overdrachten niet op de forward pass: in plaats daarvan worden KV-paginakopieën gestart zodra het model klaar is met het toevoegen van KV-cache-items aan de KV-cache voor individuele lagen. Omdat prefill-verzoeken in chunks kunnen worden opgedeeld, brengt de batch-scheduler de KV-messenger vóór uitvoering op de hoogte van de momenteel geplande chunks. Om CUDA-grafieken te ondersteunen terwijl lagen kunnen worden getraceerd, behoudt de messenger een toegewezen draad die een teller polt die wordt verhoogd na de output-projectie van attention. De teller wordt alleen bijgehouden op de lead-node in een sharded omgeving: hoewel de KV-cache-items geldig zijn na toevoeging en vóór attention, wordt de output-projectie gereduceerd over de rangen, waardoor ze impliciet worden gesynchroniseerd. Zodra een wijziging in de teller wordt waargenomen, wordt de messenger hiervan op de hoogte gesteld en roept deze fabric-lib aan om de overdracht van een laag te starten.

Vergroten

Nadat de overdracht van de laatste chunk is voltooid, worden ook alle aanvullende metadata gekopieerd: speculatief decoderen of MTP vereisen dat logits en verborgen toestanden naar de decoder worden verplaatst. Deze kopieën worden ook uitgevoerd via RDMA, van en naar vooraf toegewezen buffers.

Na voltooiing van alle in behandeling zijnde overdrachten van de laatste chunk, deallocateert de prefiller-node de KV-pagina's en voltooit het verzoek. De decoder-node wordt niet expliciet op de hoogte gesteld: in plaats daarvan gebruikt deze onmiddellijke tellers om het aantal voltooide operaties bij te houden. Het aantal RDMA-operaties aan de prefiller-kant is evenredig aan het aantal overgedragen pagina's. Na voltooiing van het bekende aantal pagina- en contextkopieën roept fabric-lib de KV-messenger aan om aan te geven dat een verzoek klaar is voor decodering. De messenger deallocateert alle context en draagt het verzoek over aan de LLM-engine.

Sharded KV-cacheoverdrachten

Als de prefiller en decoder vertrouwen op Tensor Parallelism (TP) en de KV-caches identiek sharden of repliceren, coördineert een enkele overdracht-engine meerdere apparaten om de pagina's van alle replica's te verzenden en te ontvangen. Om een enkele messenger en overdracht-engine te kunnen gebruiken, ondanks het feit dat de uitvoerder van het model over meerdere apparaten en processen is gerepliceerd, worden cuMem en cuMemImportFromShareableHandle gebruikt om het apparaatgeheugen toe te wijzen dat de KV-caches ondersteunt en dit in het hoofdproces in kaart te brengen. De overdracht-engine inspecteert de topologie van de node om de NIC's en de CPU's in de dichtstbijzijnde NUMA-node te vinden die moeten worden gebruikt voor de overdrachten van elk van de KV-cache-slices.

Als de bron en bestemming identiek sharden, zijn overdrachten triviaal aangezien er een één-op-één-toewijzing is van de apparaten en pagina's van de bron en bestemming. In deze situatie helpt sharding impliciet bij de overdrachtslatentie: door meer GPU's te gebruiken, kunnen meer bijbehorende NIC's worden ingezet, waardoor een volledige bandbreedtebenutting wordt benaderd. Als er echter een mismatch is, moet de overdracht-engine pagina's splitsen of reconstrueren, afhankelijk van de verhouding tussen bron- en bestemmingsslices.

Vergroten

Als de prefiller de KV-cache over meer apparaten verdeelt, worden volledige pagina's op de decoder gereconstrueerd door de overeenkomstige helften van de prefiller-apparaten te verzenden. Als de decoder meer shards heeft, ontvangt deze pagina's van meerdere bronnen. De decoder moet op de hoogte zijn van het sharding-schema van de prefiller om het aantal RDMA-schrijfacties te kunnen berekenen dat deze naar verwachting zal ontvangen. Als replicatie betrokken is, groepeert de prefiller de apparaten in replica-sets die de volledige KV-cache binnen zichzelf repliceren. Bestemmings-replica-sets worden willekeurig toegewezen aan een van de bron-sets om alle beschikbare apparaten te gebruiken voor het initiëren van RDMA-schrijfacties.

Vergroten

Sharded overdrachten vereisen een kleine aanpassing aan KV-caches. Standaard vertrouwt FlashInfer op de NHD-layout, die de tokens binnen een pagina binnen de heads ordent. Aangezien caches waarschijnlijk worden geshard langs het aantal attention-heads, creëert dit discontinuïteit binnen de head. RDMA-overdrachten ondersteunen impliciet geen strided schrijfacties, wat één operatie per head vereist om de overdracht uit te voeren. In plaats daarvan, om het aantal interacties met libfabric te verminderen, organiseren we KV-caches met behulp van de HND-layout die de head-dimensie vóór het aantal tokens plaatst. Dit zorgt voor continuïteit, waardoor een pagina met één schrijfactie kan worden gekopieerd.

Speculatief decoderen

Speculatief decoderen vereist kleine aanpassingen aan gedisaggregeerde prefill-decode. In onze implementatie mogen prefiller-nodes geen tokens samplen. Aangezien de Sonar-modellen van Perplexity gestructureerde uitvoer ondersteunen, willen we niet de complexiteit aangaan van het synchroniseren van de schema-processor-implementaties tussen prefillers en decoders. In de MTP- en speculatieve decode-mechanismen omvat het prefillen van het conceptmodel tot het laatste token het samplen van tokens uit het doelmodel.

Vergroten

Om deze problemen te omzeilen, bevat prefill niet het laatste token van de invoersequentie. In plaats daarvan worden verborgen toestanden of logits van prefill die voorafgaan aan het laatste token overgedragen en wordt dit in de volgende stap op de decoder behandeld als een decode-token. Hoewel dit de latentie enigszins verhoogt, aangezien een volledige decode-stap moet worden uitgevoerd na prefill om het eerste token uit te geven, wordt de complexiteit van de implementatie aanzienlijk verminderd.

Gedisaggregeerde implementaties

We hebben meerdere gedisaggregeerde configuraties met verschillende modellen ingezet of ermee geëxperimenteerd, om productieverkeer of interne evaluatieworkloads te ondersteunen. Op basis van de grootte en het attention-mechanisme van modellen, kozen we geschikte sharding-schema's voor prefiller- en decoder-nodes om GPU's optimaal te benutten.

DeepSeek-R1

Met DeepSeek overwogen we zowel Tensor-Parallel (TP)- als Data-Parallel (DP)-implementaties. Zoals besproken in eerdere blogposts, bieden TP-implementaties een betere latentie ten koste van een lagere doorvoer, waarbij meer GPU's nodig zijn om zwaar verkeer te bedienen. DP-implementaties schalen veel beter met de belasting, maar hun piekdoorvoer is lager vanwege de kosten van communicatie tussen apparaten of nodes.

DeepSeek vertrouwt op Multi-Head Latent Attention, wat KV-caches comprimeert. Aangezien alle KV-heads worden gecomprimeerd tot één latente vector, kan TP de KV-caches niet sharden, aangezien het de latente vectoren op alle rangen moet repliceren. Sharding gebeurt na decompressie, aangezien elke rang verschillende heads uit dezelfde latente representatie kan extraheren. Bijgevolg zijn alle KV-cache-shards identiek over zowel prefiller- als decoder-shards.

Met een intranode TP-setup zijn zowel prefillers als decoders identiek geshard. Overdrachten worden verzonden vanaf alle rangen om alle beschikbare NIC's volledig te benutten. Echter, bij een DP-implementatie, waarbij de TP-ranggrootte lager is of elke DP-rang aan een enkele GPU is toegewezen, kan elk prefiller-apparaat dat een gerepliceerde kopie van de KV-cache bevat, deze verzenden. Om verzoeken over alle beschikbare NIC's te balanceren, selecteren we willekeurig een GPU en een NIC om de KV-cache van de prefiller naar de decoder te sturen.

Met gemengde prefill-decode had onze R1-implementatie moeite om consistent de 50 TPS te overschrijden vanwege frequente prefill-onderbrekingen van honderden milliseconden. Daarentegen liepen we door prefill te scheiden een boete van ongeveer 100 ms op voor TTFT voor elk verzoek, maar kon een enkele prefiller-node consistente batchgroottes op 3 decoder-nodes handhaven, wat een doorvoer van meer dan 90 TPS opleverde terwijl een belasting van ongeveer 1 QPS per decoder-node werd verwerkt. Bij data-parallelle implementaties was de TPS iets lager, rond de 50, maar de instanties konden een belasting van 1 QPS per rang aan, met 8 rangen per node.

Qwen3-Coder

Dit 480B-model gebruikt Grouped-Query Attention (GQA), dus attention kan eenvoudig worden geshard en kan profiteren van tensor-parallelisme zonder geheugen voor KV-caches op te offeren. Hierdoor konden we het model over 8 GPU's sharden voor zowel prefill als decode, waarbij we ongeveer 3 decoder-nodes paarden met één prefiller-node. Omdat attention is geshard, vertrouwen we op de HND KV-cache-layout om prefiller- en decoder-KV-caches te sharden, waarbij prefiller-rangen worden gekoppeld aan decoder-rangen en alle NIC's volledig worden benut om slices parallel over te dragen.

Citering

@online{perplexity2025DisaggregatedPrefillAndDecode,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Kevin Hu and Erin McKnight and Phi Hoang},
  title = {{Disaggregated prefill and decode}},
  date = {2025-08-01},
  url = {https://perplexity.coolkeypoint.com/hub/blog/disaggregated-prefill-and-decode}
}