Hardware-Software Co-design for 3D-DRAM-based LLM Serving Accelerator

Le papier présente Helios, un accélérateur de serving pour grands modèles de langage basé sur la co-conception matériel-logiciel et l'hybridation de puces 3D-DRAM, qui surpasse les conceptions existantes en optimisant la gestion dynamique du cache KV et l'exécution distribuée de l'attention pour réduire la latence et améliorer l'efficacité énergétique.

L'essentiel

Imaginez que les grands modèles de langage (comme ceux qui font fonctionner les chatbots intelligents) sont comme des cuisiniers de génie dans un restaurant très fréquenté.

Ces cuisiniers doivent gérer deux types de tâches très différents :

1. La préparation (Prefill) : C'est le moment où ils lisent la commande du client (le prompt). C'est une tâche qui demande beaucoup de force brute et de calcul rapide.
2. La cuisson (Decoding) : C'est le moment où ils écrivent la réponse, mot par mot. Ici, le défi n'est pas la force, mais la mémoire. Ils doivent se souvenir de tout ce qui a été dit précédemment pour rester cohérents. Plus la conversation est longue, plus ils doivent se souvenir de beaucoup de choses.

Le Problème : Le "Goulot d'Étranglement" de la Mémoire

Dans les serveurs actuels (comme les puces GPU), il y a un déséquilibre. Le cuisinier est très fort pour calculer, mais sa "mémoire" (le placard où il range les ingrédients) est loin et difficile d'accès.

• Quand il doit se souvenir de 10 000 mots passés pour écrire le prochain, il passe son temps à courir chercher les ingrédients dans le placard plutôt qu'à cuisiner.
• De plus, les clients arrivent de manière imprévisible : parfois un seul client, parfois des centaines. Et certains veulent des réponses courtes, d'autres des romans entiers.

Les solutions actuelles tentent de rapprocher la mémoire du cuisinier (c'est ce qu'on appelle le Near-Memory Processing ou NMP), mais elles sont trop rigides. C'est comme si le restaurant avait des placards fixes :

• Si un client commande un petit plat, on lui donne tout un placard vide (gaspillage).
• Si un client commande un énorme plat, le placard est trop petit et on ne peut pas l'accepter, même s'il reste de la place ailleurs dans le restaurant.
• Les cuisiniers sont assignés à des placards spécifiques et ne peuvent pas s'entraider facilement.

La Solution : Helios, le Restaurant du Futur

Les chercheurs de l'Université de Pékin et de leurs partenaires ont inventé Helios. Imaginez Helios comme un restaurant ultra-moderne construit avec une technologie de "collage" 3D (le Hybrid Bonding).

Voici comment Helios change la donne, avec des analogies simples :

1. La Technologie de "Collage" (Hybrid Bonding)

Au lieu d'avoir la mémoire et le processeur séparés sur des étages différents, Helios les colle directement l'un sur l'autre, comme deux étages d'un immeuble reliés par des ascenseurs ultra-rapides et courts.

• Avantage : La distance entre le cuisinier et les ingrédients est quasi nulle. La vitesse de transfert est énorme, et la consommation d'énergie est faible.

2. La Gestion Dynamique des Placards (KV Cache)

C'est le cœur de l'innovation. Au lieu de donner un placard entier à un client, Helios découpe la mémoire en petits blocs (comme des Lego).

• L'analogie : Imaginez que vous avez un mur de Lego.
  • Les anciens systèmes : Ils réservent un mur entier pour un client, même s'il n'utilise que quelques briques. Si un autre client arrive avec un gros projet, il n'y a pas de mur libre, même si le total des murs libres suffirait.
  • Helios : Il prend les briques libres un par un. Si un client a besoin de 100 briques, on lui donne 100 briques prises ici et là. Si un client finit sa commande, ses briques retournent immédiatement dans le stock pour être réutilisées.
• Résultat : Pas de gaspillage, et le restaurant peut accueillir beaucoup plus de clients en même temps, peu importe la taille de leur commande.

3. La Cuisine Collaborative (Attention Distribuée)

Dans les anciens systèmes, un seul cuisinier devait gérer toute la mémoire d'un client. Si la mémoire était trop grande, il s'essoufflait.

