Le 1er août, le 9e Salon technique Meet AI Compiler s'est tenu à Pékin. Cet événement était consacré aux 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 points de vue 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 en langage de haut niveau à l'exécution matérielle.

dans,Guo Hui et Xiao Hang, chercheurs chez BAAI AI Compiler, ont présenté sur « FlagTree : extension de langage Triton-TLE, backend IR Tile et pratiques d'optimisation du compilateur », introduisant une série d'explorations menées par l'équipe FlagTree autour de Triton.

Face à des architectures matérielles de plus en plus complexes et à des opérateurs de modèles plus irréguliers, l'équipe a proposé l'extension de langage Triton (TLE) pour tenter d'établir un canal progressif entre le DSL de haut niveau et le contrôle matériel de bas niveau ; dans le même temps, ils ont effectué des optimisations de compilation sur le compilateur FlagTree, en se concentrant sur le backend TileIR, le réglage automatique, la disposition des données, la planification des instructions et d'autres aspects.

L'objectif principal de ce travail est de permettre aux développeurs d'approfondir les détails matériels selon leurs besoins, tout en préservant la facilité d'utilisation et l'écosystème communautaire de Triton, et de couvrir l'optimisation rapide des opérateurs, le réglage en fonction de l'architecture et l'optimisation au niveau du code natif avec un système de développement relativement unifié.

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 "...".

S’appuyant sur Triton, rééquilibrer abstraction et performance.

Ces dernières années, les compilateurs d'IA ont été confrontés à une contradiction de plus en plus marquée :Les architectures matérielles et les opérateurs de modèles deviennent de plus en plus complexes, mais les développeurs espèrent toujours utiliser des DSL de plus haut niveau pour atteindre des performances proches de celles des noyaux écrits par des experts.

Si le niveau d'abstraction est trop élevé, le compilateur risque de ne pas obtenir suffisamment d'informations pour effectuer une optimisation approfondie ; si le niveau d'abstraction est trop bas, les développeurs reviendront au mode de développement complexe de CUDA ou des langages propriétaires du fournisseur.

Le succès de Triton réside dans le passage du développement des opérateurs GPU du niveau du thread au niveau de la tuile.Les utilisateurs utilisent le DSL Python pour décrire les relations de calcul entre les blocs de données, tandis que des tâches telles que le mappage des threads, l'allocation des registres, la disposition des données, le pipeline et la synchronisation sont principalement gérées par le compilateur.

Cette approche facilite le développement d'opérateurs performants et a également favorisé l'émergence d'un vaste écosystème communautaire. Cependant, avec le développement continu des GPU de nouvelle génération, des architectures dédiées et des puces d'IA produites localement, l'abstraction initiale de Triton commence à montrer ses limites.

d'une part,Si le moteur de compilation ne prend pas encore en charge la structure de stockage, le mécanisme de communication ou l'unité de calcul du nouveau matériel, il sera difficile pour les développeurs frontaux d'utiliser ces fonctionnalités par eux-mêmes.d'autre part,De plus en plus d'opérateurs critiques exigent un contrôle précis des niveaux de stockage, de la granularité parallèle, de la collaboration CTA et du chevauchement de la communication et du calcul, et le code Triton original a parfois du mal à exprimer pleinement ces intentions.

L'émergence de nouveaux langages et DSL tels que Gluon, TLX et TileLang reflète la même tendance : le développement d'opérateurs d'IA ne se limite plus à l'écriture d'un seul noyau, mais englobe également l'expression de la disposition des données, de la hiérarchie parallèle, du pipeline, de la topologie de communication et des caractéristiques matérielles.

TLE ne cherche pas à remplacer Triton, mais plutôt à étendre sa syntaxe et son écosystème par couches.Il se compose de trois couches : TLE-Lite, TLE-Struct et TLE-Raw, qui correspondent respectivement à des indications sémantiques légères, à un contrôle prenant en compte l'architecture et à une optimisation au niveau du code natif.

TLE-Lite est conçu pour les ingénieurs en algorithmes et les scénarios d'optimisation rapide.Les développeurs n'ont pas besoin de se préoccuper directement du matériel sous-jacent ; ils peuvent plutôt compléter le compilateur avec des informations structurelles plus explicites.Par exemple, un tenseur peut nécessiter un accès par sous-tuiles, un calcul peut s'exécuter sur un maillage distribué, ou une application de traitement d'images peut utiliser un pipeline producteur-consommateur. Prenons l'exemple des opérations sur les sous-tuiles : les développeurs peuvent extraire directement un sous-bloc logique d'un tenseur plus grand, effectuer des opérations d'activation, de normalisation ou des calculs statistiques, puis le réécrire, sans avoir à calculer manuellement les décalages, à construire des masques ni à gérer les frontières.

