Raggiungere un'efficienza estrema tramite la generazione di kernel GPU specializzati

Tradizionalmente, i sistemi di inferenza in produzione si affidano a kernel generici per gestire modelli e carichi di lavoro diversi. Questo approccio è subottimale perché le forme delle operazioni GPU sono determinate da una combinazione di parametri statici del modello e fattori dinamici al momento della richiesta; ad esempio, mentre un modello definisce una delle dimensioni per la moltiplicazione di matrici, l'altra dimensione varia in base al numero specifico di token di ciascuna richiesta. C'è un crescente interesse nella generazione di kernel GPU basata su agenti, e i recenti sforzi si sono rivelati promettenti. Abbiamo esplorato una domanda fondamentale: se la generazione dei kernel può essere automatizzata, perché modelli di dimensioni estremamente diverse (da 1 miliardo a 1 trilione di parametri) dovrebbero affidarsi allo stesso kernel? Specializzando i kernel per le forme specifiche riscontrate a runtime, possiamo raggiungere un'efficienza estrema.

In questo blog condividiamo i nostri successi e le nostre scoperte derivanti dall'uso di agenti per generare kernel GPU. Abbiamo creato Proteus, un sistema progettato per ottenere una specializzazione estrema, che richiede un harness su misura per un'ottimizzazione rigorosa, la convalida e la gestione del contesto.

I tradizionali harness di programmazione spesso falliscono in questo caso perché gli agenti tendono al reward-hacking: seguono la lettera della legge piuttosto che lo spirito. Se si fornisce un benchmark a un agente, questo potrebbe ottimizzare il benchmark e non l'operazione prevista.

Per risolvere questo problema, Proteus propone dei kernel, li verifica rispetto a un'implementazione di riferimento controllata, misura i tempi di quelli corretti e migliora iterativamente i risultati migliori. Sebbene il processo sia lineare, il suo successo dipende interamente dalla risoluzione di due sfide fondamentali. La Figura 1 mostra l'architettura semplificata del nostro progetto. Utilizzando il nostro harness Proteus, abbiamo generato kernel Qwen 3.5 122B che sono risultati da 1,8 a 5,2 volte più veloci dei migliori disponibili in vLLM.

Validazione

Inizialmente consideravamo la ricerca dei kernel come la parte difficile: come esplorare un ampio spazio di programmi senza rimanere bloccati in una fase di stallo senza miglioramenti? In pratica, la prima domanda era più fondamentale. Stiamo davvero misurando ciò che pensiamo di misurare?

Un modello ottimizza il punteggio che gli viene assegnato. Non ha bisogno di un exploit esotico: potrebbe semplicemente accadere che la valutazione si basi su un presupposto errato. Un esempio è rappresentato dai kernel per i rotary position embeddings (RoPE), un passaggio comune nei livelli di attenzione (attention layers). Un candidato potrebbe riutilizzare il codice compilato rimasto da un tentativo precedente e apparire meno costoso rispetto a una corretta ricostruzione da zero. Un altro potrebbe registrare un batch di lanci GPU in un grafo (ad esempio, un grafo CUDA) e riprodurli come un'unica unità, mentre la baseline con cui effettuavamo il confronto lanciava ancora ogni pezzo separatamente, quindi le due parti non stavano svolgendo lo stesso lavoro. Un altro ancora era efficace sulle dimensioni di input inserite nel test set visibile e debole su quelle che non gli erano state mostrate.

Quindi abbiamo dedicato il lavoro di progettazione iniziale allo strumento di controllo (checker), non al prompt. Misurare il tempo di entrambe le parti nello stesso modo, anche con più di un timer (ad esempio, timer degli eventi CUDA, tempo reale e timer CUPTI) quando è necessario un controllo incrociato. Cancellare lo stato compilato residuo che non dovrebbe persistere e mantenere coerente l'ordine di configurazione e smantellamento (setup e teardown), in modo che una parte non possa saltare il lavoro che l'altra deve comunque eseguire. Misurare nuovamente i tempi dei vincitori prima di utilizzarli come punto di partenza per il round successivo. Mantenere alcuni test nascosti al candidato, in modo che non possa adattarsi solo all'esame. Per evitare "imbrogli" di valutazione con prestazioni gonfiate artificialmente, implementiamo controlli di coerenza automatizzati per segnalare accelerazioni (speedup) teoricamente impossibili (ad esempio, >100x) che superano la larghezza di banda fisica della GPU e i limiti di calcolo. Questo protegge dalle stesse insidie di reward-hacking riscontrate in passati casi del settore, in cui gli agenti ottimizzavano per le metriche dell'harness piuttosto che per reali guadagni prestazionali. Senza questi vincoli, la generazione di più kernel produceva principalmente più rumore.

