Sintesi Tecnica: KForge – Generazione di Kernel Cross-Platform guidata da LLM per Acceleratori AI

Problema

La proliferazione di applicazioni di IA agentica ha spostato i carichi di lavoro di inferenza dall'esecuzione monolitica a pipeline eterogenee. Queste pipeline intercalano ragionamento, chiamate a strumenti (tool calls) e coordinamento multi-agente, con ogni fase che possiede profili di calcolo e memoria distinti. Di conseguenza, l'implementazione ottimale richiede la mappatura di specifiche fasi della pipeline sugli acceleratori più idonei. Questa eterogeneità crea una sfida significativa a livello di sistemi: gli sviluppatori devono produrre kernel ad alte prestazioni su una crescente diversità di backend hardware (NVIDIA, AMD, Intel, Apple) e modelli di programmazione (CUDA, Triton, SYCL, Metal, ecc.).

Scrivere e portare manualmente questi kernel è dispendioso in termini di tempo, richiede una profonda competenza a basso livello e non riesce a scalare all'aumentare della complessità del kernel. Sebbene i Large Language Models (LLM) abbiano mostrato promesse nella generazione di codice, l'applicazione allo sviluppo di kernel a basso livello rimane difficile a causa della fragilità del codice di performance (dove piccole variazioni nel layout della memoria o nella precisione possono compromettere la correttezza) e della scarsità di dati di addestramento per le piattaforme non-CUDA. Inoltre, i benchmark esistenti spesso valutano i kernel isolatamente, fallendo nel catturare l'impatto dell'ottimizzazione del kernel sull'esecuzione end-to-end dell'intero modello all'interno dei runtime di produzione.

Metodologia: Il Framework KForge

KForge è un framework di sintesi di programmi autonomo, cross-platform e multi-stadio, progettato per rispecchiare il flusso di lavoro iterativo di un ingegnere dei kernel. Impiega un ciclo collaborativo tra due agenti specializzati basati su LLM:

1. Agente di Generazione (Generation Agent): Responsabile della produzione e del raffinamento progressivo dei candidati kernel. Opera all'interno di un loop guidato da prompt che incorpora:

   • Raffinamento Iterativo: Utilizzo di errori di compilazione, errori di runtime e discrepanze numeriche dalle iterazioni precedenti per correggere il codice.
   • Traduzione Cross-Platform: Sfruttamento di implementazioni funzionali esistenti (es. CUDA) come riferimenti per generare codice per piattaforme sottorappresentate (es. Metal o SYCL).
   • Guardrail dei Candidati: Applicazione di vincoli per prevenire comportamenti indesiderati come il downcasting dei tipi di dati o cambiamenti di precisione.

2. Agente di Analisi delle Performance (Performance-Analysis Agent): Un agente specializzato distinto dal generatore per gestire la scarsità di segnali di ottimizzazione all'interno di grandi dataset di profiling. Interpreta diversi input di profiling (metriche programmatiche da strumenti come Nsight Systems o output visivi da Xcode Instruments) per generare raccomandazioni concrete a livello di codice. Queste raccomandazioni mirano a metriche di utilizzo dell'hardware come la larghezza di banda della memoria, l'occupazione dei warp e l'intensità aritmetica.

Il Ciclo di Sintesi:
Il sistema alterna due fasi:

• Passaggi Funzionali: L'agente di generazione itera finché il codice compila, viene eseguito senza errori e produce output numericamente corretti.
• Passaggi di Ottimizzazione: Una volta stabilita la correttezza, l'agente di analisi delle performance fornisce feedback per guidare il round successivo di sintesi verso la chiusura del divario di performance con i baseline ottimizzati a mano.

Il framework supporta un'interfaccia uniforme per quattro vendor (NVIDIA, AMD, Intel, Apple) e sei modelli di programmazione, utilizzando un registro di modelli per rimanere agnostico rispetto all'LLM sottostante.

Contributi Chiave