Le compilateur, reconnaissant qu'il s'agit d'un accès régulier à une tuile, peut optimiser davantage l'organisation des données, la vectorisation, la gestion des conflits de banques et la réutilisation des registres. Cette approche est particulièrement adaptée à l'attention parcimonieuse, à la normalisation locale, aux statistiques de blocs et aux opérateurs de routage. Dans les environnements distribués, TLE utilise Device Mesh pour décrire différents niveaux (nœuds, GPU, clusters de blocs et blocs) et les organiser en une topologie multidimensionnelle unifiée.

Les développeurs conçoivent des programmes selon une approche basée sur un maillage, et le compilateur et l'environnement d'exécution établissent ensuite une correspondance entre la topologie logique et le matériel et les mécanismes de communication réels. Ainsi, la communication en anneau, la synchronisation des barrières et l'accès fragmenté ne se limitent plus à des groupes de communication et des rangs dispersés dans le code, mais deviennent des sémantiques structurées et analysables. Une fois les relations de communication explicitement exprimées, le compilateur peut effectuer une planification prenant en compte la topologie, la gestion du chevauchement communication-calcul, la fusion des barrières et la détection des interblocages.

TLE-Lite abstrait également la collaboration interne du CTA en un modèle producteur-consommateur via des primitives de pipeline. Les développeurs décrivent principalement qui produit les données et qui les consomme, tandis que les mécanismes sous-jacents de barrière, de réutilisation des tampons et de synchronisation sont gérés par le compilateur.Cela ne signifie pas masquer les capacités sous-jacentes, mais plutôt les transformer en une structure de programme qui puisse être analysée, vérifiée et optimisée.

De la sémantique légère au transfert natif, couvrant différents niveaux de profondeur d'optimisation.

Si TLE-Lite traite principalement de la représentation sémantique multiplateforme,TLE-Struct, en revanche, se concentre davantage sur la sensibilisation à l'architecture et le peaufinage.

Les différents GPU, DSA et accélérateurs d'IA diffèrent considérablement au niveau de leurs niveaux de stockage, de leurs unités d'exécution, de leurs mécanismes de synchronisation et de leurs réseaux sur puce. TLE-Struct expose donc aux développeurs une structure hiérarchique parallèle et de stockage, leur permettant de définir explicitement la disposition des données, le mappage des calculs et la hiérarchie de la mémoire, sans avoir à se lier directement à l'interface propriétaire d'un fournisseur.

Par exemple, une même mémoire tampon locale peut être mappée sur la mémoire partagée du GPU et sur le bloc mémoire temporaire ou la SRAM intégrée du DSA. L'utilisateur définit une intention d'accès à la mémoire structurée, et le compilateur se charge de la traduire en un espace d'adressage et des instructions d'accès à la mémoire adaptés au matériel cible.

Le comptage des experts dans MoE est un scénario typique. Cet opérateur compte essentiellement le nombre de jetons acheminés vers différents experts et est facilement affecté par des facteurs tels que la disposition de la mémoire partagée, les mises à jour simultanées, les conflits bancaires et l'agrégation inter-blocs.

Avec TLE-Struct, les développeurs peuvent organiser explicitement la disposition des compteurs, en associant différents experts ou jetons à différentes zones de stockage local ; une fois que le compilateur a obtenu ces informations structurelles, il génère les méthodes d’accès et de synchronisation appropriées.

TLE-Raw est conçu pour les experts en optimisation des performances, préservant les interfaces pour le code natif du fournisseur.

Certaines méthodes d'optimisation nécessitent l'utilisation directe de CUDA, d'assembleur ou de fonctions intrinsèques dédiées. Les contraindre à être réintégrées dans un DSL de plus haut niveau risque de limiter les performances ou d'augmenter les coûts de migration. TLE-Raw permet aux développeurs d'intégrer du code natif dans l'architecture Triton/TLE ou d'accéder directement aux pipelines de compilation du fournisseur.

