Améliorer l’Auto-tuning des Kernels GEMM sur les GPU NVIDIA avec des heuristiques et CUTLASS 4.2

TL;DR

• nvMatmulHeuristics propose des heuristiques rapides pour les GEMM et prédit un petit ensemble de configurations de kernel à fort potentiel.
• Intégré à CUTLASS 4.2, il réduit drastiquement le temps de tuning des kernels passant d’une recherche exhaustive à une évaluation ciblée de candidats.
• Résultats initiaux : en utilisant 16 configurations, on atteint environ 96% des performances maximales en environ 150 minutes, contre plus de 700 minutes pour une recherche exhaustive (Llama 3 405B sur H100 SXM).
• Pour le DeepSeek-R1 671B sur un GPU B200, 8 configurations atteignent 99% des performances de la recherche exhaustive avec plus de 5x de gain de temps de build et d’auto-tuning. Cela rend les flux de travail JIT pratiques dans les cadres DL.
• L’approche privilégie des tailles de cluster statiques connues à la compilation, ce qui peut offrir des gains de performance par rapport à des tailles de cluster dynamiques utilisées par certains kernels précompilés. Ces résultats démontrent qu’une bonne heuristique peut pousser les performances vers le kernel optimal tout en réduisant drastiquement les coûts de tuning. NVIDIA Developer Blog

Contexte et antécédents

Choisir le meilleur kernel GEMM pour un problème et un matériel donnés est un problème d’optimisation complexe. Les performances du GEMM dépendent d’un ensemble étendu de méta-paramètres: CTA, tailles de tiles par warp et niveau instruction, ordonnancement du kernel, stratégies de rasterisation, dimensions du cluster et facteurs de division. Historiquement, trouver le kernel optimal impliquait de générer des milliers de configurations, de les compiler et de les tester en tuning exhaustif. Ce flux peut prendre des heures et devient un frein pour l’adoption dans les bibliothèques hors ligne comme CUTLASS ou dans les environnements de compilation JIT (par ex. Torch Inductor, OpenAI Triton) où la compilation rapide est critique. Le billet NVIDIA présente nvMatmulHeuristics, un module d’optimisation de méta-paramètres pour kernels GPU, conçu pour délivrer un ensemble concis de configurations de kernel présentant un fort potentiel pour un GEMM donné et son matériel cible. Le module analyse les paramètres de l’opération et les capacités matérielles pour fournir un ensemble compact de configurations offrant des performances proches du maximum. L’intégration dans CUTLASS 4.2 vise à transformer le flux de génération de kernels et de tuning en un processus plus rapide et prévisible. Le feature est décrit comme une partie centrale des heuristiques cuBLAS et est disponible en accès anticipé avec une intégration dans CUTLASS. NVIDIA Developer Blog Le flux passe d’une recherche par force brute à une exploration guidée. Les utilisateurs préparent une liste de problèmes GEMM en format JSON, construisent CUTLASS avec les heuristiques activées et précisent -DCUTLASS_LIBRARY_HEURISTICS_PROBLEMS_FILE= et -DCUTLASS_LIBRARY_HEURISTICS_CONFIGS_PER_PROBLEM=N, où N est le nombre de configurations émis par nvMatmulHeuristics pour chaque GEMM. L’étape de build produit une liste CSV de tests répertoriant tous les cas à tester pour effectuer le tuning sur les configurations émises. Cette CSV peut être utilisée avec du benchmarking personnalisé ou avec cutlass_profiler pour exécuter les configurations “out of the box”. Pour des résultats de profiling cohérents, exécutez avec des horloges verrouillées. La fonctionnalité fait partie de l’écosystème CUTLASS et est compatible avec des GEMMs FP16 en entrée/sortie et FP32 pour le calcul (HSH). NVIDIA Developer Blog

Exemple FP16 et performance pratique

Pour un scénario FP16 en entrée/sortie et FP32 en calcul (HSH), nvMatmulHeuristics peut émettre huit configurations. Ces huit kernels peuvent être compilés et mesurés, et le meilleur kernel peut être choisi pour le déploiement, souvent avec des performances proches du kernel optimal découvert par une recherche exhaustive. Dans les résultats publiés, huit configurations ont atteint 104% du baseline lorsque testées sur des GEMMs avec tailles de cluster statiques, démontrant que concentrer l’effort de compilation et de profilage sur un ensemble réduit peut atteindre des performances de pointe avec un coût moindre que la recherche exhaustive. NVIDIA Developer Blog

Résultats empiriques et conclusions

