Kernel-Smith: A Unified Recipe for Evolutionary Kernel Optimization

1. Il Problema

La generazione di kernel GPU ad alte prestazioni è fondamentale per massimizzare il throughput nei sistemi di grandi modelli (LLM) e nei carichi di lavoro scientifici. Tuttavia, gli approcci attuali basati sui Large Language Models (LLM) presentano due limitazioni critiche:

• Generazione "One-Shot" Insufficiente: I modelli esistenti tendono a fallire nel mantenere ottimizzazioni affidabili oltre la prima generazione di codice. La creazione di kernel efficienti richiede spesso la ricerca di molte varianti (fusion patterns, strategie di tiling, ecc.), cosa che un singolo passaggio non riesce a coprire.
• Divario tra Correttezza e Prestazioni: La capacità di generare codice corretto non equivale alla capacità di generare codice veloce. I sistemi attuali faticano a sostenere un'ottimizzazione iterativa che sfrutti efficacemente la potenza di calcolo aggiuntiva (test-time compute) per migliorare progressivamente i candidati.
• Instabilità della Valutazione: Il rumore nelle misurazioni delle prestazioni (profiling noise) su GPU può portare a preservare kernel subottimali o scartare quelli promettenti, compromettendo la dinamica della ricerca evolutiva.

2. Metodologia: Kernel-Smith

Kernel-Smith è un framework unificato che combina un agente evolutivo guidato dalla valutazione con una ricetta di post-training orientata all'evoluzione.

A. Framework dell'Agente Evolutivo

Invece di un ciclo di dialogo multi-turno che affina una singola traccia, Kernel-Smith mantiene una popolazione di candidati eseguibili e li evolve iterativamente:

• Gestione della Popolazione: Utilizza un archivio che conserva sia i programmi ad alte prestazioni che quelli diversificati, ispirandosi agli algoritmi evolutivi "Island-based" e MAP-Elites.
• Feedback Strutturato: A ogni passo evolutivo, il sistema fornisce feedback strutturato (stato di compilazione, correttezza numerica, speedup, metadati hardware, log di errore) invece di un semplice reward scalare.
• Backend di Valutazione Robusti: Sono stati costruiti servizi di valutazione specifici per Triton (su GPU NVIDIA) e Maca (su GPU MetaX). Per mitigare il rumore temporale, il sistema utilizza:
  • Esecuzioni di warm-up.
  • Misurazioni multiple con rimozione degli outlier.
  • Tecnologia CUDAGraph per stabilizzare i tempi di esecuzione (fluttuazioni < 1%).
  • Rilevamento di "hacking" (modelli che chiamano operatori PyTorch nativi invece di scrivere kernel reali).

B. Ricetta di Training (Post-Training)

L'obiettivo non è ottimizzare il modello per la generazione one-shot, ma per agire come un miglioratore locale forte all'interno del ciclo evolutivo.

• Sintesi dei Dati: Creazione di un corpus di 59k moduli PyTorch da repository GitHub reali, integrati con dati sintetizzati tramite modelli "teacher" (es. DeepSeek-V3.2-Speciale, Gemini-3.0-pro).
• Selezione dei Passaggi Chiave: Le traiettorie evolutive a lungo termine vengono compresse. Invece di imitare ogni transizione intermedia, il modello viene addestrato solo sui passaggi di miglioramento ad alto guadagno (high-gain revisions) che preservano la correttezza e aumentano le prestazioni.
• Fasi di Addestramento:
  1. SFT (Supervised Fine-Tuning): Utilizza dati filtrati per correttezza e speedup (>1.0) per fornire un cold-start.
  2. RL (Reinforcement Learning): Utilizza l'algoritmo GRPO (Group Relative Policy Optimization). Vengono selezionati solo i "migliori passaggi" delle traiettorie evolutive come segnali di addestramento, evitando l'inquinamento da informazioni future (leakage) e focalizzandosi sull'ottimizzazione atomica.

3. Risultati Chiave

Prestazioni su KernelBench (Backend NVIDIA Triton)

Il modello Kernel-Smith-235B-RL ha raggiunto lo stato dell'arte (SOTA):

• Speedup Medio: Ha ottenuto il miglior rapporto di accelerazione medio (3.70x), superando modelli proprietari all'avanguardia come Gemini-3.0-pro (2.83x) e Claude-4.6-opus (3.33x).
• Complessità: Su task difficili (Livello 3), ha mantenuto un tasso di correttezza del 94%, superando significativamente i modelli open-weight e i competitor proprietari.
• Curva di Crescita: La curva di miglioramento nel tempo mostra che il modello beneficia in modo superiore dell'ulteriore calcolo di test-time rispetto ad altri modelli, formando l'inviluppo superiore delle prestazioni.

Adattabilità Eterogenea (Backend MetaX MACA)

Il framework è stato validato anche su GPU MetaX:

• Kernel-Smith-MACA-30B ha superato modelli su larga scala come DeepSeek-V3.2-think e Qwen3-235B-2507-think nella generazione di kernel MACA, dimostrando la capacità di adattamento senza modifiche sostanziali all'obiettivo di ottimizzazione.

Applicazioni nel Mondo Reale

Il framework ha prodotto contributi upstream reali in sistemi di produzione:

• SGLang: Un kernel Triton fuso per normal_decode_set_metadata è stato mergiato, ottenendo uno speedup locale di 4.78x.
• LMDeploy: Ottimizzazione del routing MoE, con uno speedup locale di 1.36x e un miglioramento del throughput end-to-end fino al 3.00%.
• DeepSeek Engram: Ottimizzazione di un modulo di ricerca recente, raggiungendo uno speedup impressionante di 14.59x.

4. Contributi Principali

1. Framework Unificato: Integrazione di un agente evolutivo stabile con un protocollo di training specifico per l'ottimizzazione iterativa, superando i limiti della generazione one-shot.
2. Valutazione Robusta: Sistemi di valutazione backend-specifici che riducono il rumore delle misurazioni e prevengono l'inganno (hacking), garantendo una ricerca evolutiva affidabile.
3. Ricetta di Training "Step-Centric": Una metodologia che trasforma le traiettorie evolutive complesse in segnali di supervisione focalizzati sui singoli passi di miglioramento critici, migliorando la capacità del modello di agire come ottimizzatore locale.
4. Validazione Pratica: Dimostrazione che l'ottimizzazione guidata da LLM può passare da benchmark controllati a sistemi di inferenza di produzione reali (SGLang, LMDeploy) con contributi accettati upstream.

5. Significato

Kernel-Smith dimostra che l'ottimizzazione dei kernel GPU non è solo un problema di generazione di codice, ma di ricerca strutturata e apprendimento guidato dal feedback.

• Sposta il paradigma dalla semplice "generazione di codice" all'ottimizzazione evolutiva continua, dove il modello apprende a iterare e migliorare basandosi su evidenze esecutive reali.
• Fornisce una prova concreta che i modelli LLM possono essere integrati nei flussi di lavoro di ingegneria del software per la produzione, non solo per la ricerca accademica.
• Stabilisce un nuovo standard per l'adattabilità su piattaforme hardware eterogenee (NVIDIA, MetaX), suggerendo che un'unica ricetta evolutiva può scalare su diversi ecosistemi di acceleratori.