Prenons l'exemple de All-Gather GEMM,Cet opérateur gère la communication, le calcul matriciel, la mise en mémoire tampon et la synchronisation. Les développeurs peuvent réutiliser les fonctionnalités de communication sous-jacentes tout en continuant d'utiliser des tenseurs et des tuiles structurés pour exprimer les calculs. Enfin, un pipeline de compilation unifié organise les différentes parties en opérateurs appelables.

Cette conception par couches permet aux ingénieurs algorithmiques, aux développeurs d'opérateurs et aux experts en performance de choisir différents niveaux d'optimisation au sein d'un même système, sans avoir à commencer par le niveau de programmation le plus bas.

Les tests de performance montrent également queLa surcharge d'abstraction globale de TLE est contrôlable.

Lors du test Radix Select, TLE a reproduit l'algorithme TensorRT-LLM, atteignant des performances d'environ 85% à 97% sur différentes formes. Pour les équipes qui doivent assurer une maintenance multiplateforme et itérer rapidement, obtenir des performances proches de celles des experts avec un code plus unifié représente un atout technique considérable.

Dans un scénario SparseMLA avec un contexte de 128K,TLE utilise des primitives Pipeline pour exprimer la collaboration entre différents rôles d'exécution, atteignant des performances d'environ 90%, soit la même que la référence FlashMLA.

L'équipe a également testé All-Gather sur un nœud unique équipé de huit processeurs NVIDIA H100 et a approfondi l'intégration de GEMM et d'All-Gather. L'objectif n'est pas simplement de remplacer la bibliothèque de communication, mais d'intégrer la communication aux expressions des opérateurs et aux optimisations de compilation, permettant ainsi la transmission des données, le calcul local, la synchronisation et leur utilisation ultérieure pour former un pipeline.

Pour les systèmes d'inférence, ce type de fusion communication-calcul est souvent plus pertinent que l'amélioration des performances maximales d'un seul GEMM isolé, car ce que les utilisateurs perçoivent en fin de compte, c'est la latence de bout en bout.

Pratiques d'optimisation de la compilation Flagtree

Outre les extensions de langage TLE, l'équipe FlagTree a également entrepris deux autres tâches :Premièrement, nous avons intégré le backend CUDA Tile IR ; deuxièmement, nous avons implémenté plusieurs optimisations de compilation Triton pour les charges de travail de modèles réels.

Le concept fondamental de CUDA Tile IR est de permettre aux programmes d'exprimer des tuiles, plutôt que de prédéterminer le mappage des threads sous-jacent.Ses entrées comprennent un programme de tuiles et une grille de tuiles tridimensionnelle. Au sein de Dialect, les calculs, les visualisations de données et les dépendances nécessaires sont exprimés par le biais de Tile Compute, View et Token-Ordered Operation (TKO).

Parmi eux, TensorView décrit le pointeur global, la forme et le pas ; PartitionView ajoute des capacités de mappage de blocs ; Token est utilisé pour contraindre les dépendances entre les TKO associés ; et Memory Model spécifie séparément la sémantique et la portée de la mémoire.

FlagTree ne remplace pas complètement le backend natif NVIDIA Triton. Il ajoute plutôt un chemin TileIR au système existant et redéfinit les interfaces TileIR View et Token via des primitives TLE. Les noyaux compatibles peuvent accéder au backend TileIR ; ceux qui ne le sont pas encore utiliseront le CUDA natif. Les deux chemins peuvent coexister dans le compilateur.

En ce qui concerne l'optimisation automatique,L'équipe a proposé FlagOSTune pour résoudre le conflit entre couverture et coût dans les modèles réels de Triton Autotune.

Dans l'inférence de modèles, un même opérateur peut correspondre à un grand nombre de formes différentes. Les statistiques de l'équipe montrent qu'il existe 1 994 formes MM uniques dans six modèles et quatre types de scénarios d'inférence, ce qui dépasse largement la couverture des benchmarks quotidiens. Si la configuration candidate est directement étendue, l'échelle de recherche théorique atteindra rapidement des millions d'ensembles.

FlagOSTune améliore les performances en élargissant l'espace de recherche et en réduisant les coûts grâce à une combinaison de prédiction par modèle et de tests limités en conditions réelles. Le système utilise d'abord XGBoost pour classer les configurations candidates, puis envoie seulement un petit nombre des configurations les plus prometteuses à la compilation et aux tests GPU, avant de poursuivre la recherche à l'aide d'un algorithme génétique.