L'enfasi sullo strumento di controllo cambia anche il collo di bottiglia della generazione di kernel basata su agenti. Nel solo lavoro di ricerca dei programmi (ovvero l'ottimizzazione iterativa in cui il sistema cerca tra i programmi generando ripetutamente varianti), i candidati validi sono rari, quindi la loro scrittura domina il costo. Possiamo produrre molte bozze in parallelo, ma non possiamo saltare la validazione. Dobbiamo progettare attentamente la validazione, e il controllo deve essere eseguito su GPU reali, in isolamento e più di una volta. Il sistema si muove alla velocità con cui può fidarsi di un kernel, non alla velocità con cui può scriverne uno.

Gestione del contesto

Un'altra sfida consiste nello stabilire cosa sia consentito vedere al modello di generazione dei kernel. È un compromesso. Fornendo al modello un prompt più ampio, questo disporrà di più informazioni: il miglior kernel attuale, i fallimenti recenti, i suggerimenti del profiler, le note delle esecuzioni precedenti. Questo può essere d'aiuto. Ma costa anche di più, perché pagiamo per ogni token letto dal modello. E man mano che il prompt cresce, è più facile che il tentativo successivo vada fuori strada. I segnali utili si mescolano a consigli obsoleti, suggerimenti contrastanti e dettagli che si applicano a una dimensione di input diversa o a un'operazione differente. Il modello non sa sempre di quali frasi fidarsi, quindi segue quelle più evidenti, o un po' tutte.

Fornendogli troppo poco, accade il contrario. Ogni tentativo parte da zero. Si ripresentano gli stessi vicoli ciechi. Nulla viene ereditato dall'ultima esecuzione o da un'operazione correlata, e il ciclo non avanza.

Volevamo un livello di conoscenza (knowledge layer) che aiutasse in questo: ricordare ciò che ha funzionato, riutilizzarlo in seguito e farlo senza l'intervento umano. Questo livello presenta un secondo compromesso, tra il livello di dettaglio di una lezione memorizzata e la sua ampiezza di applicazione.