• Framework KForge: Introduzione di un sistema di sintesi autonoma multi-stadio dove gli agenti di generazione e di analisi delle performance collaborano utilizzando feedback di compilazione, correttezza e profiling.
• Architettura Cross-Platform: Un design di sistema che supporta quattro vendor di acceleratori e sei modelli di programmazione attraverso un'interfaccia unificata, consentendo il trasferimento di idee algoritmiche tra piattaforme.
• Valutazione Dual-Mode: Casi studio che dimostrano l'efficacia del framework in due scenari distinti: competere contro runtime maturi ottimizzati dai vendor e servire come strumento di "bring-up" per hardware emergente privo di implementazioni di riferimento.

Risultati Sperimentali

Gli autori hanno valutato KForge su due backend con diverse disponibilità di baseline:

1. NVIDIA B200 (Competizione con Runtime Ottimizzati dai Vendor)

• Contesto: Ottimizzazione dei kernel del percorso di decodifica per il modello gpt-oss-20b contro TensorRT-LLM v1.3.0rc9.
• Micro-benchmark: I kernel generati da KForge hanno superato costantemente il baseline in vari batch size. In particolare, il kernel "MoE finalize" ha mostrato un aumento della velocità passando da 1.06× con piccoli batch size a 1.43× con batch size 128.
• Performance End-to-End: L'integrazione di questi kernel ottimizzati nella pipeline di inferenza completa ha prodotto un miglioramento del 2.12% nel throughput end-to-end (2656.61 vs. 2601.55 token/sec) e una riduzione del 2.07% nel tempo di esecuzione (wall-clock time). Gli autori osservano che, sebbene questi guadagni possano apparire modesti, sono significativi considerando gli anni di sforzo ingegneristico già investiti nel baseline del vendor.

2. Intel Arc B580 (Bring-Up di Hardware Emergente)

• Contesto: Generazione di kernel Triton per Intel Arc B580 (Battlemage) dove non esistono riferimenti ottimizzati a mano comparabili.
• Benchmark: Valutato su 37 carichi di lavoro GEMM + tail-ops da KernelBench Level 2.
• Risultati: KForge ha raggiunto un guadagno di velocità (geometric mean speedup) di 5.13× rispetto al più veloce tra PyTorch eager mode e torch.compile.
• Tecniche di Ottimizzazione: L'analisi dei kernel generati ha rivelato l'applicazione riuscita di:
  • Fusione Layer-wise: Collasso delle operazioni post-GEMM in singoli kernel per eliminare i round-trip completi dei tensori.
  • Esecuzione Mixed-Precision: Utilizzo di FP16 per le moltiplicazioni matriciali (matmuls) su unità XMX con accumulazione in FP32.
  • Riduzioni Single-Pass: Sostituzione di algoritmi multi-pass per GroupNorm e LogSumExp con algoritmi di streaming online.

Significato e Rivendicazioni

Il paper posiziona KForge come uno strumento pratico sia per l'ottimizzazione in produzione che per il rapido bring-up dell'hardware. Gli autori affermano che:

• La Sintesi Agentica è Viabile: La separazione tra generazione e analisi delle performance permette agli LLM di navigare efficacemente i complessi compromessi dell'ottimizzazione dei kernel a basso livello.
• Scalabilità dei Guadagni: Anche i miglioramenti percentuali marginali (es. ~2%) contro runtime altamente ottimizzati sono significativi quando scalati a livello di deployment, moltiplicandosi in guadagni di efficienza rilevanti.
• Colmare il Divario: Il framework affronta con successo il problema del "cold start" per l'hardware nuovo (come Intel Arc) generando kernel performanti senza richiedere riferimenti preesistenti ottimizzati a mano, pur competendo con i runtime stabiliti su hardware maturo.

Il lavoro conclude che la sintesi di kernel agentica sta transitando da concetto di ricerca a componente funzionale dello stack dei sistemi AI, capace di gestire l'eterogeneità dei carichi di lavoro di inferenza moderni.