Precompilazione e Decodifica Disaggregata

Per generare token di output da un prompt di input, l'inferenza LLM è suddivisa in due fasi: prefill e decode. Il prefill utilizza i token di input per popolare le cache KV, prima di entrare nella fase di decode che genera i token uno alla volta.

Mentre un singolo passo di decode di solito richiede decine di millisecondi, il prefill richiede molto più tempo. Se eseguiti sugli stessi dispositivi, mescolare prefill con decode degrada le prestazioni di decode. In questo articolo esploriamo una soluzione consolidata sotto forma di prefill e decode disaggregati, eseguendoli su dispositivi separati per massimizzare sia il throughput di prefill che le latenze di decode.

Prestazioni di Prefill vs Decode

In un motore di servizio LLM tipico, il batch scheduler seleziona le richieste da elaborare in ogni passo di esecuzione di un modello. Quando viene eseguito su un singolo dispositivo o nodo, sia le richieste di prefill che di decode vengono batchate insieme. Il costo dell'attenzione, che si aggrega lungo la lunghezza della sequenza, cresce sia per prefill che decode, proporzionalmente alla lunghezza delle voci nella cache KV (kv_len). Le richieste di decode di solito inoltrano un singolo token (qo_len=1), a 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 (grande qo_len).

La latenza di un passaggio in avanti è più fortemente influenzata dal numero di token indipendenti che passano attraverso gli strati densi (qo_len) rispetto al numero di token recuperati dalla cache KV durante l'attenzione (kv_len). L'attenzione può parallelizzare sia il numero di richieste che il kv_len proporzionale alle lunghezze delle sequenze, ottenendo una buona utilizzazione. Il prefill è limitato dal calcolo: essendo il qo_len elevato, i kernel GEMM possono allocare blocchi sufficienti lungo la dimensione M per utilizzare appieno le capacità computazionali delle GPU moderne. Il decode è limitato dalla memoria: a causa delle dimensioni tipicamente ridotte dei batch, il numero di input lungo M è di solito piccolo, sufficiente per un solo blocco. Mentre i kernel GEMM Split-K possono migliorare l'utilizzazione degli SM per dimensioni di batch di token ridotte, le cache e le unità di moltiplicazione delle matrici rimangono tipicamente sotto-utilizzate.

Quando mescolati insieme, i batch contenenti richieste per prefill subiscono maggiori latenze attraverso il passaggio in avanti, influendo negativamente sul throughput di decode dell'intera istanza. Mentre mescolare richieste di prefill con richieste di decode o impiegare un prefill chunked può migliorare leggermente le prestazioni di decode, è difficile mantenere un throughput di prefill sufficiente per processare abbastanza richieste su un'istanza per massimizzare il throughput di decode. Nel caso di modelli grandi, con lunghezze di output tipiche, per mantenere una grande dimensione di batch per decode, il prefill deve essere eseguito abbastanza spesso da compromettere significativamente la latenza media e causare balbuzie nell'output.

Questi problemi possono essere affrontati utilizzando un set separato di nodi per eseguire prefill e decode. Associando un nodo prefiller con più nodi decoder, possono essere pianificate richieste sufficienti per il prefill per massimizzare il throughput e mantenere un numero abbastanza alto di richieste concorrenti sui nodi decoder per massimizzare anche il throughput di decode. I nodi prefiller popolano le cache KV, che vengono poi trasferite ai nodi decoder. Poiché i decoder non devono più interrompersi per il prefill, le latenze diventano molto più deterministiche, poiché l'impatto complessivo dell'aumento del kv_len delle richieste attive è molto meno pronunciato. Il costo è pagato in un aumento del Tempo al Primo Token (TTFT), poiché il trasferimento delle cache KV attraverso la rete può richiedere decine a centinaia di millisecondi.

Messenger KV

