Accelerating GPU Inference of Large Language Models with Moderately Unstructured Sparse Weight Matrices

Sintesi Tecnica: Accelerare l'inferenza GPU di Large Language Models con matrici di pesi sparse moderatamente non strutturate

Problematica
L'implementazione dei Large Language Models (LLM) affronta un collo di bottiglia critico nei costi di inferenza. Sebbene le tecniche di pruning dei modelli (ad esempio, Wanda, RIA) possano ridurre il carico computazionale introducendo la sparsità nelle matrici di pesi, esse sono vincolate dai requisiti di qualità del modello. A differenza dei modelli di piccola scala, gli LLM sono altamente sensibili ad alti livelli di sparsità; mantenere le prestazioni limita tipicamente il pruning a una sparsità non strutturata moderata (circa il 50%).

A questo specifico livello di sparsità, gli esistenti kernel GPU per la Moltiplicazione di Matrici Sparse (SpMM) non riescono a superare le librerie di moltiplicazione di matrici dense come cuBLAS. Questo divario di prestazioni deriva da tre sfide fondamentali:

1. Incompatibilità con i Tensor Core: Le moderne GPU offrono Tensor Core sparsi ad alta efficienza, ma supportano rigorosamente la sparsità strutturata 2:4. Non possono accelerare direttamente i pattern di pruning irregolari e non strutturati.
2. Overhead dei Metadati: I formati di archiviazione tradizionali (ad esempio, CSR) richiedono la memorizzazione di metadati posizionali completi per ogni elemento non nullo. Al 50% di sparsità, questi metadati possono occupare quanto lo spazio dei valori stessi, annullando i benefici della larghezza di banda della memoria.
3. Colli di bottiglia nella decodifica on-chip: Le tecniche di compressione che riducono i metadati (ad esempio, la codifica bitmap) spostano il carico di lavoro sulla decodifica on-chip. Sulle moderne GPU con memoria ad alta larghezza di banda (HBM), l'accesso alla memoria globale non può sovrapporsi completamente al lavoro di decodifica eseguito dai più lenti core CUDA, creando un blocco delle prestazioni (stall).

Metodologia
Gli autori propongono una soluzione olistica che comprende un nuovo formato di archiviazione a tre livelli e un kernel SpMM co-ottimizzato progettato per la sparsità non strutturata moderata.

1. Formato di Archiviazione a Tre Livelli
Il formato proposto decompone la matrice sparsa in tre livelli distinti per bilanciare la compatibilità hardware e l'efficienza di archiviazione:

• Livello Sparse-TC: Questo livello estrae pattern di sparsità strutturata 2:4 dalla distribuzione irregolare. Conserva due valori non nulli da ogni quattro pesi consecutivi. Poiché il pattern è fisso, richiede solo 2 bit per valore per i metadati, garantendo la compatibilità diretta con i tensor core sparsi senza overhead di decodifica.
• Livello Slot-Filling: Gli elementi non nulli in eccesso che non rientrano nel pattern 2:4 vengono ridistribuiti negli slot vuoti lasciati dal livello Sparse-TC. Per minimizzare i metadati per questi ~17,7% degli elementi, gli autori introducono la codifica Parallel Differential Distance (PDD). Invece di memorizzare indici assoluti, la PDD codifica le distanze relative tra gli elementi ridistribuiti. La maggior parte delle distanze è sufficientemente piccola da essere codificata in 4 bit. Per distanze più grandi, vengono inseriti punti fittizi (dummy points) per interrompere l'intervallo, assicurando che nessuna singola distanza superi il limite di codifica.
• Livello Residuo: Gli altri elementi non nulli rimanenti (tipicamente <1% del totale) che non possono essere gestiti efficientemente dai primi due livelli, sono archiviati in un tradizionale formato Compressed Sparse Row (CSR). Grazie al volume minimo di dati, l'overhead dei metadati qui è trascurabile.

2. Kernel SpMM Co-Ottimizzato
Basandosi su questo formato, gli autori progettano un kernel GPU che utilizza congiuntamente tensor core sparsi e core CUDA all'interno di un modello di esecuzione in pipeline:

