Riassunto Tecnico: CUDA-L2

Definizione del Problema

La moltiplicazione di matrici (matmul) è un'operazione fondamentale nei Large Language Models (LLM), poiché rappresenta una parte significativa del tempo di computazione sia nell'addestramento che nell'inferenza. Sebbene esistano kernel CUDA ottimizzati manualmente, raggiungere prestazioni ottimali attraverso diverse dimensioni di matrici ($M, N, K$) e architetture GPU rimane una sfida significativa. Diversi tripletti di dimensioni spesso richiedono strategie di ottimizzazione distinte che non si trasferiscono bene tra l'hardware. Inoltre, anche per dimensioni identiche, le strategie ottimali variano in base ai dettagli di configurazione, come la precisione dell'accumulatore (FP16 vs. FP32).

Gli approcci automatizzati esistenti, come quelli che utilizzano i Large Language Models (LLM) per la generazione di kernel, si sono concentrati principalmente su benchmark con singole configurazioni fisse (ad esempio, KernelBench). Rimane poco chiaro come questi metodi si traducano in ambienti di produzione dove devono competere contro librerie altamente ottimizzate e a codice chiuso come cuBLAS e cuBLASLt di NVIDIA. Vi è una mancanza di sistemi automatizzati capaci di esplorare sistematicamente il vasto spazio di configurazione della moltiplicazione di matrici in mezza precisione (HGEMM) per superare i baseline ottimizzati da esperti umani.

Metodologia: CUDA-L2

Gli autori propongono CUDA-L2, un sistema che combina Large Language Models (LLM) e Reinforcement Learning (RL) per ottimizzare automaticamente i kernel CUDA per HGEMM. Il sistema è progettato per gestire 1.000 configurazioni distinte che rappresentano tutte le combinazioni di valori $M, N, K$ da $\{64, 128, 256, 512, 1024, 2048, 4096, 8192, 12288, 16384\}$.

Architettura del Sistema e Addestramento

CUDA-L2 estende il precedente framework CUDA-L1 con quattro miglioramenti tecnici chiave:

1. Pre-addestramento Continuativo (Continued Pretraining): Il modello è pre-addestrato su un corpus diversificato di codice CUDA oltre i benchmark standard, inclusi sorgenti da PyTorch, ATen, CUTLASS e tutorial NVIDIA. Le descrizioni delle istruzioni vengono generate per gli snippet di codice grezzo, e viene aggiunto un contesto di recupero aumentato (retrieval-augmented) per incorporare nuove conoscenze e caratteristiche architettoniche.
2. Addestramento RL Multi-stadio: L'addestramento procede da tipi di kernel generali ampi a domini specializzati (specificamente matmul).
   • RL per Kernel Generali: Addestra il modello su una vasta gamma di operazioni (algebra lineare, convoluzione, riduzione, ecc.) utilizzando la velocità di esecuzione come ricompensa.
   • RL per HGEMM: Si concentra specificamente sui kernel HGEMM con diverse configurazioni $M, N, K$.
3. Contesto Potenziato con Metriche di Profilazione: Il contesto RL include metriche di profilazione complete di NVIDIA Nsight Compute (NCU) (ad esempio, throughput della memoria, occupazione SM, efficienza della cache) per consentire decisioni di ottimizzazione migliori rispetto al semplice tempo di esecuzione end-to-end.
4. Design della Funzione di Ricompensa: La ricompensa per un kernel generato si basa sul punteggio medio di accelerazione (speedup) attraverso le iterazioni di test. Per garantire la robustezza, la ricompensa è penalizzata per:
   • Deviazione numerica da una verità di base (ground truth) FP32 della CPU.
   • Lunghezza eccessiva del codice (per incoraggiare la concisione).
   • Il sistema impone criteri di correttezza rigorosi, inclusi match esatti per input binari e deviazioni limitate rispetto a baseline affidabili (cuBLAS, cuBLASLt).

Protocollo di Valutazione

Il sistema è valutato in due scenari per riflettere diverse realtà di deployment:

• Modalità Offline: I kernel vengono eseguiti consecutivamente senza intervalli di tempo, simulando il throughput massimo.
• Modalità Server: I kernel vengono eseguiti a intervalli casuali per simulare l'inferenza in tempo reale, tenendo conto della dinamica termica della GPU (periodi di raffreddamento e clock boosting).

La valutazione confronta CUDA-L2 con quattro baseline:

