Amélioration de l’auto-tuning GEMM avec nvMatmulHeuristics dans CUTLASS 4.2

Aperçu

La sélection de kernels GEMM sur les GPU NVIDIA modernes est un problème d’optimisation avec de nombreux paramètres. La performance du kernel dépend d’un ensemble étendu de méta-paramètres de compilation et d’exécution, notamment les tailles CTA, warp et tiles des instructions; les plannings de kernels; les stratégies de rasterisation; les dimensions du cluster; et les facteurs de division (split-k). L’approche traditionnelle génère des milliers de configurations, les compile et exécute une tuning exhaustive pour trouver la plus rapide. Ce flux peut prendre des heures et devient un frein pour les bibliothèques hors-ligne comme CUTLASS, et est particulièrement problématique pour les environnements JIT comme Torch Inductor ou OpenAI Triton où la compilation rapide est cruciale. NVIDIA introduit nvMatmulHeuristics, un module d’optimisation des méta-paramètres qui fournit des heuristiques rapides pour prédire un petit ensemble de configurations GEMM prometteuses. Le module analyse les paramètres d’un problème GEMM et les capacités matérielles cibles pour prédire les configurations offrant les meilleures performances. L’intégration dans CUTLASS 4.2 transforme le flux de génération et de tuning des kernels d’un brute-force vers une approche ciblée et efficace. Pour un problème GEMM donné, nvMatmulHeuristics émet un ensemble concis de configurations candidates à tester. Le feature est présenté comme une composante clé des heuristiques cuBLAS et est disponible en accès anticipé pour un usage général, avec une intégration dans la bibliothèque CUTLASS. Un exemple montre un GEMM FP16 simple (A transposé ou non, B en colonne-major, et C/D selon le layout). Le flux consiste à fournir une liste de problèmes GEMM au format JSON et à construire CUTLASS avec des options CMake spécifiques afin d’émettre un ensemble de configurations par problème. Dans la pratique, cette approche réduit considérablement le temps total de tuning. Plutôt que de compiler et profiler de manière exhaustive des milliers de kernels, vous générez un ensemble restreint et ciblé et exécutez uniquement celui-ci. Cela est particulièrement utile pour les environnements comme PyTorch, où les temps de compilation JIT sont critiques, et pour les bibliothèques qui veulent des performances élevées sans longs cycles de tuning.

Caractéristiques clés

Prédit un petit ensemble top-N de configurations GEMM par problème
Réduit le temps total de tuning en concentrant la compilation et le profilage sur un ensemble de candidats
Intégré à CUTLASS 4.2 et aux heuristiques cuBLAS; disponible en accès anticipé
Accepte une entrée JSON décrivant les problèmes GEMM; facile à intégrer dans les flux existants
Génère une liste CSV de tests à partir des configurations émisses pour un profilage automatique
Compatible avec cutlass_profiler pour exécuter les configurations générées
Autorise des tailles de cluster statiques à la compilation, réduisant la variabilité en exécution
Montre des performances proches du maximum avec beaucoup moins d’effort; par exemple, 16 configurations atteignent 96% du pic avec un temps bien moindre
Aide les flux DL, compilateurs et bibliothèques de kernels à atteindre rapidement des performances élevées

Tableau: comparaison des approches de tuning

| Approche | Comment cela fonctionne | Avantages typiques | Remarques |---|---|---|---| | Recherche exhaustive | Génère et compile des milliers de kernels; tuning complet | Meilleur kernel possible; perf théorique maximale | Temps de build+tuning très long (exemples > 700 minutes) |nvMatmulHeuristics (cette fonctionnalité) | Prédit top-N kernels par problème; compile un petit ensemble et profile | Gain substantiel de temps; proche du pic | Nécessite un fichier JSON de problèmes et une intégration de build |Pré-compilé/cluster statique | Utilise des tailles de cluster fixes des kernels pré-compilés | Configuration rapide; profiling en exécution réduit | Moins de flexibilité; les performances dépendent des choix de compilation |

