Compiler des LLM en MegaKernel pour obtenir une inférence à faible latence
(zhihaojia.medium.com)- Un compilateur a été développé pour transformer automatiquement l’inférence des LLM en un megakernel unique
- L’approche MegaKernel (kernel persistant) intègre entièrement calcul et communication dans un seul kernel GPU pour l’inférence LLM, ce qui permet une latence très faible
- L’architecture distribuée des frameworks ML et des bibliothèques de kernels existants rend très difficile la transformation de tout le pipeline en un kernel unique
- Mirage Persistent Kernel (MPK) convertit automatiquement l’inférence LLM multi-GPU en un megakernel hautes performances grâce à un compilateur et un système d’exécution
- MPK transforme le graphe de calcul en graphe de tâches à grain fin, afin de maximiser le pipeline logiciel et le chevauchement calcul-communication
- Avec MPK, la latence de génération des tokens diminue par rapport aux systèmes existants, et le gain de performance devient encore plus marqué à mesure que le nombre de GPU augmente
Vue d’ensemble et avantages de l’approche MegaKernel
- Dans l’inférence des grands modèles de langage (LLM), l’un des moyens les plus efficaces de réduire la latence consiste à fusionner tous les calculs et toutes les communications dans un seul megakernel (kernel cohérent)
- Cette approche permet à un unique kernel GPU d’exécuter sans interruption tous les traitements, y compris les opérations couche par couche sur l’ensemble du modèle et les communications entre GPU
- Les principaux avantages sont les suivants
- suppression du surcoût de lancement des kernels en évitant les appels répétés
- possibilité de mettre en place un pipeline logiciel sur l’ensemble des couches
- exécution simultanée du calcul et de la communication pour masquer la latence
Limites des approches existantes et arrivée de MPK
- Les frameworks ML existants comme PyTorch, Triton et TVM ne prennent pas en charge nativement la génération automatique d’un megakernel end-to-end
- Les systèmes LLM réels reposent sur une combinaison de bibliothèques de kernels variées — NCCL/NVSHMEM (communication), FlashInfer/FlashAttention (attention), CUDA/Triton (opérations personnalisées) — ce qui rend difficile leur intégration dans un kernel unique
- Dans ce contexte, des chercheurs de CMU, UW, Berkeley, NVIDIA et Tsinghua ont développé Mirage Persistent Kernel (MPK)
- MPK associe un compilateur et un runtime pour convertir automatiquement l’ensemble du pipeline d’inférence LLM en megakernel hautes performances
Valeur clé de MPK
- MPK élimine complètement le surcoût de lancement des kernels et maximise le chevauchement entre calcul inter-couches, chargement des données et communication, afin de créer un environnement d’inférence LLM à très faible latence
- Lors de tests réels (prompt de 39 tokens, génération de 512 tokens, sans speculative decoding),
- sur une seule NVIDIA A100 40GB, la latence de décodage par token des systèmes optimisés existants comme vLLM/SGLang était de 14,5 ms, contre 12,5 ms avec MPK
- ce résultat se rapproche de la borne théorique inférieure (10 ms), sur la base d’une bande passante mémoire de 1,6 TB/s et d’un chargement de poids de 16 GB
- En environnement multi-GPU, l’intégration complète du calcul et de la communication rend l’avantage de performance de MPK encore plus visible à mesure que le nombre de GPU augmente
Structure détaillée du fonctionnement de MPK
Part 1. Compilateur – transformation du graphe de calcul LLM en graphe de tâches
- En général, le calcul d’un LLM est représenté sous forme de graphe de calcul, où chaque opération (par ex. multiplication de matrices, attention) ou opération de communication (par ex. all-reduce) constitue un nœud, et où les dépendances de données sont représentées par des arêtes
- Dans les conceptions classiques, il est courant d’exécuter un kernel distinct par opérateur, mais cette approche ne reflète que des dépendances au niveau kernel, et non au niveau réel des données dépendantes, ce qui limite les possibilités de pipeline
- Exemple : lorsqu’un all-reduce suit une multiplication de matrices, l’exécution de l’all-reduce ne peut commencer qu’une fois toute la multiplication terminée. En pratique, il est pourtant possible de découper les données et d’exploiter des dépendances partielles
- Le compilateur MPK affine le graphe de calcul et le transforme automatiquement en graphe de tâches à grain fin adapté aux unités réelles de données
- chaque tâche (rectangle) est une unité de calcul ou de communication affectée à un SM GPU individuel
- chaque événement (cercle) est un point de synchronisation entre tâches
- les arêtes entre tâches et événements représentent efficacement les dépendances de données et de contrôle
- Grâce à ce graphe de tâches, MPK permet un chevauchement plus important, partiel ou parallèle, entre calcul et communication
- Le superoptimiseur de kernels Mirage génère aussi automatiquement une implémentation CUDA hautes performances adaptée à chaque tâche
Part 2. Runtime – exécution du graphe de tâches à l’intérieur du megakernel
- Le runtime MPK exécute intégralement le graphe de tâches à l’intérieur d’un seul kernel GPU (megakernel)
- Tous les SM (Streaming Multiprocessors) du GPU sont répartis statiquement entre des rôles de workers et de schedulers
Workers
- Chaque worker fonctionne au niveau d’un SM et gère sa propre file de tâches
- En boucle, il
- récupère la tâche suivante dans la file
- l’exécute (par ex. matmul, attention, transfert de données)
- notifie un événement une fois la tâche terminée
- répète le processus
- Cela permet d’optimiser l’utilisation des ressources de chaque worker et d’exécuter des opérations inter-couches de manière asynchrone
Schedulers
- Des schedulers distribués fonctionnent au niveau d’un seul warp dans chaque SM, avec jusqu’à 4 schedulers pouvant s’exécuter simultanément
- Chaque scheduler gère une file d’événements activés et assigne aux workers les tâches dont les conditions sont remplies
- Cela permet une distribution à grande échelle des tâches sans surcoût de synchronisation centralisée
Exécution pilotée par événements
- Lorsqu’une tâche se termine, elle incrémente un compteur d’événement spécifique. Quand le compteur atteint un seuil, l’événement est activé et inséré dans la file du scheduler
- Le scheduler exécute ensuite les tâches dépendantes de cet événement
- Cela permet naturellement un pipeline logiciel à grain fin et un chevauchement calcul-communication
- par ex. le matmul d’une couche peut s’exécuter en même temps que l’attention d’une autre couche
- dès qu’un résultat partiel de matmul est disponible, la communication allreduce peut démarrer
- Comme toute la planification et tous les changements de tâche se produisent dans le contexte d’un seul kernel, le surcoût inter-tâches reste très faible, de l’ordre de 1 à 2 microsecondes (μs)
Orientations futures
-
Objectif de MPK : permettre aux développeurs de compiler facilement des LLM en megakernel avec seulement quelques dizaines de lignes de Python, tout en obtenant des performances maximales
-
Principales pistes d’évolution
- prise en charge des architectures GPU récentes : par exemple NVIDIA Blackwell, spécialisations au niveau du warp, etc.
- gestion des workloads dynamiques : recherche sur des stratégies de compilation pour les modèles nécessitant un contrôle de flux dynamique, comme les mixture-of-experts (MoE)
- ordonnancement avancé des tâches : exploration et application possible de politiques modernes fondées sur la priorité, l’optimisation du débit, etc.
-
MPK représente un tournant fondamental dans la manière de compiler et d’exécuter l’inférence LLM sur GPU, et l’équipe souhaite élargir la collaboration avec la communauté
Ressources supplémentaires
- Le code et la documentation de MPK (Mirage Persistent Kernel), ainsi que les résultats de recherche les plus récents, sont disponibles sur GitHub (https://github.com/mirage-project/mirage)
1 commentaires
Commentaire Hacker News
À l’auteur : il est intéressant de voir que l’approche par interpréteur on-GPU semble être une direction d’avenir très prometteuse. Il existe aussi d’autres travaux avec une approche presque identique, donc je recommande de jeter un œil à ce post lié. On voit de ses propres yeux à quel point le modèle de programmation fondamental de CUDA (par ex. le lancement de kernel) est contourné pour du parallélisme fondé sur des tâches très fines, et comment cette approche augmente davantage l’utilisation du matériel. Je me demande si CUDA ne nous a pas retenus à plusieurs niveaux. J’espère aussi que cette recherche de l’auteur pourra entrer comme backend expérimental de PyTorch. Et, petit détail de typo : les deux paragraphes de la première partie sont presque identiques.
J’ai travaillé de près pendant un certain temps sur vLLM et SGLang, et je suis convaincu que ce projet ressemble exactement à la suite idéale. J’ai été impressionné par l’analyse du graphe de dépendances des opérations, la fusion d’opérations et la planification plus intelligente des tâches. Félicitations à l’équipe.
J’ai parcouru l’article et le README GitHub, et je trouve le projet vraiment excellent. Je me demande si ce type d’optimisation pourrait s’appliquer non seulement à l’inférence, mais aussi à la phase d’entraînement. Je comprends en particulier que la fusion des opérations backward et des communications de gradients constitue un défi. Je crois comprendre qu’à l’heure actuelle vous ne prenez pas en charge les workloads dynamiques (par ex. MoE), mais je mentionne un article récent sur le traitement de MoE dans un seul kernel : FlashDMoE: Fast Distributed MoE in a Single Kernel.
Merci d’avoir lu l’article et même le README. La prise en charge de l’entraînement est également possible, mais comme les kernels d’entraînement sont en général plus gros, l’overhead de lancement de kernel y est moins problématique ; c’est donc l’inférence, en particulier à faible latence, qui en bénéficie le plus. J’ai aussi trouvé l’article FlashDMoE que vous avez partagé très intéressant, et je souligne que la prise en charge des modèles MoE fait partie de nos prochains objectifs.
Personnellement, je suis un peu sceptique quant au fait d’investir du temps dans l’optimisation de l’entraînement fondé sur les gradients. En pratique, beaucoup de tâches d’apprentissage ont des caractéristiques de valeurs discrètes, et je pense qu’elles se prêtent mal à un apprentissage fondé sur les gradients.
Le rêve, pour la suite, ce serait de compiler directement en Verilog et d’acheter son propre matériel LLM sur AliExpress.
Partage d’un article présentant des technologies matérielles comme Chisel. Avant l’arrivée de l’IA et des GPU, cette idée de conversion directe du logiciel vers le matériel était une approche prometteuse. Comme les progrès des CPU sont au point mort et que le désir d’optimiser davantage la couche intermédiaire entre logiciel et matériel reste constant, il est probable que le calcul parallèle de style GPU demeure la voie principale de l’accélération. Les CPU généralistes devraient au final rester comme un petit cerveau chargé de piloter les GPU. En revanche, il semble peu probable qu’une approche de conversion directe du logiciel vers le matériel devienne dominante.
Si la structure des LLM se stabilise dans 5 à 10 ans, il deviendra peut-être réaliste de les mapper directement sur du matériel. Avec la technologie actuelle, il est évoqué que même des centaines de milliards de paramètres pourraient tenir sur un seul wafer en n’utilisant que des portes logiques à ultra-basse précision, autour de 1,5 bit. Quand la précision augmente, le nombre de portes croît exponentiellement ; pour l’instant, il est donc plus efficace de conserver une mémoire de poids et de partager les unités de calcul. À l’avenir, le développement de LLM à ultra-basse précision deviendra une tâche indispensable.
Blague sur le fait que le coût de l’entraînement est déjà élevé et qu’ajouter en plus le coût du masking ne ferait qu’aggraver les choses, accompagnée du constat plus lucide que les startups de matériel IA explorent en réalité ce genre de direction depuis longtemps.
L’idée d’un LLM-in-a-box réellement existant serait très séduisante. Je vais bientôt avoir l’occasion de travailler dans un environnement hors ligne (
air-gap), et une telle solution me semblerait extrêmement utile.J’ai fait tourner le code moi-même dans un environnement GPU Modal, et les chiffres de gain de performances avancés dans la recherche se reproduisent bien en pratique. Partage du code de résultats du projet mirage. Avec la combinaison Triton + FlashInfer, la latence était d’environ 19,2 ms par token ; avec MPK, dans les mêmes conditions, j’ai observé une forte amélioration à 7,7 ms.
J’avais participé autrefois à un petit concours CUDA. C’était sur des algorithmes parallèles en image ou en vision, et j’avais voulu faire le malin en mettant en cache les résultats intermédiaires en mémoire. Quand les résultats du concours sont tombés, j’ai été surpris de voir que les autres avaient soumis du code bien plus rapide que le mien. En regardant pourquoi, j’ai découvert qu’ils ne mettaient pas du tout ces résultats intermédiaires en cache, mais les recalculaient en permanence. Le coût du calcul était bien plus faible que les allers-retours mémoire. J’imagine que ce projet suit probablement une logique similaire. En compilant en megakernel, les frontières entre couches disparaissent ; le partage de résultats intermédiaires diminue et la quantité de calcul augmente, mais au global le gain est énorme grâce à la réduction des allers-retours mémoire. Il y aura probablement un sweet spot, notamment pour les réseaux convolutionnels, mais je ne sais pas comment ce point est géré dans le megakernel.
On continue encore aujourd’hui à voir apparaître de nouvelles métaphores pour les LLM. Je me demande parfois s’il ne faudrait pas considérer les LLM comme des transistors. J’imagine qu’on est peut-être à un stade comparable à celui des ordinateurs de la taille d’une pièce qui ne faisaient que des multiplications avec des cartes perforées. C’est amusant d’imaginer ce qui se passerait si on pouvait lancer simultanément un million de requêtes o3-pro.
Ce projet vient de CMU (Carnegie Mellon). Mention du blog de Hazy Research à Stanford sur les megakernels, No Bubbles. Il est impressionnant de voir à quel point la concurrence est active dans ce domaine. (Ajout) Il existe aussi un article traitant de la vision plus large du projet "mirage", mais il ne couvre pas l’approche megakernel : lien vers l’article
L’auteur du post répond lui-même. Il confirme que les recherches avec Stanford avancent en parallèle. La principale différence est l’accent mis sur un compilateur de génération automatisée de megakernels.
Il est aussi mentionné que ThunderKittens de Hazy Research est une bibliothèque vraiment très cool. Beaucoup d’efforts se concentrent récemment sur la formalisation, le pipelining, le diviser-pour-régner, la maximisation de l’efficacité, et le développement de compilateurs/DSL dédiés afin d’exploiter au maximum les modèles récents de GPU NVIDIA.
Si les chiffres de performances sur Qwen 8B sont confirmés, ce sera extrêmement impressionnant. Cela paraît plus pratique que les approches megakernel précédentes. Ce type de kernel maintenu à raison d’un par SM rappelle l’ancien Larrabee. Je me demande à quoi ressemblerait le monde actuel si l’on avait suivi une voie plus traditionnelle processus-thread-SIMD, plutôt que CUDA.
Idée autour de la création de LLM fixes en pur ASIC plutôt que par inférence logicielle. Y aurait-il un avantage en coût ? Serait-il possible d’ajouter des couches que le logiciel pourrait encore manipuler ou affiner ? Étant donné qu’on semble presque avoir atteint un niveau « assez bon », on pourrait voir dans les 2 à 4 prochaines années des décisions consistant à figer cela dans des puces spécialisées. Je me demande à partir de quel moment les avantages d’un matériel ultra-spécialisé deviendraient vraiment déterminants.