💬 NLP

Axon: A Synthesizing Superoptimizer for Tensor Programs

Axon è un superottimizzatore di sintesi che genera automaticamente programmi tensoriali ad alte prestazioni per acceleratori AI utilizzando la sintesi di programmi e la verifica semantica basata su SMT per scoprire trasformazioni algebriche, ottimizzare il tiling e la selezione delle istruzioni, e fondere gli operatori senza fare affidamento su regole di riscrittura create manualmente.

Autori originali: Akash Kothari, Shaowei Zhu, Daniel Kroening, Chungha Sung

Pubblicato 2026-06-26

📖 5 min di lettura🧠 Approfondimento

CC BY 4.0

Autori originali: Akash Kothari, Shaowei Zhu, Daniel Kroening, Chungha Sung

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 voler preparare la torta perfetta, ma di lavorare in una cucina hi-tech con forni e impastatrici molto specifici e capricciosi. Sai cosa vuoi cucinare (la ricetta), ma far sì che la macchina lo faccia nel modo più veloce possibile richiede di sapere esattamente come caricare gli ingredienti, quando mescolare e quali pulsanti specifici premere sul forno.

Per molto tempo, scrivere il codice per i chip AI (i "forni") è stato come cercare di preparare quella torta a mano, un piccolo bricio alla volta. Devi essere uno chef maestro che comprende l'esatta meccanica del forno per ottenere un buon risultato. Se commetti un piccolo errore nel modo in cui carichi gli ingredienti, la torta uscirà lenta e asciutta.

Axon è un nuovo strumento che agisce come un sous-chef automatizzato e super intelligente che non si limita a seguire la tua ricetta, ma inventa modi migliori per cucinarla.

Ecco come funziona Axon, suddiviso in semplici passaggi:

1. Il "Cosa" vs. Il "Come"

Di solito, un programmatore scrive una ricetta dicendo: "Mescola farina e uova, poi cuoci in forno". Non si preoccupano di come il forno mescoli.

Il Problema: Il forno (il chip AI) è esigente. Ha diversi scomparti (come una grande ciotola per mescolare, una piccola frusta e un elemento riscaldante). Se dici di mescolare tutto nella ciotola grande quando avrebbe potuto usare la frusta, sprechi tempo.
Il Lavoro di Axon: Axon prende la tua semplice ricetta e dice: "Aspetta, posso farlo più velocemente". Non si limita a seguire la ricetta; riscrive le istruzioni per adattarle ai punti di forza specifici del forno.

2. La "Riorganizzazione Magica" (Trasformazioni Algebriche)

Immagina che la tua ricetta dica: "Per prima cosa, trita tutte le carote, poi mescolale nella zuppa, poi aggiungi il sale".

Il Vecchio Modo: Triti, poi mescoli, poi aggiungi il sale. Devi trasportare la ciotola avanti e indietro tra il tagliere e il fornello.
Il Modo di Axon: Axon guarda la matematica e realizza: "Ehi, posso aggiungere il sale mentre le carote vengono tritate, oppure posso mescolare le carote nella zuppa prima di tritarle, e il gusto sarà esattamente lo stesso!".
Il Risultato: Riorganizzando l'ordine delle operazioni, Axon permette al forno di fare due cose contemporaneamente (come tritare e scaldare simultaneamente) invece di una dopo l'altra. Dimostra matematicamente che la torta avrà lo stesso sapore, anche se i passaggi sono diversi.

3. Il "Puzzle dell'Imballaggio" (Tiling)

Immagina di avere una grande pila di ingredienti, ma la tua ciotola è piccola. Devi spostare gli ingredienti dalla dispensa (lenta) alla ciotola (veloce), mescolare un po', spostarli di nuovo indietro e ripetere.

Il Problema: Se sposti troppi ingredienti in una volta sola, la ciotola trabocca. Se ne sposti troppo pochi, passi tutto il tempo a camminare avanti e indietro verso la dispensa.
Il Modo di Axon: Axon è un maestro dell'imballaggio. Prova milioni di modi diversi per raggruppare gli ingredienti in "lotti" (tile). Chiede: "E se spostassi 100 carote alla volta? E se ne spostassi 50?". Testa queste diverse raggruppamenti per trovare la dimensione perfetta che mantenga la ciotola piena senza farla traboccare, minimizzando i viaggi verso la dispensa.

