Scrivere kernel GPU a mano richiede pazienza, impegno e anche la tua sanità mentale. Sto sperimentando con gli LLM moderni (Claude Opus 5, GPT 5.6 Sol) per farmi stampare letteralmente i kernel. Ti racconto cosa ho imparato facendolo.
Onestamente, scrivere kernel GPU è un compito imbarazzantemente verificabile.
Per prima cosa decidi per quale operazione (o quali operazioni) vuoi scrivere il kernel. In secondo luogo, scrivi la prima versione e ti assicuri che compili senza errori evidenti. In terzo luogo, verifichi la correttezza con un'implementazione di riferimento lenta. Se non corrisponde, provi a correggere il problema di correttezza. Una volta fatto, esegui un benchmark del tempo di esecuzione del kernel. Le versioni successive si basano su questo e continui a ottimizzare finché non raggiungi le metriche roofline o non sei soddisfatto.
Loop verificabile: sviluppo di kernel GPU
L'immagine qui sopra lo mostra come un loop verificabile:
- Il controllo di compilazione è un loop locale serrato (B <-> C) prima ancora di pensare alla correttezza.
- Il controllo di correttezza (D <-> E <-> F) è il loop centrale di ricompensa verificabile. È questa la parte che rende la scrittura di kernel un buon problema per la verifica automatica, dato che disponi di una ground truth di riferimento con cui confrontare i risultati.
- L'ottimizzazione (G -> H -> di nuovo a D) riutilizza lo stesso loop di correttezza per ogni nuova versione: un kernel veloce ma sbagliato non vale nulla, quindi la correttezza va verificata per ogni versione.
- Il controllo roofline/soddisfazione (I) è il loop esterno che decide se continuare a ottimizzare o fermarsi.
Prima versione
L'immagine qui sopra è comunque una visione ad alto livello e il diavolo si nasconde nei dettagli. Dobbiamo assicurarci che il nostro agente LLM abbia tutto il contesto necessario per iniziare a scrivere una buona prima versione.
È qui che entrano in gioco i DSL CUDA. Triton, CuTeDSL e Tilelang sono quelli con cui è molto facile iniziare, in Python. La curva di apprendimento è meno ripida rispetto a CUDA C++, tuttavia le astrazioni di questi DSL possono confondere ancora di più il nostro agente. Dobbiamo trovare un modo per passare all'agente il contesto di queste astrazioni.
Gli LLM moderni sanno già come scrivere Triton "buono". Funzionano bene con le astrazioni di Triton anche senza alcun contesto. Tuttavia, per altri DSL come CuTeDSL (che offre molto più controllo di Triton), ho scoperto che avere una directory di contesto in cui l'agente possa cercare per capire le astrazioni del DSL aiuta moltissimo.
Ad esempio, clonare il repository
NVIDIA cutlass
nella directory di contesto è un buon modo per permettere all'agente di cercare le astrazioni relative a
Layout Algebra, atomi Copy/GEMM, gerarchia di memoria, kernel di esempio
e così via mentre scrive kernel in CuTeDSL.
Nella mia esperienza, una buona prima versione del kernel compila senza errori evidenti e supera il test di correttezza di cui parlerò qui sotto.
Test, benchmark e profilazione
Dato abbastanza contesto all'agente, ora il vero collo di bottiglia si sposta sulla validazione. L'implementazione di riferimento in sé, e la validazione rispetto a essa, diventano sempre più importanti. Chiamo questa fase test di correttezza o semplicemente test. La velocità dell'implementazione di riferimento non conta quanto il suo scopo: ciò che intendi misurare e verificare è ciò su cui il tuo agente ottimizzerà.
Di solito, quando il calcolo non deve avvenire in precisione ridotta, misuro l'errore massimo assoluto/relativo (MAE), l'errore quadratico medio (MSE/RMSE) e il PSNR (rapporto segnale-rumore di picco). Quando sono coinvolte precisioni ridotte, tendo a misurare il PSNR e la similarità coseno (cossim).
Il modo in cui fai girare davvero le versioni del kernel su una GPU dipende dal fatto che la GPU sia disponibile localmente o tramite cloud. In ogni caso, il nostro agente dovrebbe avere la possibilità di accedere ai propri output in un modo o nell'altro.
Trovo che la metodologia rung qui sotto sia un buon modo per avere un numero N di funzioni di test:
1def rung(name):2 def deco(fn):3 try:4 out = fn()5 results[name] = {"ok": True, **(out or {})}6 print(f"[{name}] ok " + " ".join(f"{k}={v}" for k, v in (out or {}).items()))7 except Exception as e:8 results[name] = {"ok": False, "err": f"{type(e).__name__}: {e}"}9 print(f"[{name}] FAILED {type(e).__name__}: {e}")10 traceback.print_exc()11 return fn12 return deco
che puoi chiamare così:
1out = {}23@rung("pre-checks")4def _():5 run_pure_checks()6 run_dsl_checks()78@rung("run")9def _():10 out["o"] = custom_kernel(*inputs)11 torch.cuda.synchronize()12 return {"shape": tuple(out["o"].shape),13 "finite": bool(torch.isfinite(out["o"]).all())}
Per il rung di benchmark, puoi fare diverse cose:
- Misurare il tempo di esecuzione end-to-end del kernel per il tempo complessivo speso
- Usare il tracing intra-kernel per fare benchmark delle sezioni interne di un kernel e riversarle nell'output (usando un tracer personalizzato o CUPTI)
- Riversare l'IR, il PTX, il SASS e il CUBIN generati in una directory dumps e lasciare che l'agente li esamini
L'ultimo punto può essere approfondito un po' di più. A volte può capitare che il DSL, durante il lowering, generi un PTX non ottimale (e di conseguenza un SASS) e che tu trovi istruzioni o forme migliori da usare al suo posto. Il nostro agente può leggere i file di testo PTX/SASS e inlineare codice di livello più basso invece di lasciare che sia il DSL a gestire la parte non ottimale. Anche in questo caso, passare la documentazione PTX "ricercabile" come contesto è molto utile.
L'ultima cosa per mettere tutto insieme è la profilazione. Se il tuo agente può accedere alla CLI di NCU (Nsight Compute Systems), puoi chiedergli di profilare il tuo kernel e generare un report, come parte del ciclo di feedback verificabile descritto sopra.
Considerazioni finali
"Lo sviluppo di kernel GPU è quindi morto?"
"Beh sì, ma in realtà no"
Sì, perché la parte difficile di layout, indicizzazione, astrazioni e struttura complessiva può essere in gran parte risolta da agenti con abbastanza contesto. Puoi facilmente ridurre il lavoro da 2-3 settimane a 1-2 giorni. No, perché il vero collo di bottiglia si è spostato dai kernel alle validazioni. Ormai non esiste un unico modo di fare le cose: migliore è il contesto e l'harness, più veloce è il processo. I casi specializzati ne trarranno ancora più beneficio, e tutto ciò di cui avrai bisogno è un buon harness.
Infine, invece di trattare gli agenti come autonomi, è ancora necessario trattarli come assistenti davvero intelligenti che puoi guidare. È qui che torna utile la tua comprensione di base di GPU e kernel. La parte umana (tu) serve ancora.
È una sensazione agrodolce, lo so :)
