Le 1er août, HyperAI a organisé à Pékin le 9e Salon technique Meet AI Compiler. Cet événement a mis l'accent sur les dernières avancées en matière de compilation d'IA. Plusieurs experts issus de l'industrie et des institutions de recherche ont partagé leurs connaissances sur les langages de programmation, le développement d'opérateurs, l'optimisation de la compilation et l'exécution de l'inférence, illustrant ainsi l'évolution collaborative des compilateurs d'IA, de l'expression de haut niveau à l'exécution matérielle.

dans,Zhu Honglin, ingénieur en calcul haute performance chez Tencent, a présenté les pratiques d'optimisation des algorithmes et des opérateurs de son équipe pour la multiplication matricielle à faible complexité lors d'une présentation intitulée « FalconGEMM : Dépasser les performances matérielles grâce à la multiplication matricielle à faible complexité ».

Face au défi que représentent les bibliothèques d'opérateurs matures telles que cuBLAS, qui ont poussé les performances de la multiplication matricielle à des niveaux proches des capacités matérielles maximales, et face à la réduction croissante de l'espace d'optimisation traditionnel au niveau du noyau, l'équipe est repartie de la complexité algorithmique. S'appuyant sur des algorithmes de multiplication matricielle à faible complexité comme Strassen et AlphaTensor, elle a construit un cadre LCMA (Low-Complexity Matrix Algorithms) unifié. En combinant QDSL, la fusion d'opérateurs, le noyau persistant, l'ordonnancement fin et le modèle de coût, elle a transformé l'avantage théorique de la réduction du nombre de multiplications en gains de performance concrets sur GPU.

Dans les tests FP16 et BF16 de NVIDIA H20, FalconGEMM surpasse cuBLAS sur un grand nombre de formes de matrices, avec des améliorations de performances maximales d'environ 10%-16%, tout en maintenant une précision numérique qui est fondamentalement cohérente avec la multiplication matricielle standard dans le benchmark du modèle de langage.

HyperAI a compilé et résumé le contenu partagé sans en altérer le sens original.

Suivez le compte officiel WeChat « HyperAI » et répondez avec le mot-clé « » en arrière-plan.Compilateur IA 0801Vous pouvez obtenir la présentation PPT du conférencier autorisé en cliquant sur "...".

En partant de Strassen, nous réexaminons l'espace d'optimisation pour la multiplication matricielle.

La multiplication matricielle est l'un des opérateurs fondamentaux les plus importants en apprentissage profond, et elle représente généralement la majeure partie du temps de calcul du modèle. Les piles logicielles telles que CUDA, MKL et cuBLAS sont optimisées depuis de nombreuses années et, dans bien des cas, les performances d'un seul noyau GEMM sont déjà très proches des performances maximales du matériel. Cela signifie que si l'on continue à effectuer des optimisations locales uniquement au niveau des instructions, du pipeline et des accès mémoire, la marge de progression devient de plus en plus limitée. C'est pourquoi l'équipe s'est de nouveau intéressée à l'algorithme classique de Strassen.

L'algorithme de Strassen a été proposé par Volker Strassen en 1969. Pour la multiplication matricielle 2×2 la plus simple, la méthode traditionnelle nécessite 8 multiplications, tandis que l'algorithme de Strassen, en recombinant la matrice d'entrée, n'en effectue que 7, puis utilise des additions et des soustractions supplémentaires pour obtenir le résultat final, ce qui équivaut à réduire d'un facteur 8 le nombre de multiplications. Si les opérandes sont de simples scalaires, ce compromis est négligeable ; cependant, lorsque les opérandes deviennent des sous-matrices, la différence de complexité entre l'addition matricielle (O(N²)) et la multiplication matricielle (O(N³)) devient significative.Cela donne une valeur pratique à l'opération « faire une multiplication matricielle de moins et plusieurs additions matricielles de plus ».

