OptiML: Un approccio completo all'ottimizzazione di kernel CUDA

La generazione di kernel CUDA ad alte prestazioni è un compito complesso, che richiede l'esplorazione di un ampio spazio di trasformazioni a basso livello. OptiML affronta questa sfida con un framework end-to-end che combina modelli linguistici di grandi dimensioni (LLM) e tecniche di ricerca per migliorare le performance dei kernel CUDA.

OptiML opera in due fasi distinte. Nella prima fase, OptiML-G, un generatore basato su un modello Mixture-of-Thoughts, crea un programma eseguibile iniziale a partire da una descrizione in linguaggio naturale. Nella seconda fase, OptiML-X, un ottimizzatore basato su ricerca, affina i kernel, siano essi sintetizzati o forniti dall'utente, utilizzando la ricerca Monte Carlo Tree Search (MCTS) guidata da LLM.

Ogni trasformazione candidata viene compilata, verificata e profilata con Nsight Compute. Le performance vengono valutate tramite una funzione obiettivo composita che combina il tempo di esecuzione con proxy di colli di bottiglia hardware e meccanismi di protezione contro le regressioni. I risultati dimostrano che OptiML è in grado di scoprire miglioramenti verificati delle performance rispetto a baseline LLM consolidate e di produrre traiettorie di ottimizzazione interpretabili basate su evidenze di profilatura.

Per chi valuta deployment on-premise, esistono trade-off da considerare. AI-RADAR offre framework analitici su /llm-onpremise per valutare questi trade-off.