Gescheiden Prefill en Decode

Om outputtokens te genereren uit een invoerprompt, wordt LLM-inferentie opgesplitst in twee fasen: prefill en decode. Prefill wordt uitgevoerd op de invoertokens en vult KV-caches in voordat het decode-stadium wordt ingevoerd dat tokens één voor één genereert.

Hoewel een enkele decode-stap meestal enkele tientallen milliseconden duurt, duurt prefill aanzienlijk langer. Als ze op dezelfde apparaten worden uitgevoerd, verslechtert het mengen van prefill met decode de decode-prestaties. In dit artikel onderzoeken we een gevestigde oplossing in de vorm van gescheiden prefill en decode, waarbij deze op aparte apparaten worden uitgevoerd om zowel de prefill-doorvoer als de decode-latenties te maximaliseren.

Prefill versus Decode-prestaties

In een typische LLM-serve engine selecteert de batchplanner verzoeken om te verwerken in elke uitvoerstap van een model. Wanneer ze op een enkel apparaat of knooppunt worden uitgevoerd, worden zowel prefill- als decode-verzoeken samen gebatched. De kosten van aandacht, die zich opstapelen langs de sequentielengte, groeien voor zowel prefill als decode, evenredig met de lengte van items in de KV-cache (kv_len). Decode-verzoeken leiden meestal één token door (qo_len=1), tegen minimale kosten door andere lagen die onafhankelijk werken op de tokens van een sequentie. Prefill-verzoeken leiden duizenden of tienduizenden tokens door tegen aanzienlijke kosten door dichte lagen (grote qo_len).

De latentietijd van een voorwaartse doorgang wordt sterker beïnvloed door het aantal onafhankelijke tokens dat door dichte lagen gaat (qo_len) dan door het aantal tokens dat uit de KV-cache wordt opgehaald tijdens aandacht (kv_len). Aandacht kan parallelliseren zowel over het aantal verzoeken als de kv_len evenredig met sequentielengtes, waardoor een goede benutting wordt bereikt. Prefill is rekenkrachtgebonden: bij een hoge qo_len kunnen GEMM-kernels voldoende blokken toewijzen langs de M-dimensie om de rekencapaciteiten van moderne GPU's volledig te benutten. Decode is geheugengebonden: vanwege typisch lage batchgroottes, is het aantal ingangen langs M meestal klein, voldoende voor slechts één blok. Hoewel Split-K GEMM-kernels de SM-benutting kunnen verbeteren voor lage tokenbatchgroottes, blijven de caches en de matrixvermeerderingsunits doorgaans onderbenut.

Wanneer ze worden gemengd, veroorzaken batches met prefill-verzoeken hogere latentietijden door de voorwaartse doorgang, wat een negatief effect heeft op de decode-doorvoer van de hele instantie. 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 handhaven om genoeg verzoeken op een instantie te verwerken om de decode-doorvoer te maximaliseren. In het geval van grote modellen, met typische uitvoerlengtes, moet prefill vaak genoeg worden uitgevoerd om een grote batchgroott voor decode te behouden, waardoor het de gemiddelde latentietijd aanzienlijk verslechtert en haperingen in de uitvoer veroorzaakt.

Deze problemen kunnen worden aangepakt door een aparte set knooppunten te gebruiken om prefill en decode uit te voeren. Door een prefill-knooppunt te associëren met meerdere decoder-knooppunten kunnen voldoende verzoeken worden gepland voor prefill om doorvoer te maximaliseren en een groot genoeg aantal gelijktijdige verzoeken op de decoder-knooppunten te behouden om ook de decode-doorvoer te maximaliseren. De prefill-knooppunten vullen de KV-caches, die vervolgens worden overgedragen aan de decoder-knooppunten. Aangezien de decoders niet langer hoeven te pauzeren voor prefill, worden latenties veel meer deterministisch, omdat de algehele impact van de groei van kv_len van actieve verzoeken veel minder uitgesproken is. De kosten worden betaald in een toename van de Time to First Token (TTFT), aangezien de overdracht van KV-caches via het netwerk tientallen tot honderden milliseconden kan duren.

KV Messenger

Bij Perplexity is onze implementatie voor gescheiden prefill en decode opgebouwd rond een KV messenger die samenwerkt met de LLM-motor om KV-cacheoverdrachten van prefill-knooppunten naar decoder-knooppunten via een netwerk te orkestreren. Aan de prefill-kant accepteert de messenger verzoeken van decoder-knooppunten, geeft deze door aan de batchplanner en houdt de uitvoering van de voorwaartse doorgang bij om KV-caches met zo min mogelijk latentie te verzenden. Aan de decoder-kant, na niet-uitzetbare pagina's te hebben toegewezen, blokkeert de messenger het verzoek om gepland te worden voor decode totdat hij wordt geïnformeerd over de voltooiing van de KV-cache en decoder-contextoverdrachten.

