Accelerating GPU Inference of Large Language Models with Moderately Unstructured Sparse Weight Matrices

Résumé Technique : Accélération de l'inférence GPU des grands modèles de langage avec des matrices de poids à parcité modérément non structurée

Énoncé du problème
Le déploiement des grands modèles de langage (LLM) fait face à un goulot d'étranglement critique concernant les coûts d'inférence. Bien que les techniques d'élagage de modèles (par exemple, Wanda, RIA) puissent réduire la charge de calcul en introduisant de la parcité dans les matrices de poids, elles sont contraintes par des exigences de qualité de modèle. Contrairement aux modèles de petite échelle, les LLM sont très sensibles aux niveaux de parcité élevés ; maintenir la performance limite généralement l'élagage à une parcité non structurée modérée (environ 50 %).

À ce niveau de parcité spécifique, les noyaux GPU existants pour la multiplication de matrices creuses (SpMM) ne parviennent pas à surpasser les bibliothèques de multiplication de matrices denses comme cuBLAS. Cet écart de performance provient de trois défis fondamentaux :

1. Incompatibilité des Tensor Cores : Les GPU modernes offrent des cœurs de tenseurs creux hautement efficaces, mais ils supportent strictement la parcité structurée 2:4. Ils ne peuvent pas accélérer directement les motifs d'élagage irréguliers et non structurés.
2. Surcharge de métadonnées : Les formats de stockage traditionnels (par exemple, CSR) nécessitent le stockage de métadonnées positionnelles complètes pour chaque élément non nul. À 50 % de parcité, ces métadonnées peuvent occuper autant d'espace que les valeurs elles-mêmes, annulant les bénéfices de la bande passante mémoire.
3. Goulots d'étranglement du décodage sur puce : Les techniques de compression qui réduisent les métadonnées (par exemple, le codage par bitmap) transfèrent la charge de travail vers le décodage sur puce. Sur les GPU modernes dotés de mémoire à haute bande passante (HBM), l'accès à la mémoire globale ne peut pas pleinement se superposer au travail de décodage effectué par les cœurs CUDA plus lents, ce qui crée un blocage de performance.

Méthodologie
Les auteurs proposent une solution holistique comprenant un nouveau format de stockage de matrice à trois couches et un noyau SpMM co-optimisé conçu pour la parcité non structurée modérée.

1. Format de stockage à trois couches
Le format proposé décompose la matrice creuse en trois couches distinctes pour équilibrer la compatibilité matérielle et l'efficacité du stockage :

• Couche Sparse-TC : Cette couche extrait les motifs de parcité structurée 2:4 de la distribution irrégulière. Elle conserve deux valeurs non nulles pour chaque quatre poids consécutifs. Comme le motif est fixe, elle ne nécessite que 2 bits par valeur pour les métadonnées, assurant une compatibilité directe avec les cœurs de tenseur creux sans surcharge de décodage.
• Couche de remplissage de créneaux (Slot-Filling Layer) : Les éléments non nuls excédentaires qui ne rentrent pas dans le motif 2:4 sont redistribués dans les emplacements laissés vides par la couche Sparse-TC. Pour minimiser les métadonnées de ces ~17,7 % d'éléments, les auteurs introduisent le codage Parallèle de Distance Différentielle (PDD). Au lieu de stocker des indices absolus, le PDD encode les distances relatives entre les éléments redistribués. La plupart des distances sont suffisamment petites pour être codées en 4 bits. Pour les écarts plus importants, des points fictifs sont insérés pour briser l'étendue, garantissant qu'aucune distance unique ne dépasse la limite de codage.
• Couche Résiduelle : Les éléments non nuls restants (typiquement <1 % du total) qui ne peuvent pas être gérés efficacement par les deux premières couches sont stockés dans un format traditionnel de ligne creuse compressée (CSR). En raison du volume minimal de données, la surcharge de métadonnées ici est négligeable.

2. Noyau SpMM co-optimisé
S'appuyant sur ce format, les auteurs conçoivent un noyau GPU qui utilise conjointement des cœurs de tenseur creux et des cœurs CUDA au sein d'un modèle d'exécution pipeliné :

