Kernel Forge: An Agent Harness for LLM-based Generation and Optimization of CUDA Kernels

Riepilogo Tecnico: Kernel Forge

Problematica

I modelli di machine learning si affidano sempre più a un piccolo set di kernel computazionali (ad esempio, moltiplicazione di matrici, convoluzione, normalizzazione) per la maggior parte del loro tempo di esecuzione. Sebbene l'ottimizzazione di questi kernel sia una via diretta per ridurre la latenza e i costi, i metodi tradizionali richiedono ingegneri esperti per scrivere manualmente codice GPU a basso livello, un processo lento, costoso e che richiede elevate competenze.

Recenti sistemi agentici basati su Large Language Models (LLM) hanno mostrato promesse nella generazione e nell'ottimizzazione di kernel con meno sforzo umano. Tuttavia, gli strumenti esistenti soffrono di quattro criticità fondamentali quando applicati a scenari reali:

1. Gap di Integrazione: Spesso emettono codice CUDA standalone, lasciando agli sviluppatori il compito di reintegrare manualmente i kernel ottimizzati in modelli PyTorch non modificati.
2. Restrizione del Carico di Lavoro: Molti strumenti si concentrano esclusivamente sui carichi di lavoro degli LLM o su isolati kernel di benchmark, fallendo nel gestire la diversità di forme, layout e argomenti presenti in modelli di visione, diffusione e audio.
3. Limitazioni della Ricerca: Gli approcci attuali spesso si affidano a un raffinamento lineare o a una ricerca di tipo beam search, che possono scartare percorsi di ottimizzazione promettenti che inizialmente appaiono subottimali o che non riescono a esplorare efficacemente lo spazio delle soluzioni.
4. Osservabilità: La maggior parte dei sistemi sono prototipi di ricerca o script privi di interfacce grafiche integrate per monitorare i progressi, ispezionare i candidati o diagnosticare i fallimenti.

Inoltre, i benchmark esistenti valutano frequentemente correttezza e latenza utilizzando tensori generati casualmente e invocazioni isolate di kernel. Ciò non riesce a catturare le condizioni specifiche del modello reale, come le forme dei tensori, le distribuzioni delle attivazioni e i comportamenti degli operatori vicini, creando un divario tra le prestazioni del kernel isolato e l'esecuzione effettiva del modello.

Metodologia: Kernel Forge

Kernel Forge è un harness agentico end-to-end open-source progettato per ottimizzare le esecuzioni di reali modelli PyTorch senza richiedere agli utenti di modificare il codice del modello. Il sistema opera attraverso una pipeline multi-stadio:

1. Ingestione del Modello e Cattura degli Operatori

Kernel Forge accetta qualsiasi modello PyTorch standard e input di esempio forniti dall'utente. Esegue il modello per catturare gli operatori specifici che vengono eseguiti durante l'inferenza. Invece di trattare gli operatori come famiglie generiche, li raggruppa in varianti compatibili con il runtime basate su pattern di invocazione concreti (forme dei tensori, dtypes, layout, argomenti e posizionamento del dispositivo). Queste varianti sono incapsulate come "operator cards", che fungono da target di ottimizzazione e preservano l'esecuzione eager di PyTorch come riferimento per la correttezza.

2. Generazione di Kernel Agentica

Il sistema impiega un agente basato su LLM per generare kernel CUDA specializzati.

• Input: L'agente riceve una operator card e la configurazione hardware/runtime.
• Processo: Per il primo candidato, l'LLM genera il codice sorgente CUDA e un wrapper di lancio. Nelle iterazioni successive, l'agente riceve un kernel genitore e la sua linea di discendenza dal controller di ricerca per proporre revisioni mirate a ridurre la latenza preservando il comportamento.
• Vincolo: L'agente è istruito a produrre codice CUDA personalizzato piuttosto che delegare a librerie esistenti (es. cuDNN, cuBLAS), sebbene il sistema riconosca che il codice generato possa talvolta avvolgere efficacemente tali backend.

3. Validazione e Riparazione

I candidati generati sono sottoposti a un rigoroso ciclo di validazione:

• Compilazione: Il codice viene compilato con nvcc. Gli errori di compilazione attivano un ciclo di riparazione limitato.
• Esecuzione: I candidati con successo vengono lanciati con gli input catturati. I fallimenti a runtime (crash, accesso illegale alla memoria) vengono riportati al ciclo di riparazione.
• Correttezza Numerica: I candidati vengono confrontati con gli output eager di PyTorch catturati durante l'ingestione, utilizzando tolleranze numeriche specifiche per l'operatore. Solo i candidati che superano questo filtro specifico per il workload entrano nell'albero di ricerca.

4. Controller Monte Carlo Tree Search (MCTS)

A differenza del raffinamento lineare, Kernel Forge utilizza MCTS per esplorare molteplici percorsi di ottimizzazione.

• Struttura dell'Albero: Il controller mantiene un albero di revisione persistente in cui i nodi rappresentano i kernel candidati, memorizzando il codice sorgente, la linea di discendenza, il numero di visite e la latenza misurata.
• Strategia di Selezione: Il controller utilizza una versione di minimizzazione di Upper Confidence bounds applied to Trees (UCT) combinata con il progressive widening. Bilancia l'esplorazione di rami meno visitati con lo sfruttamento di percorsi promettenti, permettendo al sistema di rivisitare candidati che potrebbero non mostrare guadagni immediati ma che potrebbero portare a soluzioni migliori in seguito.
• Feedback: Le misurazioni della latenza vengono propagate verso l'alto nell'albero per guidare la selezione del successivo kernel genitore per il raffinamento.