L'utilisation récursive de l'algorithme de Strassen permet de réduire davantage le nombre de multiplications. Par exemple, la multiplication d'une matrice par blocs 4×4 nécessite généralement 64 multiplications par blocs, tandis que deux couches de Strassen n'en requièrent que 49. Cependant, l'augmentation du nombre de couches récursives engendre également une surcharge liée aux additions, à l'organisation des données et aux accès mémoire.Par conséquent, les systèmes pratiques n'utilisent souvent qu'un nombre limité de couches, cherchant un équilibre entre la réduction des calculs et la surcharge supplémentaire.

En 2022, AlphaTensor de DeepMind a encore élargi cet espace algorithmique. Il transforme la multiplication matricielle en un problème de décomposition tensorielle et utilise l'apprentissage par renforcement pour rechercher des méthodes de décomposition de rang inférieur, démontrant ainsi qu'en plus de la méthode classique de Strassen, il peut exister un grand nombre d'algorithmes de multiplication matricielle à faible complexité pour différentes valeurs de M, N et K.

Il subsiste toutefois un problème pratique lors du passage de la découverte d'algorithmes à leurs applications d'ingénierie :Si chaque algorithme à faible complexité nécessite un noyau GPU écrit à la main distinct, les coûts de développement et de maintenance seraient évidemment trop élevés.Pour résoudre ce problème, l'équipe a abstrait ces algorithmes en LCMA (Low-Complexity Matrix Algorithms), qui décrit de manière uniforme quels sous-blocs de la matrice d'entrée doivent être pré-combinés, combien de multiplications matricielles sont effectivement effectuées et comment les résultats intermédiaires sont finalement combinés dans la matrice de sortie. L'implémentation correspondante est ensuite générée automatiquement à l'aide de Codegen.

La question qui se pose est donc celle de « comment implémenter un noyau de Strassen ».L'attention s'est portée sur « la manière de construire un cadre unifié capable de prendre en charge une variété d'algorithmes matriciels à faible complexité tout en maintenant des performances élevées ».

Parallèlement, les algorithmes à faible complexité doivent également prendre en compte la précision numérique. Bien que la multiplication de Strassen soit algébriquement équivalente à la multiplication matricielle standard, les opérations en virgule flottante ne respectent pas strictement l'associativité, et toute modification de l'ordre de calcul peut engendrer des erreurs d'arrondi supplémentaires. Par conséquent, l'algorithme LCMA, tout en optimisant ses performances, doit également maîtriser la propagation des erreurs dans les calculs à faible précision.

Une fois la description unifiée de l'algorithme établie, l'étape suivante consistait à trouver une implémentation GPU adaptée. L'équipe a testé CUDA, Triton, TiLang et QDSL. CUDA offrait le meilleur contrôle matériel, mais face à un grand nombre d'algorithmes LCMA différents, les registres, la mémoire partagée et les structures de sommation intermédiaires nécessitaient tous des ajustements spécifiques, entraînant des coûts d'extension et de maintenance élevés.

Triton atteint des performances proches de celles de CUDA dans les scénarios Strassen de base, mais lorsque l'algorithme passe à des structures de blocs plus importantes, il nécessite une réutilisation précise des tampons de registres pour plusieurs calculs intermédiaires, ce qui le rend vulnérable aux débordements de mémoire. TiLang offre une plus grande flexibilité dans la gestion des registres et de la mémoire partagée, mais ses performances lors des tests en équipe restent inférieures d'environ 51 à 101 TP3T à celles de Triton. Pour le Strassen de base, avec un gain théorique de seulement 12,51 TP3T, cette perte est suffisamment significative pour annuler les gains de l'algorithme.

final,L'équipe a choisi QDSL comme principal moteur d'implémentation pour FalconGEMM. QDSL offre une granularité de développement proche de celle de CUDA, tout en intégrant des capacités de génération de code et la prise en charge de PTX embarqué. Ceci facilite la migration des implémentations hautes performances existantes et convient à la génération de code par lots à partir de différentes descriptions LCMA, offrant ainsi une plus grande flexibilité pour l'intégration ultérieure et l'optimisation personnalisée.

De LCMA à FalconGEMM, transfert des gains algorithmiques vers les GPU

