Prefill i decode desagregats

Per a generar tokens de sortida a partir d'un missatge d'entrada, la inferència d'un LLM es divideix en dues etapes: prefill i decode. El prefill s'executa sobre els tokens d'entrada, omplint les memòries cau KV, abans d'entrar a l'etapa de decode que genera els tokens un per un. Tot i que un sol pas de decode s'executa normalment en desenes de mil·lisegons, el prefill triga considerablement més. Si s'executa als mateixos dispositius, barrejar el prefill amb el decode degrada el rendiment del decode. En aquest article explorem una solució establerta en forma de prefill i decode desagregats, executant-los en dispositius separats per a maximitzar tant el rendiment (throughput) del prefill com les latències del decode.

Rendiment de Prefill vs Decode

En un motor de servei de LLM típic, el planificador per lots (batch scheduler) selecciona les sol·licituds a processar en cada pas d'execució d'un model. Quan s'executa en un sol dispositiu o node, tant les sol·licituds de prefill com les de decode s'agrupen per lots. El cost de l'atenció (attention), que s'agrupa al llarg de la longitud de la seqüència, creix tant per al prefill com per al decode, de manera proporcional a la longitud de les entrades a la memòria cau KV (kv_len). Les sol·licituds de decode normalment transmeten un sol token (qo_len=1), amb un cost mínim a través d'altres capes que operen de manera independent sobre els tokens d'una seqüència. Les sol·licituds de prefill transmeten milers o desenes de milers de tokens amb un cost significatiu a través de capes tridimensionals o trables (grans qo_len).

La latència d'una passada cap endavant està més fortament influenciada pel nombre de tokens independents passats per les capes denses (qo_len) que pel nombre de tokens recuperats de la memòria cau KV durant l'atenció (kv_len). L'atenció pot paral·lelitzar tant a través del nombre de sol·licituds com del kv_len proporcionalment a les longituds de les seqüències, aconseguint una bona utilització. El prefill està limitat pel càlcul: sent alt el qo_len, els nuclis GEMM poden assignar suficients blocs al llarg de la dimensió M per a utilitzar plenament les capacitats de càlcul de les GPUs modernes. El decode està limitat per la memòria: a causa de les mides de lot típicament baixes, el nombre d'entrades al llarg de M sol ser petit, suficient només per a un bloc. Tot i que els nuclis Split-K GEMM poden millorar la utilització dels SM per a mides de lot de tokens baixes, les memòries cau i les unitats de multiplicació de matrius solen continuar infrautilitzades.

Amplia

Quan es barregen, els lots que contenen sol·licituds de prefill incorren en latències més altes a través de la passada cap endavant, afectant negativament el rendiment de decode de tota la instància. Tot i que barrejar sol·licituds de prefill amb sol·licituds de decode o emprar un prefill fragmentat (chunked prefill) pot millorar lleugerament el rendiment del decode, és difícil mantenir un rendiment de prefill suficient per a processar prou sol·licituds en una instància per a maximitzar el rendiment del decode. En el cas de models grans, amb longituds de sortida típiques, per a mantenir una mida de lot gran per al decode, el prefill s'ha de dur a terme amb prou freqüència com per a degradar significativament la latència mitjana i causar interrupcions en la sortida.

Amplia

Amplia

Aquests problemes es poden solucionar utilitzant un conjunt separat de nodes per a dur a terme el prefill i el decode. En associar un node de prefill amb múltiples nodes de decode, es poden programar suficients sol·licituds per al prefill per a maximitzar el rendiment i mantenir un nombre prou gran de sol·licituds concurrents als nodes de decode per a maximitzar també el rendiment del decode. Els nodes de prefill omplen les memòries cau KV, les quals es transfereixen després als nodes de decode. Com que els decodificadors ja no han de parar per al prefill, les latències es tornen molt més deterministes, ja que l'impacte general de l'augment del kv_len de les sol·licituds actives és molt menys pronunciat. El cost es paga en un augment del temps fins al primer token (Time to First Token - TTFT), ja que la transferència de les memòries cau KV a través de la xarxa pot trigar des de desenes fins a centenars de mil·lisegons.

KV Messenger

A Perplexity, la nostra implementació per al prefill i el decode desagregats es basa entorn d'un missatger KV (KV messenger) que interactua amb el motor LLM per a orquestrar les transferències de la memòria cau KV des dels nodes de prefill fins als nodes de decode a través d'una xarxa. Pel costat del prefill, el missatger accepta sol·licituds dels nodes de decode, lliurant-les al planificador per lots i fent un seguiment de l'execució de la passada cap endavant per a enviar les memòries cau KV amb la menor latència possible. Pel costat del decode, després que s'assignin les pàgines no desallotjables (un-evictable pages), el missatger bloqueja la sol·licitud perquè no es programi per al decode fins que se l'informi de la finalització de la memòria cau KV i de les transferències de context del decodificador.

Amplia