• Décodage Parallèle : Les cœurs CUDA sont chargés de décoder les métadonnées PDD de la couche de remplissage de créneaux en parallèle. Puisque cette couche représente une faible fraction des données totales, la surcharge de décodage est gérée efficacement.
• Accélération Structurée : Les cœurs de tenseur creux traitent la couche Sparse-TC en utilisant les instructions wgmma (architecture Hopper), exploitant la structure 2:4 pour un débit élevé.
• Gestion des Résidus : Les cœurs CUDA traitent également la couche résiduelle creuse, écrivant les résultats dans un tampon distinct sur puce pour être réduits avec les sorties du cœur de tenseur à la fin du pipeline.
• Chevauchement de Pipeline : Le noyau implémente un pipeline fin qui superpose trois opérations critiques :
  • Le chargement des matrices d'entrée denses (mises en cache dans le L2) et des tuiles de poids creuses (depuis la mémoire globale).
  • Le décodage des métadonnées PDD sur les cœurs CUDA.
  • L'exécution de la multiplication de matrices sur les cœurs de tenseur creux.
    Cette stratégie « charger-en-creux, calculer-en-creux » garantit que le calcul sur puce se superpose à l'accès à la mémoire globale, maximisant l'utilisation de la HBM.

Principales Contributions

• Format de Stockage Novateur : Un format à trois couches adapté à la parcité non structurée modérée qui permet la compatibilité avec les cœurs de tenseur creux, réduit considérablement la surcharge de métadonnées via le codage PDD, et minimise les coûts de décodage sur puce.
• Noyau Co-optimisé : Un noyau SpMM qui coordonne efficacement les cœurs de tenseur creux et les cœurs CUDA, en utilisant le pipelinage pour superposer l'accès à la mémoire et le calcul.
• Premier à battre le Dense : L'implémentation est la première à démontrer qu'une SpMM de parcité modérément non structurée peut surpasser la multiplication de matrices denses sur les GPU modernes équipés de HBM.

Résultats
Les évaluations ont été menées sur des GPU NVIDIA H100 SXM5 avec les modèles OPT-30B et OPT-66B à des niveaux de parcité de 50 %, 60 % et 70 %.

• Performance au niveau du Noyau : À 50 % de parcité, la méthode proposée atteint une accélération allant jusqu'à 1,64× par rapport à SpInfer (Best Paper EuroSys'25) et surpasse significativement d'autres bases de référence comme SparTA et FlashLLM. Elle surpasse également la performance dense de cuBLAS.
• Inférence de bout en bout : La méthode délivre une accélération de bout en bout allant jusqu'à 1,41× par rapport à FlashLLM (VLDB'24) et un débit jusqu'à 41,4 % plus élevé que FlashLLM sur OPT-30B avec une taille de lot de 16.
• Efficacité Mémoire : L'approche réduit l'utilisation de la mémoire globale de pointe d'environ 21,4 % par rapport aux bases de référence denses à 50 % de parcité, bien qu'elle entraîne une utilisation de la mémoire légèrement supérieure à celle de SpInfer en raison de l'absence de codage par bitmap pour les métadonnées.
• Scalabilité : La méthode maintient ses avantages sur les niveaux de parcité de 50 % à 70 %.

Signification et Revendications
L'article affirme traiter une lacune critique dans l'optimisation de l'inférence des LLM : l'incapacité des noyaux creux existants à exploiter efficacement la parcité non structurée modérée. En introduisant un format de stockage qui comble le fossé entre les motifs d'élagage irréguliers et l'accélération matérielle structurée, les auteurs démontent que l'élagage peut produire de réelles accélérations d'inférence sans sacrifier la qualité du modèle.

Ce travail souligne que si une parcité élevée (>70 %) dégrade la qualité du modèle, une parcité modérée (~50 %) est le « point idéal » pour les LLM. La solution proposée débloque le potentiel de performance de ce point idéal, qui était auparavant inaccessible en raison des inefficacités des noyaux creux existants. Les auteurs notent que leur optimisation est spécifiquement ciblée sur la phase de décodage de l'inférence des LLM (où les tailles de lots sont faibles et limitées par la mémoire), ce qui la rend très adaptée à un déploiement évolutif dans des environnements de production où les phases de prefill et de décodage sont souvent découplées. Le code source est rendu disponible pour faciliter la recherche et l'adoption futures.