Articolo
Bassa latenza e maggiore throughput con il deployment Multi-Node DeepSeek
Nella maggior parte dei sistemi, latenza e throughput sono spesso obiettivi in conflitto che richiedono compromessi durante la progettazione e il deployment. Ad esempio, nei modelli di linguaggio di grandi dimensioni densi, aumentare la dimensione del batch può migliorare il throughput ma aumenta anche la latenza; incrementare il parallelismo dei tensori all'interno di una singola macchina può ridurre la latenza ma diminuisce il numero di repliche, portando a un throughput inferiore.
I modelli Mixture of Experts (MoE) come DeepSeek-V3/R1 hanno recentemente dimostrato eccellenti capacità modellistiche ed efficienza operativa. Ad esempio, il modello DeepSeek-V3/R1 ha un totale di 671 miliardi di parametri, ma ogni token utilizza solo 37 miliardi di parametri durante l'inferenza. Questa architettura del modello presenta sia sfide che opportunità per i sistemi di inferenza.
Questo articolo dimostra che, contrariamente ai sistemi convenzionali, i modelli MoE come DeepSeek-V3/R1 possono raggiungere simultaneamente un throughput maggiore e una latenza inferiore quando si utilizzano più GPU in deployment multi-nodo nella maggior parte degli scenari.
Architetture di Deployment
Data l'enorme quantità di piccoli esperti che il modello possiede, i deployment devono essere distribuiti su diversi dispositivi. Abbiamo considerato sia i deployment su un singolo nodo con 8xH200 GPU che quelli multi-nodo su 8xH100 GPU.
Entrambe le architetture di deployment sfruttano il Data Parallelism, orchestrato tramite il nostro schedulatore di richieste interno. L'implementazione di Data Parallelism comporta il lancio di più istanze del motore di inferenza, ognuna operante indipendentemente per servire e mantenere le richieste. Lo schedulatore di richieste, che interagisce con il motore tramite GRPC, ha la responsabilità di distribuire le richieste in modo uniforme, facilitando anche il riutilizzo di KV, inviando le richieste con prefissi parzialmente coincidenti ai server contenenti la cache. Le istanze del motore non si estendono su più nodi. Possono, opzionalmente, usare il parallelismo dei tensori per suddividere l'attenzione su più dispositivi. Le istanze sono interconnesse tramite NVLink nel caso a singolo nodo o via InfiniBand per il caso multi-nodo, spedendo e raccogliendo gli esperti.
La configurazione di deployment a singolo nodo offre una latenza superiore con piccole dimensioni di batch; tuttavia, le prestazioni decadono rapidamente sotto condizioni di carico aumentato.
Per distribuire il motore di servizio, avviamo un pod per nodo ospitante più istanze del motore. PyTorch è responsabile di impostare la comunicazione distribuita e di negoziare l'inizializzazione NVSHMEM. Per la comunicazione, ci affidiamo a kernel CUDA personalizzati descritti in un precedente post sul blog. L'implementazione delle due tipologie di deployment è praticamente identica, con il modello che sceglie i kernel corretti da usare in base al fabric che implementa il parallelismo esperto.
Tecniche di Parallelizzazione
Prima di immergerci nei nostri confronti sulle prestazioni, è essenziale comprendere le principali strategie di parallelizzazione che rendono possibile il deployment di modelli MoE di grandi dimensioni come DeepSeek-V3/R1.
Parallelismo dei Tensori
Nell'inferenza LLM, il Parallelismo dei Tensori (TP) è tipicamente usato per ridurre l'uso della memoria e la computazione per GPU, riducendo così la latenza. Solitamente, possiamo suddividere le Proiezioni Lineari in Attention e Layers MLP lungo le dimensioni di riga o colonna e suddividere le operazioni di Attention lungo la dimensione della testa di attenzione.
Con TP, l'architettura Llama-3 non ha calcoli duplicati per le operazioni di Proiezione Lineare e Attention su GPU, il che è un metodo ideale di suddivisione. Tuttavia, nei modelli DeepSeek-V3/R1, TP non può raggiungere questo obiettivo.
I modelli DeepSeek-V3/R1 usano l'Attenzione Multi-Latente (MLA). Un MLA Layer usa prima una Proiezione Lineare kv_a_proj per calcolare il vettore latente, quindi usa un'altra Proiezione Lineare kv_b_proj per trasformarlo nello spazio di ciascuna testa di attenzione. Dato che tutte le teste di attenzione condividono lo stesso vettore latente, TP non può suddividere il vettore latente, quindi tutti i ranghi TP devono replicare i parametri e il calcolo di kv_a_proj e kv_b_proj. Analogamente, poiché MLA memorizza il vettore latente nella KV Cache, ogni rango TP archivia una copia identica della KV Cache.
Nonostante qualche duplicazione in MLA, il Parallelismo dei Tensori offre comunque una riduzione parziale nelle richieste di calcolo, rendendolo prezioso per scenari che richiedono alte velocità di output.
Parallelismo degli Esperti
I modelli DeepSeek-V3/R1 sostituiscono i Layers MLP con MoE Layers. Un MoE Layer ha 256 esperti instradati e un esperto condiviso. Ogni token viene spedito a 8 diversi esperti instradati per la computazione, e i risultati vengono sommati ponderatamente. Ogni token viene anche computato nell'esperto condiviso, e il risultato viene aggiunto al risultato dagli esperti instradati.
Il Parallelismo degli Esperti (EP) serve come approccio di suddivisione tipico per i MoE Layers, con ogni GPU che gestisce 256 / EP esperti instradati mantenendo una copia dell'esperto condiviso. Rispetto a TP, il vantaggio di EP è che può distribuire la computazione su più GPU, riducendo l'uso della computazione e della memoria per GPU.
Prima di eseguire la computazione degli esperti, tutte le GPU devono eseguire una comunicazione AllToAll per inviare token alle GPU dove si trovano gli esperti corrispondenti; dopo la computazione degli esperti, è necessaria un'altra comunicazione AllToAll per raccogliere i risultati della computazione dalle varie GPU e effettuare una somma ponderata. Abbiamo implementato una versione ottimizzata di questi due Kernel di comunicazione AllToAll, Dispatch e Combine, usando NVSHMEM. Nel nostro precedente post sul blog, abbiamo dettagliato l'implementazione, e i nostri kernel sono stati open-source su GitHub.
Data Parallelism
Con EP, possiamo distribuire la computazione MoE su 128 o più GPU. Tuttavia, la computazione MLA non può essere partizionata con EP. A questo punto, possiamo introdurre il Data Parallelism (DP). Ogni Gruppo DP ha una copia completa del MLA Layer. Ogni Gruppo DP accetta ingressi diversi e esegue il calcolo del MLA Layer indipendentemente.
Il DP del layer MLA e TP possono essere combinati, con un Gruppo DP suddiviso in più ranghi TP. Il layer MoE EP può essere combinato con DP/TP del layer MLA. EP = DP * TP. Ad esempio, su 16 macchine, EP128 DP32 TP4 significa distribuire esperti instradati su 128 GPU, con ogni 4 GPU che formano un Gruppo DP, per un totale di 32 gruppi DP indipendenti.
Singolo Nodo vs Multi-Nodo
I 671 miliardi di parametri di DeepSeek superano la capacità di memoria di una singola macchina H100 a 8 GPU (80 GB * 8), ma una singola macchina H200 a 8 GPU può ospitare completamente l'intero modello (141 GB * 8). Utilizzando la configurazione EP8 DP8 TP1, il modello utilizza circa 100 GB di memoria per GPU, lasciando circa 40 GB per KV Cache e altri risultati intermedi. Un token occupa 70,272 byte di KV Cache. Assumendo che ogni richiesta abbia 5,000 token, ogni GPU può accogliere approssimativamente 100 richieste.
Volevamo comprendere le differenze di prestazioni tra deployment a singolo nodo e multi-nodo sotto diverse configurazioni. Abbiamo usato una macchina H200 per il deployment a singolo nodo e fino a 16 macchine H100 per i deployment multi-nodo. Per ogni ambiente di deployment, abbiamo utilizzato combinazioni di TP 1, 2, 4, 8 e dimensioni di batch per GPU di 1, 2, 4, 8, 16, 32, 64, 128. Abbiamo presumito che ogni richiesta avesse una lunghezza della KV Cache di 5,000 token. Abbiamo anche presumito che Multi-Token Prediction (MTP) predicesse un token aggiuntivo (cioè, la lunghezza della query per ogni richiesta è 2), e abbiamo conservativamente supposto un tasso di accettazione del 60%. La figura sottostante mostra il throughput e la velocità di output per diverse configurazioni.
L'asse orizzontale rappresenta la velocità di output per richiesta in token/s. L'asse verticale utilizza una scala logaritmica per mostrare il throughput per macchina in token/s. Abbiamo segnato la Frontiera di Pareto per ogni configurazione EP con linee di colori diversi.
In scenari con requisiti di velocità di output estremamente alti, usare un singolo nodo EP8 DP1 TP8 con una dimensione di batch di 1 può raggiungere una velocità di output superiore a 100 token/s, ma il throughput è estremamente basso, equivalente alla velocità di output. In questo scenario, l'intero batch ha solo 2 token, che possono essere inviati a un massimo di 2*8=16 esperti, attivando un totale di al massimo 57 miliardi di parametri.
Nell'intervallo di velocità di output tra 80-40 token/s, con l'aumento del throughput, la velocità di output diminuisce significativamente. Al contrario, EP128 ha circa 5 volte il throughput più alto rispetto al deployment a singolo nodo alla stessa velocità di output.
Questo fenomeno può essere spiegato esaminando il comportamento dei deployment a singolo nodo: aumentare la dimensione del batch si correla direttamente con un aumento degli esperti attivati. Quando la dimensione del batch è 1, il numero medio di esperti attivati per GPU è 2 * 8 / 8 = 2. Quando il batch è abbastanza grande, tutti gli esperti vengono attivati, il che significa che ogni GPU attiva 256 / 8 = 32 esperti. Attivare più esperti significa che la GPU deve leggere più parametri dalla memoria, aumentando significativamente la pressione della larghezza di banda della memoria. Dato che la fase di decodifica dei modelli di linguaggio di grandi dimensioni è già limitata dalla larghezza di banda della memoria piuttosto che dalla prestazione di calcolo, aumentare la dimensione del batch nei deployment a singolo nodo riduce significativamente la velocità di output.
Il confronto delle quattro configurazioni di deployment multi-nodo (EP16, EP32, EP64, e EP128) rivela che valori EP più alti spostano la Frontiera di Pareto verso miglioramenti simultanei nel throughput e nella velocità di output.
Usando un numero EP più alto significa che a ogni GPU sono assegnati meno esperti. Ad esempio, EP128 significa che ogni GPU è responsabile di 256 / 128 = 2 esperti, quindi la pressione della larghezza di banda della memoria è significativamente ridotta. In altre parole, usando un numero EP maggiore, acquisite realmente più larghezza di banda di memoria. Quando la dimensione del batch per GPU è inferiore a 64, aumentare la dimensione del batch non influisce significativamente sulla velocità di calcolo degli esperti, perché aumentare il numero di ingressi non aumenta significativamente la pressione della larghezza di banda della memoria. Pertanto, osserviamo che con EP128, aumentare la dimensione del batch non influisce sulla velocità di output in modo significativo.
Interessante, con dimensioni di batch più grandi (64 richieste per GPU), abbiamo osservato un nuovo fenomeno: il throughput del deployment a singolo nodo è leggermente superiore al deployment multi-nodo. Parte del motivo è che NVLink intra-nodo ha una larghezza di banda più alta rispetto a InfiniBand inter-nodo. Un'altra parte è dovuta a limiti della nostra implementazione. Analizzeremo questo fenomeno più in dettaglio in seguito.
A causa delle limitazioni della capacità di memoria, la configurazione EP8 DP8 TP1 non può raggiungere una dimensione di batch di 128 per GPU, quindi il deployment multi-nodo è comunque una scelta migliore negli scenari che perseguono un throughput più alto.
Sovrapposizione di Computazione e Comunicazione
Come brevemente introdotto sopra riguardo al Parallelismo degli Esperti, le GPU sono inattive durante la comunicazione del MoE Layer. Per ridurre sprechi e abbassare la latenza, dobbiamo trovare compiti di computazione indipendenti dai dati per riempire questo tempo inattivo.
La parte superiore della figura sopra mostra il flusso di computazione di un layer. La computazione MoE dipende da Dispatch, e la computazione del layer successivo dipende dal risultato di Combine.
Posizioniamo l'esperto condiviso su ogni GPU. In questo modo, la computazione dell'esperto condiviso non richiede la comunicazione AllToAll. Pertanto, possiamo eseguire la computazione dell'esperto condiviso immediatamente dopo Dispatch Send, quindi attendere il completamento di Dispatch Recv. Chiamiamo questo schema di sovrapposizione "Dispatch Overlap".
Dispatch Overlap offre un'implementazione semplice e una vasta applicabilità. Questa tecnica nasconde il tempo di computazione dell'esperto condiviso su tutte le dimensioni EP e le dimensioni dei batch.
Per aumentare ulteriormente la sovrapposizione di computazione e comunicazione, abbiamo usato micro batching menzionato nel rapporto tecnico di DeepSeek per eliminare la dipendenza dei dati. Come mostrato nella parte inferiore della figura, abbiamo diviso la computazione di uno Strato Transformer in 5 stadi:
Stadio 1: InputNorm, QKVProj, AppendKV, BMM
Stadio 2: BMM, Attn, OProj, PostNorm, Gate
Stadio 3: Dispatch Send, Esperto Condiviso
Stadio 4: Dispatch Recv, MoE, Combine Send
Stadio 5: Combine Recv
Nei primi 3 Strati Transformer Densi, usiamo l'intero batch. Nei successivi 58 Strati Transformer MoE, dividiamo equamente il batch in due micro batch. I due micro batch vengono eseguiti alternati, con uno sfasamento di 3 stadi. Poiché non vi è dipendenza dai dati tra questi due micro batch, possiamo passare al calcolo dell'altro micro batch dopo Dispatch Send e dopo Combine Send.
Scomposizione della Latenza
Successivamente, confrontiamo gli effetti della sovrapposizione attraverso un esperimento, così come confrontiamo le differenze di prestazioni tra deployment a singolo nodo EP8 e deployment multi-nodo EP128. Per facilitare il confronto, abbiamo utilizzato GPU H100 per l'esperimento successivo. Abbiamo usato TP1, una dimensione di batch di 128 per GPU, una lunghezza di Query di 2 per richiesta e una lunghezza di KV Cache di 5000.
La figura sopra mostra il tempo totale impiegato su un MoE Transformer Layer e la proporzione di latenza di diversi tipi di kernel. Ad eccezione di Dispatch, Combine e GroupGEMM, il tempo di esecuzione di altri kernel dovrebbe essere uguale nelle serie EP8, EP128 NoOverlap, ed EP128 DispatchOverlap perché la dimensione del batch è la stessa.
Sovrapposizione
Confrontiamo prima gli effetti dei tre metodi di sovrapposizione. NoOverlap ha impiegato 2667µs in totale, DispatchOverlap ha impiegato 2651µs, risparmiando 16µs o solo 0.6%. MicroBatch ha mostrato un miglioramento molto significativo, impiegando 1896µs, un incremento del 29%. Sia Dispatch che il tempo di Combine sono stati significativamente ridotti. Dispatch è diminuito da 593µs a 367µs, e Combine da 1012µs a 237µs.
Si noti che per i kernel di computazione, suddividere un batch di dimensione 128 in due batch di dimensione 64 aumenta il tempo totale di esecuzione. Pertanto, sebbene sia stato ridotto il tempo speso per la comunicazione di 1001µs, il tempo totale è stato ridotto solo di 771µs. Spiegheremo il motivo usando il modello Roofline nella sezione seguente.
Per questo motivo, il micro batching non migliora sempre le prestazioni.
La figura sopra mostra il miglioramento delle prestazioni di Microbatch rispetto a DispatchOverlap per dimensioni di batch 4-128. Quando la dimensione del batch è inferiore a 32, Microbatch riduce le prestazioni dal 5% al 40%. Quando la dimensione del batch è maggiore o uguale a 32, Microbatch può migliorare le prestazioni dal 10% al 35%.
EP8 vs EP128
Torniamo alla figura precedente e confrontiamo EP8 e EP128 Microbatch. EP8 ha impiegato 1802µs in totale, leggermente meno di EP128's 1896µs. Oltre al tempo di esecuzione del kernel aumentato portato da Microbatch menzionato sopra, le principali differenze sono nel GroupGEMM utilizzato per il calcolo MoE e i due kernel di comunicazione, Dispatch e Combine.
Il GroupGEMM di EP8 ha impiegato 555µs, mentre il GroupGEMM di EP128 ha impiegato 270µs, riducendo di metà. Questo è il vantaggio principale del deployment multi-nodo.
Purtroppo, il tempo speso per la comunicazione è aumentato di 213µs, il che ha notevolmente compensato il vantaggio di GroupGEMM. In test di prestazioni separati dei nostri kernel di comunicazione, abbiamo scoperto che possono raggiungere solo metà della larghezza di banda di InfiniBand. Continueremo a ottimizzare i nostri kernel di comunicazione.
Un altro kernel che è significativamente in ritardo è GEMM. Microbatch ha aumentato GEMM di 95µs. Analizzeremo GEMM in maggiore profondità nella sezione del modello Roofline di seguito. Crediamo che l'implementazione corrente di GEMM non abbia ancora raggiunto prestazioni ottimali.
Roofline
Il modello Roofline è un buon strumento per analizzare le prestazioni dei kernel. L'asse orizzontale è Intensità Aritmetica, il rapporto tra FLOP e byte di I/O di memoria. Il valore dell'asse orizzontale può essere calcolato direttamente dalla semantica del kernel. L'asse verticale rappresenta le prestazioni raggiunte, calcolate dividendo FLOP per la latenza di benchmark.
Il limite teorico superiore delle prestazioni del kernel è determinato direttamente dalle specifiche della GPU. Il picco di prestazioni FP8 dell'H100 è di 1979 TFLOP/s, rappresentato come una linea orizzontale nel modello Roofline. La larghezza di banda della memoria dell'H100 è di 3.35 TB/s, rappresentata come la pendenza di una linea che passa attraverso l'origine. Le due linee danno i limiti di prestazione per i kernel legati alla computazione e alla memoria, rispettivamente.
Di seguito, discutiamo le prestazioni dei kernel GroupGEMM e GEMM.
GroupGEMM
Il kernel GroupGEMM in MoE esegue la seguente computazione: ci sono un totale di g gruppi, il gruppo i-esimo ha m_i token, eseguendo una moltiplicazione di matrici di [m_i, k] x [k, n] -> [m_i, n]. Nei test di prestazioni, assumiamo che il numero di token in ogni gruppo sia lo stesso, denotato come m_i = m. Poi il conteggio FLOP per GroupGEMM è 2 * g * m * k * n, e i byte di I/O di memoria sono g * (m * k + n * k + m * n).
Nel modello DeepSeek-V3/R1, ci sono 256 esperti, e ogni token è inviato a 8 esperti per la computazione. Supponendo una dimensione di batch di 128, lunghezza della query di 2, utilizzando la configurazione EP128 DP128, il numero medio di token ricevuti da ogni esperto (cioè, m) è 128 * 2 * 8 * 128 / 256 = 1024. Analogamente, possiamo calcolare m per altre configurazioni e dimensioni di batch.
Abbiamo utilizzato l'implementazione GroupGEMM di DeepGEMM per test delle prestazioni. I punti di test hanno coperto combinazioni di configurazioni EP8, EP16, EP32, EP64, EP128 con TP1 e dimensioni di batch 1-128.
La figura sopra mostra il modello Roofline per GroupGEMM sotto diverse configurazioni EP. Diversi EP corrispondono a un diverso numero di gruppi. La figura illustra linee di prestazioni quasi sovrapposte, indicando che le prestazioni di GroupGEMM sono principalmente determinate dal conteggio totale dei token (rappresentato come g * m).
Le stelle contrassegnano i punti dati corrispondenti a una dimensione di batch di 128 per GPU per ogni configurazione EP. Confrontando questi punti dati segnati, possiamo vedere che mentre EP aumenta (e DP aumenta sincronamente), anche il numero di token per esperto m aumenta. Ad EP8, m=128, mentre ad EP128, m=2048.
Con l'aumento di m, aumenta anche l'Intensità Aritmetica. Nella maggior parte delle configurazioni, GroupGEMM è limitato dalla larghezza di banda della memoria, quindi aumentare m migliora le prestazioni.
GEMM
Il kernel GEMM corrisponde alle Proiezioni Lineari nel modello, come Q/K/V/O Projection. Per una moltiplicazione di matrici di [m, k] x [k, n] -> [m, n], il conteggio FLOP è 2 * m * k * n, e i byte di I/O di memoria sono m * k + n * k + m * n. Possiamo anche testare la latenza per dimensioni di batch 1-128.
La figura sopra mostra il modello Roofline per GEMM sotto diverse configurazioni EP. Possiamo vedere che le prestazioni di GEMM sono limitate dalla larghezza di banda della memoria. Con l'aumento della dimensione del batch, aumenta anche l'Intensità Aritmetica, migliorando così le prestazioni.
Microbatch
Quando utilizziamo il micro batching, dividiamo il batch uniformemente in due parti. Dalle due figure sopra, possiamo vedere che quando m diventa m/2, l'efficienza della moltiplicazione di matrici diminuisce. Pertanto, l'esecuzione di due moltiplicazioni di matrici di dimensione m/2 richiede più tempo dell'esecuzione di una moltiplicazione di matrici di dimensione m.
Multi-Token Prediction
Nel corso di questo articolo, abbiamo presumuto l'uso della Multi-Token Prediction (MTP) per la decodifica speculativa. MTP modifica la lunghezza della query per richiesta da 1 a 2. Per la moltiplicazione di matrici, questo equivale a modificare m in m * 2, aumentando così l'efficienza della moltiplicazione di matrici. D'altra parte, se disegniamo il modello Roofline per MLA, scopriremmo che aumentare la lunghezza della query aumenta significativamente l'efficienza del kernel MLA.
Pertanto, l'uso di MTP svolge un ruolo importante nell'efficienza del modello.
Implementazione & Ottimizzazioni
In questa sezione, introdurremo alcuni dettagli relativi all'implementazione e all'ottimizzazione del nostro modello DeepSeek-V3/R1.
Quantizzazione
DeepSeek-V3/R1 è stato allenato nativamente su FP8 usando uno schema di quantizzazione per blocco, con pesi quantificati staticamente e attivazioni quantificati al volo. Invece di calcolare un fattore di scala per canale o per matrice statica, i fattori di scala vengono calcolati su vettori di 128 elementi per le attivazioni e su tabelle di 128x128 elementi per le matrici, limitando il degrado dell'accuratezza a causa della quantizzazione.
Alla Perplexity, ci affidiamo a un mix di kernel CUDA e Triton per supportare l'inferenza, con CUDA utilizzato per i kernel più sensibili alle prestazioni e modificati meno frequentemente (come attenzione e GEMM), con Triton che implementa una vasta gamma di kernel di attivazione, normalizzazione e utilità. Triton ci ha permesso di adattare rapidamente i kernel allo schema di quantizzazione a blocchi.
Per i layer lineari e MoE, mescoliamo i kernel Deep GEMM con i nostri kernel Triton GEMM, poiché abbiamo notato che per certe dimensioni di matrice e basse dimensioni di batch, Split-K offre una minore latenza. Se il layer non quantizzato esegue una moltiplicazione (M, K) x (K, N), necessita di (M x ceil_div(K, 128)) x (ceil_div(K, 128), ceil_div(N, 128)) fattori di scala per la quantizzazione a blocchi. Per la quantizzazione a blocchi, i fattori di scala delle attivazioni sono calcolati al volo, anziché essere pre-calibrati. Dato che i fattori di scala delle attivazioni sono aggregati solo lungo la dimensione K e non lungo la dimensione M, i kernel richiedono solo lievi alterazioni per supportare lo schema.
La funzione di attivazione SiLU utilizzata da DeepSeek-V3/R1 ha richiesto sostanziali cambiamenti per supportare i grafi CUDA, la quantizzazione a blocchi e il conteggio dinamico dei token instradati. La quantizzazione a blocchi può essere problematica in quanto introduce riduzioni orizzontali, tuttavia il kernel già suddivideva le attivazioni lungo la loro dimensione nascosta in blocchi di 1024 elementi. All'interno di un blocco, il tensore da quantizzare è stato ulteriormente suddiviso in blocchi di 128 per calcolare il valore assoluto più grande, con Triton che genera riduzioni massime efficienti cross-warp, aggiungendo un overhead minimo.
Per supportare l'instradamento MoE sotto i grafi CUDA, i kernel devono essere consapevoli delle informazioni di instradamento che indicano il numero di token per esperto, anziché programmare il lavoro in base alla dimensione dei buffer che sono stati allocati per tenere il limite superiore dei conteggi dei token. Non possiamo suddividere il problema sulla base delle dimensioni del tensore di input, quindi avviamo un numero fisso di kernel persistenti che leggono le informazioni di instradamento per determinare quanti token sono popolati e distribuiscono dinamicamente il lavoro di elaborazione delle attivazioni tra loro.
Abbiamo già upstreamato alcuni dei nostri kernel al progetto FlashInfer e in futuro renderemo open-source più del nostro codice.
MLA Layer
Usiamo FlashInfer per il calcolo MLA. FlashInfer supporta impostazioni di Page Table flessibili e prestazioni estremamente elevate.
Abbiamo unito q_a_proj e kv_a_proj in un unico qkv_a_proj. La latenza è diminuita da 15.4 µs + 14.8 µs = 30.2 µs a 16.7 µs.
Abbiamo scomposto kv_b_proj in due matrici, k_b_proj e v_b_proj. Abbiamo scritto un kernel BMM quantizzato a blocchi FP8 per calcoli relativi a queste due matrici.
Cuda Graph
Cuda Graph può ridurre significativamente il sovraccarico di lancio dei kernel, il che è cruciale per le prestazioni. Creiamo un Cuda Graph per ogni dimensione di batch.
Prima di sviluppare il nostro Kernel AllToAll, abbiamo usato torch.all_to_all_single() per la comunicazione AllToAll. Questa operazione richiede che tutte le GPU usino la stessa dimensione di batch. Tuttavia, diversi Gruppi DP possono eseguire diverse dimensioni di batch.
Per garantire che all_to_all_single() sia compatibile con diversi Gruppi DP che utilizzano diverse dimensioni di batch, abbiamo utilizzato un'operazione allreduce() prima di ogni esecuzione del modello per ottenere la dimensione massima del batch tra tutti i Gruppi DP. Poi abbiamo fatto in modo che tutti i Gruppi DP utilizzassero questa dimensione di batch per eseguire.
allreduce() aggiuntiva. Secondo, i Gruppi DP con dimensioni di batch più piccole sono costretti a riempire gli spazi. Terzo, rende il nostro codice di implementazione complesso.
Dopo aver implementato il nostro Kernel AllToAll, non richiediamo più a tutte le GPU di usare la stessa dimensione di batch. Pertanto, non abbiamo più bisogno di eseguire operazioni allreduce() aggiuntive o di riempire le dimensioni di batch.
MoE Router
Il router MoE è implementato in Triton, basato su un sort modificato derivato dalla libreria standard che tiene traccia anche degli indici degli elementi ordinati. L'implementazione è condivisa tra tutti i modelli MoE, poiché l'instradamento Mixtral è un caso speciale delle rotte DeepSeek in cui il gruppo Top-K è lo stesso del gruppo di tutti gli esperti. I kernel sparsi consumano gli indici Top-K e i punteggi direttamente, mentre gli schemi di dispatch/combine densi basati su all-to-all richiedono che le informazioni di instradamento siano aggregate per esperto anziché su base per-token.
Lavoro Futuro
Nel lavoro futuro, pianifichiamo di ottimizzare ulteriormente le prestazioni del modello DeepSeek.
L'ottimizzazione più importante successiva è la Disaggregazione Prefill. La fase Prefill e la fase Decode del modello DeepSeek-V3/R1 hanno caratteristiche computazionali molto diverse. Entrambe possono utilizzare diverse strategie di ottimizzazione e schemi di deployment.
Per il MLA Layer, nella fase Decode, usiamo l'Assorbimento Matrice per ridurre il conteggio FLOP del calcolo MLA. Nella fase Prefill, prima proiettare il vettore latente nello spazio K/V e poi computare nella forma Multi-Head Attention (MHA) darebbe risultati migliori.
Se Prefill e Decode sono eseguiti sulla stessa GPU, per ridurre l'impatto di Prefill sulla velocità di output Decode, tipicamente usiamo il prefill chunked per dividere la query in più chunk per Prefill. Poiché la KV Cache archivia il vettore latente, diventa difficile convertire MLA nella forma MHA.
Per il MoE Layer, nella fase Decode, usiamo EP e DP il più grandi possibile per aumentare il numero di token di input per esperto, migliorando così le prestazioni del GroupGEMM. Nella fase Prefill, poiché il numero di token è già abbastanza grande, GroupGEMM è già legato alla computazione. Pertanto, per Prefill, possiamo usare EP e DP più piccoli.
Se Prefill e Decode sono eseguiti sulla stessa GPU, finché un qualsiasi Gruppo DP sta eseguendo Prefill, la latenza dei Layers MoE su tutte le GPU aumenterà, influenzando significativamente la velocità di output Decode.
Oltre alla Disaggregazione Prefill, pianifichiamo anche di ottimizzare i seguenti aspetti:
Prestazioni AllToAll: Il nostro kernel AllToAll attualmente può raggiungere solo 1/3 della larghezza di banda InfiniBand. Continueremo a ottimizzare questo kernel.
Decodifica speculativa stile EAGLE: Nei dati sopra, abbiamo assunto l'uso di decodifica speculativa per predire 1 token. L'EAGLE può usare una struttura ad albero per predire più token, migliorando la lunghezza di accettazione, che può aumentare significativamente la velocità di output.
Kernel GEMM: Nel modello Roofline mostrato in precedenza, possiamo trovare che l'efficienza del kernel GEMM è ancora lontana dal limite teorico. Continueremo a ottimizzare questo kernel.
GB200 NVL72: Nella soluzione GB200 NVL72 di NVIDIA più recente, 72 GPU Blackwell sono interconnesse tramite NVLink ad alta velocità. Per i modelli di architettura MoE, questa è una grande opportunità e sfida.
Conclusione
Il deployment multi-nodo dei modelli MoE DeepSeek raggiunge ciò che è tipicamente impossibile con LLM densi: miglioramento simultaneo sia del throughput che della latenza. Distribuendo gli esperti su più GPU, riduciamo la pressione della larghezza di banda della memoria per dispositivo, consentendo un'elaborazione più rapida e un throughput di sistema superiore. I nostri esperimenti mostrano configurazioni EP128 che raggiungono fino a 5x di throughput superiore a velocità di output equivalenti rispetto ai deployment a singolo nodo.
Le tecniche di sovrapposizione di computazione-comunicazione come il micro-batching riducono significativamente l'overhead di comunicazione multi-nodo, con la nostra implementazione che mostra fino al 40% di velocità migliorata. I nostri kernel di comunicazione AllToAll personalizzati e le implementazioni di kernel ottimizzate hanno permesso un deployment efficiente del modello da 671 miliardi di parametri.
Poiché le architetture MoE guadagnano popolarità per la loro capacità, queste strategie di deployment forniscono preziose intuizioni per scalare efficacemente tali modelli.