Alla Perplexity, la nostra implementazione per prefill e decode disaggregati è costruita attorno a un messenger KV che interagisce con il motore LLM per orchestrare i trasferimenti di cache KV dai nodi prefiller ai nodi decoder attraverso una rete. Sul lato prefiller, il messenger accetta richieste dai nodi decoder, consegnandole al batch scheduler e tenendo traccia dell'esecuzione del passaggio in avanti per inviare le cache KV con il minor ritardo possibile. Sul lato decoder, dopo che le pagine non evacuabili sono allocate, il messenger blocca la richiesta dall'essere pianificata per il decode fino a che non viene notificato del completamento del trasferimento della cache KV e del contesto del decoder.

Disaggregare il prefill richiede connessioni ad alto throughput e bassa latenza, quindi la nostra implementazione è adattata per RDMA, supportando sia EFA che Controller di Interfaccia di Rete ConnectX (NIC). Il Messenger KV è costruito su libfabric, utilizzando i nostri wrapper fabric-lib per fornire astrazioni di alto livello a bassa latenza sopra le primitive di Accesso Diretto alla Memoria Remoto (RDMA), implementando trasferimenti efficienti di pagine e metadati, insieme a segnalazioni a bassa latenza. In background, fabric-lib coordina una GPU e i suoi NICs direttamente connessi per copiare i dati dal nodo prefiller al nodo decoder.

Alla ricezione, il nodo prefiller alloca un set corrispondente di pagine KV sorgente e pianifica la richiesta per il prefill utilizzando il suo motore locale. Per minimizzare la latenza, i trasferimenti non aspettano il passaggio in avanti: invece, le copie delle pagine KV vengono avviate non appena il modello termina di aggiungere voci di cache KV alla cache KV per strati individuali. Poiché le richieste di prefill possono essere suddivise, il batch scheduler avvisa il messenger KV dei chunk attualmente pianificati prima dell'esecuzione. Per supportare i grafi CUDA pur essendo in grado di tracciare gli strati, il messenger mantiene un thread dedicato che poll il contatore incrementato dopo la proiezione di output dell'attenzione. Il contatore è mantenuto solo sul nodo guida in un ambiente shardato: anche se le voci di cache KV sono valide dopo l'append e prima dell'attenzione, la proiezione di output è ridotta attraverso i ranghi, sincronizzandoli implicitamente. Una volta osservato un cambiamento nel contatore, il messenger viene avvisato e chiama fabric-lib per avviare il trasferimento di uno strato.

Dopo che il trasferimento dell'ultimo chunk è completato, eventuali metadati aggiuntivi vengono anche copiati: la decodifica speculativa o MTP richiede che i logit e gli stati nascosti siano trasferiti al decoder. Queste copie sono eseguite anche tramite RDMA, da e verso buffer pre-allocati.

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

Trasferimenti di Cache KV Shardati

