Pooling Engram Conditional Memory in Large Language Models using CXL

Ce papier propose d'utiliser un pool de mémoire CXL pour stocker la mémoire conditionnelle des engrammes dans les grands modèles de langage, offrant une solution de stockage évolutive et rentable qui maintient des performances d'inférence proches de celles de la DRAM grâce à l'intégration dans SGLang.

L'essentiel

Voici une explication simple et imagée de ce papier de recherche, conçue pour être comprise par tout le monde, même sans être expert en informatique.

🧠 Le Problème : Le Cerveau qui a trop de souvenirs

Imaginez que vous essayez de construire un super-cerveau artificiel (une intelligence artificielle comme les modèles de langage que nous utilisons aujourd'hui). Pour être très intelligent, ce cerveau a besoin de deux choses :

1. De la puissance de calcul pour réfléchir et raisonner (comme faire des maths).
2. Une énorme bibliothèque de connaissances pour se souvenir des faits, des mots et des règles (comme un dictionnaire géant).

Le problème actuel, c'est que pour que ce cerveau soit rapide, on doit mettre toute cette bibliothèque directement dans sa mémoire vive (la RAM), comme si on devait empiler des milliers de livres sur le bureau de l'ordinateur pendant qu'il travaille.

• Résultat : C'est très cher (la mémoire coûte une fortune) et cela limite la taille du cerveau. Si vous voulez un cerveau plus grand, vous devez acheter plus de bureaux, ce qui devient vite impossible.

💡 La Solution : Le "Mémoriel" (Engram)

Les chercheurs ont inventé une nouvelle idée appelée Engram. C'est comme si le cerveau avait une capacité spéciale : au lieu de tout lire dans sa mémoire vive, il va chercher des informations précises au moment où il en a besoin, comme un bibliothécaire ultra-rapide.

Mais il y a un hic : ce bibliothécaire doit aller chercher des informations très rapidement, mais de manière très sporadique. Il ne prend pas un livre entier, mais parfois juste une phrase, parfois un mot, et cela arrive des milliers de fois par seconde.

🚀 L'Innovation : Le "CXL" (Le Tapis Roulant Magique)

Pour résoudre ce problème, l'équipe de chercheurs (de l'Université de Pékin, Alibaba Cloud, etc.) a proposé d'utiliser une nouvelle technologie appelée CXL (Compute Express Link).

Voici l'analogie pour comprendre la différence entre l'ancienne méthode et la nouvelle :

1. L'ancienne méthode (RDMA) : Le Camion de Livraison 🚚

Imaginez que pour récupérer un mot, vous deviez appeler un camion de livraison.

