CUCo: An Agentic Framework for Compute and Communication Co-design

CUCo est un cadre agentic sans entraînement qui automatise la génération de noyaux CUDA haute performance en co-optimisant conjointement le calcul et la communication, réduisant ainsi la latence globale jusqu'à 1,57 fois par rapport aux approches existantes.

L'essentiel

Imaginez que vous organisez un immense banquet pour des milliers d'invités (les données) dans un château rempli de chefs cuisiniers ultra-rapides (les puces GPU).

Le problème, c'est que dans la méthode traditionnelle, les chefs cuisiniers et les serveurs (qui apportent les plats d'une cuisine à l'autre) ne travaillent pas ensemble.

1. Le chef cuisine un plat.
2. Il s'arrête, lève la main, et attend qu'un serveur (le CPU, le chef d'orchestre humain) vienne lui dire : "Ok, le serveur est là, tu peux envoyer le plat".
3. Le serveur part, revient, et le chef reprend la cuisson.

Ce va-et-vient constant entre le chef et le serveur crée des temps d'attente. Le chef cuisine vite, mais il passe beaucoup de temps à attendre le serveur. C'est comme si un coureur de 100 mètres s'arrêtait à chaque mètre pour demander la permission de courir au suivant.

CUCo, c'est le nouveau système qui change tout. C'est un "architecte intelligent" (un agent) qui apprend aux chefs à cuisiner et à servir en même temps, sans avoir besoin de demander la permission à chaque fois.

Voici comment cela fonctionne, expliqué simplement :

1. Le Problème : Deux mondes séparés

Jusqu'à présent, les programmes pour les puces graphiques (CUDA) étaient écrits à la main par des experts. C'était comme essayer de composer une symphonie où les violons et les trompettes ne pouvaient jouer ensemble que si un chef d'orchestre leur donnait le signal à chaque mesure. C'était lent, difficile, et les erreurs étaient fréquentes.

2. La Solution : CUCo, le "Chef d'Orchestre Robot"

CUCo est un système automatisé qui écrit ces programmes complexes pour nous. Il ne se contente pas de copier-coller du code ; il réfléchit et expérimente. Il fonctionne en deux étapes, comme un apprenti puis un maître :

• Étape 1 : L'Apprenti Rapide (Fast-Path)
  Imaginez un apprenti qui a pour seule mission de s'assurer que le plat ne brûle pas et que le service fonctionne. Il ne cherche pas encore à être le plus rapide du monde, il veut juste que ce soit correct.
  Il prend le vieux système (où le chef attend le serveur) et le transforme en un système où le chef et le serveur travaillent dans la même cuisine. Il crée une version de base, sûre et fonctionnelle, mais qui n'est pas encore optimisée. C'est la fondation solide.

• Étape 2 : Le Maître Expérimentateur (Slow-Path)
  Une fois que l'apprenti a créé la version de base, le "Maître" prend le relais. C'est ici que la magie opère. Le Maître utilise une technique d'évolution (comme la sélection naturelle).

  • Il prend la recette de base.
  • Il crée 100 variantes : "Et si on envoyait le plat pendant qu'on coupe les légumes ?", "Et si on utilisait deux serveurs au lieu d'un ?", "Et si on changeait l'ordre des tâches ?".
  • Il teste ces variantes sur de vrais ordinateurs.
  • Il garde les meilleures, jette les mauvaises, et mélange les meilleures idées pour créer une nouvelle génération encore plus rapide.

Il répète ce processus des dizaines de fois, comme un éleveur qui sélectionne les chevaux les plus rapides, jusqu'à trouver la recette parfaite qui utilise chaque seconde de la machine.

3. Les Résultats : Une course de vitesse

Grâce à cette méthode, CUCo a réussi à faire ce que les humains peinent à faire :

• Supprimer les temps d'attente : Plus besoin d'attendre le serveur. Le chef cuisine pendant que le serveur porte le plat suivant.
• Gagner du temps : Sur les tests, CUCo a rendu les tâches jusqu'à 1,57 fois plus rapides. C'est énorme. Imaginez que votre trajet en voiture de 1 heure devienne 40 minutes.

En résumé

CUCo, c'est comme avoir un ingénieur génie qui ne dort jamais. Au lieu de vous demander de dessiner manuellement un plan de maison complexe et risqué, il dit : "Laissez-moi faire".

1. Il construit d'abord une maison solide et habitable (l'Apprenti).
2. Ensuite, il fait des milliers de rénovations, teste la solidité des murs, l'isolation, et la circulation de l'air, jusqu'à ce que la maison soit la plus confortable et la plus efficace possible (le Maître).

