Aller au contenu

4 · Mirage Persistent Kernel

L'approche compilateur. Là où Stanford écrit les instructions à la main, CMU les génère — à partir d'un modèle Hugging Face et de quelques dizaines de lignes de Python.

4.1 Le papier

Titre : Mirage Persistent Kernel: A Compiler and Runtime for Mega-Kernelizing Tensor Programs

Date : 22 décembre 2025 (arXiv:2512.22219)

Auteurs : Xinhao Cheng, Zhihao Zhang, Yu Zhou, Jianan Ji, Jinchen Jiang, Zepeng Zhao, Ziruo Xiao, Zihao Ye, Yingyi Huang, Ruihang Lai, Hongyi Jin, Bohan Hou, Mengdi Wu, Yixin Dong, Anthony Yip, Songting Wang, Wenqin Yang, Xupeng Miao, Tianqi Chen, Zhihao Jia

Affiliations : Carnegie Mellon University, Tsinghua University, NVIDIA, University of Michigan, Purdue University

Il se présente comme « le premier compilateur et système d'exécution qui transforme automatiquement l'inférence de modèle multi-GPU en un unique mega-noyau de haute performance ».

Échelle de l'implémentation : ~40 000 lignes de C++, ~84 000 lignes de CUDA, ~10 000 lignes de Python.

4.2 Le problème qu'il attaque

Le constat des auteurs : les cadres existants ne savent pas faire cela.

Les cadres d'apprentissage automatique de haut niveau comme PyTorch, Triton et TVM ne supportent pas nativement la génération de mega-noyaux de bout en bout, et les systèmes LLM modernes sont construits à partir de bibliothèques de noyaux spécialisés diverses, ce qui rend difficile la consolidation de tout le pipeline d'inférence en un seul noyau.

Autrement dit : le problème n'est pas seulement technique, il est architectural. Un moteur d'inférence assemble cuBLAS, FlashAttention, des noyaux de quantification, NCCL — chacun étant un noyau autonome.

4.3 Le ttGraph : un graphe au niveau du SM

L'innovation centrale.

Là où un graphe de calcul classique a pour nœuds des opérateurs (une GEMM, une attention), le ttGraph de MPK a pour nœuds des tâches — l'unité de travail d'un seul SM.

Graphe classique              ttGraph (niveau SM)
────────────────              ───────────────────
[GEMM QKV]                    [t0][t1][t2] ... [t147]   ← 148 tâches
    │                              │
    ▼                          (événement e0)
[Attention]                        │
    │                         [t148][t149] ...
    ▼                              │
[Proj O]                      (événement e1)
                                   │
                                  ...

Tâches et événements alternent : chaque tâche n'a que des arêtes sortantes vers des événements déclencheurs, et des arêtes entrantes depuis des événements de dépendance.

Cette granularité permet exactement ce que le découpage manuel de Hazy Research faisait à la main : une tâche produit une portion de sortie, et les tâches consommatrices de cette portion précise sont débloquées immédiatement.

4.4 Le pipeline de compilation

Cinq étapes, toutes automatiques.

1. Décomposition d'opérateurs. Partitionner la sortie de chaque opérateur entre les SM. Une GEMM \(4096 \times 4096\) devient \(N\) tâches, chacune produisant une tuile.

2. Analyse de dépendances. Énumérer les paires de tâches et introduire un événement chaque fois que la région de sortie de l'une recouvre la région d'entrée de l'autre.

3. Fusion d'événements. Beaucoup d'événements sont redondants. MPK applique une fusion par ensembles de successeurs et par ensembles de prédécesseurs : si deux événements ont exactement les mêmes tâches en aval (ou en amont), ils sont fusionnés.

4. Normalisation du graphe. Garantir que chaque tâche a au plus un événement de dépendance et un événement déclencheur. Cela simplifie radicalement le runtime.

