Préremplissage et Décodage Désagrégés

Pour générer des tokens de sortie à partir d'une invite d'entrée, l'inférence LLM est divisée en deux étapes : pré-remplissage et décodage. Le pré-remplissage fonctionne sur les tokens d'entrée, remplissant les caches KV, avant de passer à l'étape de décodage qui génère les tokens un par un.

Alors qu'une seule étape de décodage prend généralement des dizaines de millisecondes, le pré-remplissage prend beaucoup plus de temps. Si l'on utilise les mêmes appareils, mélanger le pré-remplissage avec le décodage dégrade les performances de décodage. Dans cet article, nous explorons une solution établie sous la forme de pré-remplissage et de décodage désagrégés, les exécutant sur des appareils séparés pour maximiser à la fois le débit de pré-remplissage et les latences de décodage.

Performance du Pré-remplissage vs Décodage

Dans un moteur de service LLM typique, le planificateur de lots sélectionne les requêtes à traiter à 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 ensemble. Le coût de l'attention, qui s'accumule le long de la longueur de séquence, augmente 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 transfèrent généralement un seul token (qo_len=1), à un coût minimal à travers d'autres couches qui opèrent indépendamment sur les tokens d'une séquence. Les requêtes de pré-remplissage transfèrent des milliers ou des dizaines de milliers de tokens à un coût significatif à travers des couches denses (grand qo_len).

La latence d'un passage vers l'avant est davantage influencée par le nombre de tokens indépendants passant par des couches denses (qo_len) que par le nombre de tokens récupérés du cache KV pendant l'attention (kv_len). L'attention peut se paralléliser à la fois selon le nombre de requêtes et le kv_len proportionnel aux longueurs de séquence, atteignant une bonne utilisation. Le pré-remplissage est lié au calcul : avec un qo_len élevé, les noyaux GEMM peuvent allouer suffisamment de blocs le long de la dimension M pour utiliser pleinement les capacités de calcul des GPU modernes. Le décodage est lié à la mémoire : en raison de la taille de lot généralement faible, le nombre d'entrées le long de M est généralement petit, suffisant pour un seul bloc. Bien que les noyaux Split-K GEMM puissent améliorer l'utilisation du SM pour les tailles de lot de tokens faibles, 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 subissent des latences plus élevées à travers le passage vers l'avant, affectant négativement le débit de décodage de l'ensemble de l'instance. Bien que le mélange des requêtes de pré-remplissage avec des requêtes de décodage ou l'utilisation d'un pré-remplissage en 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 un nombre suffisant 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é suffisamment souvent pour qu'il 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 séparé 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, suffisamment de requêtes peuvent être planifiées pour le pré-remplissage pour maximiser le débit et maintenir un nombre suffisant de requêtes simultanées sur les nœuds de décodage pour également maximiser le débit de décodage. Les nœuds de pré-remplissage remplissent les caches KV, qui sont ensuite transférés vers les nœuds de décodage. Puisque les décodeurs n'ont plus à interrompre 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 est payé par une augmentation du temps de premier token (TTFT), car le transfert des caches KV sur le réseau peut prendre des dizaines à des centaines de millisecondes.

Messager KV

Chez Perplexity, notre implémentation du pré-remplissage et du décodage désagrégés repose sur un messager KV qui interagit avec le moteur LLM pour orchestrer les transferts du 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 transmet au planificateur de lots et suit l'exécution du passage vers l'avant pour envoyer les caches KV avec le moins de latence possible. Du côté du décodeur, après que les pages non-éjectables ont été allouées, le messager bloque la requête pour qu'elle ne soit pas planifiée pour le décodage jusqu'à ce qu'il soit informé de l'achèvement des transferts du cache KV et du contexte du décodeur.