Le résultat ? Des intelligences artificielles (comme les modèles de langage que nous utilisons) qui apprennent et réfléchissent beaucoup plus vite, car elles ne perdent plus de temps à attendre que les données arrivent. C'est une révolution pour l'avenir de l'informatique.

Résumé technique

1. Problématique

L'entraînement et l'inférence de grands modèles de langage (LLM) distribués reposent de plus en plus sur des clusters massifs de GPU. Pour maximiser les performances, il est crucial d'optimiser non seulement les noyaux de calcul (computation), mais aussi les noyaux de communication, et surtout leur chevauchement (overlap).

Historiquement, ces deux domaines ont évolué de manière isolée :

• Orchestration côté hôte : Les noyaux de calcul et les primitives de communication (via NCCL) étaient gérés séparément par le CPU. Le CPU devait synchroniser les phases, introduisant des surcharges de synchronisation et limitant la granularité du chevauchement.
• Complexité de la fusion : Bien que de nouvelles interfaces (NVSHMEM, NCCL Device-Side) permettent désormais aux GPU d'initier directement les communications, la conception manuelle de noyaux CUDA qui fusionnent calcul et communication est extrêmement difficile. L'espace de conception est vaste, contre-intuitif et dépend fortement du matériel (topologie réseau, architecture GPU).
• Limites des outils actuels : Les compilateurs ML modernes supposent que la communication est une "effet secondaire" côté hôte. Les agents génératifs existants (basés sur LLM) échouent souvent dans ce contexte car ils ne comprennent pas les sémantiques distribuées, manquent de support pour les communications collectives multi-GPU et tendent à générer du code incorrect ou inefficace en sautant directement à la génération de code sans décomposition préalable.

2. Méthodologie : Le Framework CUCo

CUCo est un framework agentic (piloté par des agents IA) conçu pour générer automatiquement des noyaux CUDA fusionnés et performants, sans nécessiter d'entraînement préalable (training-free). Il repose sur trois composantes interdépendantes :

A. Spécification de l'Espace de Conception (Design Space Specification)

Pour guider l'agent, CUCo formalise l'espace d'optimisation en un ensemble de directives structurées, évitant ainsi les "hallucinations" de stratégies invalides. Cet espace est défini par cinq dimensions :

1. Backend (B) : Choix entre GIN (GPU-Initiated Networking, RDMA) ou LSA (Load/Store Accessible, mémoire locale).
2. Placement (P) : Aggressivité du chevauchement vs complexité de l'ordonnancement.
3. Portée de synchronisation (S) : Granularité de la synchronisation (CTA, Warp, Thread).
4. Émetteur (I) : Qui initie la communication (Thread, Warp, CTA).
5. Granularité (G) : Taille des blocs de transfert.

L'agent doit d'abord émettre une directive d'optimisation explicite avant de générer du code, ancrée dans le contexte matériel et les API NCCL.

B. Agent "Fast-Path" (Priorité à la Correction)

L'objectif de cet agent est de produire un noyau fonctionnel et correct rapidement, en ignorant temporairement l'optimisation fine des performances.

• Processus : Il analyse le code source hôte (calcul + appels NCCL), transforme les collectifs gérés par le CPU en primitives côté dispositif (Device-Side), et ajoute des annotations pour les blocs modifiables.
• Stratégie : Il utilise des paramètres conservateurs (synchronisation globale, granularité grossière) pour garantir que le noyau ne se bloque pas (deadlock) et reste cohérent.
• Résultat : Il fournit une "graine" (seed) fiable qui sert de point de départ pour l'agent suivant.