5. Linéarisation du graphe. Un parcours en largeur ordonne les tâches de sorte que celles déclenchées par un même événement occupent des indices contigus — ce qui permet au runtime de les distribuer par plages plutôt qu'une par une.

Pourquoi ces cinq étapes comptent

Les étapes 3, 4 et 5 sont ce qui distingue un compilateur d'un simple générateur. Elles réduisent le nombre d'objets de synchronisation et rendent leur manipulation régulière — donc rapide.

Sans elles, un graphe au niveau SM sur un modèle de 32 couches contiendrait des dizaines de milliers d'événements, et le coût de leur gestion annulerait le gain.

4.5 Le runtime parallèle embarqué

MPK partitionne les SM en deux rôles.

Configuration SM totaux Workers Schedulers
A100 108 104 16
H100 132 128 16
B200 148 144 16

Les workers exécutent des files de tâches. Les schedulers gèrent les dépendances : ils surveillent les événements, et quand un événement est activé (ses prérequis sont satisfaits), ils poussent les tâches en aval dans les files.

Dédier 16 SM à l'ordonnancement

C'est un choix contre-intuitif : 11 à 15 % des SM ne font aucun calcul.

La justification : le coût d'une gestion de dépendances distribuée (chaque worker sondant tous les compteurs qui le concernent) serait supérieur. En centralisant, on réduit le trafic atomique et on permet un ordonnancement plus intelligent.

Les auteurs notent d'ailleurs comme limite que « l'implémentation actuelle utilise uniquement l'état local » pour l'ordonnancement, et identifient l'ordonnancement centralisé comme une piste d'exploration.

Le lancement hybride

MPK combine deux modes :

  • juste-à-temps (JIT) : les tâches sont poussées dans les files au moment où elles deviennent prêtes → bon équilibrage de charge ;
  • anticipé (AOT) : certaines tâches sont pré-affectées → latence réduite, pas d'attente sur l'ordonnanceur.

Les workers priorisent les tâches JIT et vérifient les tâches AOT quand leur file JIT est vide.

Les optimisations du runtime

Optimisation Effet mesuré
Mémoire partagée paginée (pages de 32 Ko) partage de ressources entre tâches
Pipelining logiciel entre tâches 1,2-1,3× sur la couche linéaire finale
Préchargement des descriptions de tâches en mémoire partagée réduit la latence de dispatch
Recouvrement calcul/communication 1,1× de réduction de latence
Fusion gather-GEMM pour MoE élimine un prétraitement représentant jusqu'à 11 % du temps du MoE

4.6 Les résultats

Mono-GPU

Comparaison Résultat
Contre SGLang et vLLM 1,0 – 1,7×
Qwen3-8B sur A100, latence par jeton 14,5 ms → 12,5 ms
Borne inférieure théorique estimée ~10 ms

Modèles testés : Qwen3 (8 B et 30 B), LLaMA, et d'autres, sur A100, H100 et B200, avec des lots de 1 à 16.

Le chiffre le plus honnête du papier

« 14,5 ms → 12,5 ms, avec une borne inférieure théorique d'environ 10 ms. »

Cela signifie : sur les 4,5 ms de surcoût par rapport au plancher physique, MPK en récupère 2 sur 4,5, soit 44 %. Il reste 2,5 ms de surcoût irréductible en l'état.

C'est une présentation bien plus informative qu'un simple « 1,16× ». Elle situe le résultat par rapport à la limite physique, pas seulement par rapport à un concurrent.

Multi-GPU

Comparaison Résultat
Contre SGLang/vLLM, 8 H100 en tensor-parallèle 1,1 – 1,4×
Contre PyTorch avec CUDA Graphs jusqu'à 10×

La communication inter-GPU utilise NVSHMEM, ce qui permet de l'exécuter depuis l'intérieur du megakernel et donc de la recouvrir avec le calcul.

L'ablation

