Résumé Technique : PTStore (Prefix Tensor Store)

Problématique

Les charges de travail d'inférence de modèles de langage de grande taille (LLM) sont devenues la charge dominante dans les centres de données de calcul haute performance (HPC), dépassant l'entraînement en termes de consommation d'énergie et de demande de ressources. L'inférence de LLM se compose de deux phases : le prefill (traitement du prompt d'entrée en parallèle) et le decode (génération de jetons séquentiellement). Pour éviter le calcul redondant des mécanismes d'attention, les systèmes utilisent un cache clé-valeur (KV) pour stocker les résultats intermédiaires.

Bien que les runtimes de pointe (par exemple, vLLM) optimisent le cache KV au sein d'un seul GPU ou d'un seul nœud, ils font face à des limitations significatives à grande échelle :

1. Absence de réutilisation inter-nœuds : Les systèmes existants échouent souvent à agréger la mémoire entre les nœuds de calcul distribués. Si une requête sur un nœud partage un préfixe avec une requête sur un autre nœud, le second nœud recompute généralement le préfixe plutôt que de réutiliser les tenseurs mis en cache.
2. Goulots d'étranglement de métadonnées et de latence : Les approches tentant le cache distribué (par exemple, LMCache, EvoStore) souffrent souvent de surcoûts d'E/S élevés dus à l'accès à la mémoire distante ou à la synchronisation complexe des métadonnées (par exemple, l'extension de Radix-Attention au-delà d'un seul nœud).
3. Contraintes de mémoire : La mémoire individuelle des GPU est insuffisante pour les fenêtres de contexte étendues, et le déchargement vers la mémoire hôte ou les SSD introduit une latence qui annule les bénéfices du cache.

Le défi central consiste à permettre une réutilisation scalable et à faible latence des préfixes de cache KV à travers un grand nombre de GPU distribués sur de nombreux nœuds de calcul, sans engendrer d'imposants surcoûts d'E/S ou de métadonnées.

Méthodologie : Architecture PTStore

PTStore (Prefix Tensor Store) est un magasin de tenseurs distribué et répliqué conçu pour répondre à ces limitations en distribuant et en répliquant les préfixes de cache KV populaires. Le système emploie un modèle client-serveur où chaque nœud de calcul exécute un serveur qui agrège la mémoire locale de l'hôte et les SSD pour servir les clients GPU locaux et distants.

Principes de conception clés

1. Stockage de tenseurs incrémentiel (Structure de type Trie) :

   • Au lieu de stocker des blocs KV complets, PTStore stocke les différences incrémentielles (tenseurs) entre un nouvel objet et le plus long préfixe commun (LCP) des objets précédemment stockés.
   • Cela permet aux préfixes de croître sans redondance au fil du temps dans des directions divergentes, de manière similaire à un trie, mais implémenté avec une granularité de niveau tenseur.
   • Métadonnées consolidées : Pour éviter l'examen coûteux de tries distribués, PTStore utilise une structure de métadonnées plate. Les métadonnées de chaque objet contiennent une liste d'identifiants de tenseurs uniques. Une opération de chargement itère à travers ces identifiants pour vérifier leur existence locale dans le cache de réplication ; si l'élément est manquant, il est récupéré à distance auprès du serveur "propriétaire".

2. Mise en cache hiérarchique distribuée avec réplication :

   • Cache possédé (Owned Cache) : Stocke les tenseurs incrémentiels dont un serveur spécifique est responsable.
   • Cache de réplication (Replication Cache) : Stocke des copies des préfixes "chauds" (populaires) localement sur le serveur pour améliorer la localité d'accès.
   • Gestion des compromis : Le système gère un seuil configurable entre les caches possédés et répliqués. Il donne la priorité à l'élimination des tenseurs répliqués (qui peuvent être récupérés à nouveau) par rapport à l'éviction des tenseurs possédés (qui nécessitent un transfert vers un stockage plus lent) afin d'équilibrer la vitesse de récupération et la capacité de stockage.

3. Éviction sensible aux modèles d'accès :

   • PTStore utilise une politique d'éviction basée sur la fréquence (adaptée de GDSF) plutôt que sur le principe LRU (Least Recently Used), car les structures de préfixes impliquent que les tenseurs précoces sont accédés plus fréquemment.
   • Il prend en compte le compromis taille/fréquence, garantissant que les petits tenseurs fréquents ne déplacent pas les tenseurs plus volumineux et coûteux à récupérer.

4. Consolidation sensible à l'RDMA :

   • Pour minimiser l'éparpillement, les incréments ajoutés à un LCP sont consolidés en une région contiguë unique sur le serveur propriétaire.
   • Les opérations de chargement utilisent l'RDMA groupé (bulk RDMA) pour récupérer les segments éparpillés en parallèle via un seul appel RPC, évitant ainsi le surcoût de la copie des données dans une région contiguë avant le transfert.

Contributions clés

1. Principes de conception : Un ensemble de principes de haut niveau pour un référentiel distribué qui intègre le stockage de tenseurs incrémentiel, les métadonnées consolidées et la réplication de préfixe.
2. Prototype PTStore : Un prototype de recherche implémentant ces principes, doté d'une API de bas niveau en C++ et d'une interface Python pour une intégration transparente avec les runtimes de LLM comme vLLM.
3. Validation des performances : Des expérimentations approfondies démontrant des réductions significatives des surcoûts d'E/S et du temps d'exécution de bout en bout par rapport aux own baselines de pointe.

Résultats expérimentaux

Les auteurs ont évalué PTStore sur le banc d'essai HPC ALCF Polaris (560 nœuds, GPU A100) en utilisant deux charges de travail de questions-réponses extractives : WikiQA (contexte long) et SQUAD (volume élevé de questions). Le LLM utilisé était Mistral-7B-instruct-V2.

Baselines

• vLLM Vanilla : vLLM standard sans partage de préfixe entre les requêtes.
• vLLM Prefix : vLLM avec partage de préfixe local (au sein d'un nœud).
• EvoStore : Un magasin de tenseurs distribué utilisant le stockage incrémentiel et l'RDMA, mais manquant de réplication de préfixe locale.
• PTStore : Le système proposé avec conscience distribuée et réplication locale.

Constats

• Faible scalabilité (8–32 GPU) : PTStore a nettement surpassé EvoStore et vLLM Prefix. Alors qu'EvoStore souffrait de surcoûts élevés d'E/S RDMA lors de la récupération de préfixes distants, la réplication locale de PTStore a atténué ce problème, résultant en un "avantage détaché" sur le Temps au Premier Jeton (TTFT).
• Scalabilité de la longueur de séquence (1k–8k jetons) :
  • Pour les séquences courtes (1k), le cache local de vLLM était compétitif.
  • À mesure que la longueur de la séquence augmentait, l'avantage de PTStore grandissait. À 8k jetons, PTStore était presque 2 fois plus rapide que le cache de préfixe de vLLM et 20 % plus rapide qu'EvoStore.
  • L'écart de performance s'est élargi avec des contextes plus longs car le coût de la recomputation ou des E/S distantes l'emportait sur les bénéfices du cache local uniquement.
• Gains d'efficacité : Sur les jeux de données de questions-réponses à passages longs, PTStore a exécuté les inférences 5 à 6 fois plus efficacement que les baselines qui n'agrègent pas la mémoire entre les nœuds et nécessitent de régénérer les caches KV.

Signification et Revendications

L'article affirme que PTStore comble une lacune critique dans le service d'inférence de LLM scalable : l'incapacité des systèmes actuels à réutiliser efficacement les préfixes de cache KV à travers des nœuds distribués. En combinant le stockage incrémentiel pour minimiser la redondance, les métadonnées consolidées pour des requêtes rapides et une stratie de réplication pour optimiser la localité, PTStore permet :

• Une expansion de plusieurs ordres de grandeur de la taille effective du cache KV en agrégeant la mémoire à travers le cluster.
• Une réduction significative du TTFT, particulièrement pour les charges de travail à contexte long où la recomputation est coûteuse.
• Une scalabilité qui évite les goulots d'étranglement de communication et les problèmes de synchronisation de métadonnées qui affectent les approches distribuées précédentes.

Les auteurs positionnent PTStore comme une étape fondamentale vers une inférence IA scalable, notant que les travaux futurs se concentreront sur l'équilibrage dynamique de la mémoire, les politiques d'éviction basées sur le ML et des benchmarks plus larges contre des systèmes comme LMCache et Mooncake sur des traces réelles de conversation et de complétion de code.