Le transfert KV en mode push de vLLM réduit la latence dans l'inférence désagrégée

Amazon Web Services a introduit un mécanisme de transfert clé-valeur (KV) en mode push pour le moteur d'inférence open source vLLM. Cette mise à jour cible les architectures d'inférence désagrégées préfill-decode, visant à réduire significativement le temps jusqu'au premier jeton (TTFT) et à améliorer le débit lors de la génération de jetons. La fonctionnalité a été validée sur des instances Amazon EC2 P5en utilisant le modèle Qwen3-32B.

Ce qui s'est passé

L'inférence des grands modèles de langage se compose de deux phases distinctes : le préfill, qui traite la prompte d'entrée pour construire le cache KV, et le decode, qui génère les jetons de sortie en utilisant ce cache. Dans le cadre du service désagrégé, ces phases s'exécutent sur des instances de moteur séparées afin de permettre une mise à l'échelle et une optimisation indépendantes. Cependant, cette séparation nécessite le transfert du cache KV de l'instance de préfill vers l'instance de decode avant que la génération ne puisse commencer, rendant l'efficacité de ce transfert critique pour la latence globale.

Auparavant, vLLM utilisait une méthode en mode pull via le connecteur NVIDIA Inference Xfer Library (NIXL). Dans ce mode par défaut, l'instance de decode attend la fin du préfill, reçoit les métadonnées d'un routeur, alloue la mémoire, puis tire les données. Ces étapes se déroulent séquentiellement après la fin du préfill, ajoutant un délai important au chemin critique. Le nouveau connecteur en mode push modifie ce flux en chevauchant les tâches de coordination avec le calcul de préfill lui-même.

Les ingénieurs d'AWS ont développé et testé cette fonctionnalité sur des clusters connectés via AWS Elastic Fabric Adapter (EFA), qui prend en charge l'accès direct à la mémoire à distance (RDMA). Bien qu'elle ait été initialement conçue pour les instances AWS Trainium, la contribution est indépendante de l'accélérateur et fonctionne également avec les GPU. Les benchmarks montrent que sous une charge élevée, le mode push peut réduire le TTFT jusqu'à trois fois par rapport au mode pull, tout en améliorant le temps par jeton de sortie (TPOT) d'environ 30 pour cent.

Comment cela fonctionne

L'architecture en mode push élimine les goulots d'étranglement de sérialisation en permettant à l'instance de decode de bloquer ses blocs de mémoire cibles à l'avance. Au lieu d'attendre la fin du préfill avant de préparer la réception des données, le côté decode envoie une notification d'enregistrement à l'instance de préfill dès qu'il alloue des blocs locaux. Cela signifie que la destination est connue avant même que les données sources ne soient prêtes.

Du côté du préfill, un thread d'arrière-plan dédié gère la préparation du transfert. Une fois le calcul de préfill terminé et l'enregistrement du décodeur arrivé, la boucle principale du moteur confie les blocs KV à ce thread d'arrière-plan. Le thread établit la connexion NIXL, mappe les identifiants de blocs et lance une opération NIXL WRITE pour copier les données directement dans la mémoire enregistrée du décodeur. Cela maintient le travail lourd de préparation des descripteurs et de soumission du transfert hors des boucles principales critiques en termes de latence des deux instances.

Le système utilise un protocole de rendez-vous pour gérer les différences de timing. Si l'enregistrement arrive avant la fin du préfill, le producteur le conserve. Si le préfill se termine en premier, il retient les blocs jusqu'à l'arrivée de l'enregistrement. L'événement qui se produit en second déclenche le transfert immédiat. Le décodeur compte les notifications de complétion provenant de chaque rang de producteur et ne procède à la génération de jetons qu'une fois toutes les données attendues arrivées.

Détails clés

• Gain de performance : Réduction du TTFT jusqu'à 3x et amélioration du TPOT de 30 % pour Qwen3-32B sur des déploiements 1P1D à taux de requêtes élevés.
• Matériel testé : Instances Amazon EC2 P5en avec EFA, utilisant GPUDirect RDMA pour le transfert de mémoire direct GPU-to-GPU.
• Mécanisme : Utilise les opérations NIXL WRITE initiées par l'instance de préfill, plutôt que les NIXL READ initiées par l'instance de decode.
• Multithreading : Un thread d'arrière-plan sur le worker de préfill gère la configuration du transfert, empêchant les arrêts dans la boucle d'exécution principale du modèle.
• Compromis : Le mode push nécessite de sélectionner l'instance de decode plus tôt dans le processus, ce qui peut réduire la flexibilité pour le routage conscient de la charge par rapport au mode pull.
• Configuration : Les utilisateurs activent cela en configurant les serveurs de préfill et de decode pour utiliser NixlPushConnector au lieu du NixlConnector par défaut.

Pourquoi c'est important

Pour les ingénieurs construisant des services LLM à haut débit, le service désagrégé est une stratégie courante pour gérer efficacement l'utilisation des ressources. Cependant, la surcharge liée au déplacement de grands caches KV entre les nœuds peut annuler les avantages si elle n'est pas gérée avec soin. Le modèle pull traditionnel introduit une période d'attente synchrone où l'instance de decode est inactive, attendant les métadonnées puis tirant les données. Ce délai devient prononcé à mesure que les longueurs de contexte augmentent et que les taux de requêtes s'accélèrent, impactant directement la latence perçue par l'utilisateur.

En déplaçant le travail de coordination vers un thread d'arrière-plan et en le chevauchant avec le calcul de préfill, le modèle push garantit que l'instance de decode est prête à recevoir les données dès qu'elles sont disponibles. Cela réduit le temps d'inactivité sur le chemin critique, entraînant des temps de réponse plus rapides pour les utilisateurs. Cela empêche également le worker de decode d'être encombré par les tâches de configuration du transfert, lui permettant de maintenir des vitesses de génération de jetons constantes même sous une forte concurrence.

Ce développement souligne l'importance des optimisations d'infrastructure bas niveau dans le service IA. À mesure que les modèles deviennent plus grands et les contextes plus longs, l'efficacité du mouvement des données entre les nœuds de calcul devient aussi importante que la vitesse des multiplications matricielles elles-mêmes. Les ingénieurs déployant vLLM sur AWS ou une infrastructure similaire compatible RDMA devraient envisager cette option pour maximiser les performances.

Ce que vous pouvez faire

• Évaluez votre déploiement vLLM actuel pour voir si le service désagrégé préfill-decode est approprié pour votre charge de travail.
• Testez le NixlPushConnector dans un environnement de staging avec des motifs de trafic réalistes pour mesurer les améliorations du TTFT.
• Assurez-vous que votre infrastructure prend en charge RDMA, comme AWS EFA, pour tirer parti des fonctionnalités d'accès direct à la mémoire.
• Surveillez le compromis entre la sélection précoce de l'instance de decode et les besoins d'équilibrage de charge dynamique dans la configuration de votre routeur.
• Désactivez la mise en cache des préfixes lors des benchmarks pour isoler les performances de transfert, car la mise en cache peut masquer les coûts de transfert.
• Consultez la documentation vLLM pour les indicateurs de configuration spécifiques requis pour passer du mode pull au mode push.