4. Il "Manuale di Istruzioni" (Sintesi ISA)

Il forno ha un pannello di controllo specifico con pulsanti come "Calore Alto", "Frusta" e "Incorporare".

Il Problema: A volte la ricetta dice "Mescola", ma il forno non ha un pulsante "Mescola". Ha solo "Frusta" e "Incorpora".
Il Modo di Axon: Axon traduce "Mescola" nell'esatta sequenza di "Frusta poi Incorpora" che il forno comprende. Trova persino scorciatoie, come realizzare che "Incorporare mentre ci si sposta verso il fornello" è un singolo comando su questo specifico forno, risparmiando tempo.

5. Il "Test del Gusto" (Selezione Empirica)

Questa è la parte più importante. Axon non si limita a indovinare quale sia il modo migliore.

Il Processo: Genera migliaia di versioni diverse della ricetta (alcune con ordini diversi, altre con diverse dimensioni dei lotti, altre con diversi pulsanti da premere).
Il Test: Esegue effettivamente tutte queste versioni sul vero forno (il chip AI) e le cronometra. Sceglie quella che finisce più velocemente. È come cuocere 1.000 piccole torte per vedere quale esce dal forno in 5 minuti invece di 10.

I Risultati: Quanto è più veloce?

Il paper ha testato Axon sui chip Amazon Trainium (un tipo di forno AI) con 20 diversi tipi di "torte" (compiti AI).

Rispetto al Forno Standard: Rispetto al modo standard di cucinare (il compilatore Neuron di Amazon), Axon ha reso i singoli compiti fino a 3,7 volte più veloci e i compiti complessi fino a 19 volte più veloci.
Rispetto agli Chef Esperti: Anche confrontato con ricette scritte da esperti umani che conoscono il forno dentro e fuori, Axon era talvolta 1,35 volte più veloce.
Rispetto ad Altri Robot: Ha battuto un altro strumento automatizzato chiamato "Mirage" di circa il 10% in media, e poteva cucinare 16 tipi di torte che l'altro strumento non era in grado di gestire affatto.

Riassunto

Axon è uno strumento che prende una semplice descrizione di ciò che un'IA deve fare e trova automaticamente il modo più veloce in assoluto per farlo su un chip specifico. Lo fa dimostrando matematicamente di poter riorganizzare i passaggi, testando milioni di modi per imballare i dati e poi eseguendo effettivamente il codice per scegliere il vincitore. Elimina la necessità per gli esseri umani di essere esperti nei dettagli minuscoli e noiosi dell'hardware, permettendo loro di scrivere semplicemente la ricetta mentre Axon si occupa della cottura.

Riepilogo Tecnico: Axon: Un Superottimizzatore di Sintesi per Programmi Tensoriali

Definizione del Problema

Scrivere kernel ad alte prestazioni per acceleratori AI (come il Trainium di Amazon) richiede una profonda conoscenza dei dettagli specifici dell'hardware, tra cui le gerarchie di memoria, le strategie di tiling, la selezione delle istruzioni, il layout dei dati e la fusione degli operatori. Sebbene i linguaggi basati su tile come NKI (Neuron Kernel Interface) astraggano alcuni dettagli di basso livello, raggiungere l'utilizzo massimo dell'hardware rimane una sfida per la maggior parte degli ingegneri ML. Gli approcci di ottimizzazione esistenti soffrono di diverse limitazioni:

Regole artigianali: Molti sistemi si affidano a regole di riscrittura scritte manualmente per trasformazioni algebriche, che sono soggette a errori e difficili da scalare su nuove combinazioni di operatori.
Generazione basata su template: Alcuni approcci utilizzano template scritti a mano per la generazione del codice, limitandone la flessibilità.
Ambito limitato: Molti ottimizzatori prendono di mira solo CPU o GPU, fallendo nel rispondere alle esigenze delle architetture uniche a più motori e guidate da DMA degli acceleratori AI moderni.
Mancanza di correttezza provabile: I test probabilistici o il riordinamento basato su euristiche possono mancare configurazioni ottimali o introdurre errori semantici.