Les performances de plusieurs opérateurs ont été améliorées sur divers systèmes de puissance de calcul de NVIDIA, Moore Threads et Muxi, avec des gains de vitesse allant de 1,21 à 7,35x.Dans de multiples expériences de forme avec l'opérateur NVIDIA H20 MM, les configurations de l'espace de recherche ont été compressées de plus de 620 000 à 4 070, ce qui a permis d'accélérer de 120 fois l'efficacité du réglage sans presque aucune perte de performance.

En ce qui concerne l'organisation des données, l'équipe s'est concentrée sur l'optimisation des performances liées à la transformation de cette organisation.

Dans Triton, la fonction `convert_layout` implique généralement un réarrangement des données entre threads, nécessitant une écriture en mémoire partagée, une synchronisation, puis une relecture ; cette opération n'est donc pas sans coût. Pour les opérateurs de petite taille ou gourmands en mémoire, même un petit nombre de transformations de disposition peut constituer un goulot d'étranglement en termes de performances.

FlagTree améliore le mécanisme de suppression des conversions de mise en page en utilisant des techniques telles que la modélisation des coûts, la rétropropagation et la résolution conjointe locale pour réduire les conversions de mise en page de données inutiles.Lors de tests avec plus de 100 opérateurs, le taux d'élimination de conversion net était d'environ 68%–79%, avec une amélioration des performances allant jusqu'à 71%.

Une autre optimisation consiste à réordonner les instructions. Le déroulement de boucle se contente de copier le corps de la boucle et n'effectue pas automatiquement plusieurs chargements à l'avance. Avec le réordonnancement du compilateur activé, des chargements indépendants peuvent être exécutés en amont, entrelacant les calculs successifs et exploitant ainsi le pipeline matériel pour masquer la latence d'accès à la mémoire.Dans les trois opérateurs typiques, cette optimisation apporte une accélération moyenne d'environ 1,19 à 1,61 fois, avec une accélération maximale de 2 fois.

L'équipe a également optimisé les performances de l'opérateur MoE de Fused Marlin pour les formes réelles dans DeepSeek-V4-Flash, réduisant ainsi le temps de chargement des données grâce à la fusion des fragments. Combinées à des optimisations au niveau de l'algorithme,Comparé à vLLM CUDA, FlagGems a obtenu des gains de vitesse sur les 53 formes testées, avec une amélioration moyenne d'environ 1,208 fois pondérée par la fréquence de la forme.

Dans le test de bout en bout DeepSeek-V4-Flash sur le NVIDIA H20,Lorsque TP=4, la latence du premier jeton diminue de 20,221 TP3T et le débit total augmente de 12,631 TP3T ; lorsque TP=8, la latence du premier jeton diminue de 12,871 TP3T et le débit total augmente de 7,651 TP3T.

Perspectives d'avenir

L'étape suivante,FlagTree se concentrera sur le développement du modèle de programmation NUMA et du compilateur MegaKernel.

Pour les configurations multi-nœuds et multi-GPU, ainsi que pour les clusters internes aux GPU et la SRAM locale, l'accès aux données présente des différences significatives en termes de proximité. TLE vise à décrire explicitement la segmentation des tenseurs le long du maillage et leur transition entre différentes méthodes de distribution, permettant ainsi au compilateur de déterminer quelles données doivent être rapprochées pour le calcul et quelles communications peuvent être fusionnées ou superposées à ce dernier.

Le compilateur MegaKernel vise à étendre le champ d'optimisation d'un seul noyau à un processus d'exécution modèle. Il identifie les calculs pouvant être co-optimisés, les décompose en tâches spécifiques, puis décide de manière uniforme de leur allocation, de leur parallélisation et de leur mise en pipeline, générant ainsi un ou plusieurs MegaKernels.

TLE sert de pont entre la sémantique du front-end et l'architecture matérielle. Les interfaces Tile, Mesh, Pipeline, Memory Layout et de capacités natives permettent au compilateur de voir non seulement une série d'opérateurs indépendants, mais un graphe de tâches pouvant être planifié, placé, synchronisé et fusionné.

Des indications sémantiques légères de TLE-Lite au contrôle prenant en compte l'architecture de TLE-Struct, en passant par la transmission des capacités natives de TLE-Raw, combinées à diverses techniques d'optimisation de la compilation, FlagTree tente d'établir un système en couches plus flexible entre la facilité d'utilisation, la migration multiplateforme et les performances ultimes de Triton, et de permettre aux nouvelles capacités matérielles d'être transformées plus rapidement en performances réelles des modèles et des opérateurs.