C. Agent "Slow-Path" (Optimisation par Évolution)

Cet agent cherche à maximiser les performances en explorant l'espace de conception à partir de la graine fournie par le Fast-Path.

• Algorithme : Il utilise une recherche évolutionnaire basée sur des "îles" (Multi-Island Population) pour éviter la convergence prématurée vers des optima locaux.
• Opérateur de Mutation par LLM : Au lieu de modifications aléatoires, un LLM propose des modifications sémantiques (patches de diff, réécritures complètes, croisements) uniquement dans les zones annotées comme modifiables.
• Boucle d'évaluation en cascade :
  1. Compilation : Rejet immédiat des erreurs de syntaxe.
  2. Exécution : Vérification de la correction numérique sur le cluster.
  3. Benchmark : Mesure de la latence et calcul du score de fitness.
• Rétroaction : Les échecs et les succès sont stockés dans une base de données partagée et utilisés pour guider les mutations futures (via un "MetaSummarizer").

3. Contributions Clés

1. Premier framework agentic pour la co-conception Calcul-Communication : CUCo est la première solution capable de générer automatiquement des noyaux CUDA fusionnés pour des environnements multi-GPU distribués, en gérant la complexité des sémantiques collectives.
2. Architecture à deux voies (Fast/Slow Path) : La séparation entre la génération d'un noyau correct (Fast-Path) et l'optimisation de performance (Slow-Path) résout le problème de la complexité excessive qui fait échouer les approches "tout-en-un".
3. Espace de conception contraint et explicite : L'introduction de directives structurées (Backend, Granularité, etc.) ancre le raisonnement de l'IA dans la réalité des API matérielles, réduisant drastiquement les erreurs de conception.
4. Support natif des API Device-Side : Le framework exploite pleinement les nouvelles capacités de NVSHMEM et NCCL Device-Side pour éliminer l'orchestration CPU.

4. Résultats Expérimentaux

CUCo a été évalué sur quatre charges de travail représentatives (Flash Attention, DeepSeek MoE, Transfert de KV Cache, GEMM + AllGather) sur des clusters de GPU A100 (interne et inter-nœuds).

• Performance : CUCo réduit la latence de bout en bout de 5,3 % à 26,2 % par rapport aux baselines gérées par l'hôte (NCCL standard).
• Accélération maximale : Jusqu'à 1,57x d'accélération sur certaines configurations.
• Analyse des gains : Les améliorations proviennent principalement de :
  • L'élimination des surcharges d'appels API côté CPU.
  • La suppression des temps d'attente (bubbles) entre les phases de calcul et de communication.
  • La possibilité d'un chevauchement au niveau de la tuile (per-tile pipelining), impossible avec les approches traditionnelles.
• Études d'ablation :
  • Sans l'agent Fast-Path, l'agent évolutionnaire échoue à produire un noyau correct pendant plusieurs générations (25% du budget gaspillé).
  • Une stratégie d'évolution en deux phases (Exploration puis Exploitation) surpasse significativement une approche purement exploitante.

5. Signification et Impact

Ce travail marque un tournant dans la programmation des systèmes ML distribués :

• Démocratisation de l'optimisation : Il rend accessible l'optimisation complexe de la fusion calcul-communication, qui était auparavant réservée à des experts en noyaux CUDA.
• Adaptabilité matérielle : En étant "hardware-aware", CUCo peut s'adapter automatiquement à différentes topologies de réseau (InfiniBand, RoCE) et architectures de GPU sans réécriture manuelle.
• Avenir du ML distribué : À mesure que les modèles grandissent et que la communication devient le goulot d'étranglement principal, des frameworks comme CUCo seront essentiels pour exploiter pleinement le potentiel des nouvelles générations de matériel et d'API de communication directe.

En résumé, CUCo démontre qu'une approche agentic structurée peut automatiser la conception de systèmes distribués complexes, dépassant les limites des méthodes manuelles et des compilateurs traditionnels.