OptiML: An End-to-End Framework for Program Synthesis and CUDA Kernel Optimization

1. Il Problema

La programmazione di acceleratori moderni (come le GPU NVIDIA) rimane un collo di bottiglia critico per il calcolo scientifico e i sistemi di machine learning. Sebbene i Large Language Models (LLM) siano in grado di generare codice CUDA sintatticamente valido e funzionalmente corretto, spesso falliscono nel produrre kernel ad alte prestazioni. Le ragioni principali includono:

• Spazio di ottimizzazione combinatorio: Le trasformazioni a basso livello (es. tiling, vettorizzazione, gestione della memoria condivisa) interagiscono in modo complesso con la pressione dei registri, l'occupazione e il coalescing della memoria.
• Feedback costoso e rumoroso: La misurazione delle prestazioni richiede compilazione, esecuzione e profilazione (es. con Nsight Compute), rendendo la ricerca di ottimizzazioni lenta e soggetta a rumore.
• Separazione tra generazione e ottimizzazione: Le ricerche precedenti trattano la generazione del codice e l'ottimizzazione come problemi separati. Spesso, se il codice generato inizialmente ha una struttura debole o non compilabile, gli ottimizzatori successivi non riescono a recuperare prestazioni elevate.
• Mancanza di interpretabilità: Le strategie di ottimizzazione esistenti (es. perturbazioni casuali o regole fisse) spesso convergono lentamente o non generalizzano, senza fornire spiegazioni chiare sul perché un'ottimizzazione abbia funzionato.

2. Metodologia: Il Framework OptiML

OptiML è un framework end-to-end che unifica la sintesi del programma e l'ottimizzazione del kernel. Mappa sia l'intento in linguaggio naturale che il codice CUDA esistente in kernel ottimizzati, formulando l'ottimizzazione come un processo di ricerca guidato dalla verifica. Il sistema è composto da due fasi disaccoppiate:

A. OptiML-G (Generazione di Proposta)

Quando l'input è un'intento in linguaggio naturale, OptiML-G genera un kernel iniziale eseguibile.

• Architettura: Utilizza un framework Mixture-of-Thoughts (MoT).
• Meccanismo: Mantiene un pool di esperti pre-addestrati (es. Qwen2.5-Coder, HPC-Coder-V2, StarCoder2). Un router leggero seleziona gli esperti più rilevanti per il task.
• Collaborazione Latente: A differenza dei metodi di routing tradizionali che usano un solo esperto, OptiML-G permette una collaborazione fine-granularità tra gli esperti selezionati in uno spazio latente condiviso. Gli stati nascosti degli esperti vengono proiettati e fusi tramite meccanismi di cross-attention, permettendo all'esperto principale di aggregare "pensieri" complementari.
• Obiettivo: Produrre un kernel iniziale strutturale forte, compilabile e pronto per l'ottimizzazione, riducendo lo spazio di ricerca per la fase successiva.

B. OptiML-X (Ottimizzazione Basata sulla Ricerca)

