The Pulse
Meta illustra TLX, un’estensione di Triton per l’orchestrazione delle GPU
Meta ha descritto TLX, un’estensione di Triton che offre agli sviluppatori un controllo più diretto sull’organizzazione del lavoro delle GPU. L’ha utilizzata in carichi di lavoro di produzione, tra cui l’addestramento e l’inferenza dei sist

AI.info Team · Scritto con l’IA · verificato ·
Un livello di controllo per il lavoro delle GPU
TLX è un’estensione di Triton che offre agli ingegneri un controllo esplicito sull’esecuzione delle GPU, mantenendo il modello di programmazione tensoriale di Triton. È un livello che permette di specificare come viene orchestrato il lavoro, affiancandosi alle operazioni di Triton che esprimono il calcolo stesso.
Il modello di esecuzione si chiama Multiple Instructions, Multiple Warps, o MIMW. I gruppi di warp sono trattati come agenti che collaborano, così i programmatori possono specificare come comunicano, si sincronizzano e condividono le risorse, anziché lavorare solo a livello dei singoli thread o di interi blocchi.
Una risposta alle nuove funzionalità delle GPU
Secondo Meta, l’hardware recente ha dato maggiore importanza al coordinamento dei gruppi di esecuzione e delle operazioni asincrone. Tra gli esempi figurano le funzionalità di Hopper di NVIDIA, come la specializzazione dei warp e il Tensor Memory Accelerator, e le capacità di Blackwell, tra cui Tensor Memory e l’esecuzione a livello di cluster.
In un esempio fornito dall’azienda, un kernel FlashAttention assegna a un gruppo di warp il trasferimento dei dati, mentre un altro si occupa del calcolo sui tensor core. CUDA può esprimere configurazioni di questo tipo, ma spesso richiede un considerevole lavoro di orchestrazione specifico per l’architettura, afferma Meta; Triton offre un modello più produttivo, anche se alcune relazioni di esecuzione sono nascoste dalle astrazioni del compilatore.
Come funziona TLX
Il ciclo sincrono di caricamento e calcolo di un kernel Triton standard viene messo a confronto con la pipeline esplicita di TLX. Nella versione TLX, i gruppi di warp produttori spostano i dati in modo asincrono, i gruppi consumatori eseguono i calcoli sui tensor core e le barriere coordinano i passaggi di consegne; il calcolo resta espresso in Triton.
L’estensione espone l’orchestrazione dei gruppi di warp, la sincronizzazione, le operazioni di memoria, tra cui TMA e TMEM, e l’esecuzione su cluster. Gli sviluppatori possono inizialmente lasciare i layout non specificati: il compilatore fornisce valori predefiniti, mentre i programmatori possono scegliere manualmente i layout quando vogliono ottimizzare le prestazioni. Secondo Meta, questo può facilitare la fusione delle operazioni senza dover riconciliare manualmente i layout tra le diverse fasi.
L’utilizzo in produzione dichiarato
TLX è in esecuzione in produzione nell’infrastruttura di machine learning di Meta da diversi mesi, con implementazioni nell’addestramento e nell’inferenza dei sistemi di raccomandazione e nel pre-addestramento di LLM su larga scala. Per l’addestramento del suo Generative Ads Model, l’ultima versione ha raggiunto O(10^24) operazioni in virgola mobile complessive su diverse migliaia di GPU NVIDIA Blackwell, con un utilizzo end-to-end dei FLOPs del modello del 20-25%, riferisce Meta.
Meta riferisce inoltre che TLX ha accelerato fino a 2 volte il passaggio forward di un kernel Generalized Dot Product Attention in GEM. Per l’inferenza dei sistemi di raccomandazione, la In-Kernel Broadcast Optimization sposta la replica delle caratteristiche degli utenti all’interno del kernel; la tabella dell’azienda riporta, per i modelli progettati congiuntamente, una riduzione fino a 2/3 della latenza netta nelle operazioni a elevata intensità di calcolo.
Secondo Meta, una terza implementazione, la grouped GEMM distribuita per il pre-addestramento di LLM, è stata validata in produzione su una vasta flotta per un periodo prolungato.
Fonte
Notizie correlate
- Meta illustra i controlli basati sull’IA per le raccomandazioni su Facebook
- Meta illustra le reti per i cluster di IA su scala di gigawatt
- Meta illustra Prometheus, il suo cluster IA su scala gigawatt
- Un team di RadixArk mette in servizio Muse Spark Open di Meta su un solo host TPU
- Meta lancia Muse Ember a 0,10 dollari per milione di token in input