Deux charges de travail illustrent l’impact de nvMatmulHeuristics. D’abord, une formation Llama 3 405B sur un GPU NVIDIA H100 SXM a montré qu’une recherche exhaustive peut prendre plus de 700 minutes pour identifier le kernel optimal, alors qu’une sélection restreinte à 16 configurations atteint 96% des performances de pic en environ 150 minutes. Cela montre une réduction importante du temps total de tuning avec une perte minime par rapport au pic. Ensuite, une charge DeepSeek-R1 671B sur un GPU NVIDIA B200 a démontré que l’utilisation d’un petit ensemble de candidats peut atteindre 99% des performances de la recherche exhaustive avec plus de 5x de vitesse dans le temps de construction et d’auto-tuning. Dans ce cas, la baseline utilisait des tailles de cluster dynamiques, tandis que les configurations proposées par nvMatmulHeuristics furent construites avec des tailles de cluster statiques connues à la compilation, produisant 104% des performances de baseline pour ces GEMMs. Ces résultats illustrent comment des heuristiques bien conçues peuvent conduire à des performances de pointe tout en réduisant le coût de tuning et en facilitant l’intégration dans des frameworks, compilateurs et bibliothèques de kernels. NVIDIA Developer Blog

Implications pour les développeurs et les entreprises

En concentrant le tuning sur un petit ensemble de kernels à fort potentiel, nvMatmulHeuristics réduit considérablement le temps nécessaire pour déployer GEMMs haute performance dans des environnements de production. Cela facilite l’adoption de kernels optimisés dans les stacks JIT et les bibliothèques hors ligne, accélérant les cycles de développement pour les frameworks DL, les compilateurs et les bibliothèques de kernels. L’intégration à CUTLASS 4.2 simplifie l’adoption pour les développeurs utilisant déjà CUTLASS et les heuristiques cuBLAS comme partie de leur stack. NVIDIA Developer Blog

Détails techniques ou mise en œuvre

nvMatmulHeuristics opère sur deux couches: un module de heuristiques prédictives qui analyse les paramètres du GEMM et les capacités du matériel, et une intégration CUTLASS qui utilise les configurations prédites pour construire et évaluer un petit ensemble de kernels. Étapes clés:

• Préparer une liste de problèmes GEMM en format JSON décrivant les transpositions et les dispositions de A, B et C/D selon les besoins (par exemple GEMMs FP16 avec entrée/sortie FP16 et calcul FP32, HSH).
• Compiler CUTLASS avec l’intégration des heuristiques et spécifier -DCUTLASS_LIBRARY_HEURISTICS_PROBLEMS_FILE et -DCUTLASS_LIBRARY_HEURISTICS_CONFIGS_PER_PROBLEM pour contrôler le nombre de configurations émis par problème.
• L’étape de build produit une liste CSV de tests énumérant tous les cas à tester pour exécuter le tuning sur les configurations émis. On peut utiliser du benchmarking personnalisé ou cutlass_profiler pour exécuter les configurations automatiquement.
• Pour des résultats de profilage cohérents, exécutez avec des horloges verrouillées. La documentation fournit des conseils sur l’utilisation et l’intégration. L’exemple des huit configurations pour FP16/FP32 (HSH) illustre le côté pratique: en compilant et testant ces huit kernels, il est possible de choisir le meilleur kernel sans le coût d’une recherche exhaustive, tout en conservant des performances proches du meilleur obtenu. NVIDIA Developer Blog

Points clés

• Les heuristiques livrent un ensemble petit et ciblé de configurations GEMM à fort potentiel, plutôt qu’une recherche exhaustive.
• L’intégration avec CUTLASS 4.2 facilite un flux de travail avec des listes JSON de problèmes, des CSV de tests et l’option cutlass_profiler.
• Les résultats montrent des gains significatifs de temps tout en préservant des performances proches de l’optimum.
• Les tailles de cluster statiques peuvent offrir des bénéfices supplémentaires selon le contexte par rapport aux dynamiques.
• nvMatmulHeuristics est une composante centrale des heuristiques cuBLAS et est disponible en accès anticipé avec intégration CUTLASS. NVIDIA Developer Blog

FAQ

• Q : Que fait nvMatmulHeuristics ? A : Il analyse les paramètres GEMM et les capacités matérielles pour prédire un petit ensemble de configurations de kernel à fort potentiel.
• Q : Comment l’utiliser avec CUTLASS ? A : Compilez CUTLASS avec les heuristiques activées et fournissez une liste de problèmes en JSON ainsi que le nombre de configurations par problème; le build génère un CSV utilisé par profiling ou cutlass_profiler.
• Q : Quels gains de performance attendre ? A : Des résultats montrent que 16 configurations peuvent atteindre environ 96% du pic en ~150 minutes, contre >700 minutes pour une recherche exhaustive. D’autres workloads montrent des gains similaires et des économies de temps de build. NVIDIA Developer Blog
• Q : nvMatmulHeuristics est-il disponible publiquement ? A : Oui, il est décrit comme une partie des heuristiques cuBLAS et est disponible en accès anticipé avec intégration CUTLASS. NVIDIA Developer Blog

Références

• https://developer.nvidia.com/blog/improving-gemm-kernel-auto-tuning-efficiency-on-nvidia-gpus-with-heuristics-and-cutlass-4-2/
• La publication NVIDIA citée constitue la source principale pour nvMatmulHeuristics, son intégration à CUTLASS 4.2 et les résultats présentés. NVIDIA Developer Blog