Questa fase raffina i kernel (sia generati da OptiML-G che forniti dall'utente) utilizzando una ricerca strutturata.

• Motore di Ricerca: Utilizza Monte Carlo Tree Search (MCTS) per esplorare traiettorie di ottimizzazione multi-step.
• LLM-as-a-Judge: Un LLM funge da "giudice" che:
  1. Diagnostica: Analizza i segnali di profilazione (da Nsight Compute) per identificare i colli di bottiglia (es. bound alla memoria, bound al calcolo).
  2. Propone: Suggerisce patch specifiche (es. tiling, unrolling, vettorizzazione) per affrontare il collo di bottiglia diagnosticato.
  3. Valuta: Assegna un punteggio di validità e coerenza.
• Reward Model (Funzione di Ricompensa): La ricerca è guidata da una ricompensa composita che combina:
  • Tempo di esecuzione: Riduzione della latenza.
  • Metriche Proxy Hardware: Utilizzo delle risorse (SOL - SM, DRAM, Texture) e riduzione del lavoro (byte DRAM, settori L1, istruzioni eseguite).
  • Guardrail: Penalità severe per candidati che violano vincoli di compilazione, correttezza funzionale (test L0/L1) o causano regressioni (es. aumento eccessivo dei registri).
  • Adattività: I pesi delle metriche proxy si adattano dinamicamente in base al tipo di collo di bottiglia rilevato (memoria-bound vs compute-bound).

3. Contributi Chiave

1. Formulazione del Problema: Definizione dell'ottimizzazione di kernel CUDA condizionata all'intento, che unifica generazione e ottimizzazione sotto vincoli di correttezza e hardware.
2. Architettura Unificata: Introduzione di OptiML, che integra una politica di generazione MoT (OptiML-G) con un ottimizzatore basato su MCTS (OptiML-X).
3. Ottimizzazione Guidata dai Profiler: OptiML-X non tratta l'ottimizzazione come un "tuning a scatola nera", ma utilizza metriche derivate dai profiler per indirizzare la ricerca verso specifici colli di bottiglia, rendendo il processo più efficiente in termini di campioni e interpretabile.
4. Pipeline di Valutazione Interpretabile: Il sistema non solo migliora le prestazioni, ma spiega perché un kernel è migliorato (es. riduzione del traffico di memoria, migliore utilizzo degli SM), fornendo traiettorie di ottimizzazione basate su evidenze.

4. Risultati Sperimentali

Il framework è stato valutato su una suite diversificata di kernel CUDA (es. Moltiplicazione di Matrici, Max Pooling 3D, Self-Attention, ReLU, Cross-Entropy) su GPU NVIDIA A100.

• Prestazioni: OptiML supera costantemente i baselines LLM-only (come GPT-5.1, Qwen2.5-Coder) e le varianti di ottimizzazione post-hoc.
  • Nel task di Moltiplicazione di Matrici, OptiML ha ottenuto un speedup di 1.64x rispetto al miglior baseline LLM, riducendo significativamente le istruzioni eseguite e i byte DRAM.
  • In task complessi come Max Pooling 3D e Cross-Entropy, molti baselines LLM fallivano la compilazione (CF), mentre OptiML ha prodotto kernel funzionanti e ottimizzati.
• Qualità della Generazione (OptiML-G): Su benchmark come ParEval, OptiML-G ha raggiunto un Pass@1 del 12.17%, superando i modelli base e dimostrando che la collaborazione tra esperti eterogenei produce kernel strutturalmente superiori.
• Efficienza della Ricerca: L'analisi "round-to-best" mostra che OptiML (G+X) converge al kernel migliore in meno iterazioni rispetto all'uso di OptiML-X su un punto di partenza debole, grazie alla migliore inizializzazione fornita da OptiML-G.
• Analisi dei Colli di Bottiglia: I risultati mostrano che OptiML-X adatta la strategia: riduce il traffico di memoria per kernel bound alla memoria (es. Pooling) e riduce l'overhead delle istruzioni per kernel bound al calcolo (es. ReLU).

5. Significato e Impatto

Il lavoro di OptiML rappresenta un passo significativo verso l'automazione dello sviluppo di software per acceleratori:

• Superamento del divario Generazione-Ottimizzazione: Dimostra che un approccio end-to-end è necessario perché un codice generato male non può essere facilmente "aggiustato" dopo la generazione; la struttura iniziale è cruciale.
• Interpretabilità e Affidabilità: L'uso di metriche hardware come segnali di guida rende il processo di ottimizzazione trasparente e riduce il rischio di "reward hacking" (ottimizzare metriche proxy a scapito delle prestazioni reali).
• Scalabilità: Il framework è progettato per adattarsi a diversi tipi di kernel e architetture hardware, offrendo una soluzione robusta per colmare il divario tra il codice generato dagli LLM e le prestazioni di livello esperto richieste nell'HPC e nel Deep Learning.

In sintesi, OptiML trasforma l'ottimizzazione dei kernel CUDA da un processo manuale o euristico in un ciclo di ricerca guidato dai dati, dove l'intelligenza artificiale non solo scrive il codice, ma ne comprende e risolve le inefficienze hardware specifiche.