1. torch.matmul (il dispatch predefinito di PyTorch).
2. cuBLAS (utilizzando la selezione del layout ottimale tra NN e TN).
3. cuBLASLt-heuristic (selezionando l'algoritmo con il ranking più alto dall'API euristica di NVIDIA).
4. cuBLASLt-AutoTuning (testando esaustivamente fino a 100 candidati per trovare il più veloce).

Contributi Chiave e Tecniche Scoperte

CUDA-L2 dimostra che l'RL guidato da LLM può scoprire implementazioni superiori esplorando sistematicamente spazi di configurazione impraticabili per la calibrazione manuale. Il sistema scopre e applica autonomamente diverse tecniche di ottimizzazione avanzata:

• Selezione dell'Astrazione: Il modello seleziona automaticamente tra le primitive raw WMMA per matrici piccole e le astrazioni CuTe per operazioni tiled più grandi e complesse.
• Padding degli Input: Il sistema scopre strategie in cui le matrici di input vengono caricate con zeri per dimensioni non divisibili (ad esempio, padding di $M=8192$ a $8320$) per abilitare dimensioni di tile ($B_M$) più efficienti che sarebbero altrimenti invalide, bilanciando l'overhead di calcolo extra con i guadagni di prestazione.
• Sovrapposizione Avanzata di Memoria e Calcolo:
  • Frammenti di Registro a Doppio Buffer: Alternanza tra due set di frammenti di registro per sovrapporre il movimento dei dati con il calcolo (esecuzione ping-pong).
  • Prefetching Aggressivo: Oltre al prefetching a singolo step, il sistema passa al prefetching a più step (caricamento di più iterazioni in anticipo) quando la pressione dei registri lo consente.
  • Scheduling Staggered del Prefetching: Intercalare il prefetching delle matrici A e B attorno all'operazione di multiply-accumulate (MMA) per massimizzare il parallelismo a livello di istruzione.
  • Copia Diretta da Registro a Memoria Condivisa: Eliminazione di tensori temporanei intermedi nell'epilogo quando i layout dei registri corrispondono a quelli della memoria condivisa, utilizzando tipi di dati ampi (ad esempio, uint128_t) per ridurre le operazioni di copia.
  • Block Swizzling: Riordinamento dell'esecuzione dei thread block per migliorare i tassi di hit della cache L2, particolarmente per dimensioni di problemi elevate.

Risultati

CUDA-L2 supera sistematicamente tutti i principali baseline su 1.000 configurazioni su GPU A100.

Prestazioni in Modalità Offline (Accelerazione Media rispetto ai Baseline):

• +22.0% rispetto a torch.matmul.
• +19.2% rispetto a cuBLAS (usando il layout ottimale).
• +16.8% rispetto a cuBLASLt-heuristic.
• +11.4% rispetto a cuBLASLt-AutoTuning (il baseline più forte, che testa esaustivamente 100 candidati).

Prestazioni in Modalità Server:
Le accelerazioni aumentano nello scenario server, probabilmente a causa della capacità del sistema di gestire meglio la dinamica termica e gli effetti di raffreddamento della cache:

• +28.7% rispetto a torch.matmul.
• +26.0% rispetto a cuBLAS.
• +22.4% rispetto a cuBLASLt-heuristic.
• +15.9% rispetto a cuBLASLt-AutoTuning.

I tassi di vittoria (win rates) variano dal 79.3% al 95.7% tra tutti i baseline, indicando che i miglioramenti sono sistematici piuttosto che guidati da outlier. Il sistema mostra anche che le accelerazioni sono generalmente più alte per dimensioni di matrici minori dove l'utilizzo della GPU è inferiore, mentre le matrici più grandi saturano il throughput della GPU, restringendo il divario con i baseline.

Significato

L'articolo sostiene che CUDA-L2 dimostra come anche i kernel più critici per le prestazioni e pesantemente ottimizzati, come l'HGEMM, possano essere migliorati attraverso l'automazione dell'RL guidata da LLM. Scalando l'ottimizzazione su spazi di configurazione difficili da calibrare manualmente, il sistema raggiunge guadagni di prestazione che superano le librerie a codice chiuso allo stato dell'arte.

Sebbene l'attuale versione sia limitata alle architetture A100 (Ampere), il framework è progettato per una vasta applicabilità. Gli autori segnalano il lavoro in corso per estendere il sistema ad altre architetture GPU, incluse Ada Lovelace, Hopper e Blackwell. Il lavoro suggerisce che i metodi automatizzati possono esplorare sistematicamente vasti spazi di configurazione per trovare ottimizzazioni che la calibrazione manuale o le librerie basate su euristiche potrebbero mancare.