🤖 machine learning

Accelerating Sparse Transformer Inference on GPU

Ce papier présente STOF, un framework GPU qui accélère l'inférence de Transformers clairsemés en exploitant une modélisation analytique pour une cartographie efficace de l'attention multi-têtes et une stratégie de recherche en deux étapes pour optimiser dynamiquement la fusion d'opérateurs, réalisant des accélérations allant jusqu'à 1,6x et 1,4x respectivement dans le calcul de l'attention multi-têtes et l'inférence de bout en bout.

Auteurs originaux : Wenhao Dai, Haodong Deng, Mengfei Rong, Xinyu Yang, Hongyu Liu, Fangxin Liu, Hailong Yang, Qianwen Cao, Qingxiao Sun

Publié 2026-05-20

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Wenhao Dai, Haodong Deng, Mengfei Rong, Xinyu Yang, Hongyu Liu, Fangxin Liu, Hailong Yang, Qianwen Cao, Qingxiao Sun

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous essayez de lire une immense bibliothèque de livres (un Modèle de Langage de Grande Taille) pour répondre à une question. La bibliothèque est organisée en salles appelées Transformers, et dans chaque salle, il y a un bibliothécaire (le mécanisme d'Attention Multi-Têtes) qui doit parcourir des milliers de pages pour trouver les phrases spécifiques pertinentes pour votre question.

Le problème est que pour de nombreuses questions, la plupart des pages sont sans rapport. Le bibliothécaire perd du temps à feuilleter des pages vides ou des pages qui n'ont pas d'importance. C'est ici que la sparsité intervient : c'est comme coller des autocollants « Ne pas lire » sur les pages non pertinentes.

Cependant, les bibliothécaires actuels (les logiciels existants) sont mauvais pour utiliser ces autocollants. Ils continuent de passer devant les pages « Ne pas lire », ou ils se confondent lorsque les autocollants sont placés selon des motifs étranges et aléatoires. De plus, la bibliothèque a d'autres tâches (comme la synthèse ou la mise en forme) qui sont généralement effectuées séparément, ajoutant plus de temps de déplacement entre les tâches.

Voici STOF, un nouveau système proposé par les chercheurs. Considérez STOF comme un système de gestion de bibliothèque super efficace et intelligent, conçu spécifiquement pour ces bibliothèques « éparses ». Voici comment il fonctionne, décomposé en parties simples :

1. Le Bibliothécaire Intelligent (Kernels MHA Unifiés)

Les chercheurs ont réalisé que différents motifs d'autocollants « Ne pas lire » nécessitent des stratégies différentes.

Le Problème : Certains motifs sont des rangées nettes d'autocollants (comme une fenêtre glissante), tandis que d'autres sont dispersés aléatoirement (comme un billet de loterie). Les anciens systèmes tentaient d'utiliser une méthode « unique pour tous », ce qui était lent.
La Solution STOF : STOF agit comme un bibliothécaire intelligent qui choisit le meilleur outil pour le travail.
- Si les autocollants forment un petit groupe net, le bibliothécaire utilise une approche « par ligne » : il saisit une rangée entière de livres d'un coup et la parcourt rapidement.
- Si les autocollants sont dispersés ou si la bibliothèque est immense, ils utilisent une approche « par blocs » : ils divisent les livres en petits morceaux gérables et n'ouvrent que les morceaux spécifiques qui ont des autocollants valides.
Le Résultat : En sautant complètement les pages « Ne pas lire » au lieu de simplement les ignorer, le bibliothécaire travaille beaucoup plus vite.

2. La Chaîne de Montage (Fusion d'Opérateurs)

Dans une bibliothèque normale, le bibliothécaire peut finir de lire, puis marcher vers un autre bureau pour résumer le texte, puis marcher vers un autre bureau pour mettre en forme la réponse. Cette marche (déplacement des données entre la mémoire et le processeur) est lente.

Le Problème : Les systèmes actuels ne combinent souvent que des tâches simples. Ils laissent le travail lourd (comme les mathématiques complexes) pour des étapes séparées, provoquant des embouteillages.
La Solution STOF : STOF construit une chaîne de montage personnalisée. Il examine l'ensemble du processus et demande : « Pouvons-nous combiner ces étapes ? »
- Il ne se contente pas de coller deux tâches simples ensemble ; il trouve le moyen parfait de combiner des tâches mathématiques complexes avec des tâches de mise en forme.
- Il utilise un « moteur de recherche » pour tester différentes façons de combiner ces tâches (comme essayer différents agencements de chaînes de montage) afin de trouver celle qui se déplace le plus rapidement pour la taille spécifique de la bibliothèque que vous lisez.

3. Le Pilote Automatique (Recherche Hiérarchique)

On ne peut pas concevoir manuellement la chaîne de montage parfaite pour chaque taille de livre et chaque type de question ; il y a trop de combinaisons.

La Solution STOF : STOF possède un Pilote Automatique qui apprend en temps réel.
- Phase 1 (La Carte) : Il examine la structure de la bibliothèque et dessine une carte approximative de l'emplacement des autocollants « Ne pas lire ».
- Phase 2 (L'Optimisation) : Il exécute une recherche en deux étapes. D'abord, il élargit les limites de la chaîne de montage pour voir jusqu'où il peut aller. Ensuite, il affine la vitesse des travailleurs (paramètres) en fonction de la réussite des tentatives précédentes.
- Il se souvient de ce qui a fonctionné (mise en cache) pour ne pas perdre de temps à retester les mêmes idées lentes.

Les Résultats : Combien de fois plus rapide ?

Les chercheurs ont testé STOF sur des cartes graphiques puissantes (GPU) en utilisant des modèles d'IA populaires (comme BERT, GPT et LLaMA).

Vitesse : Par rapport aux meilleures méthodes existantes, STOF a rendu la tâche de lecture principale (MHA) jusqu'à 1,6 fois plus rapide.
Vitesse Globale : En regardant l'ensemble du processus de réponse à une question (de bout en bout), il était jusqu'à 1,4 fois plus rapide.
Grandes Bibliothèques : Plus la bibliothèque est grande (séquences de texte plus longues), plus STOF brillait, car il sautait tant de travail inutile.

Résumé

Pensez à STOF comme à un système qui empêche l'IA de perdre du temps à lire des pages qu'elle n'a pas besoin de lire, et qui l'empêche de faire des allers-retours entre les bureaux. Il utilise une stratégie intelligente et adaptative pour sauter les déchets et combiner les étapes utiles en un seul mouvement fluide et rapide. Cela permet aux modèles d'IA de fonctionner significativement plus vite, en particulier lorsqu'ils traitent des textes longs ou complexes.

Résumé technique : Accélération de l'inférence des Transformers clairsemés sur GPU

Énoncé du problème

Les grands modèles de langage (LLM) reposent fortement sur l'architecture Transformer, où l'attention multi-têtes (MHA) constitue le composant de calcul central. Bien que les couches de masques induisant de la parcimonie (par exemple, causales, fenêtre glissante, aléatoires) soient largement utilisées pour réduire le volume de calcul, elles introduisent des défis de performance significatifs sur les GPU :

Calcul fragmenté : Les masques clairsemés créent des motifs d'accès aux données fragmentés, exacerbant les goulots d'étranglement de la bande passante mémoire et empêchant une utilisation efficace des ressources GPU.
Optimisation rigide : Les implémentations MHA haute performance existantes (par exemple, FlashAttention) ne prennent souvent en charge que des motifs de masquage spécifiques et continus (comme causaux ou fenêtre glissante). Elles peinent avec la parcimonie discrète ou non structurée (par exemple, Bigbird, attention aléatoire) ou recourent à des méthodes inefficaces telles que la réinitialisation des matrices de scores après le calcul.
Fusion d'opérateurs statique : Les schémas actuels de fusion d'opérateurs dans les frameworks d'apprentissage profond sont largement statiques ou pilotés par des règles. Ils échouent à s'adapter à des scénarios d'application divers, à des longueurs de séquence variables et à différents hyperparamètres de modèle. De plus, la configuration optimale pour les opérateurs fusionnés diffère souvent considérablement de la somme des opérateurs ajustés individuellement, rendant l'ajustement séquentiel sous-optimal.

Méthodologie : Le cadre STOF

Les auteurs proposent STOF, un cadre conçu pour accélérer l'inférence des Transformers clairsemés sur GPU en intégrant un support de masquage flexible avec une fusion d'opérateurs adaptative. STOF se compose de deux modules principaux :

1. Module MHA unifié

Ce module répond au défi du calcul efficace de l'attention clairsemée à travers divers motifs de masquage.

Format de stockage à deux niveaux : Pour gérer des motifs de masques arbitraires (à la fois structurés et non structurés), STOF introduit un nouveau format de stockage combinant Row Compressed Sparse Block (BSR) et bitmap.
- Tuiles extérieures (OT) : Représentent les blocs sautés globaux. Les OT « Pleins » (tous les éléments valides) et les OT « Partiels » (mélange d'éléments valides/invalides) sont suivis via des tableaux de pointeurs.
- Tuiles intérieures (IT) : Représentent la distribution des éléments au sein d'un bloc. Les éléments valides au sein des OT « Partiels » sont encodés à l'aide de bitmap 64 bits (un uint64 par bloc 8x8), permettant un masquage précis sans matérialiser la matrice clairsemée complète.
Stratégie de noyau dual : STOF implémente deux types de noyaux distincts basés sur un modèle analytique :
- Noyau ligne par ligne : Optimisé pour les scénarios hautement clairsemés avec des séquences courtes. Il découpe la requête ( $Q$ ) en lignes, utilise des mélanges au niveau du warp (warp-level shuffles) et élimine la synchronisation inter-warp pour maximiser la localité des données.
- Noyau par bloc : Un noyau polyvalent pour des scénarios plus larges. Il partitionne $Q$ , $K$ et $V$ en sous-blocs, exploitant la mémoire partagée (SMEM) et la copie de données asynchrone pour chevaucher les transferts mémoire avec le calcul.
Sélection de noyau : Un modèle analytique détermine s'il faut utiliser le noyau ligne par ligne ou le noyau par bloc en fonction de la parcimonie du masque et de la longueur de la séquence, garantissant le chemin d'exécution le plus efficace.

2. Module de fusion d'opérateurs

Ce module répond au besoin d'une fusion adaptative au-delà de la couche MHA.

Conversion de schéma de fusion : STOF convertit les schémas de fusion en une représentation numérique utilisant un encodage par hachage. Les opérateurs adjacents dans le graphe de calcul sont mappés vers des tableaux binaires (représentant les limites de fusion), qui sont ensuite décodés en modèles de compilation spécifiques (implémentés via Triton et TileLang).
Moteur de recherche hiérarchique : Au lieu de règles statiques, STOF emploie un mécanisme de recherche en deux étapes pour déterminer les limites de fusion optimales et les paramètres de noyau :
1. Initialisation : Utilise le hachage neuronal et des règles prédéfinies pour identifier les motifs de sous-graphes fréquents et établir un schéma de fusion initial.
2. Ajustement en deux étapes :
  - Étape 1 (Expansion de fusion) : Utilise un algorithme basé sur la récompense (expansion, saisie, compétition) pour étendre itérativement les limites de fusion, testant les gains de performance par rapport à un cache de résultats précédents.
  - Étape 2 (Échantillonnage de paramètres) : Une fois les limites fixées, il échantillonne les paramètres de noyau (par exemple, tailles de blocs, étapes de pipeline) pour trouver la configuration optimale pour le schéma de fusion spécifique.

Contributions clés

Analyse complète : Les auteurs analysent l'impact de divers motifs de masquage et configurations d'inférence, exposant des opportunités d'optimisation tant dans le calcul MHA que dans la fusion d'opérateurs en aval.
Module MHA unifié : Un module novel implémentant des noyaux ligne par ligne et par bloc avec des formats de stockage uniques (BSR + bitmap) et des optimisations, régit par un modèle de sélection analytique.
Fusion d'opérateurs adaptative : Un module qui traduit les schémas de fusion en modèles de compilation via un décodage numérique et utilise un moteur de recherche en deux étapes pour déterminer dynamiquement les paramètres optimaux.
Cadre STOF : Un système d'inférence complet permettant des motifs de masquage flexibles et une fusion d'opérateurs optimale sur GPU.

Résultats expérimentaux

Les auteurs ont évalué STOF sur des GPU NVIDIA RTX 4090 (Ada) et A100 (Ampere) par rapport à des références de l'état de l'art incluant PyTorch Native, FlashAttention-2 (FA2), FlexAttention, ByteTransformer, MCFuser et Bolt.

Performance MHA : STOF a atteint des accélérations maximales de 1,6× par rapport aux meilleures implémentations MHA existantes (par exemple, FlexAttention) dans le calcul MHA. Il a démontré une supériorité constante sur divers motifs de masquage (causal, fenêtre glissante, Longformer, Bigbird) et longueurs de séquence (jusqu'à 4 096).
Inférence de bout en bout : STOF a atteint des accélérations maximales de 1,4× dans l'inférence de bout en bout sur des modèles tels que BERT, GPT-2, LLaMA, T5 et ViT.
Efficacité de l'ajustement : Le processus d'ajustement de STOF était significativement plus rapide que celui des concurrents (MCFuser et Bolt), étant 6,7 à 6,9 fois plus rapide dans les scénarios à grande échelle grâce à ses mécanismes d'échantillonnage basés sur la récompense et de mise en cache.
Évolutivité : Dans les expériences avec des séquences ultra-longues (jusqu'à 16k), STOF a montré des gains de performance significatifs (jusqu'à 16,8× par rapport à PyTorch Compile) et une meilleure efficacité mémoire, évitant les erreurs de mémoire insuffisante (OOM) à des longueurs de séquence où les références échouaient.

Importance et revendications

L'article revendique que STOF est le premier système à permettre simultanément des motifs de masquage flexibles et des schémas de fusion d'opérateurs divers pour les scénarios de Transformers clairsemés. Son importance réside dans le comblement du fossé entre les noyaux spécifiques à la parcimonie, ajustés manuellement, et la fusion d'opérateurs denses automatique. En fournissant une pile complète intégrant des noyaux MHA personnalisés avec des modèles de compilation génératifs, STOF établit des opportunités d'optimisation plus larges qui étaient auparavant inaccessibles avec des approches statiques ou pilotées par des règles. Les auteurs soulignent que leur approche est extensible aux architectures émergentes de LLM (par exemple, Mixture-of-Experts) et aux nouvelles générations de GPU, offrant une voie vers des améliorations de performance durables dans l'inférence de modèles à grande échelle.