Se il prefiller e il decoder si affidano al Parallelismo di Tensore (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. Al fine di poter utilizzare un singolo messenger e motore di trasferimento nonostante il fatto che l'esecutore del modello sia replicato su più dispositivi e processi, cuMem e cuMemImportFromShareableHandle vengono utilizzati per allocare la memoria del dispositivo che supporta le cache KV e per mappare nel processo principale. Il motore di trasferimento ispeziona la topologia del nodo per trovare i NIC e le CPU nel nodo NUMA più vicino da utilizzare per i trasferimenti di ciascuna delle sezioni della cache KV.

Se la shard sorgente e destinazione sono identiche, i trasferimenti sono banali poiché c'è una mappatura uno a uno dai dispositivi e dalle pagine della sorgente alla destinazione. In questa situazione, la shardatura aiuta implicitamente le latenze del trasferimento: utilizzando più GPU, più NIC associate possono essere impiegate, raggiungendo una più vicina piena utilizzazione della larghezza di banda. Tuttavia, se c'è una discrepanza, il motore di trasferimento deve dividere o ricostruire le pagine a seconda del rapporto tra sezioni sorgenti e di destinazione.

Se il prefiller divide la cache KV su più dispositivi, le pagine complete vengono ricostruite sul decoder inviando le metà corrispondenti dai dispositivi prefiller. Se il decoder ha più shard, riceve pagine da più sorgenti. Il decoder deve conoscere lo schema di shard del prefiller per essere in grado di calcolare il numero di scritture RDMA che ci si aspetta di ricevere. Se è coinvolta la replica, il prefiller raggruppa i dispositivi in ​​set replica che replicano la cache KV completa all'interno di essi. I set di replica di destinazione vengono assegnati casualmente a uno dei set sorgente per utilizzare tutti i dispositivi disponibili per avviare scritture RDMA.

I trasferimenti shardati richiedono una leggera modifica alle cache KV. Di default, FlashInfer si affida al 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 a strisce, 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 pone la dimensione della testa prima del numero di token. Ciò garantisce continuità, permettendo di copiare una pagina con una sola scrittura.

Decodifica Speculativa

La decodifica speculativa richiede lievi modifiche al prefill-decode disaggregato. Nella nostra implementazione, ai nodi prefiller non è permesso campionare token. Poiché i modelli Sonar di Perplexity supportano output strutturati, 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 dei token dal modello target.

Per aggirare questi problemi, il prefill non include l'ultimo token della sequenza di input. Invece, gli stati nascosti o i logit dal prefill che precedono l'ultimo token vengono trasferiti e vengono trattati come un token di decoding nel passo successivo sul decoder. Mentre ciò aumenta leggermente le latenze, poiché deve essere eseguito un intero passo di decoding dopo il prefill per emettere il primo token, la complessità dell'implementazione è notevolmente ridotta.

Distribuzioni Disaggregate

Abbiamo distribuito o sperimentato con più configurazioni disaggregate con diversi modelli, per supportare sia il traffico di produzione che i carichi di lavoro di valutazione interna. Basandoci sulla dimensione e sul meccanismo di attenzione dei modelli, abbiamo scelto schemi di shard adatti per i nodi prefiller e decoder per utilizzare al meglio le GPU.

DeepSeek-R1

Con DeepSeek, abbiamo considerato sia deployment Tensor-Parallel (TP) che Data-Parallel (DP). Come discusso nei post precedenti del blog, i deployment TP offrono migliori latenze e al costo di un throughput inferiore, richiedendo più GPU per servire traffico pesante. I deployment DP scalano molto meglio con il carico, tuttavia il loro picco di throughput è inferiore a causa del costo della comunicazione inter-dispositivo o inter-nodo.

DeepSeek si basa sull'Attenzione Latente Multi-Testa, comprimendo le cache KV. Poiché tutte le teste KV sono compresse in un unico vettore latente, TP non può shardare le cache KV, poiché invece deve replicare i vettori latenti su tutti i ranghi. La shardatura avviene dopo la decompressione, poiché ogni rango può estrarre teste diverse dalla stessa rappresentazione latente. Di conseguenza, tutte le shard di cache KV sono identiche sia nei shard prefiller che decoder.

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

Con prefill-decode misto, il nostro deployment R1 ha lottato per superare costantemente 50 TPS a causa di frequenti interruzioni di prefill dell'ordine di centinaia di millisecondi. Al contrario, separando il prefill, abbiamo subito una penalità di circa 100 ms a TTFT per ogni richiesta, ma un singolo nodo prefiller poteva mantenere dimensioni batch costanti su 3 nodi decoder, fornendo un throughput superiore a 90 TPS gestendo un carico di circa 1 QPS per nodo decoder. Con deployment data-parallel, il TPS era leggermente inferiore a circa 50, tuttavia le istanze potevano gestire un carico di 1 QPS per rango, con 8 ranghi in un singolo nodo.

Qwen3-Coder

Questo modello da 480B utilizza l'Attenzione a Richieste Raggruppate (GQA), quindi l'attenzione può essere facilmente shardata e può beneficiare del parallelismo tensoriale senza sacrificare memoria per le cache KV. Di conseguenza, abbiamo potuto shardare il modello su 8 GPU sia per prefill che decode, accoppiando circa 3 nodi decoder con un singolo nodo prefiller. Poiché l'attenzione è shardata, ci affidiamo al layout di cache KV HND per shardare le cache KV prefiller e decoder, accoppiando i ranghi prefiller con i ranghi decoder e utilizzando appieno tutti i NIC per trasferire le sezioni in parallelo.