5. Dispatch Protetto e Packaging

L'output finale è un pacchetto protetto (.cast) che integra i kernel ottimizzati nel modello.

• Logica di Guardia: Un kernel generato viene distribuito solo se è valido, compatibile con il pattern di chiamata catturato, numericamente corretto e più veloce della latenza PyTorch eager misurata per quella specifica variante.
• Fallback: Se un candidato fallisce qualsiasi controllo o è più lento, il sistema torna automaticamente al percorso Py torch eager originale. Ciò garantisce la sicurezza e previene regressioni delle prestazioni.
• Auditabilità: Il sistema distingue tra codice CUDA custom generato, wrapper di backend e percorsi di fallback, fornendo una chiara traccia di audit di ciò che è stato effettivamente ottimizzato.

Interfaccia Utente

Kernel Forge fornisce sia un'interfaccia grafica (GUI) che un'interfaccia a riga di comando (CLI). Questi strumenti consentono agli utenti di monitorare il progresso dell'ottimizzazione, ispezionare i kernel candidati, visualizzare l'albero di revisione MCTS, debuggare i fallimenti e confrontare i risultati tra diversi provider di modelli.

Contributi Chiave

1. Sistema Agentico End-to-End: L'introduzione di Kernel Forge, uno strumento che ottimizza le esecuzioni di reali modelli PyTorch catturando gli operatori, generando kernel CUDA specializzati, esplorando molteplici percorsi tramite MCTS e integrando automaticamente i risultati nel modello.
2. Osservabilità e Usabilità: La fornitura di una GUI e di una CLI per il monitoraggio, l'ispezione e il debug del processo di ottimizzazione, colmando il divario tra prototipi di ricerca e strumenti pratici.
3. Valutazione su Scenari Reali: Una valutazione su quattro diversi modelli PyTorch (ResNet-50, Stable Diffusion 3.5 Medium, Gemma 4 E2B e Qwen 3.5 35B-A3B) che rappresentano carichi di lavoro di visione, diffusione e LLM.
4. Rilascio Open-Source: Il rilascio dello strumento per supportare la ricerca riproducibile e l'adozione pratica dell'ottimizzazione di kernel tramite agenti.

Risultati

Valutato su un NVIDIA DGX Spark con GPU GB10, Kernel Forge è stato testato con 50 iterazioni di ottimizzazione per kernel. I risultati evidenziano sia il potenziale che le sfide dell'ottimizzazione guidata da LLM in carichi di lavoro reali:

• Guadagni di Prestazione: Kernel Forge ha ottimizzato con successo 14 kernel rendendoli superiori alla modalità PyTorch eager. Notevoli accelerazioni includono:
  • 1.52× su adaptive_avgpool2d in ResNet-50.
  • 1.70× su group_norm in Stable Diffusion 3.5 Medium.
  • 2.83× su softmax in Gemma 4 E2B.
  • 1.54× su softmax in Qwen 3.5 35B-A3B.
• Specificità dell'Operatore: Accelerazioni significative sono state trovate prevalentemente nei percorsi open-source/nativi di PyTorch (es. tensor_iadd, relu, softmax). Al contrario, gli operatori supportati da implementazioni proprietarie o di vendor mature (es. conv2d, linear, scaled_dot_product_attention) sono spesso rimasti al livello o al di sotto delle prestazioni PyTorch eager.
• Impatto sul Runtime: Gli operatori che hanno ottenuto i più alti speedup locali spesso rappresentavano una piccola frazione del tempo di esecuzione totale catturato. Al contrario, i contributori dominanti del runtime erano frequentemente i percorsi supportati da vendor altamente ottimizzati, difficili da migliorare.
• Costo di Ricerca: La valutazione ha notato che i costi di ricerca (utilizzo API LLM) aumentano significativamente ad alti conteggi di iterazione (opt50), ma ciò non si è sempre correlato a guadagni proporzionali nelle prestazioni aggregate della regione degli operatori, specialmente per gli operatori dominanti.

Significato e Claim

Il paper sostiene che Kernel Forge dimostra come l'ottimizzazione dei kernel guidata da LLM possa essere applicata efficacemente oltre i benchmark isolati agli operatori catturati da interi modelli PyTorch.

Gli autori sottolineano l'importanza di una valutazione consapevole del deployment. Argomentano che, sebbene gli LLM possano generare kernel standalone veloci, il vero valore risiede nel modo in cui questi kernel si comportano all'interno di esecuzioni di modelli reali. Il meccanismo di "guarded dispatch" è presentato come un'innovazione critica, che consente al sistema di applicare i kernel generati solo dove vi è evidenza di un beneficio, preservando al contempo l'affidabilità delle implementazioni esistenti dei framework.

Il paper conclude che l'ottimizzazione dei kernel tramite agenti dovrebbe essere valutata non solo per la sua capacità di generare codice veloce, ma per la sua capacità di navigare i compromessi tra profondità di ottimizzazione, costo di ricerca e le caratteristiche specifiche dei carichi di lavoro reali, in particolare la prevalenza di implementazioni backend mature per gli operatori dominanti.