Les bugs de synchronisation dans les kernels GPU sont particulièrement difficiles à identifier. Lorsqu’une barrière provoque un interblocage (deadlock), le matériel ne fournit ni trace d’exécution, ni code d’erreur avant que le test n’atteigne un délai d’attente, ce qui peut allonger chaque cycle de débogage de plusieurs dizaines de minutes. Pour remédier à cette difficulté, un ingénieur spécialisé en vérification formelle a proposé une approche alternative : encoder le modèle de synchronisation des kernels dans le langage Promela et le soumettre au model checker SPIN. Cette méthode permet non seulement de réduire le temps de débogage, mais aussi de démontrer de manière déterministe l’existence d’un contre-exemple à une propriété souhaitée (par exemple un deadlock) ou de prouver qu’aucun entrelacement problématique n’existe.
Contexte : un kernel MoE pour DeepSeek-V3 sur B200
Ce travail a été motivé par la participation au défi FlashInfer MLSYS, dans le cadre duquel l’équipe a développé un kernel persistant et fusionné de mélange d’experts (MoE) pour le modèle DeepSeek-V3, exécuté sur un GPU NVIDIA B200. Le kernel, écrit en CuTe DSL (le langage dédié de la bibliothèque CUTLASS de NVIDIA), est assemblé à partir de trois étapes — FF1, SwiGLU et FF2 — réparties sur des clusters de blocs de threads (CTA). La coordination entre ces étapes repose sur des mbarriers en mémoire partagée, des mbarriers atteignables entre CTA via des pointeurs DSMEM traduits par l’instruction mapa, et des compteurs atomiques en mémoire globale (GMEM) entre clusters. Une erreur dans l’un de ces mécanismes peut produire un interblocage difficile à diagnostiquer, d’autant plus que les GPU n’étaient accessibles que via un fournisseur de cloud à la demande.
Primitives de synchronisation Blackwell modélisées
Le model checker SPIN a été choisi pour sa capacité à analyser exhaustivement les entrelacements possibles. Pour cela, il a fallu modéliser les primitives de synchronisation spécifiques à l’architecture Blackwell.
L’mbarrier est un objet matériel de 64 bits logé en mémoire partagée. Il contient un compteur d’arrivées, un compteur attendu prédéclaré et un bit de phase. Une fois initialisé, il fonctionne comme une machine à états avec deux transitions : l’arrivée (mbarrier_arrive), qui incrémente atomiquement le compteur – si celui-ci atteint la valeur attendue, la barrière se complète en réinitialisant le compteur et en basculant le bit de phase – et l’attente (mbarrier_wait(phase=P)), qui bloque le thread jusqu’à ce que la phase diffère de P. Ce mécanisme permet de réutiliser la barrière pour plusieurs itérations, à condition de gérer correctement le basculement du bit de phase entre chaque cycle. Une erreur dans cette gestion peut entraîner un deadlock.
En CuTe DSL, chaque opération sur mbarrier correspond à un petit ensemble d’appels : initialisation (mbarrier_init, mbarrier_init_fence), arrivée locale, arrivée avec attente de transaction asynchrone (TMA), attente de phase, et arrivée sur une barrière appartenant à un CTA voisin via mapa_shared_cluster. Les barrières entre clusters ne peuvent pas être partagées ; la communication utilise alors des compteurs entiers en mémoire globale, avec des opérations atomiques et des barrières mémoire (sem="acquire"/"release", scope="gpu"). Le choix du scope est crucial : un scope trop petit (comme cta au lieu de gpu) peut entraîner des lectures périmées de la mémoire globale, même après qu’une opération atomique ait été effectuée.
L’outil de traduction : CuTe vers Promela
L’ingénieur a développé un outil de traduction nommé cute2promela, disponible sur GitHub, qui abaisse automatiquement les kernels CuTe en code Promela. Le modèle Promela inclut des canaux pour les mbarriers et des processus pour les threads, ce qui permet à SPIN de vérifier l’absence de deadlock, de violation d’ordre mémoire ou d’autres propriétés de sécurité. L’outil a été testé sur un exemple simple : un lisseur de Jacobi à deux CTA, qui échange des valeurs aux frontières via mémoire partagée DSMEM. Ce noyau, bien que simplifié, reproduit le squelette de synchronisation du kernel MoE complet.
Résultats et implications
La méthode proposée permet de raccourcir considérablement le cycle de débogage : au lieu de lancer un kernel sur GPU et d’attendre un timeout, le développeur peut exécuter SPIN sur son poste de travail et obtenir un contre-exemple en quelques secondes. De plus, SPIN offre une preuve formelle d’absence de deadlock lorsque la propriété est vérifiée, ce qui n’est pas possible avec des tests classiques. Cette approche pourrait être étendue à d’autres types de bugs de synchronisation, comme les conditions de course liées à l’ordre mémoire.
L’auteur du projet souligne que l’encodage des barrières mémoire et des fences (comme fence_view_async_shared) est indispensable pour capturer les problèmes de visibilité, sous peine de ne pas détecter certains bugs. L’outil cute2promela est présenté comme une preuve de concept, mais il pourrait être généralisé pour vérifier d’autres kernels CuTe ou même d’autres langages de programmation GPU.