Résumé technique : Xe-Forge

Énoncé du problème

Le portage d'algorithmes d'apprentissage profond vers de nouveaux accélérateurs matériels, spécifiquement les GPU Intel, oblige les développeurs à appliquer manuellement des optimisations de bas niveau répétitives (quantification, fusion de mémoire, réglage des tuiles, contournements spécifiques à l'architecture) à chaque noyau Triton. Cet effort manuel crée un goulot d'étranglement significatif : bien que les motifs d'optimisation soient cohérents, les contraintes matérielles (nombre de warps, modes GRF, dimensionnement SLM) varient d'un dispositif à l'autre et sont souvent absentes des données d'entraînement des grands modèles de langage (LLM). Les systèmes existants basés sur les LLM se concentrent principalement sur la génération de noyaux à partir de zéro pour le matériel NVIDIA ou AMD, laissant une lacune critique dans l'optimisation des noyaux Triton existants et fonctionnellement corrects pour l'architecture spécifique d'Intel (séries Arc/Arc Pro).

Méthodologie

Les auteurs présentent Xe-Forge, un pipeline multi-étapes alimenté par un LLM, conçu pour automatiser l'optimisation de noyaux Triton existants pour les GPU Intel. Le système fonctionne via les composants principaux suivants :

1. Architecture de pipeline multi-étapes

Xe-Forge traite un noyau Triton fonctionnellement correct à travers jusqu'à neuf étapes d'optimisation séquentielles. Un planificateur basé sur un LLM détermine l'ordre d'exécution en fonction des problèmes détectés, sous réserve de contraintes de dépendance strictes (par exemple, la restructuration algorithmique doit précéder la conversion de type de données). Les étapes incluent :

1. Analyse : Détecte les opportunités d'optimisation.
2. Optimisation algorithmique : Restructure le calcul pour réduire les FLOPs et les accès mémoire (par exemple, élimination des GEMM redondants).
3. Découverte : Applique des optimisations novatrices et ouvertes non couvertes par les étapes nommées.
4. Correction de type de données (Dtype Fix) : Convertit les précisions inutiles (par exemple, de float64 à float32/16).
5. Fusion : Fusionne des noyaux séparés ou des chaînes élémentaires.
6. Accès mémoire : Optimise la fusion, la prélecture et la disposition des données.
7. Pointeurs de bloc : Remplace l'arithmétique de pointeurs manuelle par tl.make_block_ptr.
8. Noyau persistant : Restructure les boucles pour la réutilisation des données entre itérations.
9. Réglage spécifique au GPU : Applique des contraintes spécifiques à Intel (nombre de warps, modes GRF, dimensionnement des tuiles, mélange GROUP_SIZE_M).
10. Autotuning : Génère une grille de configuration @triton.autotune.

Les étapes sans problèmes associés sont ignorées pour réduire les coûts et les modifications de code.

2. Chaîne de vérification et de raffinement (CoVeR)

Chaque étape est exécutée par un agent CoVeR (implémenté dans DSPy) qui génère itérativement des candidats, les valide et les affine en fonction des retours. L'agent fonctionne en boucle :

• Générer : Propose des modifications de code.
• Vérifier : Exécute une cascade à quatre niveaux :
  1. Syntaxe : Analyse AST Python.
  2. Structure : Vérifie les contraintes matérielles (par exemple, warps puissants de deux, modes GRF valides).
  3. Correction : Compare les sorties à une référence PyTorch en utilisant torch.allclose.
  4. Performance : Effectue des tests de performance sur du matériel réel ; rejette si plus lent que la référence.
• Raffiner : Si la vérification échoue, les diagnostics d'erreur sont réinjectés dans le contexte pour la prochaine itération.

3. Base de connaissances curatée

Pour pallier le manque de données spécifiques à Intel dans les corpus d'entraînement des LLM, Xe-Forge utilise une base de connaissances YAML lisible par machine. Celle-ci encode des contraintes strictes (par exemple, nombres de warps valides, limites de capacité SLM) et des motifs d'optimisation (transformations de code avant/après) spécifiques aux architectures Xe d'Intel. Cela garantit que le modèle reste dans des limites architecturalement valides.

4. Infrastructure de validation (AI Bench)

Le système repose sur AI Bench, un framework autonome qui découple la mesure de la logique d'optimisation. Il assure un chronométrage reproductible (en utilisant le chronométrage d'événements matériels avec vidage du cache L2) et la vérification de la correction. Crucialement, le pipeline impose une séparation du harnais de noyau : le LLM ne peut modifier que le corps du noyau, tandis que l'encapsulation Model, la génération d'entrées et le chemin d'exécution restent fixes. Cela empêche les optimisations adverses où le modèle contournerait la logique de benchmark.