Optimisation retirée Impact
Pipelining entre tâches −1,2 à 1,3× sur la couche linéaire finale
Recouvrement calcul/communication −1,1× de latence
Fusion gather-GEMM (MoE) +11 % de temps sur le MoE

4.7 L'usage

C'est l'argument le plus fort de MPK : la facilité.

Mirage permet de compiler des LLM du zoo Hugging Face en un mega-noyau en quelques dizaines de lignes de Python — principalement pour définir les entrées et les sorties du noyau.

Le dépôt est mirage-project/mirage, et le projet est référencé par le groupe Catalyst de CMU.

C'est la différence fondamentale avec l'approche de Stanford : là où il faut écrire sept instructions CUDA à la main, MPK les génère.

4.8 Les limites reconnues

Le papier les énonce explicitement, ce qui est appréciable.

Limite Détail
Surcoût de compilation optimisation par taille de lot ; plusieurs ttGraphs nécessaires pour des lots dynamiques
Surcoût mémoire la normalisation ajoute des tâches et des événements (annoncé < 1 %)
Ordonnancement décentralisé état local uniquement ; l'ordonnancement centralisé reste à explorer
Charges dynamiques il faut pré-générer des ttGraphs pour des tailles de lot représentatives
Taille des descriptions 352 octets par tâche en mémoire du périphérique

La limite du dynamisme est la plus structurante, et c'est précisément celle qu'attaque Event Tensor — dont plusieurs auteurs sont les mêmes.

4.9 Ada-MK : le megakernel en production

Un travail complémentaire, et le premier déploiement industriel documenté.

Titre : Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference

Date : 12 mai 2026 (arXiv:2605.11581)

Auteurs : Wenxin Dong, Mingqing Hu, Guanghui Yu, Qiang Fu, Peng Xu, Hui Xu, Yue Xing, Xuewu Jiao, Shuanglong Li, Lin Liu — Baidu

Le contexte

Le résumé pose le cadre :

Quand les grands modèles de langage servent l'inférence en temps réel dans les systèmes publicitaires commerciaux en ligne, la latence de bout en bout doit être strictement bornée à l'ordre de la milliseconde. Or chaque jeton généré pendant la phase de décodage déclenche des milliers de lancements de noyaux, et le surcoût de lancement seul peut représenter 14,6 % du temps d'inférence de bout en bout.

Le problème spécifique

Les megakernels existants font face à « une tension fondamentale entre portabilité et efficacité sur des GPU à ressources contraintes comme NVIDIA Ada » — c'est-à-dire les cartes de la classe L4, L40S, RTX 4090, qui ont bien moins de mémoire partagée par SM que Hopper (100 Ko contre 227) et pas de TMA.

La méthode

  • une recherche sur DAG hors ligne, combinée à MLIR, qui détermine les chemins d'exécution optimaux à la compilation, éliminant le branchement à l'exécution ;
  • un découpage selon la dimension K pour tenir dans la mémoire partagée contrainte.

Les résultats

Comparaison Gain
Contre TensorRT-LLM (débit à lot 1) +23,6 %
Contre vLLM +50,2 %
Réduction du pic de mémoire partagée −50 %
Réduction de latence de bout en bout (vLLM, SGLang, TensorRT-LLM) 10 à 50 %

Le déploiement

C'est le point remarquable : Ada-MK est déployé en production dans le système publicitaire en ligne commercial de Baidu, via un « moteur hybride hétérogène qui intègre le MegaKernel comme greffon de TensorRT-LLM ».

L'architecture est instructive : le système utilise TensorRT-LLM pour le préremplissage (haut débit) et le megakernel pour le décodage (basse latence), « sans coût de refonte métier supplémentaire ».

Ce que ce déploiement démontre

Les megakernels sont sortis du laboratoire. Et l'architecture retenue — megakernel pour le décodage, moteur classique pour le préremplissage — est cohérente avec toute l'analyse de ce document : le megakernel est rentable exactement là où le travail par passe est faible.