L'implémentation GPU la plus simple de l'algorithme de Strassen peut être divisée en plusieurs étapes : la combinaison des sous-matrices A et B pour générer 7 nouvelles paires d'entrées ; l'exécution de 7 GEMM par lots ; et enfin la combinaison des 7 ensembles de résultats intermédiaires dans la matrice finale C. Comparée à un GEMM classique, la partie de multiplication matricielle véritablement intensive en calcul ne représente que 7/8 de l'originale.Par conséquent, tant que le temps supplémentaire consacré au prétraitement et au post-traitement est inférieur à 1/8 de la charge de calcul économisée, il existe une possibilité globale d'en tirer des avantages.

L'équipe a d'abord effectué des tests sur la carte graphique NVIDIA H20. Grâce à sa large bande passante mémoire et à sa faible charge de calcul de pointe, la H20 est parfaitement adaptée à cette approche qui consiste à « augmenter le traitement de certaines données en échange d'une réduction de la charge de calcul ». Pour des matrices de taille 2048³ et plus, l'implémentation de base a déjà démontré des gains stables. Cependant, pour des matrices plus petites, la proportion de combinaisons d'entrées, d'écritures de résultats intermédiaires et de combinaisons de sorties augmente rapidement, consommant facilement la charge de calcul économisée.

Par conséquent, l'optimisation ultérieure s'est concentrée non plus sur GEMM lui-même, mais sur l'accès à la mémoire intermédiaire. L'approche la plus directe consiste en la fusion d'opérateurs.Essayez de conserver autant que possible les résultats intermédiaires sur la puce, plutôt que de les réécrire sans cesse dans la mémoire globale.Cependant, l'intégration directe des entrées A et B dans GEMM n'est pas appropriée car un même sous-bloc peut être utilisé par plusieurs SM, ce qui risque d'entraîner des chargements et des sommations redondants. En revanche, la fusion en post-traitement de Batched GEMM et de Combine H est plus pertinente.

La véritable difficulté réside dans le fait que les sept résultats intermédiaires de Strassen contribuent de différentes manières aux quatre sous-matrices de sortie finales. Si H est utilisé comme unité parallèle, plusieurs SM peuvent simultanément écrire dans la même C, ce qui entraîne de graves conflits atomiques ; si C est utilisé comme unité parallèle, certaines valeurs de H seront calculées à plusieurs reprises par différents SM. Ces deux approches annulent les avantages de la réduction de la multiplication.

L'équipe a finalement abandonné l'organisation des tâches en fonction des résultats intermédiaires de Strassen. Elle a plutôt opté pour un regroupement des tâches selon les coordonnées spatiales de la matrice : sept tuiles de multiplication situées à la même position dans sept GEMM par lots ont été regroupées et exécutées sur le même SM. Ainsi, une fois les calculs d'un groupe terminés, les résultats pouvaient être directement accumulés dans le code C final sur la puce, évitant ainsi la nécessité de réécrire les résultats intermédiaires dans la mémoire globale et les conflits d'écriture inter-SM importants.

Cette approche de fusion réduit considérablement les accès mémoire supplémentaires introduits par Strassen, mais la granularité plus fine des groupes entraîne un déséquilibre de charge. Par exemple, avec une multiplication matricielle de 4096³, une planification à gros grains peut engendrer un gaspillage de ressources d'environ 211 TP3T, dépassant même la réduction de calcul de 12,51 TP3T obtenue avec Strassen.

à cette fin,L'équipe s'est inspirée de Stream-K pour diviser un groupe en deux SM pour l'exécution, si nécessaire.La couche de planification utilise toujours le groupe comme unité de base, mais l'exécution réelle peut être affinée au niveau de la tuile, réduisant ainsi le temps d'inactivité du SM de queue et améliorant l'utilisation du matériel tout en conservant les avantages de la réutilisation des données au niveau du groupe.

