Prefill e decode disaggregati

Per generare token di output da un prompt di input, l'inferenza LLM è suddivisa in due fasi: prefill e decode. Il prefill viene eseguito sui token di input, popolando le cache KV, prima di entrare nella fase di decode che genera i token uno alla volta. Mentre un singolo passaggio di decode richiede solitamente decine di millisecondi, il prefill richiede molto più tempo. Se eseguiti sugli stessi dispositivi, combinare prefill e decode degrada le prestazioni di decode. In questo articolo esploriamo una soluzione consolidata basata sul disaccoppiamento di prefill e decode, eseguendoli su dispositivi separati per massimizzare sia il throughput del prefill che le latenze di decode.

Prestazioni di Prefill vs Decode

In un tipico motore di servizio LLM, lo scheduler dei batch seleziona le richieste da elaborare in ogni passaggio di esecuzione di un modello. Quando si esegue su un singolo dispositivo o nodo, sia le richieste di prefill che quelle di decode vengono raggruppate in batch. Il costo dell'attention, che si aggrega lungo la lunghezza della sequenza, aumenta sia per il prefill che per il decode, proporzionalmente alla lunghezza delle voci nella cache KV (kv_len). Le richieste di decode solitamente inoltrano un singolo token (qo_len=1), con un costo minimo attraverso altri livelli che operano indipendentemente sui token di una sequenza. Le richieste di prefill inoltrano migliaia o decine di migliaia di token a un costo significativo attraverso strati densi (qo_len elevato).

La latenza di un passaggio forward è influenzata più fortemente dal numero di token indipendenti passati attraverso strati densi (qo_len) che dal numero di token recuperati dalla cache KV durante l'attention (kv_len). L'attention può parallelizzare sia tra il numero di richieste che su kv_len in proporzione alla lunghezza delle sequenze, ottenendo una buona utilizzazione. Il prefill è vincolato al calcolo (compute bound): con un qo_len elevato, i kernel GEMM possono allocare blocchi sufficienti lungo la dimensione M per utilizzare appieno le capacità di calcolo delle moderne GPU. Il decode è vincolato alla memoria (memory bound): a causa di dimensioni dei batch solitamente ridotte, il numero di input lungo M è solitamente piccolo, sufficiente per un solo blocco. Sebbene i kernel Split-K GEMM possano migliorare l'utilizzazione SM per batch di token ridotti, le cache e le unità di moltiplicazione matriciale rimangono solitamente sottoutilizzate.

Espandi

Quando mescolati, i batch contenenti richieste di prefill subiscono latenze più elevate durante il passaggio forward, influenzando negativamente il throughput di decode dell'intera istanza. Sebbene combinare richieste di prefill con richieste di decode o impiegare il prefill chunked possa migliorare leggermente le prestazioni di decode, è difficile mantenere un throughput di prefill sufficiente per elaborare abbastanza richieste su un'istanza al fine di massimizzare il throughput di decode. Nel caso di modelli di grandi dimensioni, con lunghezze di output tipiche, per mantenere un batch size elevato per il decode, il prefill deve essere eseguito abbastanza spesso da degradare significativamente la latenza media e causare interruzioni nell'output.

Espandi

Espandi

Questi problemi possono essere risolti utilizzando un insieme separato di nodi per eseguire prefill e decode. Associando un nodo di prefill a più nodi di decode, è possibile pianificare un numero sufficiente di richieste per il prefill per massimizzare il throughput e mantenere un numero abbastanza elevato di richieste simultanee sui nodi di decode per massimizzare anche il throughput di decode. I nodi di prefill popolano le cache KV, che vengono poi trasferite ai nodi di decode. Poiché i decoder non devono più interrompersi per il prefill, le latenze diventano molto più deterministiche, poiché l'impatto complessivo della crescita di kv_len delle richieste attive è molto meno pronunciato. Il costo si paga con un aumento del Tempo al Primo Token (TTFT), poiché il trasferimento delle cache KV sulla rete può richiedere da decine a centinaia di millisecondi.

KV Messenger

In Perplexity, la nostra implementazione per il prefill e decode disaggregati è costruita attorno a un KV messenger che interagisce con il motore LLM per orchestrare i trasferimenti della cache KV dai nodi di prefill ai nodi di decode attraverso una rete. Sul lato prefill, il messenger accetta richieste dai nodi di decode, passandole allo scheduler dei batch e tenendo traccia dell'esecuzione del passaggio forward per inviare le cache KV con la minor latenza possibile. Sul lato decoder, dopo che le pagine non espellibili sono state allocate, il messenger blocca la richiesta affinché non venga programmata per il decode finché non viene notificata la conclusione dei trasferimenti della cache KV e del contesto del decoder.

Espandi