• L'analogie Helios : Quand un client a une très longue conversation, Helios ne donne pas tout le travail à un seul cuisinier. Il répartit les briques de mémoire sur plusieurs cuisiniers qui travaillent ensemble en temps réel.
• Ils se passent les informations très vite grâce à leurs ascenseurs ultra-rapides. C'est comme une équipe de relais où chaque coureur court une petite distance, mais ensemble, ils vont très vite.

Les Résultats : Pourquoi c'est impressionnant ?

Grâce à cette conception intelligente (co-conception matériel/logiciel), Helios bat les meilleurs serveurs actuels (comme les puces NVIDIA A100) de manière spectaculaire :

• Vitesse : Il est 3,25 fois plus rapide en moyenne. C'est comme passer d'une voiture de ville à une Formule 1.
• Économie d'énergie : Il consomme 3,36 fois moins d'énergie pour faire le même travail. C'est comme avoir une voiture qui roule plus vite mais qui consomme moins d'essence.
• Stabilité : Même quand le restaurant est bondé et que les commandes arrivent en rafale, Helios maintient un service fluide. Les temps d'attente (le temps entre deux mots générés) restent très bas, même pour les clients exigeants.

En Résumé

Helios est une révolution pour l'intelligence artificielle générative. Il résout le problème de la "mémoire" en transformant un système rigide et gaspilleur en un système flexible, collaboratif et ultra-rapide.

C'est comme passer d'un entrepôt où les camions doivent faire de longs détours pour charger des marchandises, à un entrepôt où les robots sont collés aux étagères et peuvent se passer les colis instantanément, peu importe la taille de la commande. Cela permet de servir des millions de conversations complexes, rapidement et à moindre coût.

Résumé technique

Voici un résumé technique détaillé du papier de recherche "Hardware-Software Co-design for 3D-DRAM-based LLM Serving Accelerator" (Co-conception matériel-logiciel pour un accélérateur de service LLM basé sur la DRAM 3D), présentant le système HeLioS.

1. Problématique et Contexte

Les modèles de langage de grande taille (LLM) sont de plus en plus déployés sous forme de services en ligne, générant des charges de travail hautement dynamiques caractérisées par :

• Des taux d'arrivée de requêtes fluctuants (pic vs creux).
• Des longueurs de séquences variables (de quelques tokens à 10 000+ tokens).
• Des besoins mixtes en calcul (étape de prefill) et en accès mémoire (étape de décodage).

Limites des solutions existantes (NMP - Near-Memory Processing) :
Les accélérateurs NMP actuels, qui intègrent des unités de traitement (PE) dans les modules de mémoire (ex: HBM), souffrent de deux limitations majeures face à cette dynamique :

1. Gestion grossière du Cache KV : Les architectures existantes allouent le cache Key-Value (KV) entier d'une tête d'attention à un module NMP fixe. Cela entraîne un déséquilibre de charge de calcul (surtout avec des variantes d'attention réduisant le nombre de têtes) et une fragmentation mémoire sévère (gaspillage d'espace car les modules doivent réserver la taille maximale de contexte, même si la requête est courte).
2. Flux d'exécution rigide : L'exécution de l'attention est souvent statique et non divisée (non tiled), empêchant une gestion fine du cache KV nécessaire pour les requêtes de longueurs variables.

2. Méthodologie : HeLioS

HeLioS est un accélérateur de service LLM basé sur la liaison hybride (Hybrid Bonding) qui vise à combler le fossé entre la nature dynamique du service LLM et l'abstraction mémoire non uniforme des architectures NMP.

A. Architecture Matérielle (Hardware)

• Technologie de Liaison Hybride (HB) : Utilisation d'une technologie de liaison 3D (Cu-Cu) pour empiler plusieurs puces DRAM sur une puce logique. Cela offre une bande passante massive (110 000 I/O/mm²) et une faible latence énergétique (0.66 pJ/bit).
• Réseau de Processeurs (PE) : La puce logique contient une grille de PE (4x4 dans la conception finale). Chaque PE a un accès direct à sa mémoire locale DRAM via des contrôleurs HB.
• Unités Spécialisées : Chaque PE intègre des unités pour les GEMM (multiplication matricielle), le softmax en ligne, la réduction et les opérations vectorielles.
• Communication Inter-PE : Utilisation de réseaux sur puce (NoC) maillés pour gérer les communications collectives (All-Reduce, All-Gather, Scatter) nécessaires à l'attention distribuée.