Cependant, après la résolution du problème d'équilibrage de charge, un nouveau souci est apparu : la saturation du cache L2. Suite à la séparation des groupes, différents types de multiplications intermédiaires pouvaient être mélangés au sein d'une même vague, et les données accédées étaient indépendantes, ce qui entraînait une baisse significative du taux d'accès au cache L2. Parallèlement, GEMM utilisait déjà intensivement les cœurs Tensor, et lorsque la pression sur l'accès à la mémoire approchait sa pleine capacité, le H20 atteignait sa limite de puissance. Lors des tests, la fréquence du cœur a chuté d'environ 1,8 GHz à 1,6 GHz, ce qui a entraîné une baisse des performances de calcul et une diminution des avantages de la fusion.

Pour remédier à la saturation du cache L2, l'équipe a ajusté l'ordre de répartition des groupes, en veillant à ce que les résultats intermédiaires de même type soient traités au sein d'une même vague, le traitement mixte n'intervenant que dans quelques vagues finales. Ceci a permis de préserver l'équilibrage de charge assuré par une planification fine tout en restaurant une meilleure localité des données L2, éliminant ainsi le problème de réduction significative de la fréquence.

Il convient de mentionner queLe noyau persistant constitue l'élément fondamental de ces optimisations d'ordonnancement.Contrairement à un noyau classique où le gestionnaire de tâches (CTA) se termine après l'exécution d'un bloc, un noyau persistant permet au CTA de résider sur le service de flux (SM) pendant une période prolongée, acceptant en continu les tâches suivantes. Ceci offre aux développeurs une plus grande flexibilité dans le contrôle de l'ordre d'exécution des groupes et des tuiles, et permet la réutilisation des ressources sur la puce. Le fractionnement des tâches, le réarrangement de l'ordonnancement et l'optimisation du cache peuvent ainsi être effectués au sein du même noyau.

Du réordonnancement du cache au modèle de coût, l'amélioration maximale des performances est de 10% à 16%.

Après fusion, équilibrage de charge et réorganisation du cache, FalconGEMM peut exploiter les avantages de calcul de LCMA sur un plus grand nombre de formes. Cependant, LCMA n'est pas supérieur à GEMM classique dans tous les cas.Son principe reste le même : privilégier le traitement de données supplémentaire au détriment de la réduction des calculs de multiplication.Si l'algorithme GEMM d'origine est déjà limité par l'accès à la mémoire, une réduction supplémentaire des calculs n'apportera pas d'avantages suffisants ; les algorithmes à faible complexité ne sont plus avantageux que lorsque la densité de calcul est élevée.

donc,L'équipe a également conçu un modèle de coûts similaire à Roofline pour déterminer quand utiliser l'analyse du coût du cycle de vie (LCMA) et laquelle des différentes LCMA choisir.L'objectif étant de « choisir le bon algorithme » plutôt que de prédire avec précision le temps d'exécution, le modèle analyse principalement les coûts de calcul et d'accès à la mémoire des différents schémas, et estime la plage de goulot d'étranglement en matière de calcul/accès à la mémoire du GPU cible en combinant sa puissance de calcul et sa bande passante.

Dans ce modèle, la réduction du nombre de multiplications dans l'algorithme à faible complexité se traduit par des gains de calcul, tandis que les combinaisons de données supplémentaires et les accès mémoire répétés engendrent une surcharge mémoire ; l'optimisation de fusion mentionnée précédemment réduit encore ce coût d'accès mémoire. Par conséquent, FalconGEMM peut déterminer les limites d'efficacité de la méthode GEMM conventionnelle et de différents schémas LCMA en fonction des valeurs de M, N et K, et sélectionner automatiquement l'implémentation la plus adaptée.

En tirant parti des capacités de génération de code de QDSL, l'ensemble du système forme un flux d'exécution relativement complet : tout d'abord, un noyau persistant fusionné est généré à partir de la description LCMA, réduisant ainsi les accès mémoire intermédiaires grâce à la fusion ; ensuite, le modèle de coût sélectionne un algorithme de multiplication matricielle adapté à la forme spécifique ; enfin, QDSL génère et compile automatiquement le code cible. De cette manière, LCMA n'est plus une simple implémentation de Strassen, mais constitue un espace d'algorithmes pouvant être sélectionné dynamiquement en fonction de la charge de travail.