La disaggregazione del prefill richiede connessioni ad alto throughput e bassa latenza, pertanto la nostra implementazione è adattata per RDMA, supportando sia EFA che i controller di interfaccia di rete (NIC) ConnectX. Il KV Messenger è basato su libfabric, utilizzando i nostri wrapper fabric-lib per fornire astrazioni di alto livello a bassa latenza sulle primitive di Remote Direct Memory Access (RDMA), implementando trasferimenti efficienti di pagine e metadati, insieme a segnalazioni a bassa latenza. In background, fabric-lib coordina una GPU e le sue NIC direttamente collegate per copiare i dati dal nodo di prefill al nodo di decode.

Al ricevimento, il nodo di prefill alloca un insieme corrispondente di pagine KV di origine e programma la richiesta di prefill utilizzando il suo motore locale. Per minimizzare la latenza, i trasferimenti non attendono il passaggio forward: invece, le copie delle pagine KV vengono avviate non appena il modello termina di aggiungere voci alla cache KV per i singoli strati. Poiché le richieste di prefill possono essere suddivise in chunk, lo scheduler dei batch notifica al KV messenger i chunk attualmente programmati prima dell'esecuzione. Per supportare i grafi CUDA pur essendo in grado di tracciare gli strati, il messenger mantiene un thread dedicato che interroga un contatore incrementato dopo la proiezione di output dell'attention. Il contatore viene mantenuto solo sul nodo principale in un ambiente sharded: anche se le voci della cache KV sono valide dopo l'aggiunta e prima dell'attention, la proiezione di output viene ridotta tra i ranghi, sincronizzandoli implicitamente. Una volta osservato un cambiamento nel contatore, il messenger viene notificato e chiama fabric-lib per avviare il trasferimento di uno strato.

Espandi

Dopo il completamento del trasferimento dell'ultimo chunk, vengono copiati anche eventuali metadati aggiuntivi: la decodifica speculativa o MTP richiedono che i logit e gli stati nascosti vengano spostati nel decoder. Queste copie vengono eseguite anch'esse tramite RDMA, verso e da buffer pre-allocati.

Al completamento di tutti i trasferimenti in sospeso dell'ultimo chunk, il nodo di prefill de-alloca le pagine KV e completa la richiesta. Il nodo di decode non viene notificato esplicitamente: utilizza invece contatori immediati per tenere traccia del numero di operazioni completate. Il numero di operazioni RDMA sul lato prefill è proporzionale al numero di pagine trasferite. Al completamento del numero noto di copie di pagine e contesti, fabric-lib chiama il KV messenger per indicare che una richiesta è pronta per il decoding. Il messenger de-alloca qualsiasi contesto e passa la richiesta al motore LLM.

Trasferimenti di Cache KV Sharded

Se il prefill e il decoder si basano sul parallelismo tensoriale (TP) e shardano o replicano le cache KV in modo identico, un singolo motore di trasferimento coordina più dispositivi per inviare e ricevere le pagine di tutte le repliche. Per poter utilizzare un unico messenger e motore di trasferimento nonostante l'executor del modello sia replicato su più dispositivi e processi, vengono utilizzati cuMem e cuMemImportFromShareableHandle per allocare la memoria del dispositivo che supporta le cache KV e per mapparla nel processo principale. Il motore di trasferimento esamina la topologia del nodo per trovare le NIC e le CPU nel nodo NUMA più vicino da utilizzare per i trasferimenti di ciascuna delle slice della cache KV.

Se l'origine e la destinazione shardano in modo identico, i trasferimenti sono banali poiché esiste una mappatura uno-a-uno dai dispositivi e dalle pagine dell'origine alla destinazione. In questa situazione, lo sharding aiuta implicitamente le latenze di trasferimento: utilizzando più GPU, è possibile impiegare più NIC associate, avvicinandosi alla piena utilizzazione della larghezza di banda. Tuttavia, se vi è una discrepanza, il motore di trasferimento deve suddividere o ricostruire le pagine a seconda del rapporto tra le slice di origine e di destinazione.

Espandi

Se il prefill suddivide la cache KV su più dispositivi, le pagine complete vengono ricostruite sul decoder inviando le metà corrispondenti dai dispositivi di prefill. Se il decoder ha più shard, riceve pagine da più fonti. Il decoder deve conoscere lo schema di sharding del prefill per poter calcolare il numero di scritture RDMA che ci si aspetta riceva. Se è coinvolta la replica, il prefill raggruppa i dispositivi in set di repliche che replicano l'intera cache KV al loro interno. Ai set di repliche di destinazione viene assegnato casualmente uno dei set di origine per utilizzare tutti i dispositivi disponibili per avviare le scritture RDMA.

Espandi