Le pré-remplissage désagrégé nécessite des connexions à haut débit et à faible latence, et notre implémentation est donc conçue pour le RDMA, supportant à la fois les contrôleurs d'interface de réseau (NIC) EFA et ConnectX. Le messager KV est construit sur libfabric, utilisant nos wrappers fabric-lib pour fournir des abstractions de haut niveau à faible latence sur les primitives d'accès direct à la mémoire à distance (RDMA), implémentant des transferts de pages et de métadonnées efficaces, ainsi qu'une signalisation à faible latence. En arrière-plan, fabric-lib coordonne un GPU et ses NICs 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 planifie la requête de pré-remplissage à l'aide de son moteur local. Pour minimiser la latence, les transferts n'attendent pas la fin du passage vers l'avant : à la place, les copies de pages KV sont initiées dès que le modèle termine d'ajouter les entrées du cache KV au cache KV pour des couches individuelles. Étant donné que les requêtes de pré-remplissage peuvent être divisées, le planificateur de lots informe le messager KV des morceaux actuellement planifiés avant l'exécution. Pour prendre en charge les graphiques CUDA tout en pouvant suivre les couches, le messager maintient un fil dédié qui sonde 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. Une fois que le compteur change, le messager est notifié et appelle fabric-lib pour initier le transfert d'une couche.

Après le transfert du dernier morceau, des métadonnées supplémentaires sont également copiées : le décodage spéculatif ou le MTP nécessitent que les logits et les états cachés soient déplacés vers le décodeur. Ces copies sont également effectuées via le RDMA, vers et depuis des tampons pré-alloués.

Une fois tous les transferts en attente du dernier morceau terminés, le nœud de pré-remplissage désalloue les pages KV et complète la requête. Le nœud décodeur n'est pas explicitement informé : à 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é pré-remplissage est proportionnel au nombre de pages transférées. Une fois que le nombre connu de copies de pages et de contextes est terminé, fabric-lib appelle le messager KV pour indiquer qu'une requête est prête à être décodée. Le messager désalloue tout contexte et transmet la requête au moteur LLM.

Transferts de Cache KV Fragmentés

Si le pré-remplissage et le décodeur reposent sur le parallélisme tensoriel (TP) et que les caches KV sont fragmentés ou dupliqués de manière identique, un seul moteur de transfert coordonne plusieurs appareils pour envoyer et recevoir les pages de toutes les répliques. Afin de pouvoir utiliser un messager et un moteur de transfert uniques 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 alimentant les caches KV et pour la mapper dans le processus principal. Le moteur de transfert examine la topologie du nœud pour trouver les NICs et les CPU dans le nœud NUMA le plus proche à utiliser pour les transferts de chaque tranche 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 un-à-un entre les appareils et les pages de la source et de la destination. Dans cette situation, le fractionnement aide implicitement à réduire les latences de transfert : en utilisant plus de GPU, plus de NICs associés peuvent être employés, approchant ainsi de la pleine utilisation de la bande passante. Cependant, s'il existe un désaccord, le moteur de transfert doit diviser ou reconstruire les pages en fonction du ratio entre les tranches source et destination.

Si le pré-remplissage divise le cache KV entre plus d'appareils, les pages complètes sont reconstruites sur le décodeur en envoyant les moitiés correspondantes à partir des appareils de pré-remplissage. Si le décodeur a plus de fragments, il reçoit des pages de plusieurs sources. Le décodeur doit être informé du schéma de fragmentation du pré-remplisseur afin de pouvoir calculer le nombre d'écritures RDMA qu'il est censé recevoir. Si la réplication est impliquée, le pré-remplisseur regroupe les appareils dans des ensembles répliques qui répliquent le cache KV complet en leur sein. Les ensembles de réplication de destination se voient attribuer de manière aléatoire un des ensembles sources pour utiliser tous les appareils disponibles pour initier les écritures RDMA.

Les transferts fragmentés nécessitent un léger ajustement aux caches KV. Par défaut, FlashInfer s'appuie sur la disposition NHD, qui ordonne les tokens à l'intérieur d'une page dans les têtes. Étant donné que les caches sont très probablement fragmentés en fonction du nombre de têtes d'attention, cela crée une discontinuité dans la tête. Les transferts RDMA ne prennent pas implicitement en charge les écritures étagées, nécessitant une opération par tête pour effectuer le transfert. À la place, pour réduire le nombre d'interactions avec libfabric, nous organisons les caches KV en utilisant la disposition HND qui place la dimension de tête avant le nombre de tokens. Cela assure la continuité, permettant de copier une page avec une seule écriture.

Décodage Spéculatif