Het scheiden van prefill vereist hoge doorvoersnelheden, verbindingen met lage latentie, dus onze implementatie is op maat gemaakt voor RDMA en ondersteunt zowel EFA als ConnectX Network Interface Controllers (NICs). De KV Messenger is opgebouwd met libfabric, waarbij onze fabric-lib wrappers worden gebruikt om hogere abstracties met lage latentie te bieden over de Remote Direct Memory Access (RDMA)-primitieven, voor het implementeren van efficiënte pagina- en metadataoverdrachten, samen met signalering met lage latentie. Op de achtergrond coördineert fabric-lib een GPU en zijn direct verbonden NICs om gegevens van het prefill-knooppunt naar het decoder-knooppunt te kopiëren.

Na ontvangst, wijst het prefill-knooppunt een overeenkomstige set bron-KV-pagina's toe en plant het verzoek voor prefill met zijn lokale motor. Om de latentie te minimaliseren, wachten overdrachten niet op de voorwaartse doorgang: 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. Aangezien prefill-verzoeken kunnen worden opgedeeld, informeert de batchplanner de KV-messenger over de huidige geplande stukken vóór de uitvoering. Om CUDA-grafieken te ondersteunen terwijl het mogelijk is om lagen te volgen, houdt de messenger een toegewijde thread bij die een teller pollt die wordt verhoogd na de uitvoerprojectie van aandacht. De teller wordt alleen bijgehouden op het hoofdknooppunt in een geshard omgeving: hoewel de KV-cache-items geldig zijn na het toevoegen en vóór aandacht, wordt de uitvoerprojectie verminderd over rangen, waardoor ze impliciet worden gesynchroniseerd. Wanneer een wijziging in de teller wordt waargenomen, wordt de messenger geïnformeerd en roept hij fabric-lib op om de overdracht van een laag te initiëren.

Na de voltooiing van de laatste chunk-overdracht worden eventuele extra metadata ook gekopieerd: speculatieve decode of MTP vereisen dat logits en verborgen toestanden worden verplaatst naar de decoder. Deze kopieën worden ook uitgevoerd via RDMA, naar en vanaf vooraf toegewezen buffers.

Na de voltooiing van alle uitstaande overdrachten van de laatste chunk, dealloceert het prefill-knooppunt de KV-pagina's en voltooit het verzoek. Het decoder-knooppunt wordt niet expliciet geïnformeerd: in plaats daarvan gebruikt het directe telmethode om het aantal voltooide operaties bij te houden. Het aantal RDMA-operaties aan de prefill-kant is evenredig met het aantal overgedragen pagina's. Na de voltooiing van het bekende aantal pagina- en contextkopieën, roept fabric-lib de KV messenger op om aan te geven dat een verzoek klaar is voor decode. De messenger dealloceert eventuele context en geeft het verzoek door aan de LLM-motor.

Gesharcde KV Cache Transfers

Als de prefill en decoder afhankelijk zijn van Tensor Parallelism (TP) en de KV-caches identiek fragmenteren of repliceren, coördineert een enkele overdrachtmotor meerdere apparaten om de pagina's van alle replica's te verzenden en te ontvangen. Om een enkele messenger en overdrachtmotor te kunnen gebruiken, ondanks het feit dat de executie 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 om het in het hoofdproces te mappen. De overdrachtmotor inspecteert de topologie van het knooppunt om de NICs en de CPU's in het dichtstbijzijnde NUMA-knooppunt te vinden voor de overdrachten van elk van de KV-cache-snedes.

Als de bron en bestemming identiek sharden, zijn transfers triviaal doordat er een een-op-een mapping is van de apparaten en pagina's van de bron en bestemming. In deze situatie helpt het sharden impliciet bij overdrachtlatenties: door meer GPUs te gebruiken, kunnen meer geassocieerde NICs worden ingezet voor een bijna volledige bandbreedte-benoeming. Als er echter een mismatch is, moet de overdrachtmotor pagina's splitsen of reconstrueren, afhankelijk van de verhouding tussen bron- en bestemmingsneden.

Als de prefill de KV-cache over meer apparaten verdeelt, worden volledige pagina's gereconstrueerd op de decoder door de corresponderende helften uit de prefill-apparaten te sturen. Als de decoder meer shards heeft, ontvangt hij pagina's van meerdere bronnen. De decoder moet op de hoogte zijn van het shardingschema van de prefill om het aantal RDMA-schrijvingen dat hij verwacht te ontvangen te kunnen berekenen. Als replicatie betrokken is, groepeert de prefill de apparaten in replicasets die de volledige KV-cache binnen zichzelf repliceren. Bestemmingsreplicasets worden willekeurig toegewezen aan een van de bronsets om alle beschikbare apparaten te gebruiken voor het initiëren van RDMA-schrijfopdrachten.