C'est aussi une réponse pratique au problème de la généralité : plutôt que de tout mega-noyauter, on l'insère comme composant.

Résumé du chapitre

À retenir

  • MPK est le premier compilateur de megakernel : il génère automatiquement ce que Stanford écrit à la main.
  • Son innovation est le ttGraph, un graphe au niveau du SM où tâches et événements alternent.
  • Cinq étapes de compilation : décomposition, analyse de dépendances, fusion d'événements, normalisation, linéarisation.
  • Le runtime dédie 16 SM à l'ordonnancement et le reste au calcul, avec un lancement hybride JIT/AOT.
  • Résultats : 1,0-1,7× contre SGLang/vLLM en mono-GPU (Qwen3-8B/A100 : 14,5 → 12,5 ms, plancher théorique ~10 ms), 1,1-1,4× sur 8 H100.
  • Utilisation : « quelques dizaines de lignes de Python » depuis Hugging Face.
  • Limites reconnues : recompilation par taille de lot, dynamisme non traité, ordonnancement décentralisé.
  • Ada-MK (Baidu, mai 2026) est déployé en production, comme greffon TensorRT-LLM : décodage par megakernel, préremplissage par le moteur classique. +23,6 % sur TensorRT-LLM, +50,2 % sur vLLM.

Vérifiez que vous avez compris

Pourquoi dédier 16 SM à l'ordonnancement plutôt que de laisser chaque worker gérer ses dépendances ?

Parce que la gestion décentralisée coûterait plus cher qu'elle ne rapporte.

Si chaque worker sonde lui-même les compteurs qui le concernent, on obtient 132 blocs faisant des lectures atomiques répétées sur des adresses partagées. Les atomiques s'exécutent dans le L2 et se sérialisent sous contention : le trafic devient prohibitif.

En centralisant sur 16 SM, on réduit le nombre de sondeurs d'un facteur ~8, et surtout on permet un ordonnancement intelligent : les schedulers voient l'état global et peuvent équilibrer.

Le coût est de 11 % des SM. Compte tenu du fait que le noyau est limité par la mémoire (les SM ne sont pas la ressource rare), c'est un bon échange.

MPK annonce « 1,0-1,7× ». Pourquoi la borne basse est-elle 1,0 ?

Parce que sur certaines configurations, MPK n'apporte rien.

C'est cohérent avec l'analyse du chapitre 1 : le gain est proportionnel à la part du surcoût dans le temps total. Sur un grand modèle à lot élevé, cette part est faible et le gain aussi.

Publier la borne basse est un signe d'honnêteté méthodologique. Beaucoup de travaux ne rapportent que le maximum, ce qui donne « 6,7× » ou « 10× » sans préciser la configuration.

Le chiffre de 10× contre PyTorch avec CUDA Graphs est du même ordre : il est vrai, et la ligne de base n'est pas un système d'inférence optimisé.

Ada-MK utilise TensorRT-LLM pour le préremplissage et le megakernel pour le décodage. Pourquoi ne pas tout faire en megakernel ?

Trois raisons, toutes pratiques.

  1. Le gain n'y est pas. Le préremplissage est limité par le calcul (\(I \approx bS\)), et le surcoût de lancement y est négligeable. Un megakernel n'y apporterait rien.
  2. Le préremplissage a besoin de GEMM de pointe. TensorRT-LLM et cuBLAS y sont excellents ; réimplémenter dans un megakernel serait un recul.
  3. Le coût d'intégration. Le greffon TensorRT-LLM permet de conserver toute l'infrastructure existante — gestion du cache KV, ordonnancement, quantification, API — et de n'y insérer que le composant qui apporte quelque chose.

C'est un bon modèle d'adoption : le megakernel comme composant, pas comme remplacement. C'est probablement ainsi que la technique se diffusera.

Chapitre suivant : 5 · Event Tensor et le dynamisme

Sources de ce chapitre