Le décodage spéculatif nécessite de légères modifications 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 tokens. Étant donné que les modèles Sonar de Perplexity prennent en charge la sortie structurée, nous ne voulons pas ajouter la complexité de synchroniser les implémentations du processeur de schéma entre le pré-remplissage et le décodage. Dans les mécanismes MTP et de décodage spéculatif, le pré-remplissage du modèle brouillon jusqu'au dernier token implique l'échantillonnage de tokens du modèle cible.

Pour contourner ces problèmes, le pré-remplissage n'inclut pas le dernier token de la séquence d'entrée. À la place, les états cachés ou les logits du pré-remplissage précédant le dernier token sont transférés et traités comme un token de décodage à l'étape suivante sur le décodeur. Bien que cela augmente légèrement les latences, car une étape complète de décodage doit être effectuée après le pré-remplissage pour émettre le premier token, la complexité de l'implémentation est considérablement réduite.

Déploiements Désagrégés

Nous avons déployé ou expérimenté de multiples configurations désagrégées avec différents modèles, pour supporter soit le trafic en production soit les charges de travail d'évaluation en interne. 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 de mieux utiliser les GPU.

DeepSeek-R1

Avec DeepSeek, nous avons envisagé à la fois des déploiements en parallélisme tensoriel (TP) et en parallélisme de données (DP). Comme discuté dans les articles de blog précédents, les déploiements TP offrent une meilleure latence au coût d'un débit plus bas, nécessitant plus de GPU pour desservir un trafic important. Les déploiements DP s'élaborent beaucoup mieux avec la charge, mais leur débit maximal est plus faible en raison du coût de la communication inter-appareils ou inter-nœuds.

DeepSeek repose sur l'attention latente multi-tête, compressant les caches KV. Étant donné que toutes les têtes KV sont compressées en un seul vecteur latent, TP ne peut pas fragmenter les caches KV, car il doit à la place répliquer les vecteurs latents sur tous les rangs. La fragmentation intervient 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 de cache KV sont identiques à la fois pour les fragments de pré-remplissage et de décodage.

Avec une configuration TP intra-nœud, les pré-remplisseurs et les décodeurs sont fragmentés de manière identique. Les transferts sont envoyés depuis tous les rangs pour utiliser pleinement tous les NICs disponibles. Cependant, avec un déploiement DP, où la taille de rang TP est inférieure ou chaque rang DP est assigné à un seul GPU, tout appareil pré-remplisseur qui détient une copie répliquée du cache KV peut l'envoyer. Pour équilibrer les requêtes à travers tous les NICs disponibles, nous sélectionnons de manière aléatoire un GPU et un NIC pour envoyer le cache KV du pré-remplisseur au décodeur.

Avec le pré-remplissage-décodage mixte, notre déploiement R1 avait du mal à dépasser régulièrement 50 TPS en raison d'interruptions fréquentes du pré-remplissage de l'ordre de centaines de millisecondes. En revanche, en séparant le pré-remplissage, nous avons engendré une pénalité d'environ 100 ms pour le TTFT de chaque requête, mais un seul nœud de pré-remplissage pouvait maintenir des tailles de lot constantes sur 3 nœuds de décodeurs, délivrant un débit supérieur à 90 TPS tout en gérant une charge d'environ 1 QPS par nœud de décodeur. Avec des déploiements en parallèle de données, TPS était légèrement inférieur à environ 50, mais les instances pouvaient gérer une charge de 1 QPS par rang, avec 8 rangs pour un seul nœud.

Qwen3-Coder

Ce modèle de 480B utilise une attention à requêtes groupées (GQA), donc l'attention peut être facilement fragmentée et bénéficie du parallélisme tensoriel sans sacrifier la mémoire pour les caches KV. Par conséquent, nous avons pu fragmenter le modèle sur 8 GPU pour le pré-remplissage et le décodage, en associant environ 3 nœuds de décodeurs à un seul nœud de pré-remplissage. Étant donné que l'attention est fragmentée, nous nous appuyons sur la disposition des caches KV HND pour fragmenter les caches KV de pré-remplissage et de décodeur, associant les rangs de pré-remplissage aux rangs de décodeur et utilisant pleinement tous les NICs pour transférer les fragments en parallèle.