Préremplissage et décodage désagrégés
La séparation du préremplissage et du décodage sur plusieurs appareils accélère l'inférence LLM et améliore les performances
Afin de générer des jetons de sortie à partir d'une invite d'entrée, l'inférence par LLM est divisée en deux étapes : le préremplissage (prefill) et le décodage (decode). Le préremplissage s'exécute sur les jetons d'entrée, remplissant les caches KV, avant d'entrer dans l'étape de décodage qui génère les jetons un par un. Alors qu'une seule étape de décodage s'exécute généralement en dizaines de millisecondes, le préremplissage prend beaucoup plus de temps. S'ils sont exécutés sur les mêmes appareils, le mélange du préremplissage et du décodage dégrade les performances de décodage. Dans cet article, nous explorons une solution établie sous la forme d'un préremplissage et d'un décodage désagrégés, en les exécutant sur des appareils séparés pour maximiser à la fois le débit de préremplissage et les temps de latence du décodage.
Performances : préremplissage vs décodage
Dans un moteur de service LLM typique, l'ordonnanceur de lots sélectionne les requêtes à traiter lors de chaque étape d'exécution d'un modèle. Lors de l'exécution sur un seul appareil ou nœud, les requêtes de préremplissage et de décodage sont regroupées par lots. Le coût de l'attention, qui s'agrège le long de la longueur de la séquence, augmente à la fois pour le préremplissage et le décodage, proportionnellement à la longueur des entrées dans le cache KV (kv_len). Les requêtes de décodage transmettent généralement un seul jeton (qo_len=1), à un coût minimal à travers les autres couches qui fonctionnent indépendamment sur les jetons d'une séquence. Les requêtes de préremplissage transmettent des milliers ou des dizaines de milliers de jetons à un coût significatif à travers les couches denses (grand qo_len).
La latence d’une passe avant est davantage influencée par le nombre de jetons indépendants transmis par les couches denses (qo_len) que par le nombre de jetons récupérés du cache KV pendant l’attention (kv_len). L’attention peut paralléliser à la fois le nombre de requêtes et le kv_len proportionnel aux longueurs de séquence, ce qui permet une bonne utilisation. Le préremplissage est lié au calcul : qo_len étant élevé, les noyaux GEMM peuvent allouer suffisamment de blocs le long de la dimension M pour exploiter pleinement les capacités de calcul des GPU modernes. Le décodage est limité par la mémoire : en raison de tailles de lot généralement faibles, le nombre d’entrées le long de M est habituellement petit, suffisant pour un seul bloc. Bien que les noyaux GEMM Split-K puissent améliorer l’utilisation des SM pour de faibles tailles de lot de jetons, les caches et les unités de multiplication matricielle restent généralement sous-utilisés.
Lorsqu'ils sont mélangés, les lots contenant des requêtes de préremplissage entraînent des latences plus élevées lors de la passe avant, affectant négativement le débit de décodage de l'ensemble de l'instance. Bien que le mélange de requêtes de préremplissage avec des requêtes de décodage ou l'utilisation d'un préremplissage par morceaux puisse légèrement améliorer les performances de décodage, il est difficile de maintenir un débit de préremplissage suffisant pour traiter assez de requêtes sur une instance afin de maximiser le débit de décodage. Dans le cas de grands modèles, avec des longueurs de sortie typiques, pour maintenir une grande taille de lot pour le décodage, le préremplissage doit être effectué assez souvent pour que cela dégrade considérablement la latence moyenne et provoque des saccades dans la sortie.
Ces problèmes peuvent être résolus en utilisant un ensemble distinct de nœuds pour effectuer le préremplissage et le décodage. En associant un nœud de préremplissage à plusieurs nœuds de décodage, il est possible d'ordonnancer suffisamment de requêtes pour le préremplissage afin de maximiser le débit et de maintenir un nombre suffisant de requêtes concurrentes sur les nœuds de décodage pour maximiser également le débit de décodage. Les nœuds de préremplissage peuplent les caches KV, qui sont ensuite transférés vers les nœuds de décodage. Étant donné que les décodeurs n'ont plus à s'interrompre pour le préremplissage, les latences deviennent beaucoup plus déterministes, car l'impact global de l'augmentation du kv_len des requêtes actives est beaucoup moins prononcé. Le coût se traduit par une augmentation du temps jusqu'au premier jeton (TTFT), car le transfert des caches KV sur le réseau peut prendre de dizaines à centaines de millisecondes.
Messager KV
Chez Perplexity, notre implémentation pour le préremplissage et le décodage désagrégés est construite autour d'un messager KV qui interagit avec le moteur LLM pour orchestrer les transferts de cache KV des nœuds de préremplissage vers les nœuds de décodage via un réseau. Du côté du préremplissage, le messager accepte les requêtes des nœuds de décodage, les transmettant à l'ordonnanceur de lots et suivant l'exécution de la passe avant pour distribuer les caches KV avec le moins de latence possible. Du côté du décodage, après l'allocation des pages non expulsables, le messager bloque la requête pour qu'elle ne soit pas ordonnancée pour le décodage jusqu'à ce qu'il soit notifié de la fin des transferts du cache KV et du contexte du décodeur.
La désagrégation du préremplissage nécessite des connexions à haut débit et à faible latence, c'est pourquoi notre implémentation est conçue pour l'architecture RDMA, prenant en charge à la fois les contrôleurs d'interface réseau (NIC) EFA et ConnectX. Le messager KV est construit sur libfabric, en utilisant nos wrappers fabric-lib pour fournir des abstractions de plus haut niveau et à faible latence sur les primitives d'accès direct à la mémoire à distance (RDMA), mettant en œuvre des transferts efficaces de pages et de métadonnées ainsi qu'une signalisation à faible latence. En arrière-plan, fabric-lib coordonne un GPU et ses NIC directement connectés pour copier les données du nœud de préremplissage vers le nœud de décodage.
À la réception, le nœud de préremplissage alloue un ensemble correspondant de pages KV source et ordonnance la requête de préremplissage à l'aide de son moteur local. Pour minimiser la latence, les transferts n'attendent pas la passe avant : au lieu de cela, les copies de pages KV sont initiées dès que le modèle finit d'ajouter les entrées du cache KV au cache KV pour les couches individuelles. Étant donné que les requêtes de préremplissage peuvent être découpées en morceaux, l'ordonnanceur de lots informe le messager KV des morceaux actuellement ordonnancés avant l'exécution. Pour prendre en charge les graphes CUDA tout en étant capable de suivre les couches, le messager maintient un thread dédié qui interroge un compteur incrémenté après la projection de sortie de l'attention. Le compteur est maintenu uniquement sur le nœud principal dans un environnement fragmenté : même si les entrées du cache KV sont valides après l'ajout et avant l'attention, la projection de sortie est réduite à travers les rangs, les synchronisant implicitement. Dès qu'un changement dans le compteur est observé, le messager est notifié et appelle fabric-lib pour initier le transfert d'une couche.
Une fois le transfert du dernier morceau terminé, toute métadonnée supplémentaire est également copiée : le décodage spéculatif ou le mécanisme MTP nécessitent le déplacement des logits et des états cachés vers le décodeur. Ces copies sont également effectuées via RDMA, vers et depuis des tampons pré-alloués.
À la fin de tous les transferts en attente du dernier morceau, le nœud de préremplissage désalloue les pages KV et termine la requête. Le nœud de décodage n'est pas explicitement notifié : à la place, il utilise des compteurs immédiats pour suivre le nombre d'opérations terminées. Le nombre d'opérations RDMA du côté du préremplissage est proportionnel au nombre de pages transférées. À l'achèvement du nombre connu de copies de pages et de contexte, fabric-lib appelle le messager KV pour indiquer qu'une requête est prête pour le décodage. Le messager désalloue tout contexte et remet la requête au moteur LLM.
Transferts de cache KV fragmentés
Si le prérempleur et le décodeur s'appuient sur le parallélisme de tenseurs (TP) et fragmentent ou répliquent les caches KV de manière identique, un moteur de transfert unique coordonne plusieurs appareils pour envoyer et recevoir les pages de toutes les réplicas. Afin de pouvoir utiliser un seul messager et moteur de transfert malgré le fait que l'exécuteur du modèle soit répliqué sur plusieurs appareils et processus, cuMem et cuMemImportFromShareableHandle sont utilisés pour allouer la mémoire de l'appareil prenant en charge les caches KV et pour la mapper dans le processus principal. Le moteur de transfert inspecte la topologie du nœud pour trouver les NIC et les processeurs (CPU) dans le nœud NUMA le plus proche à utiliser pour les transferts de chacune des tranches du cache KV.
Si la source et la destination sont fragmentées de manière identique, les transferts sont triviaux car il existe une correspondance biunivoque entre les appareils et les pages de la source et de la destination. Dans cette situation, la fragmentation aide implicitement les latences de transfert : en utilisant plus de GPU, il est possible d'employer davantage de NIC associés, se rapprochant d'une utilisation complète de la bande passante. Cependant, en cas de non-correspondance, le moteur de transfert doit diviser ou reconstruire les pages en fonction du rapport entre les tranches source et destination.
Si le prérempleur divise le cache KV sur un plus grand nombre d'appareils, les pages complètes sont reconstruites sur le décodeur en envoyant les moitiés correspondantes depuis les appareils du prérempleur. Si le décodeur possède plus de fragments, il reçoit des pages de plusieurs sources. Le décodeur doit connaître le schéma de fragmentation du prérempleur afin de pouvoir calculer le nombre d'écritures RDMA qu'il est censé recevoir. En cas de réplication, le prérempleur groupe les appareils en ensembles de réplicas qui répliquent l'ensemble du cache KV en leur sein. Les ensembles de réplicas de destination se voient assigner de manière aléatoire l'un des ensembles source afin d'utiliser tous les appareils disponibles pour initier des écritures RDMA.
Les transferts fragmentés nécessitent un léger ajustement des caches KV. Par défaut, FlashInfer s'appuie sur la disposition NHD, qui ordonne les jetons au sein d'une page à l'intérieur des têtes. Étant donné que les caches sont très probablement fragmentés le long du nombre de têtes d'attention, cela crée une discontinuité au sein de la tête. Les transferts RDMA ne prennent pas implicitement en charge les écritures échelonnées (strided), nécessitant une opération par tête pour effectuer le transfert. À la place, afin de réduire le nombre d'interactions avec libfabric, nous organisons les caches KV en utilisant la disposition HND qui place la dimension de la tête avant le nombre de jetons. Cela garantit la continuité, permettant de copier une page avec une seule écriture.
Décodage spéculatif
Le décodage spéculatif nécessite de légers ajustements au préremplissage-décodage désagrégé. Dans notre implémentation, les nœuds de préremplissage ne sont pas autorisés à échantillonner des jetons. Étant donné que les modèles Sonar de Perplexity prennent en charge la sortie structurée, nous ne souhaitons pas encourir la complexité de la synchronisation des implémentations de processeur de schéma entre les prérempleurs et les décodeurs. Dans les mécanismes MTP et de décodage spéculatif, le préremplissage du modèle brouillon jusqu'au dernier jeton implique l'échantillonnage de jetons à partir du modèle cible.
Pour contourner ces problèmes, le préremplissage n'inclut pas le dernier jeton de la séquence d'entrée. À la place, les états cachés ou les logits du préremplissage précédant le dernier jeton sont transférés et ce dernier est traité comme un jeton de décodage lors de l'étape suivante sur le décodeur. Bien que cela augmente légèrement les latences, dans la mesure où une étape de décodage complète doit être exécutée après le préremplissage pour émettre le premier jeton, la complexité de l'implémentation est grandement réduite.
Déploiements désagrégés
Nous avons déployé ou expérimenté plusieurs configurations désagrégées avec différents modèles, afin de prendre en charge soit le trafic de production, soit les charges de travail d'évaluation internes. En fonction de la taille et du mécanisme d'attention des modèles, nous avons choisi des schémas de fragmentation adaptés pour les nœuds de préremplissage et de décodage afin d'utiliser au mieux les GPU.
DeepSeek-R1
Avec DeepSeek, nous avons envisagé des déploiements à la fois en parallélisme de tenseurs (TP) et en parallélisme de données (DP). Comme indiqué dans de précédents articles de blog, les déploiements TP offrent une meilleure latence au détriment d'un débit plus faible, nécessitant davantage de GPU pour servir un trafic lourd. Les déploiements DP s'échelonnent beaucoup mieux avec la charge, mais leur débit de pointe est inférieur en raison du coût de la communication inter-appareil ou inter-nœud.
DeepSeek repose sur l'attention latente multi-tête (MLA), compressant les caches KV. Étant donné que toutes les têtes KV sont compressées en un seul vecteur latent, le TP ne peut pas fragmenter les caches KV, car il doit à la place répliquer les vecteurs latents sur tous les rangs. La fragmentation se produit après la décompression, car chaque rang peut extraire différentes têtes de la même représentation latente. Par conséquent, tous les fragments du cache KV sont identiques à la fois sur les fragments de préremplissage et de décodage.
Avec une configuration TP intranœud, les prérempleurs et les décodeurs sont fragmentés de manière identique. Les transferts sont envoyés depuis tous les rangs afin d'utiliser pleinement tous les NIC disponibles. Cependant, avec un déploiement DP, où la taille du rang TP est inférieure ou où chaque rang DP est assigné à un seul GPU, n'importe quel appareil de préremplissage qui détient une copie répliquée du cache KV peut l'expédier. Pour équilibrer les requêtes sur tous les NIC disponibles, nous sélectionnons au hasard un GPU et un NIC pour envoyer le cache KV du prérempleur au décodeur.
Avec un préremplissage-décodage mixte, notre déploiement R1 avait du mal à dépasser de manière cohérente 50 TPS en raison d'interruptions fréquentes de préremplissage de l'ordre de centaines de millisecondes. En revanche, en séparant le préremplissage, nous avons subi une pénalité d'environ 100 ms sur le TTFT pour chaque requête, mais un seul nœud de préremplissage pouvait maintenir des tailles de lots cohérentes sur 3 nœuds de décodage, offrant un débit supérieur à 90 TPS tout en gérant une charge d'environ 1 QPS par nœud de décodage. Avec les déploiements en parallélisme de données, le TPS était légèrement inférieur, autour de 50, mais les instances pouvaient gérer une charge de 1 QPS par rang, avec 8 rangs par nœud.
Qwen3-Coder
Ce modèle de 480B utilise l'attention à requête groupée (GQA), de sorte que l'attention peut être facilement fragmentée et peut bénéficier du parallélisme de tenseurs sans sacrifier de mémoire pour les caches KV. Par conséquent, nous avons pu fragmenter le modèle sur 8 GPU à la fois pour le préremplissage et le décodage, en associant environ 3 nœuds de décodage à un seul nœud de préremplissage. L'attention étant fragmentée, nous nous appuyons sur la disposition de cache KV HND pour fragmenter les caches KV du prérempleur et du décodeur, en associant les rangs du prérempleur aux rangs du décodeur et en utilisant pleinement tous les NIC pour transférer les tranches en parallèle.
Citation
@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}
}