TileLang è un linguaggio specifico di dominio (DSL) di alto livello in Python per la progettazione e la compilazione di kernel GPU orientati alle prestazioni tramite TVM. Poiché i modelli linguistici di grandi dimensioni e altre architetture transformer spingono verso richieste computazionali sempre maggiori, la capacità di scrivere kernel GPU personalizzati che sfruttino appieno i tensor core — gli acceleratori dedicati alla moltiplicazione matriciale di NVIDIA — è diventata sempre più critica per l'efficienza di training e inferenza. Tuttavia, la scrittura di tali kernel richiede tradizionalmente competenze approfondite in CUDA C++, programmazione a livello di warp e gerarchie di memoria specifiche dell'hardware. TileLang colma questa lacuna permettendo agli sviluppatori di esprimere i calcoli a livello di tile mentre il compilatore gestisce i dettagli di basso livello. Questo tutorial offre una panoramica completa delle capacità di TileLang, a partire dalla validazione dell'ambiente e dall'istituzione di utility riutilizzabili per il benchmarking e la verifica numerica. Da lì, implementa progressivamente l'addizione vettoriale, la moltiplicazione matriciale a tile su tensor core, l'esplorazione degli schedule, gli epiloghi GEMM fusi, la softmax per righe e FlashAttention.

Durante tutto il tutorial, gli sviluppatori lavorano direttamente con i tile in shared memory, i frammenti di registro, i cicli pipelined, le primitive di iterazione parallela, le riduzioni e gli operatori GEMM su tensor core di TileLang. Il compilatore gestisce automaticamente la mappatura dei thread, i layout di memoria, la sincronizzazione, la vettorizzazione e la generazione di istruzioni CUDA di basso livello. I kernel vengono confrontati con i baseline PyTorch e cuBLAS, con ispezione del codice sorgente CUDA generato, valutazione del throughput di memoria e calcolo, e autotuning per identificare configurazioni del kernel dipendenti dall'architettura. TileLang è disponibile su GitHub.

Il tutorial inizia configurando l'ambiente CUDA di Google Colab, installando TileLang con un fallback nightly e importando i moduli PyTorch e TileLang necessari. Vengono definite utility riutilizzabili per il benchmarking, la validazione e il reporting per misurare la latenza del kernel e confrontare i risultati numerici utilizzando l'errore relativo. Viene quindi implementato un kernel TileLang di addizione vettoriale, eseguito sulla GPU, confrontato in termini di bandwidth con PyTorch e ispezionato per il codice sorgente CUDA generato dal compilatore.

Successivamente, viene implementato un kernel di moltiplicazione matriciale a tile su tensor core che sposta i tile di input attraverso la memoria globale, la shared memory e i frammenti di registro. Le dimensioni dei tile, gli stadi di pipeline, i conteggi dei thread e il L2 swizzling sono controllati manualmente, mentre TileLang genera le istruzioni per i tensor core, la sincronizzazione e la logica di trasferimento della memoria. Diverse configurazioni di schedule vengono sottoposte a benchmark, ne viene verificata l'accuratezza numerica e viene identificata la configurazione del kernel con le migliori prestazioni dipendenti dall'architettura. Il codice completo del tutorial è disponibile qui.

Il kernel di moltiplicazione matriciale viene quindi esteso fondendo l'addizione del bias e l'attivazione GELU direttamente nell'accumulatore residente in registro. Questo approccio riduce il traffico intermedio in memoria globale completando l'epilogo prima di scrivere il tensore di output finale. La fusione dei kernel è una tecnica consolidata per ridurre i colli di bottiglia della bandwidth di memoria, che spesso dominano la latenza nei carichi di lavoro di reti neurali su larga scala. L'implementazione fusa viene confrontata con l'esecuzione eager di PyTorch. Inoltre, viene implementato un kernel softmax per righe utilizzando riduzioni di massimo e somma a livello di frammento, mantenendo il processo di normalizzazione in gran parte all'interno dei registri.

Viene implementato un kernel FlashAttention forward fused che elabora i tile di query, key e value senza materializzare la matrice completa degli attention score in memoria globale. Gli aggiornamenti online della softmax vengono applicati utilizzando i massimi progressivi, le somme di normalizzazione, i fattori di riscalamento e le moltiplicazioni matriciali a tile su tensor core. FlashAttention, originariamente introdotto da ricercatori di Stanford, è diventato una tecnica ampiamente adottata per ridurre l'impronta di memoria della self-attention da uno spazio quadratico a uno lineare, ed è ora integrato nei principali framework, inclusa l'API scaled dot-product attention di PyTorch. Sia l'attention causale che non causale vengono validate rispetto alla scaled dot-product attention di PyTorch, con confronto di latenza e throughput computazionale.

Viene definito uno spazio di ricerca per l'autotuning su dimensioni dei tile matriciali, dimensioni dei blocchi K, profondità di pipeline e conteggi di thread, filtrando le configurazioni che superano il budget di shared memory. Il decoratore di autotuning di TileLang viene utilizzato per compilare, sottoporre a benchmark, validare e memorizzare in cache molteplici schedule di kernel per lo stesso carico di lavoro di moltiplicazione matriciale. Il kernel selezionato viene eseguito, verificato rispetto a PyTorch e valutato in termini di latenza raggiunta e throughput su tensor core. Questa ricerca automatizzata affronta una sfida pratica nello sviluppo di kernel GPU: le dimensioni ottimali dei tile e le profondità di pipeline variano tra architetture GPU come Ampere e Hopper, rendendo il tuning manuale fragile quando si mira a molteplici generazioni di hardware.

Il flusso di lavoro di debug e introspezione di TileLang viene introdotto tramite la stampa lato device, l'ispezione del codice CUDA generato e il profiler di kernel integrato. Vengono esaminati i landmark emessi dal compilatore, tra cui le operazioni su tensor core, le copie asincrone, le barriere di sincronizzazione e le istruzioni di caricamento matriciale. Tutte le sezioni del tutorial sono organizzate in un runner fault-tolerant che registra lo stato di esecuzione, riporta le informazioni sui tempi e stampa una guida di riferimento compatta per la programmazione in TileLang.

Il tutorial dimostra come TileLang traduca programmi Python a livello di tile in kernel GPU ottimizzati senza richiedere la gestione manuale degli indici di thread, dei layout di dati a livello di warp, delle istruzioni per i tensor core o delle barriere di memoria asincrona. I kernel implementati e validati coprono operazioni elementwise vincolate dalla bandwidth, carichi di lavoro GEMM ad alta intensità di calcolo, epiloghi di reti neurali fusi, riduzioni residenti in registro e attention con online softmax. L'esplorazione mostra inoltre come le dimensioni dei blocchi, il consumo di shared memory, la profondità di pipeline, il conteggio dei thread, la forma dei tile e il L2 swizzling influenzino le prestazioni tra diverse architetture GPU. L'ispezione del codice generato, il debug lato device, il profiling e la ricerca automatizzata degli schedule stabiliscono insieme un flusso di lavoro completo per lo sviluppo, la verifica, il benchmarking e il perfezionamento di kernel TileLang personalizzati.