• Decodifica Parallela: I core CUDA hanno il compito di decodificare i metadati PDD dal livello Slot-Filling in parallelo. Poiché questo livello rappresenta una piccola frazione del totale, l'overhead di decodifica è gestito efficacemente.
• Accelerazione Strutturata: I tensor core sparsi elaborano il livello Sparse-TC utilizzando istruzioni wgmma (architettura Hopper), sfruttando la struttura 2:4 per un'alta produttività.
• Gestione del Residuo: I core CUDA elaborano anche il livello sparso Residuo, scrivendo i risultati in un buffer separato on-chip per essere ridotti con gli output dei tensor core alla fine della pipeline.
• Sovrapposizione della Pipeline: Il kernel implementa una pipeline a grana fine che sovrappone tre operazioni critiche:
  • Caricamento delle matrici di input dense (cache in L2) e dei tile di pesi sparsi (da Global Memory).
  • Decodifica dei metadati PDD sui core CUDA.
  • Esecuzione della moltiplicazione di matrici sui tensor core sparsi.
    Questa strategia "carica come sparso, computa come sparso" assicura che il calcolo on-chip si sovrapponga all'accesso alla memoria globale, massimizzando l'utilizzo dell'HBM.

Contributi Chiave

• Formato di Archiviazione Innovativo: Un formato a tre livelli adattato per la sparsità non strutturata moderata che consente la compatibilità con i tensor core sparsi, riduce drasticamente l'overhead dei metadati tramite la codifica PDD e minimizza i costi di decodifica on-chip.
• Kernel Co-Ottimizzato: Un kernel SpMM che coordina efficientemente i tensor core sparsi e i core CUDA, utilizzando la pipelining per sovrapporre l'accesso alla memoria con il calcolo.
• Primo a Superare il Dense: L'implementazione è la prima a dimostrare che la SpMM con sparsità non strutturata moderata può superare la moltiplicazione di matrici densa sulle moderne GPU dotate di HBM.

Risultati
Le valutazioni sono state condotte su GPU NVIDIA H100 SXM5 con i modelli OPT-30B e OPT-66B con livelli di sparsità del 50%, 60% e 70%.

• Prestazioni a livello di Kernel: Al 50% di sparsità, il metodo proposto ottiene un acceleramento fino a 1,64× rispetto a SpInfer (EuroSys'25 Best Paper) e supera significativamente altri baseline come SparTA e FlashLLM. Supera anche le prestazioni dense di cuBLAS.
• Inferenza End-to-End: Il metodo fornisce un acceleramento end-to-end fino a 1,41× rispetto a FlashLLM (VLDB'24) e uno throughput fino al 41,4% superiore rispetto a FlashLLM su OPT-30B con un batch size di 16.
• Efficienza della Memoria: L'approccio riduce l'uso del picco di memoria globale di circa il 21,4% rispetto ai baseline densi al 50% di sparsità, sebbene comporti un uso di memoria leggermente superiore rispetto a SpInfer a causa della mancanza di codifica bitmap per i metadati.
• Scalabilità: Il metodo mantiene i vantaggi attraverso i livelli di sparsità dal 50% al 70%.

Significatività e Rivendicazioni
Il documento sostiene di affrontare un vuoto critico nell'ottimizzazione dell'inferenza degli LLM: l'incapacità degli esistenti kernel sparsi di sfruttare efficacementmente la sparsità non strutturata moderata. Introducendo un formato di archiviazione che colma il divario tra pattern di pruning irregolari e accelerazione hardware strutturata, gli autori dimostrano che il pruning può produrre reali accelerazioni di inferenza senza sacrificare la qualità del modello.

Il lavoro evidenzia come, mentre un'alta sparsità (>70%) degrada la qualità del modello, la sparsità moderata (~50%) sia il "punto ottimale" (sweet spot) per gli LLM. La soluzione proposta sblocca il potenziale prestazionale di questo punto ottimale, precedentemente inaccessibile a causa delle inefficienze dei kernel sparsi esistenti. Gli autori sottolineano che la loro ottimizzazione è specificamente mirata alla fase di decode dell'inferenza degli LLM (dove i batch size sono piccoli e limitati dalla memoria), rendendola altamente adatta allo spiegamento scalabile in ambienti di produzione dove le fasi di prefill e decode sono spesso disaccoppiate. Il codice sorgente è reso disponibile per facilitare ulteriore ricerca e adozione.