L'addestramento del proxy CUDA per LLM scrive kernel GPU, superando torch.compile in velocità
Il team Seed di ByteDance, in collaborazione con l'Istituto di Ricerca sull'Intelligenza Artificiale (AIR) dell'Università di Tsinghua, ha lanciato CUDA Agent, un sistema di apprendimento per rinforzo basato su agenti su larga scala progettato per insegnare ai modelli linguistici a generare kernel CUDA ad alte prestazioni.
CUDA Agent addestra modelli linguistici di grandi dimensioni a scrivere kernel GPU più veloci di torch.compile
Introduzione
Il team Seed di ByteDance, in collaborazione con l'Istituto di Ricerca sull'Intelligenza Artificiale (AIR) dell'Università di Tsinghua, ha lanciato CUDA Agent, un sistema di apprendimento per rinforzo basato su agenti su larga scala progettato per insegnare ai modelli linguistici a generare kernel CUDA ad alte prestazioni.
Questo lavoro affronta una debolezza di lunga data nella generazione di codice di basso livello da parte dell'IA. I modelli linguistici all'avanguardia spesso riescono a generare codice CUDA che compila e funziona correttamente, ma per i carichi di lavoro GPU di livello produttivo la sola correttezza non è sufficiente. I kernel generati devono anche competere con le implementazioni prodotte da compilatori altamente ottimizzati.
CUDA Agent è costruito proprio attorno a questo secondo problema: non solo generare codice CUDA valido, ma imparare a profilare, verificare e ottimizzare iterativamente i kernel fino a ottenere miglioramenti significativi delle prestazioni in fase di esecuzione.
Perché generare solo codice CUDA corretto non basta
Prima dell'introduzione dell'apprendimento per rinforzo basato su agenti, il modello base Seed1.6 mostrava già discrete capacità di programmazione CUDA su KernelBench.
Nei 250 task KernelBench di livello 1–3 utilizzati dai ricercatori, Seed1.6 ha raggiunto un tasso di superamento del 74,0%. In altre parole, era in grado di generare soluzioni funzionalmente corrette per la maggior parte dei task di riferimento del modello base.
Ma le prestazioni erano un'altra storia.
Solo il 27,2% dei kernel generati dal modello base era più veloce di torch.compile, con una velocità media geometrica rispetto alla baseline del compilatore di appena 0,69×.
| Modello | Tasso di superamento | Più veloce di torch.compile |
Velocità media geometrica rispetto a torch.compile |
|---|---|---|---|
| Modello base Seed1.6 | 74,0% | 27,2% | 0,69× |
| CUDA Agent | 98,8% | 96,8% | 2,11× |
Questo divario evidenzia la sfida principale della generazione automatica di kernel GPU.
I modelli linguistici possono comprendere la sintassi CUDA, ma possono comunque produrre accessi alla memoria inefficienti, troppi avvii di kernel, scelte di tiling inadeguate, sincronizzazioni non necessarie o una gestione subottimale dei registri, tutti fattori che hanno un impatto enorme sulle prestazioni reali.
Apprendimento per rinforzo basato su agenti per la generazione di kernel
I ricercatori sostengono che il problema non sia la comprensione di CUDA, ma l'ottimizzazione. I modelli linguistici raramente ricevono un feedback diretto sul fatto che il codice generato sia più veloce o più lento di una soluzione di riferimento.
CUDA Agent affronta questo problema utilizzando l'apprendimento per rinforzo basato su agenti.
Il processo inizia con un'analisi iniziale del task da parte dell'agente, che poi genera codice CUDA, lo compila e lo esegue in un ambiente sandbox. Il sistema raccoglie metriche di profiling e feedback di verifica, che vengono restituiti al modello come osservazioni.
L'agente può quindi riflettere sui risultati, modificare il codice e provare di nuovo.
Questo ciclo può continuare per molti round. Ad ogni iterazione, il modello impara ad associare scelte di implementazione specifiche a risultati di prestazioni osservabili, costruendo gradualmente una strategia di ottimizzazione.
Il file SKILL.md fornisce competenze di ottimizzazione specializzate
L'addestramento del modello si basa su un file SKILL.md che codifica competenze specialistiche di ottimizzazione CUDA.
Il file include linee guida per il workflow e strategie strutturate come:
- Ottimizzazione iterativa del kernel
- Guida all'ottimizzazione specifica per CUDA
- Ambiente sandbox limitato
- Segnali di ricompensa legati ai risultati di esecuzione effettivi
Il file SKILL.md limita anche le scorciatoie che potrebbero distorcere i risultati dei benchmark. Ad esempio, l'agente non può semplicemente ripiegare su operazioni PyTorch arbitrarie nell'implementazione del kernel personalizzato.
Questo è importante perché i sistemi di apprendimento per rinforzo potrebbero trovare modi per massimizzare la ricompensa senza risolvere il problema di ottimizzazione previsto.
L'addestramento PPO porta l'agente a traiettorie di ottimizzazione lunghe
I ricercatori hanno addestrato CUDA Agent utilizzando Proximal Policy Optimization (PPO).
Una sfida principale è che l'ottimizzazione GPU è intrinsecamente un compito a lungo orizzonte temporale. Il modello può dover passare attraverso più cicli di modifica del codice, compilazione, debug, profiling e ulteriore ottimizzazione prima di raggiungere il risultato desiderato.
Il documento adotta quindi diverse lunghezze di contesto e fasi di addestramento per stabilizzare l'apprendimento.
Il modello base è Seed1.6, un modello mixture-of-experts con 230 miliardi di parametri totali e 23 miliardi di parametri attivi.
Per l'apprendimento per rinforzo basato su agenti:
- La finestra di contesto è di 131.072 token.
- Il replay di addestramento consente fino a 150 round di interazione dell'agente.
- La valutazione consente fino a 200 round di interazione dell'agente.
- Il modello viene addestrato per 150 passi di RL.
- I modelli actor e critic adottano una strategia di warm-up multi-fase prima dell'ottimizzazione completa a lungo orizzonte.
Questo è più complesso che semplicemente fare fine-tuning su una coppia di prompt e risposte CUDA.
L'obiettivo è insegnare al modello come migliorare i propri kernel attraverso un'interazione ripetuta con il feedback di esecuzione.
La sandbox separa la compilazione dal profiling GPU
La misurazione affidabile delle prestazioni è fondamentale, poiché la velocità di esecuzione fa parte della ricompensa.
I ricercatori hanno quindi costruito un'architettura sandbox con risorse CPU-GPU disaccoppiate.
La sandbox terminale basata su Docker gestisce i task orientati alla CPU (come la compilazione), mentre la verifica e il profiling vengono assegnati a un pool dedicato di sandbox GPU contenente 128 GPU NVIDIA H20.
Questa separazione ha un duplice scopo.
Primo, isola la compilazione e l'interazione dell'agente dai benchmark GPU. Secondo, l'allocazione GPU dedicata riduce l'interferenza tra carichi di lavoro concorrenti, rendendo le misurazioni di latenza più stabili.
Questo è importante per l'apprendimento per rinforzo: se le misurazioni temporali sono influenzate dal rumore, il modello riceve segnali di ricompensa inaffidabili e potrebbe apprendere comportamenti di ottimizzazione errati.
CUDA Agent raggiunge il 98,8% di tasso di superamento e supera il compilatore nel 96,8% dei task
Il sistema finale è stato valutato su 250 task dei livelli 1–3 di KernelBench.
CUDA Agent ha raggiunto:
- 98,8% di tasso di superamento complessivo
- 98,4% più veloce dell'esecuzione immediata di PyTorch
- 96,8% più veloce di
torch.compile - Accelerazione media geometrica di 2,60× rispetto all'esecuzione immediata
- Accelerazione media geometrica di 2,11× rispetto a
torch.compile
Questi risultati rappresentano un miglioramento significativo rispetto al punto di partenza di Seed1.6.
Il modello non solo è diventato migliore nel generare codice che supera i test di correttezza, ma i kernel che genera hanno anche maggiori probabilità di superare le baseline dei compilatori.
Risultati per livello di difficoltà KernelBench
I miglioramenti delle prestazioni rimangono solidi in tutti e tre i livelli di difficoltà di KernelBench.
| Difficoltà KernelBench | Tasso di superamento | Percentuale più veloce di torch.compile |
Accelerazione media geometrica rispetto a torch.compile |
|---|---|---|---|
| Livello 1 | 100,0% | 97,0% | 1,87× |
| Livello 2 | 100,0% | 100,0% | 2,80× |
| Livello 3 | 94,0% | 90,0% | 1,52× |
Le prestazioni al Livello 2 sono particolarmente notevoli.
Questi task coinvolgono sequenze di operatori in cui le opportunità di ottimizzazione derivanti da fusione e spostamento della memoria non sono sempre sfruttate efficacemente dalle euristiche statiche dei compilatori.
Gli autori del documento ritengono che un ottimizzatore iterativo basato sull'apprendimento possa esplorare uno spazio più ampio di strategie implementative, inclusi accessi alla memoria specifici per l'hardware, scelte di tiling e fusioni.
I dati di addestramento saranno presto resi pubblici
Attualmente, i pesi completi del modello addestrato non sono inclusi nella versione pubblica del progetto.
Tuttavia, i ricercatori hanno rilasciato diverse componenti importanti dello stack di addestramento.
CUDA-Agent-Ops-6K
Il dataset CUDA-Agent-Ops-6K contiene 6.000 task di addestramento sintetici a livello di operatore.
Il processo di costruzione comprende:
- Raccolta di operatori seed da librerie come
torch,transformers, ecc. - Utilizzo di LLM per combinare più operatori in task di fusione più complessi.
- Filtraggio dei task generati tramite controlli basati sui risultati di esecuzione.
La fase di filtraggio rimuove casi con randomicità, troppo semplici, non validi o troppo simili ai task di valutazione.
Il dataset è stato pubblicato su Hugging Face con licenza CC BY 4.0.
SKILL.md e ambiente agente
Il repository GitHub include anche il file di istruzioni SKILL.md specifico per CUDA e l'ambiente di lavoro dell'agente.
Questi materiali documentano il workflow di ottimizzazione, gli strumenti disponibili, i vincoli e la configurazione dell'ambiente di esecuzione utilizzata per guidare l'agente.
Meccanismo di ricompensa e schema di addestramento
Il documento e il repository descrivono la progettazione delle ricompense, le fasi di warm-up, l'inizializzazione del critic e lo schema di addestramento basato su PPO per stabilizzare il processo di ottimizzazione a lungo ciclo.
Questo è particolarmente utile per i ricercatori interessati all'addestramento di agenti per la programmazione di sistema, non solo per riprodurre i punteggi finali dei benchmark.
Perché questo va oltre KernelBench
L'ottimizzazione dei kernel CUDA supporta gran parte dell'infrastruttura AI moderna.
Kernel più veloci possono migliorare:
- Il throughput dell'addestramento dei modelli
- Le prestazioni di latenza dell'inferenza LLM
- L'utilizzo della GPU
- L'efficienza dei costi di servizio
- Le pipeline AI in tempo reale
- I carichi di lavoro di calcolo numerico ad alte prestazioni
Applicazioni potenziali in settori come infrastrutture, servizi di inferenza per modelli di grandi dimensioni, guida autonoma e trading quantitativo — ambiti in cui anche minime differenze di latenza possono essere cruciali.
I risultati di CUDA Agent non implicano che i compilatori tradizionali siano diventati obsoleti.
torch.compile rimane una solida baseline automatica, e la valutazione di CUDA Agent si basa su ambienti di benchmark controllati e su specifiche configurazioni GPU.
Le conclusioni mostrate in questa ricerca, sebbene più ristrette, restano importanti: con un adeguato feedback di esecuzione e un apprendimento per rinforzo dedicato, i modelli linguistici possono apprendere strategie di ottimizzazione che superano la baseline del compilatore statico nella stragrande maggioranza dei task di kernel GPU testati.
Ciò significa trasformare l'ingegneria delle prestazioni di basso livello in un ambito ragionevole per l'apprendimento agentico, non solo per la generazione di codice in un'unica passata.
Domande frequenti
Cos'è CUDA Agent?
CUDA Agent è un sistema di apprendimento per rinforzo agentico su larga scala sviluppato da ByteDance Seed, dall'Istituto di Ricerca sull'Industria Intelligente (AIR) dell'Università Tsinghua e dal loro SIA-Lab congiunto. Addestra modelli linguistici a scrivere, validare, analizzare e ottimizzare kernel CUDA in modo iterativo.
Su quale modello si basa CUDA Agent?
La ricerca utilizza Seed1.6 come modello di base. L'articolo lo descrive come un modello a esperti misti con 230 miliardi di parametri totali e 23 miliardi di parametri attivi.
Cos'è KernelBench?
KernelBench è un benchmark aperto per valutare se i modelli linguistici possono generare kernel GPU corretti ed efficienti per programmi PyTorch. CUDA Agent è stato valutato su 250 task che coprono i livelli da 1 a 3 di KernelBench.
Quanto è più veloce CUDA Agent rispetto a torch.compile?
Sull'intero benchmark, CUDA Agent ottiene un'accelerazione media geometrica di 2,11 volte rispetto a torch.compile. I kernel generati sono più veloci della baseline del compilatore nel 96,8% dei task valutati.
CUDA Agent usa l'apprendimento per rinforzo?
Sì. Il sistema utilizza PPO, combinato con warm-up in più fasi, pre-addestramento del modello di valore, progettazione robusta delle ricompense e traiettorie agentiche lunghe e multi-round.
Quanto può durare una traiettoria di ottimizzazione di CUDA Agent?
Il rollout di addestramento consente fino a 150 round agentici, mentre la valutazione consente fino a 200 round. La finestra di contesto di addestramento dell'agente è di 131.072 token.
CUDA Agent è open source?
Il progetto ha pubblicato il repository GitHub, l'ambiente agentico, SKILL.md, la ricetta di addestramento e il dataset CUDA-Agent-Ops-6K. I pesi completi del modello addestrato non sono inclusi nei materiali attualmente rilasciati pubblicamente.
Cos'è CUDA-Agent-Ops-6K?
CUDA-Agent-Ops-6K è un dataset con 6.000 task di addestramento CUDA sintetici a livello di operatore. È progettato specificamente per l'apprendimento per rinforzo agentico su larga scala e include passaggi di filtraggio per garantire che i task siano eseguibili, deterministici, non banali e separati dal set di valutazione di KernelBench.
Strumenti correlati
- CUDA Agent: pagina ufficiale del progetto del sistema di generazione di kernel CUDA di ByteDance Seed e dell'Istituto di Ricerca sull'Industria Intelligente dell'Università Tsinghua.
- CUDA Agent GitHub: repository ufficiale con ambiente agentico,
SKILL.mde materiali di progetto pubblicati. - CUDA-Agent-Ops-6K: dataset di addestramento ufficiale con 6.000 campioni rilasciato con il progetto.
- KernelBench: benchmark aperto per valutare i kernel GPU generati da modelli linguistici di grandi dimensioni.
- PyTorch: framework di deep learning la cui modalità eager e l'esecuzione con
torch.compilesono baseline chiave nella ricerca. - NVIDIA CUDA: documentazione ufficiale NVIDIA per la programmazione CUDA e lo sviluppo di kernel GPU.
Collegamenti correlati
- Pagina del progetto CUDA Agent: panoramica ufficiale, risultati dei benchmark, diagrammi del progetto, articolo, codice e collegamenti ai dataset.
- Articolo di CUDA Agent (arXiv): articolo di ricerca completo, "Apprendimento per rinforzo agentico su larga scala per la generazione di kernel CUDA ad alte prestazioni."
- Repository GitHub di CUDA Agent: repository del codice sorgente con ambiente e file di progetto rilasciati.
- Dataset CUDA-Agent-Ops-6K: dataset ufficiale Hugging Face con 6.000 task di operatori sintetici.
- Repository GitHub di KernelBench: ambiente di benchmark e valutazione per la generazione efficiente di kernel GPU.
- PyTorch
torch.compile: documentazione ufficiale PyTorch della baseline del compilatore utilizzata nella ricerca. - Guida alla programmazione CUDA C++ di NVIDIA: guida ufficiale NVIDIA per la programmazione e l'ottimizzazione CUDA.
Riepilogo
CUDA Agent affronta una debolezza specifica del codice di sistema generato dai modelli linguistici: il codice CUDA generato non deve solo essere corretto, ma deve anche essere realmente più veloce delle alternative prodotte dal compilatore.
Partendo da Seed1.6, i ricercatori hanno utilizzato un ambiente agentico dedicato a CUDA, feedback di esecuzione, una progettazione robusta delle ricompense e un addestramento PPO a lungo orizzonte, portando il tasso di successo dal 74,0% al 98,8% e la percentuale di kernel più veloci di torch.compile dal 27,2% al 96,8%.
Il progetto ha inoltre rilasciato un dataset di addestramento con 6.000 task, SKILL.md, materiali per l'ambiente agentico e la ricetta di addestramento, fornendo una base concreta per ulteriori ricerche sull'ottimizzazione GPU appresa.
Il principale contributo di CUDA Agent è dimostrare che l'ingegneria delle prestazioni può essere appresa attraverso cicli agentici iterativi — non solo approssimata tramite la generazione di codice in un'unica passata.