Riepilogo Tecnico: MPK - Un Compilatore e Runtime per Programmi Tensoriali Mega-Kernelizzati

1. Definizione del Problema

Gli attuali sistemi di inferenza su GPU ad alte prestazioni per i grandi modelli linguistici (LLM) e altri programmi tensoriali si affidano prevalentemente a un modello di esecuzione basato su kernel-per-operatore. In questo paradigma, ogni operatore (ad esempio, moltiplicazione di matrici, attention, all-reduce) viene eseguito come un kernel GPU distinto. Sebbene questo approccio semplifichi la programmazione, introduce tre limitazioni critiche che ne ostacolano le prestazioni:

1. Barriere Implicite tra i Kernel: L'esecuzione GPU convenzionale impone una barriera tra lanci di kernel consecutivi sullo stesso stream. Ciò impedisce il pipelining software cross-operatore, costringendo gli operatori dipendenti a eseguire strettamente in sequenza e creando "bolle di pipeline" in cui le risorse hardware (Tensor Cores, TMA) rimangono inattive.
2. Sovrapposizione a Grana Grossa: Le dipendenze sono imposte alla granularità di interi operatori. Di conseguenza, la sovrapposizione fine-grained tra computazione e comunicazione è impossibile. Ad esempio, un'operazione di AllReduce deve attendere la fine dell'intero MatMul precedente, anche se frammenti specifici della riduzione dipendono solo da frammenti specifici dell'output della moltiplicazione.
3. Overhead di Lancio e Flessibilità Statica: L'esecuzione di centinaia o migliaia di kernel per iterazione di inferenza comporta un overhead di lancio significativo. Sebbene i CUDA Graph mitigino il problema, sono largamente statici; cambiamenti dinamici nel controllo di flusso, nelle forme dei tensori o nelle dipendenze dei dati spesso richiedono la ri-instanziazione del grafo, limitando la flessibilità per carichi di lavoro dinamici come il serving di LLM.

Le soluzioni esistenti, come i mega-kernel manuali (es. FlashDMoE, Spector et al.), offrono vantaggi prestazionali ma richiedono un enorme sforzo ingegneristico e una profonda competenza GPU, fallendo nel generalizzarsi su modelli o hardware diversi.

2. Metodologia

Gli autori presentano Mirage Persistent Kernel (MPK), il primo sistema di compilazione e runtime che trasforma automaticamente l'inferenza di modelli multi-GPU in un singolo mega-kernel ad alte prestazioni. MPK raggiunge questo obiettivo attraverso due componenti primarie:

A. Rappresentazione del Grafo a Livello di SM (tGraph)

MPK introduce una nuova rappresentazione di grafo chiamata tGraph che cattura computazione e comunicazione alla granularità dei singoli Streaming Multiprocessors (SM), piuttosto che dell'intera GPU.

• Nodi: Rappresentano task (computazione o comunicazione) in esecuzione su un singolo SM.
• Archi: Codificano dipendenze fine-grained tra i task tramite eventi.
• Vantaggio: Questa granularità espone un parallelismo altrimenti inaccessibile ai modelli a livello di kernel, consentendo il pipelining software cross-operatore e la sovrapposizione fine-grained di computazione e comunicazione.

Ottimizzazioni del Compilatore:
Il compilatore MPK trasforma un programma tensoriale standard in un tGraph ottimizzato attraverso diverse fasi:

1. Decomposizione degli Operatori: Suddivide gli operatori in insiemi disgiunti di task basati sulla partizione dei tensori di output (ad esempio, il tiling delle righe/colonne della moltiplicazione matriciale) per corrispondere al numero di SM disponibili.
2. Analisi delle Dipendenze: Elenca le coppie produttore-consumatore tra i task per inserire eventi di sincronizzazione solo dove esistono sovrapposizioni di dati.
3. Fusione degli Eventi: Unisce punti di sincronizzazione ridondanti usando la fusione dell'insieme dei successori (unendo eventi con set di consumatori identici) e la fusione dell'insieme dei predecessori (unendo eventi con set di produttori identici) per ridurre la complessità del grafo.
4. Normalizzazione: Trasforma il grafo in modo che ogni task abbia al massimo un evento dipendente e un evento di attivazione, ottenuto inserendo task ed eventi dummy per gestire il fan-in/fan-out, garantendo una rappresentazione uniforme e a basso overhead.
5. Linearizzazione: Utilizza un algoritmo basato su BFS per assegnare indici contigui ai task attivati dallo stesso evento. Ciò consente al runtime di codificare il fan-out degli eventi come un semplice intervallo di indici (primo/ultimo ID del task) invece di una lista esplicita, riducendo significativamente l'impronta di memoria del dispositivo.
6. Generazione dell'Implementazione del Task: Utilizza la superottimizzazione del compilatore (Mirage) per generare automaticamente implementazioni CUDA ad alte prestazioni per ogni singolo task, incluse ottimizzazioni intra-SM come il riutilizzo dei registri e il layout della memoria condivisa.

B. Runtime Parallelo In-Kernel

MPK esegue il tGraph all'interno di un singolo mega-kernel persistente utilizzando un'architettura decentrata ed event-driven:

