Surprise : Des noyaux AI générés extrêmement rapides chez Stanford CRFM Introduction Le Centre de Recherche en Intelligence Artificielle de Stanford (Stanford CRFM) a réussi à créer des noyaux de calcul générés par IA, écrits en pur CUDA-C, sans recourir à des bibliothèques ou des DSLs comme CUTLASS et Triton. Ces noyaux ont surpassé ou égalé les performances des noyaux optimisés par des experts dans PyTorch. Les résultats sur un GPU Nvidia L40S sont notables : Multiplication de matrices (FP32) : 101.3% de la performance de torch.matmul. Conv2D : 179.9% de la performance de torch.nn.Conv2D. Softmax : 111.8% de la performance de torch.softmax. LayerNorm : 484.4% de la performance de torch.nn.LayerNorm. Conv2D + ReLU + MaxPool : 290.1% de la performance de torch et 189.0% de la performance de torch.compile(). Méthode L'équipe a utilisé le cadre de référence KernelBench, publié en décembre 2024, avec quelques modifications. Le processus se déroule en plusieurs étapes : Initialisation : Le modèle de langage (LLM) reçoit du code PyTorch et génère un noyau CUDA afin d'améliorer les performances. Optimisation par cache en lecture seule : Utilisation de __ldg pour charger des tenseurs invariants. Conversion en FP16 GEMM : Transformation de la convolution en une multiplication matricielle implicite sur les Tensor Cores. Double tamponnage pour le pipelinage : Chevauchement des charges mémoire globale avec les calculs des Tensor Cores pour réduire le temps d'exécution. Précomputation d'indices : Stockage des indices de décomposition des noyaux/entrées dans la mémoire partagée pour éviter des calculs redondants. Écriture vectorisée en mémoire partagée : Utilisation de données de type half2 pour optimiser les écritures en mémoire partagée. Parallélisation des sorties : Élimination de la sérialisation entre les files de calcul lors des écritures en mémoire globale. Pipelinage logiciel des charges en mémoire : Chevauchement des lectures en mémoire partagée avec les calculs WMMA pour améliorer les performances. Raffinement des noyaux : Itérations successives pour améliorer la performance et vérifier la correction. Exemple d'Optimisation d'un Noyau de Conv2D Voici une trajectoire d'optimisation pour un noyau de Conv2D, comparé à la référence PyTorch qui prend 1.41 ms : Round 0 : Durée : 7.02 ms (20.1% de la référence). Idée : Remplacement de l'opération par un noyau CUDA. Round 1 : Durée : 7.54 ms (18.8% de la référence). Idée : Utilisation du cache en lecture seule pour les tenseurs invariants. Round 2 : Durée : 3.46 ms (41.0% de la référence). Idée : Conversion de la convolution en une multiplication matricielle implicite (GEMM) en FP16. Round 3 : Durée : 3.67 ms (38.7% de la référence). Idée : Pipelinage avec double tamponnage pour chevaucher les charges mémoire globale avec les calculs de Tensor Core. Round 4 : Durée : 3.46 ms (41.0% de la référence). Idée : Utilisation d'une multiplication matricielle implicite (GEMM) avec un noyau préalablement générée. Round 5 : Durée : 1.91 ms (74.9% de la référence). Idée : Précomputation et réutilisation des indices du noyau et de l'entrée en mémoire partagée. Round 6 : Durée : 1.37 ms (103.6% de la référence). Idée : Précomputation et mise en cache des indices de la GEMM en multidimension en mémoire partagée. Round 7 : Durée : 1.38 ms (102.9% de la référence). Idée : Parallélisation de l'écriture du noyau CUDA en utilisant des tampons de mémoire partagée par file. Round 8 : Durée : 1.37 ms (103.6% de la référence). Idée : Précomputation et mise en cache des coordonnées d'entrée de base en mémoire partagée. Round 9 : Durée : 1.36 ms (105.1% de la référence). Idée : Pipelinage logiciel du chargement des fragments B pour chevaucher les lectures en mémoire partagée avec les calculs WMMA. Round 10 : Durée : 1.07 ms (133.6% de la référence). Idée : Réutilisation des indices de décomposition N-dimension de la GEMM pour le calcul des adresses de sortie. Round 11 : Durée : 1.21 ms (117.4% de la référence). Idée : Suppression de la décomposition hi/lo dans les opérations half WMMA, en tenant compte de l'accumulateur FP16 standard si la précision résultante le permet. Round 12 : Durée : 1.01 ms (141.2% de la référence). Idée : Chevauchement des charges globales de Asub (poids) et Bsub (entrées) avec le calcul WMMA en utilisant un double tamponnage. Round 13 : Durée : 0.795 ms (179.9% de la référence). Idée : Implémentation des écritures vectorisées en mémoire partagée en utilisant des types de données plus larges comme half2. Réflexions et Bilan Recherche raisonnée et exploration parallèle : La combinaison d'une forte capacité de raisonnement et d'une exploration parallèle de différentes hypothèses a conduit à des améliorations significatives. Cela résonne avec d'autres recherches en IA, où des stratégies simples mais bien conçues peuvent mener à des avancées importantes. Amélioration continue : Malgré les limites actuelles, comme la génération d'idées d'optimisation plus performantes et de meilleurs noyaux, cette méthode offre un potentiel prometteur pour l'amélioration future des noyaux de calcul. Synthèse de données : La méthode proposée aide également à générer de meilleures données synthétiques pour l'amélioration future des modèles de génération de noyaux. Futur : L'équipe est enthousiaste à l'idée de continuer à explorer et perfectionner cette approche, en espérant atteindre l'éventuelle création de systèmes IA capables de s'auto-améliorer. Importance de l'Exploitation de la Parallélisation La principale différence avec les approches traditionnelles de révision séquentielle repose sur l'exploration parallèle et la génération d'idées en langage naturel, ce qui permet d'explorer diverses stratégies d'optimisation plutôt que de stagner autour de transformations locales. Problèmes Courants Les problèmes couramment rencontrés dans la génération de noyaux sont souvent liés à des optimisations spécifiques aux tailles, car chaque problème a des dimensions précises sur lesquelles il faut travailler. Cela nécessite des ajustements fins et la prise en compte des particularités de chacun. Conclusions Cette recherche montre que même des méthodes relativement simples, combinées à une exploration parallèle et à un raisonnement en langage naturel, peuvent produire des noyaux de calcul très performants. En dépit des défis actuels, les résultats obtenus donnent confiance dans le potentiel d'amélioration continuous et de développement futur de l'IA dans ce domaine. Remerciements Nous tenons à remercier Christopher Rinard, Saman Amarasinghe et Allen Nie pour leurs discussions enrichissantes, ainsi que Standard Kernel Co. et Prime Intellect pour leur soutien. Standard Kernel Co. est une entreprise de pointe spécialisée dans l'optimisation des noyaux de calcul pour les applications de deep learning et de machine learning. Ils se concentrent sur le développement de solutions innovantes pour accroître l'efficacité des modèles de calcul sur GPU et CPU, en mettant particulièrement l'accent sur la réductions des temps d'exécution et l'optimisation des ressources.