La desagregació del prefill requereix connexions d'alt rendiment i baixa latència, per la qual cosa la nostra implementació està adaptada per a RDMA, donant suport tant als controladors d'interfície de xarxa (NIC) EFA com ConnectX. El missatger KV està construït sobre libfabric, utilitzant els nostres encaselladors (wrappers) fabric-lib per a proporcionar abstraccions de nivell superior i baixa latència sobre les primitives d'accés directe a memòria remota (RDMA), implementant transferències eficients de pàgines i metadades, juntament amb senyalització de baixa latència. En segon pla, fabric-lib coordina una GPU i les seves NIC connectades directament per a copiar dades des del node de prefill fins al node de decode.

En rebre-ho, el node de prefill assigna un conjunt corresponent de pàgines KV d'origen i programa la sol·licitud de prefill utilitzant el seu motor local. Per a minimitzar la latència, les transferències no esperen la passada cap endavant: en comptes d'això, les còpies de pàgines KV s'inicien tan bon punt el model acaba d'afegir les entrades de la memòria cau KV a la memòria cau KV per a les capes individuals. Com que les sol·licituds de prefill es poden fragmentar, el planificador per lots notifica al missatger KV els fragments programats actualment abans de l'execució. Per a admetre els gràfics CUDA alhora que es pot fer un seguiment de les capes, el missatger manté un fil dedicat que avalua (polling) un comptador incrementat després de la projecció de sortida de l'atenció. El comptador es manté només al node principal (lead node) en un entorn fragmentat (sharded): tot i que les entrades de la memòria cau KV són vàlides després de l'addició (append) i abans de l'atenció, la projecció de sortida es redueix entre rangs, sincronitzant-los implícitament. Un poc s'observa un canvi en el comptador, es notifica al missatger i aquest crida a fabric-lib per a iniciar la transferència d'una capa.

Amplia

Un poc completada la transferència de l'últim fragment, qualsevol metadada addicional també es copia: el descodificador especulatiu (speculative decoding) o MTP requereixen que els logits i els estats ocults es moguin al decodificador. Aquestes còpies també es realitzen a través de RDMA, cap a i des de memòries intermèdies (buffers) preassignades.

Un cop finalitzades totes les transferències pendents de l'últim fragment, el node de prefill desassigna les pàgines KV i completa la sol·licitud. El node de decode no es notifica explícitament: en canvi, utilitza comptadors immediats per a fer un seguiment del nombre d'operacions completades. El nombre d'operacions RDMA pel costat del prefill és proporcional al nombre de pàgines transferides. En completar-se el nombre conegut de còpies de pàgines i de context, fabric-lib crida al missatger KV per a indicar que una sol·licitud està llesta per a la descodificació. El missatger desassigna qualsevol context i lliura la sol·licitud al motor LLM.

Transferències de memòria cau KV fragmentades

Si el prefill i el decode depenen del paral·lelisme de tensors (Tensor Parallelism - TP) i fragmenten o repliquen les memòries cau KV de manera idèntica, un sol motor de transferència coordina múltiples dispositius per a enviar i rebre les pàgines de totes les rèpliques. Per a poder utilitzar un sol missatger i motor de transferència malgrat el fet que l'executant (executor) del model està replicat en múltiples dispositius i processos, s'utilitzen cuMem i cuMemImportFromShareableHandle per a assignar la memòria de dispositiu que dóna suport a les memòries cau KV i per a assignar-la (map) al procés principal. El motor de transferència inspecciona la topologia del node per a trobar les NIC i les CPU al node NUMA més proper per a utilitzar-les per a les transferències de cadascuna de les llesques (slices) de la memòria cau KV.

Si l'origen i la destinació es fragmenten de manera idèntica, les transferències són trivials ja que hi ha una assignació un a un (one-to-one) a partir dels dispositius i les pàgines de l'origen i la destinació. en aquesta situació, la fragmentació ajuda implícitament a les latències de transferència: utilitzant més GPUs, es poden emprar més NIC associades, apropant-se a la utilització de l'amplada de banda completa. Tanmateix, si hi ha una desajustada (mismatch), el motor de transferència ha de dividir o reconstruir les pàgines depenent de la ràtio entre les llesques d'origen i destinació.

Amplia

Si el prefill divideix la memòria cau KV entre més dispositius, les pàgines completes es reconstrueixen al decodificador enviant les meitats corresponents des dels dispositius de prefill. Si el decodificador té més fragments, rep pàgines de múltiples orígens. El decodificador necessita conèixer l'esquema de fragmentació del prefill per a poder calcular el nombre d'escriptures RDMA que s'espera que rebi. Si hi ha replicació, el prefill agrupa els dispositius en conjunts de rèpliques que repliquen la memòria cau KV completa dins d'ells mateixos. Als conjunts de rèpliques de destinació se'ls assigna aleatòriament un dels conjunts d'origen per a utilitzar tots els dispositius disponibles per a iniciar escriptures RDMA.

Amplia

