Sintesi Tecnica: GPU Forecasters: Modelli di Linguaggio come Surrogati Selettivi per l'Ottimizzazione del Runtime dei Kernel

Definizione del Problema

I kernel GPU sono i motori fondamentali del moderno deep learning, dove anche piccoli guadagni di efficienza si traducono in riduzioni significative dei costi di addestramento e inferenza. L'ottimizzazione di questi kernel implica tipicamente un ciclo di generazione-valutazione, in cui i Large Language Models (LLM) propongono candidati kernel, che vengono poi compilati ed eseguiti su hardware fisico per misurarne le prestazioni. Tuttavia, questo processo di valutazione rappresenta un collo di bottiglia critico. Ogni valutazione richiede la compilazione (che impiega dai 60 ai 120 secondi per CUDA) ed l'esecuzione sul dispositivo, rendendo il throughput di valutazione ordini di grandezza più lento rispetto alla generazione. Man mano che i budget di ricerca scalano, il costo delle ripetute misurazioni fisiche limita la capacità di esplorare il vasto spazio di potenziali implementazioni dei kernel.

Mentre il lavoro precedente si è concentrato sull'uso degli LLM per generare migliori kernel, questo articolo affronta il problema complementare di ridurre il costo di valutazione. Gli autori propongono di utilizzare gli LLM non solo come generatori, ma come surrogati selettivi della GPU fisica. L'obiettivo è prevedere le prestazioni relative di un kernel candidato rispetto a un'implementazione di riferimento e identificare quando il modello è incerto, rimandando così solo i casi incerti alla costosa GPU fisica.

Metodologia

Formulazione del Task: Forecasting Selettivo

Gli autori formulano il problema come la previsione dello speedup relativo ($S_i = T_{ref} / T_i$) di un kernel candidato rispetto a un kernel di riferimento, date le specifiche dell'hardware. Prevedere il runtime assoluto in secondi è considerato intrattabile a causa della complessità della modellazione del comportamento del compilatore e dei dettagli micro-architetturali. Inveov, il target viene discretizzato in otto bin ordinali di speedup (che vanno da "Severe Slowdown" a "Extreme Speedup").

Il surrogato (un LLM) riceve in input il kernel di riferimento, il kernel candidato e le specifiche dell'hardware. Genera una catena di ragionamento e produce una distribuzione di probabilità sugli otto bin. La confidenza dichiarata dal modello è definita come la massima probabilità assegnata a un singolo bin.

Strategia di Predizione Selettiva

Il sistema opera sotto un framework di predizione selettiva. Il surrogato viene utilizzato per classificare i candidati. Se la confidenza del surrogato supera una certa soglia, la previsione viene accettata e il candidato non viene misurato sulla GPU. Se la confidenza è bassa (indicando incertezza), il candidato viene deferito alla GPU fisica per la misurazione del valore reale (ground-truth). Ciò consente al processo di ricerca di considerare molti più candidati entro un budget fisso di misurazioni GPU.

Apprendimento per Rinforzo (RL) per la Calibrazione

Per migliorare l'accuratezza del surrogato e, soprattutto, la sua calibrazione (l'allineamento tra la confidenza dichiarata e l'accuratezza effettiva), gli autori applicano l'Apprendimento per Rinforzo (Reinforcement Learning) a un modello open-weights (GPT-OSS-20B). Definiscono ricompense basate su:

1. Correttezza: Se il bin predetto corrisponde al bin misurato.
2. Calibrazione Categorica (Brier Score): Penalizza la distribuzione in base alla distanza dal target reale one-hot.
3. Calibrazione Sensibile alla Distanza (CRPS): Penalizza la distribuzione cumulativa, tenendo conto della natura ordinale dei bin di speedup (ad esempio, predire uno speedup di 2× quando la realtà è 1.5× è meno grave rispetto a predire 4×).

I dati di addestramento sono costruiti come sottoprodotto della ricerca stessa dei kernel, utilizzando i candidati di successo di precedenti sessioni di ricerca come esempi di addestramento.

Contributi Chiave

