Technische Samenvatting: MPK - Een Compiler en Runtime voor Mega-Kernelizing Tensorprogramma's

1. Probleemstelling

Huidige high-performance GPU-inferentiesystemen voor grote taalmodellen (LLM's) en andere tensorprogramma's vertrouwen voornamelijk op een kernel-per-operator executiemodel. In dit paradigma wordt elke operator (bijv. matrixvermenigvuldiging, aandacht, all-reduce) uitgevoerd als een afzonderlijke GPU-kernel. Hoewel deze aanpak het programmeren vereenvoudigt, introduceert het drie kritieke beperkingen die de prestaties belemmeren:

1. Impliciete Kernelbarrières: Conventionele GPU-executie dwingt een barrière af tussen opeenvolgende kernel-launches op dezelfde stream. Dit voorkomt cross-operator software pipelining, waardoor opeenvolgende operators strikt sequentieel moeten worden uitgevoerd en "pipeline bubbles" ontstaan waarbij hardwarebronnen (Tensor Cores, TMAs) ongebruikt blijven.
2. Grofmazige Overlap: Afhankelijkheden worden afgedwongen op het niveau van volledige operators. Bijgevolg is fine-grained overlap tussen berekening en communicatie onmogelijk. Bijvoorbeeld, een AllReduce-operatie moet wachten op de volledige voltooiing van een voorafgaande MatMul, zelfs als specifieke fragmenten van de reductie alleen afhankelijk zijn van specifieke fragmenten van de output van de vermenigvuldiging.
3. Launch Overhead en Statische Flexibiliteit: Het uitvoeren van honderden of duizenden kernels per inferentie-iteratie brengt aanzienlijke launch-overhead met zich mee. Hoewel CUDA Graphs dit verzachten, zijn ze grotendeels statisch; dynamische veranderingen in controleflow, tensorvormen of datadependenties vereisen vaak het opnieuw instantiëren van de grafiek, wat de flexibiliteit beperkt voor dynamische workloads zoals LLM-serving.

Bestaande oplossingen, zoals handmatige mega-kernels (bijv. FlashDMoE, Spector et al.), bieden prestatievoordelen maar vereisen uitgebreide engineeringinspanning en diepgaande GPU-expertise, waardoor ze niet generaliseerbaar zijn over modellen of hardware.

2. Methodologie

De auteurs presenteren Mirage Persistent Kernel (MPK), de eerste compiler- en runtimesysteem dat multi-GPU modelinferentie automatisch transformeert naar een enkele, high-performance mega-kernel. MPK bereikt dit door twee primaire componenten:

A. SM-Level Grafische Representatie (tGraph)

MPK introduceert een nieuwe grafische representatie genaamd tGraph die berekening en communicatie vastlegt op het niveau van individuele Streaming Multiprocessors (SM's), in plaats van de gehele GPU.

• Nodes: Vertegenwoordigen taken (berekening of communicatie) die op een enkele SM draaien.
• Edges: Codificeren fine-grained afhankelijkheden tussen taken via events.
• Voordeel: Deze granulariteit legt parallellisme bloot dat ontoegankelijk is voor kernel-niveau modellen, wat cross-operator software pipelining en fine-grained overlap van berekening en communicatie mogelijk maakt.

Compiler Optimalisaties:
De MPK-compiler transformeert een standaard tensorprogramma naar een geoptimaliseerde tGraph via verschillende stadia:

1. Operator Decompositie: Splitst operators op in disjuncte sets van taken op basis van output-tensorpartitionering (bijv. het partitioneren van matrixvermenigvuldigingsrijen/kolommen) om overeen te komen met het aantal beschikbare SM's.
2. Dependency Analyse: Inventariseert producer-consumer paren tussen taken om synchronisatie-events in te voegen waar alleen dat overlappen bestaat.
3. Event Fusie: Voegt redundante synchronisatiepunten samen met behulp van successor-set fusion (het samenvoegen van events met identieke consumer-sets) en predecessor-set fusion (het samenvoegen van events met identieke producer-sets) om de grafiekcomplexiteit te verminderen.
4. Normalisatie: Transformeert de grafiek zodat elke taak maximaal één afhankelijke event en één triggerende event heeft. Dit wordt bereikt door dummy-taken en -events in te voegen om fan-in/fan-out af te handelen, wat zorgt voor een uniforme, low-overhead representatie.
5. Linearisatie: Gebruikt een BFS-gebaseerd algoritme om opeenvolgende indices toe te wijzen aan taken die door dezelfde event worden getriggerd. Dit stelt de runtime in staat om event fan-out te coderen als een eenvoudige indexrange (eerste/laatste taak ID) in plaats van een expliciete lijst, wat de device memory footprint aanzienlijk verkleint.
6. Taak Implementatie Generatie: Gebruikt compiler superoptimalisatie (Mirage) om automatisch high-performance CUDA-implementaties voor elke individuele taak te genereren, inclussief intra-SM optimalisaties zoals registerhergebruik en shared memory layout.

B. In-Kernel Parallel Runtime

MPK voert de tGraph uit binnen een enkele persistente mega-kernel met behulp van een gedecentraliseerde, event-gestuurde architectuur:

• Worker/Scheduler Partitionering: De SM's van de GPU worden statisch verdeeld in Workers (die taken uitvoeren) en Schedulers (die event-queues beheren).
• Event-Gestuurde Executie: Schedulers pollen op geactiveerde events (waar aan alle vereisten is voldaan) en dispatchen de bijbehorende taken naar workers. Dit elimineert de noodzaak voor host-side synchronisatie of herhaalde kernel-launches.
• Hybride Taak Launch: Om een balans te vinden tussen latency en load balancing, gebruikt MPK twee modi:
  • Just-In-Time (JIT): Taken worden pas gelanceerd nadat hun afhankelijkheden zijn voldaan. Dit past zich aan dynamische workloads aan (bijv. variabele sequentielengtes in attention) om load imbalance te voorkomen.
  • Ahead-Of-Time (AOT): Taken worden vooraf in worker-queues geplaatst voordat hun afhankelijkheden zijn voldaan. Dit vermindert de dispatch-latency voor voorspelbare operators.
• Paged Shared Memory: Om cross-task software pipelining mogelijk te maken, introduceert MPK een gepaged abstractie voor shared memory. Taken verkrijgen en geven fixed-size pages vrij, waardoor de volgende taak data kan prefetchen naar shared memory terwijl de huidige taak berekent, zonder resourceconflicten.

3. Belangrijkste Bijdragen

1. Eerste Geautomatiseerde Mega-Kernel Systeem: MPK is het eerste systeem dat willekeurige multi-GPU tensorprogramma's automatisch transformeert naar een enkele mega-kernel, waardoor de noodzaak voor handmatige kernel fusie verdwijnt.
2. SM-Level Abstractie: De introductie van de tGraph-representatie verschuift dependency tracking van het operator-niveau naar het SM-taak niveau, wat optimalisaties (pipelining, fine-grained overlap) mogelijk maakt die voorheen onhaalbaar waren onder het kernel-per-operator model.
3. End-to-End Compiler en Runtime: Het systeem integreert een compiler die taakgrafieken optimaliseert en CUDA-code genereert met een runtime die deze taken asynchroon binnen een enkele kernel uitvoert, waarbij ondersteuning wordt geboden voor dynamische workloads zoals continuous batching en paged attention.
4. Hybride Executiestrategie: De combinatie van JIT en AOT taak-launching stelt het systeem in staat om zowel de dynamische aard van LLM-inferentie (variabele batch/sequentielengtes) als de statische efficiëntie van vaste operators te verwerken.

4. Evaluatieresultaten

De auteurs hebben MPK geëvalueerd op vijf veelgebruikte LLM's (inclusief varianten van Qwen3 en Llama-3) over drie generaties NVIDIA GPU's (A100, H100, B200), waarbij het werd vergeleken met state-of-the-art kernel-per-operator systemen (SGLang, vLLM) en vanilla PyTorch.

• Prestatiewinst: MPK bereikt 1.0× tot 1.7× lagere end-to-end inferentie-latency vergeleken met SGLang en vLLM. De verbeteringen zijn het meest uitgesproken voor kleinere modellen en nieuwere hardwaregeneraties (bijv. B200), waar kernel launch-overhead en pipeline bubbles relatief significanter zijn ten opzichte van de rekentijd.
• Doorvoersnelheid (Throughput): In multi-GPU settings (tot 8 H100's) verbetert MPK de throughput met 1.1× tot 1.4× ten opzichte van geoptimaliseerde baselines.
• Latency Limieten: Op een A100 verminderde MPK de per-token decoding latency voor Qwen3-8B van 14,5 ms (vLLM/SGLang) naar 12,5 ms, waarmee de geschatte hardware lower bound van ~10 ms wordt benaderd.
• Ablatie Studies:
  • Cross-Task Pipelining: Verminderde de taak runtime met 1.2–1.3×.
  • Compute-Communication Overlap: Verminderde de per-iteratie latency met 1.1×.
  • Kernel Launch Reductie: Elimineerde de ~1.1 ms per-token overhead geassocieerd met 293 kernel launches per token in baseline systemen.
• Schaalbaarheid: MPK schaalt effectief over 2, 4 en 8 GPU's, waarbij de snelheidswinst ten opzichte van baselines behouden blijft.

5. Betekenis en Claims

Het artikel claimt dat MPK een paradigmaverschuiving vertegenwoordigt in GPU-inferentiesystemen door de executie-pipeline te verenigen in een enkele mega-kernel. De betekenis ligt in:

• Het Doorbreken van de Kernelbarrière: Door synchronisatie naar het SM-niveau te verplaatsen, ontsluit MPK het potentieel voor hardware-utilisatie dat momenteel ontoegankelijk is voor kernel-per-operator systemen.
• Automatisering vs. Handmatige Inspanning: In tegen tegenoverde vorige mega-kernel benaderingen die expert handmatig ontwerp vereisten, automatiseert MPK het proces, waardoor high-performance mega-kernel executie toegankelijk wordt voor standaard PyTorch developers met minimale code-wijzigingen.
• Hardware Nabijheid: Het systeem duwt de LLM-inferentieprestaties dicht bij de theoretische limieten van de onderliggende hardware door software-overheads (launch, scheduling, synchronisatie) te minimaliseren en hardware overlap te maximaliseren.

De auteurs concluderen dat MPK een nieuw pad opent naar high-performance, compiler-gestuurde inferentiesystemen die de flexibiliteit van bestaande programmeermodellen behouden terwijl ze prestaties leveren die nabij de hardwarelimieten liggen.