Les transferències fragmentades requereixen un petit ajust de les memòries cau KV. Per defecte, FlashInfer es basa en el disseny (layout) NHD, que ordena els tokens dins d'una pàgina dins dels caps (heads). Atès que les memòries cau molt probablement estan fragmentades al llarg del nombre de caps d'atenció, això crea una discontinuïtat dins del cap. Les transferències RDMA no admeten implícitament escriptures amb salts (strided writes), requerint una operació per cap per a realitzar la transferència. En canvi, per a reduir el nombre d'interaccions amb libfabric, organitzem les memòries cau KV utilitzant el disseny HND que col·loca la dimensió del cap abans del nombre de tokens. Això garanteix la continuïtat, permetent que una pàgina sigui copiada amb una sola escriptura.

Descodificació especulativa

La descodificació especulativa requereix petits ajustos al prefill-decode desagregat. En la nostra implementació, no es permet als nodes de prefill mostrejar (sample) tokens. Com que els models Sonar de Perplexity admeten sortida estructurada, no volem incórrer en la complexitat de sincronitzar les implementacions del processador d'esquemes entre els nodes de prefill i els decodificadors. En els mecanismes MTP i de descodificació especulativa, preomplir (prefilling) el model esborrany (draft model) fins a l'últim token implica mostrejar tokens del model de destí.

Amplia

Per a solucionar aquests problemes, el prefill no inclou l'últim token de la seqüència d'entrada. En canvi, els estats ocults o logits del prefill que precedeixen l'últim token es transfereixen i es tracten com un token de decode en el pas següent al decodificador. Tot i que això augmenta lleugerament les latències, ja que s'ha de dur a terme un pas de decode complet després del prefill per a emetre el primer token, la complexitat de la implementació es redueix considerablement.

Despleiaments desagregats

Hem desplegat o experimentat amb múltiples configuracions desagregades amb diferents models, per a donar suport tant al trànsit de producció com a càrregues de treball d'avaluació pròpies. Segons la mida i el mecanisme d'atenció dels models, vam triar esquemes de fragmentació adequats per als nodes de prefill i decode per a utilitzar de la millor manera les GPUs.

DeepSeek-R1

Amb DeepSeek, vam considerar tant despleiaments de paral·lelisme de tensors (TP) com de paral·lelisme de dades (DP). Tal com es va tractar en publicacions anteriors del blog, els despleiaments TP proporcionen una millor latència a costa d'un menor rendiment, requerint més GPUs per a servir trànsit intens. Els despleiaments DP s'escalen molt millor amb la càrrega, tot i que el seu rendiment màxim és inferior a causa del cost de la comunicació entre dispositius o entre nodes.

DeepSeek es basa en l'atenció latent multicap (Multi-Head Latent Attention), comprimint les memòries cau KV. Com que tots els caps KV es comprimen en un sol vector latent, el TP no pot fragmentar les memòries cau KV, ja que ha de replicar en el seu lloc els vectors latents en tots els rangs. La fragmentació es produeix després de la descompressió, ja que cada rang pot extreure diferents caps de la mateixa representació latent. Per consegüent, tots els fragments de la memòria cau KV són idèntics tant a través dels fragments de prefill com de decode.

Amb una configuració TP intranode, tant els nodes de prefill com els decodificadors es fragmenten de manera idèntica. Les transferències s'envien des de tots els rangs per a utilitzar plenament totes les NIC disponibles. No obstant això, amb un desplegament DP, on la mida del rang TP és menor o a cada rang DP se li assigna una sola GPU, qualsevol dispositiu de prefill que contingui una còpia replicada de la memòria cau KV pot enviar-la. Per a equilibrar les sol·licituds entre totes les NIC disponibles, seleccionem aleatòriament una GPU i una NIC per a enviar la memòria cau KV des del node de prefill fins al decodificador.

Amb el prefill-decode barrejat, el nostre desplegament R1 estava lluitant per a superar de manera consistent els 50 TPS a causa de freqüents interrupcions de prefill de l'ordre de centenars de mil·lisegons. En canvi, en separar el prefill, vam incórrer en una penalització d'aproximadament 100 ms per al TTFT per a cada sol·licitud, però un sol node de prefill podia mantenir mides de lot consistents en 3 nodes de decode, oferint un rendiment superior a 90 TPS mentre gestionava una càrrega d'aproximadament 1 QPS per node de decode. Amb els despleiaments de paral·lelisme de dades, el TPS era lleugerament inferior, al voltant de 50, tot i que les instàncies podien gestionar una càrrega d'1 QPS per rang, amb 8 rangs per a un sol node.

Qwen3-Coder

Aquest model de 480B utilitza l'atenció de consultes agrupades (Grouped-Query Attention - GQA), de manera que l'atenció es pot fragmentar fàcilment i pot beneficiar-se del paral·lelisme de tensors sense sacrificar memòria per a les memòries cau KV. Per consegüent, vam poder fragmentar el model entre 8 GPUs tant per al prefill com per al decode, aparellant alrededor de 3 nodes de decode amb un sol node de prefill. Atès que l'atenció està fragmentada, ens basem en el disseny de memòria cau KV HND per a fragmentar les memòries cau KV de prefill i decode, aparellant els rangs de prefill amb els rangs de decode i utilitzant plenament totes les NIC per a transferir llesques en paral·lel.

Citació

@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}
}