B. Conception Logicielle et Algorithmique (Software)

• Attention Tillée Distribuée (Distributed Tiled Attention) : Contrairement aux approches statiques, HeLioS divise le cache KV en blocs fins (ex: 64 tokens). Ces blocs sont distribués dynamiquement sur tous les PE disponibles.
• Allocation de Cache KV Sensible à l'Espace (Spatially-Aware Allocation) : Un algorithme d'allocation intelligent place les blocs de cache KV en tenant compte de :
  1. L'équilibre de charge de calcul (nombre de tokens par PE).
  2. La topologie du réseau (distance Manhattan vers les PE centraux) pour minimiser le trafic de communication inter-PE.
• Flux d'Exécution Optimisé :
  • Fusion des opérations de softmax en ligne avec les GEMM pour masquer la latence.
  • Gestion dynamique des transferts de cache KV (tampon de transfert) pour éviter les conflits entre écriture mémoire et calcul.
  • Support natif des variantes d'attention (MQA, GQA, MLA) et des modèles MoE (Mixture of Experts).

3. Contributions Clés

1. Analyse des Limites NMP : Identification précise des goulots d'étranglement liés à la gestion statique du cache KV et à l'absence de flexibilité dans les flux d'attention des designs NMP existants.
2. Architecture HeLioS : Proposition d'une architecture matérielle basée sur la liaison hybride permettant une attention distribuée et native au NMP, avec une granularité de blocs KV fine.
3. Stratégie d'Allocation Spatiale : Développement d'un mécanisme d'allocation qui équilibre la charge de calcul et minimise les transferts de données inter-PE, résolvant le problème de fragmentation mémoire.
4. Co-conception Complète : Intégration de la gestion du cache, de l'exécution des opérateurs et de la communication collective dans un système unifié.

4. Résultats Expérimentaux

Les évaluations ont été menées sur plusieurs modèles (OPT, LLaMA3, Mixtral, Qwen, DeepSeek) avec des charges de travail réalistes (traces MoonCake).

• Performance de Décodage :
  • Accélération : HeLioS offre un speedup géométrique moyen de 3,25x par rapport aux conceptions GPU/NMP existantes.
  • Efficacité Énergétique : Amélioration de 3,36x par rapport aux meilleures solutions de base.
  • Latence (TBT - Time Between Tokens) : Réduction de la latence P50/P99 de 72% à 76% par rapport aux solutions de pointe.
• Capacité de Service :
  • HeLioS maintient des performances élevées même sous forte charge, évitant la saturation prématurée observée chez les solutions NMP traditionnelles.
  • Réduction drastique de la fragmentation mémoire (facteur de 2,21x à 23,10x moins de fragmentation que les baselines).
• Comparaison avec WaferLLM : HeLioS surpasse l'approche récente WaferLLM (également basée sur des PE maillés) grâce à une meilleure efficacité de calcul (pas de partitionnement excessif des dimensions courtes) et une communication inter-PE optimisée.

5. Signification et Impact

Ce travail démontre que la simple intégration de processeurs dans la mémoire (NMP) ne suffit pas pour les services LLM dynamiques. La clé réside dans une co-conception matérielle-logicielle qui :

• Abandonne l'allocation statique au profit d'une gestion fine et dynamique du cache KV.
• Exploite pleinement la bande passante et la capacité de calcul offertes par la liaison hybride 3D.
• Permet aux accélérateurs NMP de rivaliser, voire de surpasser, les GPU massifs pour les phases de décodage, qui sont souvent limitées par la bande passante mémoire.

HeLioS ouvre la voie à des infrastructures de service LLM plus efficaces, capables de gérer des charges de travail imprévisibles avec une latence réduite et un coût énergétique inférieur, ce qui est crucial pour le déploiement économique des modèles génératifs à grande échelle.