PyTorch ajoute le backend CuteDSL, mise sur Python plutôt que C++ pour les kernels GPU

Meta a intégré CuteDSL de NVIDIA comme quatrième backend d'autotuning dans TorchInductor, rejoignant Triton, CUTLASS C++ et cuBLAS pour l'optimisation des multiplications matricielles. Ce langage spécialisé basé sur Python offre des performances comparables aux kernels C++ optimisés à la main tout en maintenant des temps de compilation équivalents aux backends existants—une amélioration significative par rapport à CUTLASS C++ qui nécessite des invocations nvcc complètes. Les benchmarks internes montrent que les kernels softmax CuteDSL et Triton approchent tous deux la bande passante terminale sur le matériel GB200, mais la vraie cible sont les GEMMs où les écarts de performance comptent le plus.

Cette intégration représente plus qu'une mise à niveau technique—c'est un pari stratégique sur l'avenir du développement de kernels GPU. Alors que Triton excelle dans les opérations limitées par la mémoire comme les calculs élément par élément et les réductions, les GEMMs qui dominent les charges de travail de transformers exigent un contrôle de plus bas niveau sur les hiérarchies de threads et de mémoire. CuteDSL fournit ce contrôle à travers les mêmes abstractions que CUTLASS C++, qui s'est révélé efficace pour les GEMMs FP8 et la fusion d'épilogue, mais l'enveloppe dans la syntaxe conviviale de Python. Meta positionne explicitement CuteDSL comme un "remplacement éventuel" pour CUTLASS C++ sur les nouvelles générations de matériel.

Le timing s'aligne avec l'élan plus large de l'industrie vers les DSLs de kernels basés sur Python, avec l'adoption par des chercheurs comme Tri Dao (bibliothèque Quack) et Jay Shah chez Colfax International. Meta a appliqué trois critères pour l'intégration de backend : charge de maintenance minimale, aucune régression du temps de compilation, et performance supérieure sur les charges de travail cibles. L'engagement de développement actif de NVIDIA et les templates de kernels optimisés satisfont la première exigence, tandis que les résultats de performance valident la troisième. Pour les développeurs qui construisent une infrastructure AI de production, cela signifie des opérations GEMM potentiellement plus rapides sans la complexité de la maintenance des kernels C++, bien que les vrais gains de performance dépendront des architectures de modèles et des configurations matérielles spécifiques.