I trasferimenti sharded richiedono una leggera modifica alle cache KV. Per impostazione predefinita, FlashInfer si basa sul layout NHD, che ordina i token all'interno di una pagina all'interno delle teste. Poiché le cache sono molto probabilmente shardate lungo il numero di teste di attenzione, ciò crea discontinuità all'interno della testa. I trasferimenti RDMA non supportano implicitamente le scritture strided, richiedendo un'operazione per testa per eseguire il trasferimento. Invece, per ridurre il numero di interazioni con libfabric, organizziamo le cache KV utilizzando il layout HND che posiziona la dimensione della testa prima del numero di token. Ciò garantisce la continuità, consentendo di copiare una pagina con una singola scrittura.

Decodifica Speculativa

La decodifica speculativa richiede leggere modifiche al prefill-decode disaggregato. Nella nostra implementazione, ai nodi di prefill non è consentito campionare i token. Poiché i modelli Sonar di Perplexity supportano l'output strutturato, non vogliamo incorrere nella complessità di sincronizzare le implementazioni del processore di schema tra prefiller e decoder. Nei meccanismi MTP e di decodifica speculativa, il prefill del modello di bozza fino all'ultimo token comporta il campionamento di token dal modello di destinazione.

Espandi

Per aggirare questi problemi, il prefill non include l'ultimo token della sequenza di input. Invece, gli stati nascosti o i logit derivanti dal prefill precedente all'ultimo token vengono trasferiti e trattati come un token di decode nel passaggio successivo sul decoder. Sebbene ciò aumenti leggermente le latenze, poiché un passaggio di decode completo deve essere eseguito dopo il prefill per emettere il primo token, la complessità dell'implementazione è notevolmente ridotta.

Distribuzioni Disaggregate

Abbiamo implementato o sperimentato molteplici configurazioni disaggregate con modelli diversi, per supportare sia il traffico di produzione che i carichi di lavoro di valutazione interna. In base alla dimensione e al meccanismo di attenzione dei modelli, abbiamo scelto schemi di sharding idonei per i nodi di prefill e decoder per utilizzare al meglio le GPU.

DeepSeek-R1

Con DeepSeek, abbiamo considerato distribuzioni sia Tensor-Parallel (TP) che Data-Parallel (DP). Come discusso in post precedenti, le distribuzioni TP offrono una latenza migliore a costo di un throughput inferiore, richiedendo più GPU per servire traffico intenso. Le distribuzioni DP scalano molto meglio con il carico, tuttavia il loro throughput di picco è inferiore a causa del costo della comunicazione inter-dispositivo o inter-nodo.

DeepSeek si basa sulla Multi-Head Latent Attention, comprimendo le cache KV. Poiché tutte le teste KV sono compresse in un singolo vettore latente, il TP non può shardare le cache KV, poiché deve invece replicare i vettori latenti su tutti i ranghi. Lo sharding avviene dopo la decompressione, poiché ogni rango può estrarre teste diverse dalla stessa rappresentazione latente. Di conseguenza, tutti gli shard della cache KV sono identici tra gli shard di prefill e decoder.

Con una configurazione TP intranodo, sia i prefiller che i decoder sono shardati in modo identico. I trasferimenti vengono inviati da tutti i ranghi per utilizzare appieno tutte le NIC disponibili. Tuttavia, con una distribuzione DP, dove la dimensione del rango TP è inferiore o ogni rango DP è assegnato a una singola GPU, qualsiasi dispositivo di prefill che detiene una copia replicata della cache KV può inviarla. Per bilanciare le richieste su tutte le NIC disponibili, selezioniamo casualmente una GPU e una NIC per inviare la cache KV dal prefill al decoder.

Con il prefill-decode misto, la nostra distribuzione R1 faticava a superare costantemente i 50 TPS a causa di frequenti interruzioni di prefill nell'ordine delle centinaia di millisecondi. Al contrario, separando il prefill, abbiamo sostenuto una penalità di circa 100ms sul TTFT per ogni richiesta, ma un singolo nodo di prefill poteva mantenere batch size costanti su 3 nodi decoder, offrendo un throughput superiore a 90 TPS gestendo un carico di circa 1 QPS per nodo decoder. Con le distribuzioni data-parallel, il TPS era leggermente inferiore, intorno a 50, tuttavia le istanze potevano gestire un carico di 1 QPS per rango, con 8 ranghi per singolo nodo.

Qwen3-Coder

Questo modello da 480B utilizza la Grouped-Query Attention (GQA), quindi l'attention può essere facilmente shardata e può beneficiare del parallelismo tensoriale senza sacrificare la memoria per le cache KV. Di conseguenza, abbiamo potuto shardare il modello su 8 GPU sia per il prefill che per il decode, accoppiando circa 3 nodi decoder con un singolo nodo di prefill. Poiché l'attention è shardata, ci affidiamo al layout della cache KV HND per shardare le cache KV di prefill e decoder, accoppiando i ranghi di prefill con i ranghi dei decoder e utilizzando appieno tutte le NIC per trasferire le slice in parallelo.

Citazione

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