🤖 machine learning

AutoMegaKernel: A Statically-Checked Agent Harness for Self-Retargeting Megakernel Synthesis

AutoMegaKernel è un harness di agenti con controllo statico che sintetizza automaticamente un singolo megakernel CUDA persistente per i modelli della famiglia HuggingFace Llama, garantendo l'esecuzione priva di deadlock su diverse architetture NVIDIA e raggiungendo fino a un incremento di velocità di 1,72x rispetto ai baseline su GPU di classe inferenza attraverso loop di agenti auto-miglioranti.

Autori originali: Jaber Jaber, Osama Jaber

Pubblicato 2026-06-09

📖 5 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Jaber Jaber, Osama Jaber

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di dover gestire una fabbrica enorme e complessa (un Large Language Model) per produrre un singolo prodotto (una parola di testo).

Il Vecchio Modo (IA Standard):
Attualmente, gestire questa fabbrica è come inviare un operaio nel magazzino per prendere uno strumento specifico, riportarlo sulla linea di montaggio, fare un piccolo lavoro, rimandare l'operaio nel magazzino per lo strumento successivo, e ripetere questo processo decine di volte per ogni singola parola.

Il Problema: L'operaio passa la maggior parte del tempo a camminare avanti e indietro (aspettando i dati dalla memoria) piuttosto che costruire. Questo è chiamato "limitato dalla larghezza di banda" (bandwidth-bound). La fabbrica è limitata dalla velocità con cui l'operaio cammina, non dalla velocità con cui costruisce.
Il Costo: Ogni volta che l'operaio si ferma per prendere uno strumento, c'è un piccolo ritardo (latenza di lancio). Fare questo per ogni passaggio si accumula, rendendo il processo lento.

Il Nuovo Modo (AutoMegaKernel):
Gli autori di questo articolo hanno costruito un sistema chiamato AutoMegaKernel (AMK). Pensa ad AMK come a un rivoluzionario manager di fabbrica che cambia completamente le regole.

1. La Fabbrica "One-Shot"

Invece di inviare lavoratori avanti e indietro, AMK organizza l'intera fabbrica in modo che un singolo lancio faccia tutto il lavoro.

L'Analogia: Immagina che invece di un operaio che va a prendere gli strumenti uno alla volta, tu abbia un team di lavoratori stazionati presso ogni singola macchina. Sono tutti collegati da una pipeline diretta e interna. Quando viene premuto il tasto "Start", l'intera linea di montaggio si avvia dall'inizio alla fine in un unico movimento continuo, senza che nessuno debba mai lasciare il pavimento della fabbrica per andare in magazzino.
Il Risultato: Questo elimina i "tempi di camminata" e i ritardi di "stop-and-start", rendendo teoricamente il processo molto più veloce.

2. L'Ispettore della Sicurezza (L'Agente Verificato Staticamente)

Costruire una fabbrica dove tutti lavorano simultaneamente è pericoloso. Se i lavoratori non sono perfettamente coordinati, potrebbero scontrarsi tra loro (race conditions) o aspettare per sempre uno strumento che non arriverà mai (deadlock).

L'Innovazione: Di solito, per far sì che questo funzioni, serve un esperto umano che scriva il codice manualmente, il che è difficile e soggetto a errori. AMK utilizza un Agente IA per progettare automaticamente il layout della fabbrica.
La Rete di Sicurezza: Prima che la fabbrica venga accesa, un Ispettore della Sicurezza congelato e immutabile (un validatore) controlla il progetto dell'IA. Esso dimostra matematicamente che:
- Nessuno rimarrà bloccato ad aspettare per sempre.
- Nessun paio di lavoratori cercherà di usare lo stesso strumento contemporaneamente.
- L'ordine delle operazioni è perfetto.
L'Affermazione: Gli autori hanno testato questo ispettore contro 7.160 progetti "insidiosi" o difettosi. L'ispettore ha individuato tutti quelli non sicuri e li ha rifiutati prima che la fabbrica partisse. Non ha mai lasciato passare un piano pericoloso.

3. L'Adattatore Universale (Self-Retargeting)

Di solito, una fabbrica progettata per un edificio specifico (un chip specifico) non funzionerà in un edificio diverso.