Gesharcde overdrachten vereisen een kleine aanpassing van KV-caches. Standaard vertrouwt FlashInfer op de NHD-indeling, die de tokens binnen een pagina rangschikt binnen de hoofden. Aangezien caches waarschijnlijk worden geshard op het aantal aandachtshoofden, veroorzaakt dit discontinuïteit binnen het hoofd. RDMA-overdrachten ondersteunen niet impliciet gestreepte schrijvingen, waardoor één operatie per hoofd benodigd is 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-indeling, waarin de hoofddimensie vóór het aantal tokens wordt geplaatst. Dit zorgt voor continuïteit, waardoor een pagina met een enkele schrijfbewerking kan worden gekopieerd.

Speculatieve Decoding

Speculatieve decodering vereist enkele aanpassingen aan gescheiden prefill-decode. In onze uitvoering zijn prefill-knooppunten niet toegestaan om tokens te nemen. Aangezien de Sonar-modellen van Perplexity gestructureerde uitvoer ondersteunen, willen we de complexiteit vermijden van het synchroniseren van de schema-verwerkersimplementaties tussen prefillers en decoders. In het MTP- en speculatieve decoderingmechanismen wordt het concept om het conceptmodel te prefabuen tot de laatste token omvat met het bemonsteren van tokens uit het doelmodel.

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

Gescheiden Implementaties

We hebben meerdere gescheiden configuraties geïmplementeerd of uitgeprobeerd met verschillende modellen, om ofwel productieverkeer te ondersteunen of in-house evaluatiewerkbelastingen. Gebaseerd op de grootte en het aandachtmechanisme van modellen, kozen we geschikte shardingschema’s voor prefillers en decoders om GPU's optimaal te benutten.

DeepSeek-R1

Met DeepSeek hebben we zowel Tensor-Parallel (TP) als Data-Parallel (DP) implementaties overwogen. Zoals besproken in eerdere blogberichten, bieden TP-implementaties een betere latentietijd ten koste van een lagere doorvoer, wat meer GPU's vereist om zware belasting te kunnen verwerken. DP-implementaties schalen veel beter met belasting, maar hun piekdoorvoer is lager vanwege de kosten van communicatie tussen apparaten of knooppunten.

DeepSeek vertrouwt op Multi-Head Latent Attention om KV-caches te comprimeren. Aangezien alle KV-hoofden worden samengeperst in een enkele latente vector, kan TP de KV-caches niet parceleren, aangezien het in plaats daarvan de latente vectoren moet repliceren op alle rangen. Het sharden gebeurt na decomprimeren, aangezien elke rang afzonderlijke hoofden kan extraheren uit dezelfde latente representatie. Bijgevolg zijn alle KV-cache-shards identiek over zowel prefill- als decoder-shards.

Bij een intranode TP-opstelling, worden zowel prefill- als decoders gelijk geshard. Overdrachten worden vanaf alle rangen verzonden om optimaal gebruik te maken van alle beschikbare NIC's. Echter, bij een DP-implementatie, waarbij de TP-ranggrootte lager is of elke DP-rang aan een enkele GPU toegewezen is, kan elk prefill-apparaat dat een gerepliceerd exemplaar van de KV-cache behoudt, deze verzenden. Om verzoeken te balanceren over alle beschikbare NIC's, selecteren we willekeurig een GPU en een NIC om de KV-cache van de prefill naar de decoder te verzenden.

Met gemengde prefill-decode had onze R1-implementatie moeilijkheden om consistent meer dan 50 TPS te overtreffen vanwege frequente prefill-onderbrekingen in de orde van honderden milliseconden. Daarentegen, door prefill te scheiden, kregen we een TTFT-boete van ongeveer 100ms voor elk verzoek, maar een enkele prefill-knooppunt kon consistente batchgroottes behouden op 3 decoder-knooppunten, dat een doorvoer leverde van meer dan 90 TPS terwijl een belasting van ongeveer 1 QPS per decoder-knooppunt werd afgehandeld. Met data-parallelle implementaties was TPS iets lager op ongeveer 50, maar de instanties konden een belasting van 1 QPS per rang aan, met 8 rangen naar een enkel knooppunt.

Qwen3-Coder

Dit 480B-model gebruikt Gerelateerde-Query Aandacht (GQA), zodat aandacht gemakkelijk kan worden geshard en kan profiteren van tensor-parallelisme zonder geheugen te offeren voor KV-caches. Bijgevolg konden we het model scharden over 8 GPU's voor zowel prefill als decode, waarbij we ongeveer 3 decoder-knooppunten koppelen met een enkele prefill-knooppunt. Aangezien aandacht wordt geshard, vertrouwen we op de HND KV-cache-indeling om prefill- en decoder-KV-caches te sharden, waarbij prefill-rangen worden gekoppeld met decoder-rangen en alle NIC's optimaal worden benut om gelijktijdig snedes over te dragen.