• Partizionamento Worker/Scheduler: Gli SM della GPU sono partizionati staticamente in Worker (che eseguono i task) e Scheduler (che gestiscono le code degli eventi).
• Esecuzione Event-Driven: Gli scheduler monitorano l'attivazione degli eventi (dove tutti i prerequisiti sono soddisfatti) e dispacciano i relativi task ai worker. Ciò elimina la necessità di sincronizzazione lato host o lanci ripetuti di kernel.
• Lancio dei Task Ibrido: Per bilanciare latenza e bilanciamento del carico, MPK impiega due modalità:
  • Just-In-Time (JIT): I task vengono lanciati solo dopo che le loro dipendenze sono state soddisfatte. Questo si adatta a carichi di lavoro dinamici (es. lunghezze di sequenza variabili nell'attention) per prevenire squilibri di carico.
  • Ahead-Of-Time (AOT): I task vengono pre-inseriti nelle code dei worker prima che le dipendenze siano soddisfatte. Ciò riduce la latenza di dispatch per operatori prevedibili.
• Memoria Condivisa a Pagine (Paged Shared Memory): Per abilitare il pipelining software cross-task, MPK introduce un'astrazione a pagine per la memoria condivisa. I task acquisiscono e rilasciano pagine di dimensione fissa, permettendo al task successivo di iniziare il prefetching dei dati nella memoria condivisa mentre il task corrente computa, senza conflitti di risorse.

3. Contributi Chiave

1. Primo Sistema di Mega-Kernel Automatizzato: MPK è il primo sistema che trasforma automaticamente arbitrari programmi tensoriali multi-GPU in un singolo mega-kernel, eliminando la necessità di fusione manuale dei kernel.
2. Astrazione a Livello di SM: L'introduzione della rappresentazione tGraph sposta il tracciamento delle dipendenze dal livello di operatore al livello di SM-task, abilitando ottimizzazioni (pipelining, sovrapposizione fine-grained) precedentemente impossibili sotto il modello kernel-per-operatore.
3. Compilatore e Runtime End-to-End: Il sistema integra un compilatore che ottimizza i grafi dei task e genera codice CUDA con un runtime che esegue questi task in modo asincrono all'interno di un singolo kernel, supportando carichi di lavoro dinamici come il continuous batching e il paged attention.
4. Strategia di Esecuzione Ibrida: La combinazione di lancio dei task JIT e AOT permette al sistema di gestire sia la natura dinamica dell'inferenza LLM (variabilità di batch/lunghezza della sequenza) sia l'efficienza statica degli operatori fissi.

4. Risultati della Valutazione

Gli autori hanno valutato MPK su cinque LLM ampiamente utilizzati (inclusi varianti di Qwen3 e Llama-3) su tre generazioni di GPU NVIDIA (A100, H100, B200), confrontandolo con sistemi kernel-per-operatore allo stato dell'arte (SGLang, vLLM) e PyTorch vanilla.

• Guadagni di Prestazione: MPK raggiunge una latenza di inferenza end-to-end da 1.0× a 1.7× inferiore rispetto a SGLang e vLLM. I miglioramenti sono più pronunciati per i modelli più piccoli e le generazioni di hardware più recenti (es. B200), dove l'overhead di lancio del kernel e le bolle di pipeline diventano più significativi rispetto al tempo di computazione.
• Throughput: In configurazioni multi-GPU (fino a 8 H100), MPK migliora il throughput di 1.1× a 1.4× rispetto ai baseline ottimizzati.
• Limiti di Latenza: Su un A100, MPK ha ridotto la latenza di decoding per token di Qwen3-8B da 14.5 ms (vLLM/SGLang) a 12.5 ms, avvicinandosi al limite hardware stimato di ~10 ms.
• Studi di Ablazione:
  • Pipelining Cross-Task: Ha ridotto il tempo di esecuzione del task di 1.2–1.3×.
  • Sovrapposizione Compute-Communication: Ha ridotto la latenza per iterazione di 1.1×.
  • Riduzione del Lancio dei Kernel: Ha eliminato l'overhead di circa 1.1 ms per token associato ai 293 lanci di kernel per token nei sistemi baseline.
• Scalabilità: MPK scala efficacemente su 2, 4 e 8 GPU, mantenendo velocità superiori rispetto ai baseline.

5. Significato e Rivendicazioni

Il paper afferma che MPK rappresenta un cambio di paradigma nei sistemi di inferenza GPU unificando la pipeline di esecuzione in un singolo mega-kernel. La sua importanza risiede nel:

• Rompere la Barriera del Kernel: Spostando la sincronizzazione al livello di SM, MPK sblocca il potenziale di utilizzo dell'hardware che è attualmente inaccessibile ai sistemi kernel-per-operatore.
• Automazione vs. Sforzo Manuale: A differenza dei precedenti approcci a mega-kernel che richiedevano un design manuale esperto, MPK automatizza il processo, rendendo l'esecuzione di mega-kernel ad alte prestazioni accessibile agli sviluppatori standard di PyTorch con modifiche minime al codice.
• Prossimità all'Hardware: Il sistema spinge le prestazioni dell'inferenza LLM vicino ai limiti teorici dell'hardware sottostante minimizzando gli overhead software (lancio, scheduling, sincronizzazione) e massimizzando la sovrapposzione hardware.

Gli autori concludono che MPK apre una nuova strada verso sistemi di inferenza ad alte prestazioni guidati dal compilatore che preservano la flessibilità dei modelli di programmazione esistenti pur fornendo prestazioni vicine ai limiti dell'hardware.