La Magia: AMK è "self-retargeting". Scrivi il progetto una volta sola e il sistema lo adatta automaticamente per funzionare perfettamente su diversi tipi di chip per computer (come l'RTX 5090, l'A100 o l'H100) senza la necessità che un essere umano riscriva il codice per ciascuno di essi.

4. La Versione "Leggera" (Quantizzazione)

Il sistema ha anche scoperto come far funzionare la fabbrica utilizzando strumenti "più leggeri".

L'Analogia: Invece di trasportare strumenti pesanti e a grandezza naturale (calcoli ad alta precisione), gli operai usano strumenti leggeri e compatti (precisione int8 o int4).
Il Beneficio: Poiché gli strumenti sono più leggeri, possono trasportarne di più alla volta e gli operai si muovono più velocemente.
- Int8: Funziona altrettanto bene degli strumenti pesanti (senza perdite/lossless) ma è più veloce del 12%.
- Int4: Si muove ancora più velocemente (tagliando il "tempo di camminata" di oltre 2 volte), ma il prodotto finale potrebbe essere leggermente meno perfetto (con perdite/lossy), sebbene sia comunque leggibile.

5. I Risultati Onesti (La "Realtà dei Fatti")

Gli autori sono molto trasparenti su dove questo sistema vince e dove perde.

Dove Vince: Sui chip di classe "Inference" (come l'L4, L40S, A10G e il consumer RTX 5090), il sistema AMK è più veloce dello standard attuale del settore (cuBLAS) nella generazione di una parola alla volta. Batte la concorrenza di 1,1x a 1,3x.
- Perché? Ha eliminato con successo i ritardi di "andare e venire" che affliggono altri sistemi.
Dove Perde: Sui massicci e velocissimi chip di classe "Training" (come l'A100 e l'H100), la versione attuale è più lenta del sistema standard.
- Perché? La "coordinazione della sicurezza" (sincronizzazione) tra i lavoratori richiede un po' di tempo. Sui chip più veloci, questo piccolo ritardo diventa il collo di bottiglia. Gli autori lo ammettono apertamente: "Abbiamo localizzato il collo di bottiglia, e siamo onesti a riguardo".

Riassunto

AutoMegaKernel è un sistema che utilizza un agente IA per progettare una fabbrica "one-shot" per eseguire modelli di IA. Include un rigoroso ispettore della sicurezza per garantire che il design non vada mai in crash. Funziona con successo su molti diversi chip per computer automaticamente.

La Grande Vittoria: È attualmente il modo più veloce per generare testo una parola alla volta su molti popolari chip consumer e da datacenter, superando gli strumenti standard eliminando i ritardi non necessari.
Il Problema: Non è ancora perfetto. Sui chip assoluti più veloci, l'overhead di coordinazione rende ancora il sistema leggermente più lento rispetto ai giganti stabiliti, ma gli autori hanno dimostrato che il sistema è sicuro, corretto e capace di auto-miglioramento.

Riepilogo Tecnico: AutoMegaKernel (AMK)

Problema

La decodifica autoregressiva a singolo stream con batch size 1 è fondamentalmente limitata dalla larghezza di banda. Il limite inferiore della latenza è determinato dal rapporto tra il numero totale di byte dei pesi e la larghezza di banda dell'HBM ( $t_{min} = \text{weight\_bytes} / \text{HBM\_bandwidth}$ ). Gli attuali framework di esecuzione (ad es. PyTorch, cuBLAS) non riescono ad avvicinarsi a questo limite perché lanciano un kernel separato per ogni operatore, incorrendo in una significativa latenza di lancio della CPU e costringendo le attivazioni a compiere viaggi di andata e ritorno attraverso l'HBM tra le operazioni. Sebbene i CUDA Graphs amortizzino l'overhead di lancio, essi mantengono i confini tra i kernel e i viaggi di andata e ritorno nell'HBM.

Gli approcci esistenti a "megakernel" (ad es. Mirage Persistent Kernel) o i kernel ottimizzati a mano affrontano il limite della latenza fondendo l'intero passaggio forward in un unico kernel persistente. Tuttavia, questi sistemi mancano di fiducia e portabilità: spesso omettono gate di sicurezza statici (rischiando deadlock o race condition), mirano manualmente a modelli/architetture specifiche e non forniscono un'interfaccia strutturata affinché agenti autonomi possano ottimizzare il sistema in modo sicuro.

Metodologia

AutoMegaKernel (AMK) è un sistema autonomo, guidato da agenti, che compila modelli della famiglia Llama di HuggingFace in un singolo kernel cooperativo persistente senza la necessità di scrivere codice CUDA manualmente per ogni modello. Il sistema è strutturato in quattro livelli e due loop di ricerca automatica:

Livello 0 (La VM): Un megakernel persistente scritto a mano, congelato e per architettura (vm/scheduler.cu). Lancia un threadblock cooperativo per ogni Streaming Multiprocessor (SM). Ogni blocco possiede una coda locale e percorre un ordine topologico globale di task, eseguendo istruzioni wait → dispatch → signal. Si affida a contatori monotoni per la sincronizzazione, evitando i lock.
Livello 1 (Istruzioni): Micro-kernel (GEMV, attention, RMSNorm, ecc.) conformi all'ABI che sono controllati per correttezza in isolamento.
Livello 2 (Scheduler): La superficie di modifica accessibile all'agente. Importa un modello, lo abbassa a una Schedule IR tipizzata (un DAG di task sincronizzato tramite contatori monotoni) e propone configurazioni strutturali (tiling, fusione, assegnazione SM, profondità di pipelining) tramite un oggetto ScheduleConfig.
Livello 3 (Roadmap): Segnaposto per la futura dinamicità (ad es. continuous batching).

Il Gate di Sicurezza:
L'innovazione principale è un validatore a controllo statico (validate() in schedule/ir.py). Prima di qualsiasi lancio su GPU, il validatore esegue un insieme di controlli grafici statici sulla Schedule IR per certificare:

Assenza di Deadlock: Garantire che il DAG produttore-consumatore sia aciclico, che le soglie di attesa siano soddisfacibili ( $1 \le t \le \#producers$ ) e che gli ordini delle code per ogni SM siano validi.
Assenza di Race: Garantire che i contatori condivisi siano uniti correttamente (tutti i produttori devono terminare prima che un consumatore attenda un contatore multi-produttore) e che gli accessi alla memoria rispettino le relazioni di "happens-before".
Correttezza: Verificare la raggiungibilità dell'output e l'integrità referenziale.

Se uno schedule fallisce questi controlli, viene rifiutato immediatamente. Ciò impedisce che schedule proposti in modo insicuro dagli agenti causino il blocco delle GPU.

Loop di Ricerca Automatica:
AMK include un loop non presidiato e guidabile da agenti che propone iterativamente modifiche alla ScheduleConfig, valida tali modifiche, esegue i gate di correttezza rispetto a un riferimento su CPU, misura la latenza e mantiene o annulla le modifiche. Questo loop migliora autonomamente il megakernel rispetto al proprio baseline.

Contributi Chiave

Schedule IR a Controllo Statico: Un validatore congelato che certifica l'assenza di deadlock e race tramite controlli grafici statici (non prove meccanizzate) su una base affidabile. Ha ottenuto zero falsi accettati su 7.160 schedule costruite avversarialmente (di cui 6.091 non sicure) pur accettando tutti i 360 abbassamenti di modelli reali.
Copertura di Generazione Automatica: AMK genera automaticamente megakernel corretti per 10 modelli su 10 supportati (inclusi checkpoint reali come SmolLM2-135M e TinyLlama-1.1B) senza scrivere codice CUDA manualmente. Rifiuta con successo varianti non supportate (ad es. proiezioni con bias) o documenta onestamente i propri punti ciechi.
Auto-Retargeting: Un unico codice sorgente genera megakernel corretti per sm_80 (A100), sm_90 (H100) e sm_120 (RTX 5090), corrispondendo al bit-per-bit o entro strette tolleranze fp32 ( $\le 3.8 \times 10^{-5}$ ) rispetto all'esecuzione eager.
Megakernel Quantizzati: Il sistema genera automaticamente kernel con pesi quantizzati int8 (greedy-lossless, 1.12× più veloce) e int4 (lossy, riduzione del traffico pesi di 2.42×) partendo dallo stesso percorso IR.
Compilatore Guidabile da Agenti: Il sistema tratta la correttezza come un invariante strutturale, permettendo a un agente di esplorare lo spazio degli schedule in modo sicuro senza rischiare blocchi della GPU.
Studio di Performance Onesto: Il paper riporta i risultati su GPU di classe inferenza e di classe training, dichiarando esplicitamente dove il sistema eccelle e dove invece resta indietro rispetto ai baseline stabiliti.

Risultati Sperimentali

Correttezza e Portabilità:

Su RTX 5090, A100 e H100, AMK riproduce la decodifica greedy di HuggingFace token per token e corrisponde alla perplexity teacher-forced fino a $2.5 \times 10^{-7}$ sul checkpoint SmolLM2-135M.
Il validatore ha rifiutato tutti i 6.091 schedule non sicuri durante il test di stress, dimostrando un tasso di falsi accettati dello 0%.

Performance su GPU di Classe Inferenza (Il Successo):

RTX 5090 (Consumer): Un megakernel auto-ottimizzato int8 weight-only (W8A16) batte cuBLAS con CUDA-graph (bf16) di 1.19–1.23× nella decodifica a batch-1. Questa è una comparazione asimmetrica di precisione (int8 vs. bf16) dove il vantaggio deriva dal flusso di meno byte di peso (0.61× i byte), non da un kernel più veloce per byte.
Flotta di Inferenza Datacenter: Il megakernel int8 batte cuBLAS bf16 con CUDA-graph in tutta la flotta di classe inferenza:
- NVIDIA L4: Fino a 1.33× più veloce (il vantaggio cresce con la dimensione del modello).
- L40S: 1.25–1.27× più veloce.
- A10G: Fino a 1.08× più veloce su scala.
Auto-Miglioramento: Il loop di ricerca automatica ha migliorato autonomamente il megakernel rispetto al proprio default di 1.25–1.72× su RTX 5090.

Performance su GPU di Classe Training (La Sconfitta):

Su GPU ad alta larghezza di banda di classe training (A100, H100), il megakernel generato automaticamente resta indietro rispetto a cuBLAS con CUDA-graph e vLLM.
- Su H100, AMK è stato 1.65× più lento del percorso CUDA-graph predefinito di vLLM.
- Su A100, l'ottimizzato GEMV in bf16 ha raggiunto solo circa il 63% del picco HBM misurato (contro il ~90% di cuBLAS).
Diagnosi: Il collo di bottiglia è il costo di sincronizzazione cross-SM per tile (un costo fisso che il megakernel paga e che cuBLAS non paga). Su silicio ad alta larghezza di banda, questo overhead di sincronizzazione domina, mentre sulle GPU di classe inferenza, il risparmio di banda derivante dalla fusione e dalla quantizzazione ammortizza tale costo.

Significato e Rivendicazioni

Il paper sostiene che AutoMegaKernel rappresenti un passaggio dalla "velocità pura" alla sicurezza sistematica e portabilità nella sintesi di megakernel. Il suo contributo primario è un'architettura di compilazione a controllo statico che permette ad agenti autonomi di esplorare in sicurezza lo spazio degli schedule di interi modelli.

Fiducia: Sostituisce il rischio di bloccare le GPU con un gate di sicurezza statico che rifiuta gli schedule non sicuri prima del lancio.
Generalità: Raggiunge l'auto-retargeting su diverse architetture (sm_80/90/120) e famiglie di modelli senza tuning manuale di CUDA.
Onestà: Gli autori inquadrano esplicitamente i risultati di performance in base al regime di utilizzo. Il sistema vince sulle GPU di classe inferenza fondendo le operazioni e sfruttando la quantizzazione per ridurre il traffico di memoria, ma attualmente perde sulle GPU di classe training a causa degli overhead di sincronizzazione che richiedono una programmazione più grossolana (lavoro futuro) piuttosto che solo migliori micro-kernel.

Il paper conclude che, sebbene la qualità attuale dei kernel non eguagli ancora i baseline ottimizzati a mano su hardware di alto livello per il training, l'impalcatura a controllo statico fornisce la base necessaria per automatizzare in sicurezza i miglioramenti futuri (ad es. sincronizzazione più grossolana, pipelining più profondo) necessari per colmare questo divario.