Riepilogo Tecnico: Sovrapposizione delle Attivazioni a Livello di Tile per l'Inferenza Efficiente di LLM

Definizione del Problema

I moderni Large Language Models (LLM), inclusi Qwen-2.5, Llama, Mistral e Gemma, utilizzano prevalentemente la funzione di attivazione SwiGLU. Questa struttura MLP a gate calcola $Y = \text{SiLU}(A \times W_1) \odot (A \times W_2)$, richiedendo due moltiplicazioni matriciali indipendenti (GeMM) seguite da un'attivazione a gate elemento per elemento.

Nelle implementazioni standard, questo processo materializza due tensori intermedi completi ($Gate$e$Up$) nella memoria ad alta larghezza di banda (HBM) tra le fasi di GeMM e di attivazione. Il profiling su GPU NVIDIA H100 rivela che questa materializzazione intermedia e il relativo calcolo dell'attivazione consumano dal 9% al 37% del tempo totale di esecuzione del MLP. L'overhead è più grave per i modelli più piccoli (es. Qwen-2.5 0.5B), dove SwiGLU incide per oltre il 30% del tempo del MLP. Man mano che la densità aritmetica dei Tensor Core aumenta tramite la quantizzazione (FP8, INT4), il costo relativo di queste operazioni limitate dalla memoria (memory-bound) cresce, creando un collo di bottiglia significativo.

L'attuale infrastruttura dei compilatori, specificamente torch.compile di PyTorch, non riesce a risolvere il problema. Non può fondere (fuse) le operazioni attraverso due GeMM separate con matrici di peso differenti a causa di limiti fondamentali nei passaggi di fusione a livello di grafo. Gli esperimenti dimostrano che torch.compile raggiunge solo il 34–94% delle prestazioni di PyTorch in modalità eager per questo pattern, con suggerimenti di fusione espliciti che forniscono miglioramenti trascurabili.

Metodologia

Gli autori presentano due complementari kernel basati su CUTLASS SM90, progettati artigianalmente, che fondono SwiGLU in GeMM a livello di tile, eliminando la materializzazione intermedia in HBM. Entrambi i kernel sfruttano il Tensor Memory Accelerator (TMA) e lo scheduling specializzato per warp disponibile su H100 (SM90).

1. Kernel-1: Sync SwiGLU tramite Sovrapposizione dello Schema Pingpong