• Le camion arrive, charge une palette entière (même si vous ne voulez qu'un seul mot).
• Il fait le tour du pâté de maisons (le réseau).
• Il décharge, puis repart.
• Problème : C'est trop lent pour récupérer un seul mot. Le temps perdu à faire venir le camion est plus long que le temps de lire le mot. C'est comme si vous deviez attendre un train pour acheter un journal à la gare.

2. La nouvelle méthode (CXL) : Le Tapis Roulant Magique 🎢

Le CXL, c'est comme un tapis roulant ultra-rapide qui relie directement le cerveau de l'ordinateur à une immense bibliothèque partagée située juste à côté.

• Vous tendez la main, et le mot arrive instantanément sur le tapis.
• Pas de camion, pas de détour, pas de temps d'attente.
• Avantage : C'est aussi rapide que si le livre était déjà sur votre bureau, mais vous n'avez pas besoin d'acheter un bureau géant pour tout stocker.

🏗️ Comment ça marche en pratique ?

L'équipe a construit un prototype avec deux choses principales :

1. Une Bibliothèque Partagée (Memory Pool) : Au lieu que chaque ordinateur ait sa propre bibliothèque, plusieurs ordinateurs partagent une seule et même bibliothèque géante (stockée sur des disques CXL).
2. Des Routiers Optimisés : Ils ont écrit des logiciels spéciaux pour que le "tapis roulant" (CXL) fonctionne à la vitesse de l'éclair, même quand des milliers de mots doivent être récupérés en même temps.

📊 Les Résultats : Moins cher, aussi rapide !

Les chercheurs ont testé leur système avec un modèle de langage (Qwen) :

• Vitesse : Le système fonctionne presque aussi vite que si la mémoire était directement dans l'ordinateur (presque pas de ralentissement).
• Coût : C'est là que c'est magique.
  • Avant : Si vous aviez 16 ordinateurs, vous deviez acheter 16 bibliothèques complètes.
  • Maintenant : Vous achetez une seule bibliothèque géante que les 16 ordinateurs partagent.
  • Économie : Plus vous avez d'ordinateurs, plus vous économisez d'argent. Pour un grand système, vous pouvez économiser des dizaines de milliers de dollars.

🎯 En Résumé

Imaginez que vous voulez construire une ville de géants (les IA de demain).

• Avant : Chaque géant devait avoir sa propre maison remplie de livres. C'était cher et encombrant.
• Aujourd'hui (avec ce papier) : Les géants partagent une immense bibliothèque centrale reliée à leurs maisons par un tapis roulant magique (CXL).
• Le résultat : Ils peuvent lire aussi vite que s'ils étaient chez eux, mais la ville coûte beaucoup moins cher à construire et peut accueillir beaucoup plus de géants.

C'est une étape clé pour rendre les futures intelligences artificielles plus puissantes, plus intelligentes et beaucoup moins chères à faire tourner.

Résumé technique

1. Problématique

Les Grands Modèles de Langage (LLM) modernes reposent souvent sur des architectures de type Mixture-of-Experts (MoE) pour le calcul conditionnel. Cependant, ces modèles manquent de mécanismes natifs efficaces pour la recherche de connaissances statiques (vocabulaire), obligeant le modèle à simuler cette recherche par le calcul, ce qui est inefficace.

Pour résoudre ce problème, l'architecture Engram a été introduite. Elle découple le stockage des connaissances statiques (sous forme de tables d'embeddings N-Gram) du calcul dynamique. Bien que prometteuse, l'Engram présente deux défis majeurs pour le déploiement à grande échelle :

• Surcharge mémoire : Les tables d'embeddings peuvent atteindre plusieurs centaines de Go, dépassant souvent la capacité de la mémoire DRAM locale des GPU.
• Motif d'accès spécifique : L'accès à la mémoire Engram est sparsé (peu de données par token), discret (de nombreux petits segments non contigus) et en lecture seule.

Les solutions de mémoire distribuée actuelles, basées sur le RDMA (Remote Direct Memory Access), souffrent de latences élevées et d'une inefficacité lors du transfert de petits paquets de données, ce qui les rend inadaptées aux motifs d'accès fins et fréquents de l'Engram.

2. Méthodologie

Les auteurs proposent d'utiliser la technologie CXL (Compute Express Link) pour créer un pool de mémoire partagé dédié au stockage des paramètres Engram.

Architecture et Infrastructure

• Pool Mémoire CXL : Utilisation d'un commutateur CXL (CXL Switch) pour connecter plusieurs nœuds de calcul à un pool de mémoire centralisé (ex: 256 Go ou 4 To).
• Accès Direct : Contrairement au RDMA qui nécessite une orchestration par le CPU et des tampons intermédiaires (bounce buffers), CXL offre une interface de chargement/stockage (load/store) gérée matériellement avec une granularité au niveau de la ligne de cache. Cela permet un accès à la mémoire distante avec une latence proche de celle de la mémoire locale (DRAM).

Implémentation Technique

L'implémentation a été intégrée dans le framework d'inférence SGLang avec les optimisations suivantes :