1. LLM come Surrogati GPU: Il documento dimostra che gli LLM possono raggiungere un'accuratezza non banale nel prevedere come una GPU fisica risponda ai kernel candidati, recuperando una frazione sostanziale dello speedup disponibile classificando i candidati in base alle previsioni piuttosto che alle misurazioni.
2. Framework di Predizione Selettiva: Gli autori inquadrano il forecasting GPU come un problema di predizione selettiva, dimostrando che i surrogati possono riportare un'incertezza calibrata. Ciò consente una strategia di ricerca che si fida delle previsioni ad alta confidenza e riserva le misurazioni fisiche per i casi incerti.
3. Impatto dell'RL sul Comportamento del Surrogato: Attraverso esperimenti, il documento caratterizza come l'addestramento RL sposti le proprietà del surrogato. Nello specifico, l'aggiunta di ricompense di calibrazione (Brier o CRPS) migliora l'affidabilità delle stime di confidenza del modello, mentre l'addestramento basato solo sulla correttezza può degradare la relazione tra confidenza e accuratezza.
4. Miglioramento della Ricerca End-to-End: Lo studio mostra che integrare il surrogato in un ciclo di ricerca di kernel (specificamente una ricerca basata su PUCT) permette alla ricerca di esplorare significativamente più candidati per ogni misurazione GPU. Nella maggior parte dei task, ciò porta a trovare kernel più veloci rispetto a un baseline che misura ogni proposta sotto lo stesso budget di GPU.
5. Rilascio del Dataset: Gli autori rilasciano una collezione riutilizzabile di 12.388 kernel GPU generati da LLM con runtime misurati, che copre 118 problemi, molteplici tipi di GPU (A100, L40S, H100) e varie metodologie di generazione. Questo dataset supporta la ricerca più ampia nella previsione delle prestazioni GPU e nella ricerca di kernel.

Risultati

• Accuratezza della Previsione: Gli LLM pronti all'uso (es. GPT-OSS-20B) recuperano circa l'83% del miglior speedup disponibile quando viene misurato su GPU solo l'1% dei migliori candidati. Questa prestazione migliora al 94% quando si misurano il top 50%.
• Calibrazione e RL: L'addestramento RL altera significativamente il comportamento del surrogato.
  • L'addestramento basato solo sulla correttezza migliora l'accuratezza del bin esatto ma degrada la correlazione tra confidenza ed errore (il modello diventa eccessivamente sicuro di sé nelle previsioni errate).
  • L'addestramento alla calibrazione (Brier/CRPS) migliora la monotonicità tra errore di previsione e confidenza. Notevolmente, l'addestramento calibrato con Brier migliora la metrica di "speedup recuperato" di circa 2,1 punti percentuali rispetto al modello base, suggerendo che distribuzioni meglio calibrate portano segnali di classificazione migliori per l'algoritmo di ricerca.
• Efficienza della Ricerca: Nella ricerca di kernel end-to-end, la ricerca abilitata dal surrogato trova kernel più veloci rispetto al baseline in quattro casi su sei, spesso raggiungendo il miglior speedup del baseline con un numero sostanzialmente inferiore di misurazioni GPU. Nei task in cui il baseline vince, la ricerca sembra saturarsi precocemente, indicando che il surrogato è più vantaggioso quando lo spazio di ricerca è ampio e inesplorato.
• Momenti di Scoperta: Il surrogato può identificare i "momenti di scoperta" (mutazioni in cui un kernel figlio è significativamente più veloce del suo genitore) meglio del caso casuale, sebbene la precisione rimanga una sfida, rendendo necessaria la conferma della GPU per decisioni ad alto rischio.

Significato e Rivendicazioni

Il documento sostiene che gli LLM possano giocare un ruolo più ampio nell'ottimizzazione dei kernel oltre la semplice generazione di codice. Agendo come modelli virtuali di una GPU, possono servire come filtri ad alto throughput e basso costo che gestiscono la scarsità delle risorse hardware fisiche.

Gli autori sottolineano che il valore primario di questo approccio non è solo la velocità, ma la capacità di eseguire un'ottimizzazione basata su modello senza interrogare il sistema originale ad ogni passaggio. Il surrogato offre un throughput più elevato e un costo per query inferiore rispetto a una GPU, scambiando il rumore della previsione che può essere gestito attraverso stime di incertezza calibrate. Il lavoro collega l'ottimizzazione dei kernel al campo più ampio della costruzione di simulatori neurali per sistemi fisici complessi, suggerendo una strada verso ottimizzatori che trattano la misurazione dell'hardware come una risorsa scarsa da utilizzare selettivamente.