Una nota molto specifica (“su questo kernel, con questa dimensione di input, esegui l'unrolling di questo ciclo”) può essere esattamente ciò di cui ha bisogno il tentativo successivo. È anche facile che venga utilizzata in modo errato sull'operazione successiva, sulla GPU successiva o su una dimensione di input diversa. Una nota molto generale (“fai un uso migliore della memoria on-chip”) si applica quasi ovunque e non dice quasi nulla al modello su cosa fare. Abbiamo riscontrato entrambe le modalità di errore. Quando le lezioni erano troppo generali, ribadivano un fallimento senza indicare un'azione. Quando memorizzavamo più dettagli, questi erano spesso troppo legati a una singola esecuzione per poter essere d'aiuto a quella successiva. In un'esecuzione prolungata, la maggior parte di ciò che il modello leggeva e scriveva veniva spesa per recuperare e instradare quella memoria piuttosto che per scrivere kernel. Il livello di memoria svolgeva molto lavoro. Ma non migliorava il candidato successivo. La Figura 2 mostra la scomposizione dei costi dei token di un sistema di questo tipo. Il costo è dominato dai livelli di conoscenza.

La versione di conoscenza che vale la pena conservare è più limitata e trova un equilibrio tra generalità e specificità. Quando il modello sta per scrivere un kernel, il suo prompt dovrebbe includere solo un contesto ad alta affidabilità: indicazioni pratiche che associano situazioni specifiche ad azioni (estrapolate dalla mappatura passata delle modifiche rispetto all'impatto) e sintetiche note di errore provenienti da esecuzioni madri strettamente correlate. Recuperate tramite un filtraggio gerarchico dei tag combinato con una ricerca ibrida (parole chiave + semantica), le lezioni dovrebbero essere sufficientemente specifiche da poter essere messe in pratica e abbastanza delimitate da chiarire dove non sono applicabili. Operazioni più profonde, come la riorganizzazione e l'ulteriore distillazione dell'archivio delle lezioni, appartengono a processi in background, non a un attraversamento sincrono multi-hop delle esecuzioni passate a ogni tentativo. Se un'indicazione non è in grado di definire la situazione e l'azione, non vale la pena inserirla nel prompt. La Figura 3 mostra la scomposizione dei costi dei token dopo aver corretto il livello di conoscenza; dopo questa correzione, la maggior parte dei token viene spesa per la generazione dei candidati.

Caso di studio: Gated DeltaNet packed decode

Un esempio concreto è il kernel packed decode sul percorso Gated DeltaNet in Qwen 3.5 122B. L'operazione aggiorna uno stato ricorrente e scrive l'output di decodifica (decode output) a partire da input QKV compressi (packed), parametri di gate e indici di stato. Abbiamo utilizzato questo task per testare l'intero ciclo di Proteus su GPU NVIDIA B200 con un backend Triton: convalidare il contratto del task, misurare l'implementazione di riferimento, richiedere agli agenti i kernel candidati, eseguire controlli statici e build, verificare la correttezza rispetto al riferimento controllato, eseguire il benchmark solo dei candidati verificati e quindi misurare nuovamente i candidati migliori.

La Figura 4 si legge da sinistra a destra. Il nodo baseline fissa il benchmark a 0.025 ms. Il Candidato 0000 è il seed sicuro: ha riprodotto la struttura packed-decode e ha superato la validazione, ma era più lento del riferimento, quindi Proteus lo ha mantenuto come nodo padre misurato anziché considerarlo una vittoria. Da quel momento, Proteus ha smesso di ottimizzare un unico kernel generico per ogni forma e ha suddiviso la ricerca in percorsi specifici per ciascuna forma.

Il percorso di ripristino Batch-1 ha prodotto un kernel specifico per la forma con il Candidato 012, raggiungendo un'accelerazione di 1,5x sulla forma di decodifica a batch singolo. I risultati migliori sono stati ottenuti sul percorso serving-decode: il Candidato 030 ha registrato la latenza del kernel misurata più bassa a 0,018 ms, e il Candidato 036 ha prodotto la migliore accelerazione della forma a 1,6x. Quel frammento di serving vincente si è specializzato per il layout Batch=4, Key=128, Value=128 e ha elaborato la dimensione del valore in blocchi (chunk) di ampiezza 64, quindi si tratta di un kernel sicuro per quella forma specifica piuttosto che di un sostituto universale.

L'ultima deviazione nella timeline mostra perché la traccia è importante. I tentativi successivi di generazione in C++ (anziché Triton) hanno riscontrato errori di compilazione e generazione, e l'esecuzione prolungata è terminata dopo aver esaurito il budget di tentativi di quel ramo. L'artefatto utile non è quindi solo il candidato più veloce. È l'intero percorso mostrato nella figura: gli errori semantici sono stati rifiutati, i kernel corretti ma più lenti sono stati misurati e i reali miglioramenti delle prestazioni sono rimasti associati alla forma che li rende sicuri da comporre in un kernel di produzione.

Su cosa lavoreremo in seguito

Il loop evolutivo che abbiamo creato è molto rigido come scrittore. Spesso chiama il modello secondo uno schema fisso: prendi il miglior kernel attuale, prova una piccola modifica, controlla, ripeti. Questo toglie l'autonomia di cui l'agente ha bisogno. Non può cambiare facilmente struttura, cambiare linguaggio o abbandonare un design senza via d'uscita.

Il loop è ancora necessario. Non per scrivere il kernel, ma per fornire all'agente un suggerimento successivo affidabile. Questo suggerimento deve provenire da due fonti.

In primo luogo, dalla comunicazione con il livello di conoscenza: alcune conclusioni che siano abbastanza specifiche da poter essere messe in pratica e delimitate in modo da sapere dove non si applicano. Senza questo, ogni tentativo parte da zero.

In secondo luogo, dai risultati di un verificatore attendibile: correttezza e tempistiche che l'agente non ha misurato da solo. Questi numeri sono suggerimenti per il tentativo successivo. Sono anche gli unici punteggi a cui dovremmo credere. Se l'agente misura i propri tempi di esecuzione, torniamo a cache residue, confronti non corrispondenti e test visibili all'agente stesso.

Quindi la distinzione che vogliamo è più sottile di "agente contro loop". Diamo all'agente l'autonomia su come viene scritto un kernel. Manteniamo il loop come canale per la memoria e per la valutazione. L'agente propone. Il loop restituisce ciò che gli è consentito vedere e se l'ultima proposta ha effettivamente avuto successo.

Proteus ha creato kernel specializzati per parti di Qwen 3.5 122B sul suo percorso Gated DeltaNet (un blocco in stile linear-attention) in esecuzione su GPU NVIDIA B200. Le accelerazioni sui singoli kernel sono state comprese tra 1,8x e 5,2x.

La lezione è che la generazione è la fase economica. La validazione e la gestione del contesto sono la parte difficile. È qui che sono necessari un attento tempo di progettazione e innovazioni.

La generazione agentica di kernel GPU ha sbloccato l'incredibile potenziale dell'estrema specializzazione, ma la creazione di harness affidabili e pronti per la produzione rimane una frontiera complessa. Stiamo affrontando le sfide più difficili all'intersezione tra AI e sistemi, e cerchiamo ingegneri audaci che si uniscano a noi per plasmare il futuro dell'inferenza efficiente. Se desideri superare i limiti di ciò che è possibile, stiamo assumendo!

(Questo post sul blog è stato tradotto utilizzando strumenti basati sull'intelligenza artificiale) Post originale