• Meccanismo: Questo kernel utilizza lo schema specializzato per warp "Pingpong" di CUTLASS, che partiziona i warp in produttori (che caricano i dati tramite TMA) e consumatori (che eseguono l'MMA).
• Strategia: Crea una griglia di threadblock $[M, N]$. All'interno della fase di consumo, dopo che il tile di Gate è stato calcolato ma prima che arrivino i dati del tile di Up, esiste una breve finestra in cui i warp consumatori sono inattivi. Il kernel sfrutta questa finestra per calcolare l'attivazione Swish ($\text{SiLU}(Gate)$) sull'accumulatore di Gate residente nei registri.
• Sovrapposizione: Il calcolo Swish si sovrappone temporalmente ai carichi TMA dei tile di peso Up del produttore. La moltiplicazione elemento per elemento finale ($\text{SiLU}(Gate) \odot Up$) e lo store avvengono nell'epilogo.
• Ottimizzazione: Questo design è ottimizzato per batch size elevate ($M \ge 2048$), dove la griglia $[M, N]$ satura sufficientemente i multiprocessori di streaming (SM) della GPU.

2. Kernel-2: Interleave SwiGLU tramite Custom Epilogue Visitor Tree

• Meccanismo: Questo kernel introduce un nodo di Epilogue Visitor Tree (EVT) personalizzato chiamato PairMulStore.
• Strategia: Concatena le matrici di peso $W_1$ e $W_2$ in una singola matrice fusa $W_{fused} \in \mathbb{R}^{K \times 2N}$ con colonne intercalate (colonne pari per Gate, dispari per Up). Il kernel lancia un GeMM standard su $[M, 2N, K]$, creando una griglia di threadblock $[M, 2N]$.
• Fusione: Il nodo PairMulStore opera sull'accumulatore $[M, 2N]$ all'interno di ogni threadblock. Legge coppie di colonne adiacenti $(D[m, 2n], D[m, 2n+1])$, applica $\text{SiLU}(D[m, 2n]) \odot D[m, 2n+1]$ e memorizza il risultato.
• Sovrapposizione: Questo design intercala il calcolo SwiGLU con gli store dei tile. L'aritmetica del tile corrente si sovrappone agli store TMA asincroni del tile precedente, fornendo una sovrapposizione tra calcolo e store.
• Ottimizzazione: L'aumento di $2\times$ del numero di threadblock migliora significativamente l'occupazione degli SM per batch size piccole ($M \le 2048$), rendendolo il default preferito per la maggior parte delle configurazioni.

Contributi Chiave

1. Fusione a Livello di Tile: La prima implementazione che fonde SwiGLU in GeMM a livello di tile, eliminando 4 operazioni HBM intermedie (lettura/scrittura di Gate e Up) per ogni step di inferenza.
2. Strategia a Doppio Kernel: Un approccio complementare che utilizza il Kernel-1 per batch grandi (sequenza Gate→Swish→Up) e il Kernel-2 per batch piccoli (pesi intercalati con EVT personalizzato), affrontando diversamente i vincoli di occupazione e scheduling.
3. Superiorità Numerica: I kernel fusi raggiungono zero discrepanze (0% di elementi con errore relativo >1%) rispetto a un riferimento FP32, mentre l'implementazione baseline di cuBLAS presenta discrepanze del 4,5–11% a causa dell'accumulo parziale aggressivo in BF16 attraverso i tile.
4. Validazione dei Limiti del Compilatore: Evidenza empirica che dimostra come torch.compile non possa replicare questa fusione, rimanendo da 3 a 7 volte più lento dei kernel hand-crafted.

Risultati

Valutati su modelli Qwen-2.5 (da 0.5B a 72B) su NVIDIA H100:

• Accelerazione (Speedup): I kernel ottengono un'accelerazione fino a 2.47× rispetto a PyTorch eager (cuBLAS + SwiGLU separato). Il Kernel-2 vince in 13 configurazioni su 20 testate, mentre il Kernel-1 vince in 7 su 20 (principalmente per modelli grandi ad alte batch size).
• Spostamento di Regime: L'analisi Roofline conferma che la fusione sposta il carico di lavoro da memory-bound a compute-bound. Per i modelli piccoli, l'intensità aritmetica aumenta fino al 247%, spostando il punto operativo attraverso la cresta del roofline.
• Utilizzo: L'utilizzo di picco di BF16 raggiunge il 79.5% su H100, quasi raddoppiando l'utilizzo della baseline per i modelli piccoli.
• Efficienza: L'efficienza di fusione (velocità effettiva vs risparmio teorico di memoria) è vicina al 100% per i modelli piccoli. Per il Kernel-2, l'efficienza supera il 100% in diverse configurazioni grazie ai benefici della sovrapposizione tra calcolo e store.
• Scalabilità: L'accelerazione diminuisce con la dimensione del modello (2.47× per 0.5B vs ~1.0× per 72B ad alte batch size) perché il calcolo GeMM domina il tempo totale e la frazione relativa del traffico di memoria intermedio diminuisce.

Significato

Il paper sostiene che per l'inferenza degli LLM moderni, specialmente per i modelli di dimensioni edge o mobile (0.5B–14B), il collo di bottiglia si è spostato dal calcolo puro al traffico di memoria causato dalla materializzazione dei tensori intermedi. Gli autori dimostrano che un design a livello di tile, hand-crafted e consapevole dell'hardware, è necessario per superare i limiti dell'attuale infrastruttura dei compilatori. Eliminando gli accessi intermedi in HBM e sovrapponendo il calcolo dell'attivazione con il movimento dei dati, questi kernel riducono significativamente la latenza di inferenza e migliorano l'utilizzo della computazione, offrendo un'ottimizzazione critica per i sistemi di serving in produzione dove latenza e throughput sono vincolati.