Les tests de performance ont été principalement réalisés sur la carte graphique NVIDIA H20. Les résultats montrent que dans divers scénarios de multiplication matricielle à faible précision,FalconGEMM surpasse cuBLAS sur un grand nombre de formes, avec des améliorations de performance maximales d'environ 10%–16%.Après avoir effectué la division du groupe et le réarrangement du cache, la grande forme peut obtenir de manière stable les avantages de calcul apportés par l'algorithme à faible complexité, les performances de la petite forme sont également améliorées et la réduction de fréquence causée par le mur de puissance déclenché par la gigue du cache L2 est évitée.

Sur la base des résultats de la sélection du modèle de coût, le modèle a pu sélectionner l'implémentation la plus performante pour la plupart des formes qui franchissent le seuil de retour LCMA, indiquant que l'approche « sélection d'algorithme + noyau de fusion » peut couvrir efficacement les scénarios de multiplication matricielle avec différentes densités de calcul.

Outre les performances, la précision numérique est également un point crucial que FalconGEMM doit vérifier. L'équipe a rencontré des erreurs importantes lors des premières expériences à faible précision, principalement dues à des changements dans l'ordre des additions en virgule flottante. Par exemple, A+C+B−C est algébriquement égal à A+B, mais dans les calculs en virgule flottante à précision finie, ces deux termes ne sont pas nécessairement strictement égaux.

Une analyse plus poussée a révélé que les erreurs réellement significatives provenaient principalement de la faible précision des tirages, et non de l'accumulation FP32 elle-même. Dans les scénarios d'entrée FP16/BF16 courants, la multiplication matricielle est généralement effectuée en premier lieu à l'aide de l'accumulation FP32, puis reconvertie en une précision inférieure ; si des résultats intermédiaires sont fréquemment tirés, les informations de la mantisse FP32 sont continuellement ignorées.

L'approche par fusion atténue ce problème. La sortie de WGMMA conserve la précision FP32, tandis que FalconGEMM combine et accumule directement les valeurs C finales sur la puce en FP32, avant de les reconvertir à la précision cible une fois le calcul terminé. Comparée à la réécriture répétée de résultats intermédiaires de faible précision entre plusieurs noyaux indépendants, cette méthode réduit d'une ou plusieurs conversions de précision, ce qui permet aux erreurs dues aux changements d'ordre de calcul de rester davantage dans les bits de poids faible de FP32.

Dans le test de modélisation du langage, les scores finaux obtenus à l'aide de FalconGEMM et de la multiplication matricielle standard sont presque identiques, avec seulement de très légères différences, ce qui indique que l'implémentation actuelle n'entraîne pas de diminution significative de la précision du modèle.

Dans la prochaine phase, l'équipe prévoit de poursuivre ses efforts dans deux directions : premièrement, intégrer davantage la combinaison A/B en ajustant l'ordre des groupes de traitement par lots et des boucles K afin de réduire encore les accès mémoire pour les résultats d'entrée intermédiaires ; deuxièmement, étendre LCMA à l'attention. L'attention Flash présente également un rapport calcul/mémoire élevé et, si des algorithmes matriciels à faible complexité peuvent être combinés à sa décomposition par blocs et à son pipeline, elle pourrait également offrir un potentiel de performance accru.

De Strassen et AlphaTensor à LCMA et FalconGEMM, l'importance de ces travaux dépasse la simple amélioration de quelques points de pourcentage des performances déjà hautement optimisées de GEMM. Ils proposent une autre approche : lorsque le noyau lui-même approche les limites matérielles, l'optimisation des performances peut non seulement s'étendre aux instructions et aux pipelines, mais aussi explorer de nouvelles pistes d'amélioration de la complexité algorithmique. Ensuite, grâce à la compilation, la fusion et l'ordonnancement, la réduction théorique du temps de calcul peut se traduire concrètement par des gains de performance à l'exécution.