Contributions clés

1. Formulation d'optimisation Triton vers Triton : Contrairement aux approches existantes de génération PyTorch vers Triton, Xe-Forge optimise les noyaux existants tout en préservant la sémantique.
2. Premier système ciblant les GPU Intel : Le premier système basé sur un LLM ciblant les GPU Intel, doté d'une base de connaissances consciente du matériel pour les contraintes absentes des données d'entraînement.
3. Pipeline CoVeR multi-étapes : Une approche décomposée où chaque étape possède des motifs et des critères de vérification dédiés, équilibrant flexibilité et reproductibilité.
4. Framework AI Bench : Un outil de benchmarking qui découple l'exécution/chronométrage de la logique d'optimisation, assurant une comparaison équitable entre les backends.
5. Atténuation des attaques adverses : Une conception qui fixe le harnais d'exécution pour empêcher le LLM de « tricher » lors du benchmark en court-circuitant le modèle.

Résultats expérimentaux

Le système a été évalué sur 97 noyaux Level-2 de KernelBench et des configurations Flash Attention sur le GPU Intel Arc Pro B70.

• Performance globale : Atteint un accélération géométrique moyenne de 1,17× par rapport à PyTorch eager, avec 67 % des noyaux montrant une amélioration.
• Accélérations extrêmes : Neuf noyaux ont dépassé une accélération de 5×, le plus élevé atteignant 82× (Matmul_Min_Subtract). Ces gains ont été motivés par une restructuration algorithmique éliminant les calculs redondants, réduisant effectivement le nombre de FLOPs en dessous du comptage original.
• Flash Attention : A atteint des accélérations de 2× à 13,3× sur les 16 configurations testées (variant les tailles de modèle, les longueurs de séquence et les tailles de lot) sans régression. Notamment, le système a permis l'inférence de séquences de 32k sur un GPU de 32 Go en utilisant des noyaux persistants et un tuilage SLM, ce qui était irréalisable avec le noyau naïf.
• Correction : 100 % des noyaux optimisés ont passé les tests de correction par rapport à la référence PyTorch.
• Analyse Roofline : Pour les noyaux limités par le calcul (GEMM/MatMul), le pipeline a élevé les performances jusqu'à une efficacité proche du pic. Pour les noyaux limités par la bande passante, les performances étaient généralement compétitives, bien que certains aient montré de légères régressions (0,5–0,8×) dues aux limites de débit mémoire.

Signification et affirmations

L'article affirme que la connaissance de domaine structurée combinée à une vérification en boucle matérielle est plus efficace que l'échelle du modèle seule pour l'optimisation de noyaux spécifiques au matériel.

• Combler le fossé des connaissances : En codant des contraintes spécifiques à Intel (nombres de warps, modes GRF) dans une base de connaissances curatée, le système compense l'absence de données d'entraînement du LLM sur ces architectures spécifiques.
• Élimination systématique de l'effort de portage : Xe-Forge démontre que le travail répétitif de portage peut être automatisé, éliminant systématiquement le goulot d'étranglement qui limite actuellement le déploiement d'algorithmes sur de nouveaux accélérateurs.
• Robustesse : La séparation du code de noyau modifiable du code de harnais immuable empêche efficacement la classe de défaillances d'« évasion » observée dans d'autres systèmes de génération, où les modèles contournent les contraintes pour passer les tests de correction.

Les auteurs concluent que, bien que le pipeline ne soit pas parfait (une minorité de noyaux montrent de légères régressions et les motifs algorithmiques novateurs restent difficiles), il automatise avec succès l'optimisation des opérateurs standard et des compositions fusionnées qui constituent la majeure partie des charges de travail de production.