1. Mappage DAX (Direct Access) : La mémoire CXL est exposée comme un dispositif DAX dans le système d'exploitation, permettant un accès direct via mmap sans copie de données inutiles.
2. Optimisation CPU : Utilisation de stratégies de lecture parallèle multi-threadées (OpenMP) pour saturer la bande passante PCIe lors du transfert des embeddings vers le CPU.
3. Optimisation GPU (P2P) : Développement de noyaux CUDA personnalisés pour permettre des transferts directs Peer-to-Peer (P2P) depuis la mémoire CXL vers la mémoire GPU (VRAM), contournant le CPU. Cela évite la surcharge de lancement de milliers de petits transferts cudaMemcpy.
4. Préchargement Asynchrone : Le mécanisme de récupération des embeddings est lancé de manière asynchrone au début de l'étape de décodage, permettant de masquer la latence d'accès mémoire par le calcul des blocs transformeurs précédents.

3. Contributions Clés

1. Premier système CXL pour Engram : Proposition et mise en œuvre du premier système déchargeant les paramètres Engram sur un pool mémoire CXL.
2. Analyse comparative RDMA vs CXL : Démonstration que CXL est supérieur au RDMA pour les motifs d'accès Engram. Les tests montrent que le RDMA subit une pénalité de latence d'ordres de grandeur pour les petits paquets, tandis que CXL maintient une latence proche de la DRAM locale.
3. Intégration SGLang et Performance : Implémentation fonctionnelle dans SGLang atteignant des performances end-to-end quasi-identiques à celles d'une configuration entièrement DRAM.
4. Analyse de Coût : Preuve que le pooling CXL devient économiquement avantageux à mesure que la taille du cluster et du modèle augmente, réduisant considérablement le coût total de possession (TCO) par rapport à l'achat de DRAM locale pour chaque nœud.

4. Résultats Expérimentaux

Les expériences ont été menées sur un testbed avec des GPU NVIDIA L20 et un commutateur CXL XC50256.

• Latence de Lecture :

  • Pour des tailles de lot (batch) variables, la latence de lecture CXL vers CPU est comparable à celle de la DRAM locale.
  • La latence CXL vers GPU, bien que légèrement supérieure, reste dans une plage acceptable et ne dégrade pas le débit global.
  • Le RDMA montre une latence nettement plus élevée, incompatible avec les fenêtres de préchargement strictes de l'Engram (notamment pour les premières couches du modèle).

• Débit d'Inférence (Throughput) :

  • Sur des modèles Qwen3-4B et Qwen3-8B, l'ajout de l'Engram via CXL entraîne une réduction de débit marginale par rapport à la DRAM locale (ex: ~1,2% de perte pour Qwen3-4B).
  • Le système conserve un débit élevé même lors de l'augmentation de la parallélisation de données (DP) et du nombre de nœuds, prouvant une scalabilité robuste.

• Analyse de Coût :

  • Pour un Engram de 100B de paramètres sur 16 nœuds, le coût du stockage CXL est d'environ 16 960 $**, contre **48 000$ pour une solution DRAM locale.
  • Les économies deviennent exponentielles avec l'augmentation de la taille du modèle (ex: 166 040 $ d'économies pour un Engram de 400B sur 16 nœuds).

5. Signification et Impact

Ce travail démontre que le CXL est une infrastructure idéale pour les futures architectures de LLM enrichies en mémoire (Memory-Augmented LLM).

• Évolutivité : Il permet de surmonter le « mur de la mémoire » (memory wall) des GPU en permettant le partage de tables d'embeddings massives sans sacrifier les performances d'inférence.
• Efficacité Économique : Il offre une solution rentable pour déployer des modèles avec des capacités de connaissances statiques exponentielles, rendant viable l'utilisation de tables N-Gram géantes.
• Futur de l'Inférence : En validant que les motifs d'accès sparses et discrets de l'Engram peuvent être servis efficacement via CXL, cette recherche ouvre la voie à une nouvelle génération de modèles plus intelligents et moins coûteux à déployer, tout en suggérant des pistes pour l'évolution des solutions RDMA (via des interfaces P2P directes et du caching).

En résumé, l'article établit que le pooling mémoire via CXL est non seulement techniquement faisable pour l'Engram, mais qu'il constitue la voie la plus prometteuse pour l'infrastructure de stockage des connaissances dans les LLM de prochaine génération.