Cas d’utilisation courants

Intégration dans des frameworks DL: accélérer les chemins JIT (par ex., PyTorch) en réduisant le temps de tuning tout en conservant de hautes performances pour les charges GEMM.
Bibliothèques hors-ligne: permettre des ensembles de kernels bien tunés avec un minimum d’effort de tuning durant le déploiement.
Charges de modèles avec d’importants GEMM (p. ex., Llama, DeepSeek) où le tuning exhaustif est prohibitif.
Scénarios nécessitant des profils reproductibles; l’approche privilégie un ensemble restreint de candidats pour des builds reproductibles.

Configuration & installation (commandes exactes)

Le flux repose sur la préparation d’une liste de problèmes GEMM en JSON, la construction de CUTLASS avec nvMatmulHeuristics activé, puis l’exécution du profiling via le CSV généré. Les commandes CMake explicitement montrées dans le texte sont:

cmake .. \
-DCUTLASS_LIBRARY_HEURISTICS_PROBLEMS_FILE= \
-DCUTLASS_LIBRARY_HEURISTICS_CONFIGS_PER_PROBLEM=N

Après configuration, une construction standard émettra un CSV de tests lors du processus. Vous pouvez utiliser cutlass_profiler pour exécuter les configurations générées via le CSV; pour des résultats de profiling cohérents, exécutez avec des horloges verrouillées. Remarques:

Construisez CUTLASS comme d’habitude, mais avec les deux flags CUTLASS_LIBRARY_HEURISTICS montrés ci-dessus.
Le format exact des problèmes JSON n’est pas spécifié ici; le flux se réfère à une liste de problèmes GEMM exprimés en JSON.

Démarrage rapide (exemple minimal)

Préparez une liste de problèmes GEMM en JSON (l’exemple dans la source mentionne un GEMM FP16 unique, A transposé ou non, et les layouts de B et C/D).
Construisez CUTLASS avec les heuristiques activées en utilisant les commandes de la section précédente.
Le système de build émettra un CSV décrivant les configurations générées pour chaque problème GEMM.
Exécutez le flux de profiling avec cutlass_profiler sur le CSV généré pour lancer les configurations. Pour des résultats cohérents, utilisez des horloges verrouillées pendant le profiling.
Comparez les performances avec une baseline exhaustive, si disponible; dans les cas publiés, 16 configurations atteignent 96% du pic avec bien moins de temps, et davantage de candidats se rapprochent de l’exhaustif tout en conservant des gains de temps significatifs.

Avantages et inconvénients

Avantages
Réduction énorme du temps de build+tuning par rapport à la recherche exhaustive
Performance proche du pic avec un petit ensemble de candidats
Meilleure expérience pour les pipelines JIT et les déploiements sur GPU
Options de cluster statique peuvent améliorer la prévisibilité du compile-time
Inconvénients
Nécessite une structure JSON pour la liste de problèmes et l’intégration correspondante
Fonctionnalité en accès anticipé; API et flux peuvent évoluer
Dépend des capacités matérielles et de l’exactitude de la caractérisation des problèmes

Alternatives (comparaisons rapides)

Recherche exhaustive: garantit le meilleur kernel mais demande un temps et des ressources importants (exemples > 700 minutes).
Heuristiques avec cluster statique: donne de bonnes performances mais dépend des choix de compilation et peut limiter la flexibilité en exécution.
Kernels précompilés/dynamiques du type Blackwell: offrent une adaptation en temps réel mais peuvent nécessiter du tuning à l’exécution; nvMatmulHeuristics cherche à déplacer plus de décisions vers le compile-time.

Prix ou licence

nvMatmulHeuristics est décrit comme une composante centrale des heuristiques cuBLAS et est disponible en accès anticipé pour un usage général, avec intégration dans CUTLASS. Aucune information sur le prix n’est fournie.