Metodologia

Axon è un superottimizzatore di sintesi progettato per programmi di acceleratori AI basati su tile. Accetta programmi tensoriali di alto livello (scritti con un'interfaccia simile a NumPy) e genera automaticamente istruzioni target ottimizzate (NKI) attraverso una pipeline multi-stadio. L'innovazione principale è la $\nu$ Graph ("nu-graph"), una rappresentazione interna persistente che traccia simultaneamente tutte le varianti semanticamente equivalenti del programma durante l'intero processo di sintesi, differendo la scelta dell'ottimizzazione alla selezione empirica.

La pipeline consiste nei seguenti stadi:

Trasformazioni Algebriche tramite Propagazione di Operatori:
- Axon scopre grafi di computazione semanticamente equivalenti propagando gli operatori attraverso il grafo di computazione.
- Tenta di riordinare gli operatori adiacenti (ad esempio, spostando una moltiplicazione elemento per elemento oltre una moltiplicazione di matrici) per esporre parallelismo e opportunità di fusione.
- Garanzia di Correttezza: Invece di utilizzare regole di riscrittura artigianali, Axon decompone le operazioni tensoriali in computazioni granulari e utilizza un solver SMT (Satisfiability Modulo Theories) su tensori illimitati per verificare in modo provabile che il riordinamento preservi la semantica. Questo gestisce forme di tensori arbitrarie e non richiede definizioni manuali di proprietà.
Tiling Simbolico:
- Il sistema effettua il tiling di ogni operazione tensoriale utilizzando parametri simbolici per le dimensioni di strip, block e tile.
- Questi parametri sono vincolati dalle descrizioni hardware (ad esempio, capacità dei buffer on-chip e limiti specifici delle istruzioni).
- Axon genera molteplici versioni con diverse configurazioni di tiling, consentendo l'ottimizzazione indipendente dei rami paralleli (ad esempio, separando i rami RMSNorm e MatMul per eseguirli su diversi motori di calcolo).
Sintesi ISA (Selezione delle Istruzioni):
- Utilizzando un approccio di sintesi di programmi basato su sketch, Axon abbassa le operazioni tensoriali di alto livello alle istruzioni ISA target (ad esempio, nc_matmul e dma_copy di Trainium).
- Esplora sequenze di istruzioni valide e trasformazioni del layout dei dati (come le trasposizioni) richieste da specifiche istruzioni hardware.
- L'equivalenza tra l'operazione di alto livello e la sequenza di istruzioni sintetizzata è nuovamente verificata tramite SMT.
Fusione di ISA e Operatori:
- Un secondo round di trasformazioni algebriche viene eseguito sul grafo ISA sintetizzato per esporre nuove opportunità di fusione create dal processo di sintesi (ad esempio, fondere una trasposizione con una moltiplicazione di matrici).
- Fusione di Nodi: Unisce operazioni consecutive con dimensioni di tiling compatibili per mantenere i risultati intermedi nei buffer on-chip (SBUF/PSUM), riducendo il traffico HBM.
- Fusione di Istruzioni: Sostituisce sequenze di istruzioni ISA con varianti fuse (ad esempio, dma_copy + nc_transpose $\rightarrow$ dma_transpose).
Emissione del Codice e Selezione Empirica:
- Axon istanzia i parametri simbolici con valori concreti per emettere programmi NKI concreti.
- Ogni candidato viene compilato dal compilatore Neuron ed eseguito sull'hardware target.
- Il kernel con le migliori prestazioni viene selezionato in base alle misurazioni empiriche della latenza e alla validazione della correttezza numerica rispetto a un'implementazione di riferimento.

Contributi Chiave

Il documento presenta tre contributi primari:

Trasformazioni Algebriche Provabili: Una metodologia per scoprire automaticamente grafi di computazione semanticamente equivalenti propagando gli operatori e verificando l'equivalenza tramite SMT su tensori illimitati. Ciò elimina la necessità di regole di riscrittura artigianali.
Pipeline Basata sulla Sintesi: Una pipeline completa per acceleratori AI basati su tile che esplora congiuntamente trasformazioni algebriche, selezione delle istruzioni, configurazioni di tiling e fusione. Mantiene tutte le varianti valide in una $\nu$ Graph e seleziona il kernel ottimale empiricamente.
Valutazione Esaustiva: Una valutazione su Amazon Trainium attraverso 20 benchmark, inclusi singoli operatori e complessi kernel multi-operatore per LLM (ad esempio, Group Query Attention, Gated MLP).

Risultati Sperimentali

Axon è stato valutato rispetto a tre baseline: il compilatore Neuron di Amazon, kernel NKI ottimizzati a mano e Mirage (un compilatore allo stato dell'arte basato sulla ricerca).

Vs il Compilatore Neuron:
- Singoli Operatori: Axon ottiene accelerazioni fino a 3.7x (ad esempio, SiLU con grandi dimensioni) con un aumento della media geometrica di 1.23x–1.90x attraverso vari operatori.
- Kernel Multi-Operatore: Axon ottiene accelerazioni fino a 19x su kernel come Transpose+MatMul e fino a 1.47x su RMSNorm+MatMul. Gli aumenti della media geometrica variano da 1.07x a 1.32x.
- Ragionamento: I guadagni sono attribuiti a migliori strategie di tiling, un utilizzo efficiente del DMA e alla capacità di fondere le operazioni tra diversi motori di calcolo (ad esempio, eseguire RMSNorm sui motori Scalar/Vector mentre il MatMul gira sul Tensor Engine in parallelo).
Vs NKI Ottimizzati a Mano:
- Axon eguaglia o supera i kernel scritti da esperti fino a 1.35x in configurazioni specifiche (ad esempio, Cumsum e RoPE), dimostrando che l'esplorazione automatizzata può competere con la sintonizzazione manuale senza richiedere una profonda esperienza hardware.
Vs Mirage:
- Sui benchmark supportati da entrambi, Axon ottiene aumenti della media geometrica del 2–10%.
- Fondamentalmente, Axon sintetizza con successo kernel per altri 16 benchmark che Mirage non può gestire a causa della sua dipendenza da template scritti a mano.
Overhead di Compilazione:
- Lo spazio di ricerca può generare circa 30.000 candidati per kernel complessi.
- La Fase 1 (sintesi simbolica) è veloce (<82 secondi). La Fase 2 (compilazione ed esecuzione) domina il tempo (fino a circa 1.7 ore per complessi layer LLM). Il documento nota che questo è un costo una tantum, paragonabile ad altri compilatori basati su schedule, e i kernel ottimizzati possono essere memorizzati nella cache per il riutilizzo.

Significato e Rivendicazioni

Il documento afferma che Axon rappresenta un passaggio dall'ottimizzazione basata su euristiche o template a un approccio di sintesi guidato dalla semantica. Sfruttando i solver SMT per garantire l'equivalenza semantica, Axon può scoprire automaticamente trasformazioni algebriche complesse (come il riordinamento delle operazioni tra diversi motori di calcolo) che sono difficili da codificare manualmente.

Il significato risiede in:

Automazione: Riduzione del carico sui programmatori per la sintonizzazione manuale dei kernel per hardware specifici.
Correttezza: Fornire garanzie provabili per le trasformazioni, a differenza dei metodi di test probabilistici.
Prestazioni: Raggiungere accelerazioni significative su moderni acceleratori AI ottimizzando congiuntamente tiling, selezione delle istruzioni e fusione in uno spazio di ricerca unificato.
Generalità: Capacità di gestire complessi kernel multi-operatore per LLM che gli esistenti compilatori basati sulla ricerca (come Mirage) non possono sintetizzare.

Gli autori sottolineano che, sebbene l'overhead di compilazione non sia trascurabile, si tratta di un costo standard per i compilatori di ottimizzazione, e i guadagni di prestazione risultanti giustificano l'investimento, in particolare per carichi di lavoro LLM su larga scala.