Compiler-Grounded Hierarchical Diagnosis for LLM-Based Triton Kernel Optimization

Sintesi Tecnica: Diagnosi Gerarchica con Grounding del Compilatore per l'Ottimizzazione di Kernel Triton basata su LLM

1. Definizione del Problema

Sebbene i Large Language Models (LLM) abbiano fatto avanzare la generazione e l'ottimizzazione automatizzata dei kernel, gli approcci esistenti si affidano prevalentemente a "segnali superficiali" come errori di compilazione, metriche di runtime e dati di profiling. Questi segnali indicano che un kernel è lento, ma raramente spiegano il perché il compilatore backend non riesca a realizzare un'ottimizzazione profittevole. Ciò è particolarmente acuto sugli emergenti acceleratori come le NPU (es. Ascend), dove lo stack del compilatore è in evoluzione, specifico per il target e opaco agli sviluppatori.

La sfida principale identificata è un gap semantico tra il codice sorgente a livello di operatore e il comportamento delle prestazioni del backend. I fallimenti delle prestazioni emergono spesso solo dopo che il codice è passato attraverso gli stadi di rappresentazione intermedia (IR), le decisioni di scheduling e i passaggi specifici del compilatore backend. Di conseguenza, gli agenti LLM, tipicamente pre-addestrati su pattern centrici su GPU, si comportano come ottimizzatori trial-and-error: propongono modifiche plausibili al codice sorgente ma mancano delle evidenze per mappare i sintomi del runtime a cause specifiche a livello di compilatore.

2. Metodologia

Il documento propone un framework di ottimizzazione gerarchico con grounding del compilatore che tratta l'ottimizzazione dei kernel come un problema di diagnosi progressiva cross-layer. Inveve di trattare il grounding del compilatore come una funzione sempre attiva, il sistema scala da un'evidenza leggera a un'analisi profonda del compilatore solo quando necessario.

Il Workflow Gerarchico:
Il sistema opera attraverso una pipeline a ciclo chiuso organizzata in quattro livelli crescenti:

• L1: Triage dei Pattern: Il sistema controlla innanzitutto ipotesi forti supportate da pattern utilizzando un indice curato di pattern di ottimizzazione, euristiche e regole esperte (incluse euristiche interne distillate da esecuzioni precedenti). Se una riscrittura a basso rischio è giustificata dalle evidenze esistenti, il processo si ferma qui.
• L2: Diagnosi tramite Profiling: Se i pattern falliscono, il sistema raccoglie evidenze di runtime (latenza, throughput, contatori hardware) per localizzare i colli di bottiglia (es. movimento di memoria, attività scalare, sincronizzazione). Questo identifica il sintomo, ma non necessariamente il meccanismo.
• L3: Attribuzione IR: Quando il profiling è insufficiente, il sistema ispeziona le Rappresentazioni Intermedie (IR) per tradurre i sintomi del runtime in diagnosi strutturali (es. carichi non unit-stride, fusion mancata, fallimenti di vettorizzazione indotti dal layout). Ricostruisce la traiettoria di lowering dell'operatore per identificare le trasformazioni stage-to-stage.
• L4: Escalation Compilatore-Sorgente: Come ultima risorsa, il sistema invoca l'analisi con grounding del compilatore per spiegare perché un pattern diagnosticato blocca l'ottimizzazione del backend. Consulta la conoscenza del compilatore (comportamenti dei pass, regole di lowering, vincoli) per recuperare le condizioni a livello di sorgente che una riscrittura di successo deve soddisfare.

Architettura di Sistema e Gestione dello Stato:

• Design Centrico sull'Agente: Implementato come un ambiente centrato sull'agente (integrato nel repository di skill di CannBot) piuttosto che come un servizio di compilazione bespoke. Utilizza un harness di esecuzione sottile per gestire validazione, correttezza e raccolta di evidenze, mentre "skill" specifiche del dominio gestiscono la logica di ottimizzazione.
• Round Stateful: Il processo di ottimizzazione è organizzato in round espliciti e auditabili. Ogni round registra la propria ipotesi, le fonti di evidenza e l'esito della validazione. Questo design stateful permette al sistema di essere riprendibile e assicura che le analisi più profonde (IR/Compilatore) siano allegate come evidenze riutilizzabili piuttosto che osservazioni effimere.
• Sintesi Retrospettiva: Oltre al loop online, il sistema esegue un passaggio di sintesi cross-operatore più lento. Mappa le storie di ottimizzazione validate verso famiglie di pattern esistenti o crea nuovi pattern riutilizzabili, permettendo al livello più "economico" (Pattern Triage) di migliorare nel tempo.

3. Contributi Chiave

1. Formulazione del Problema: Gli autori formulano l'ottimizzazione dei kernel basata su LLM come un problema di diagnosi progressiva cross-layer, identificando la mancanza di evidenze con grounding del compilatore come la ragione primaria per cui gli agenti basati solo sul profiling falliscono nelle ottimizzazioni specifiche del backend.
2. Workflow Gerarchico: Descrivono un workflow che scala dal triage dei pattern al profiling, all'attribuzione IR e all'escalation compilatore-sorgente, proponendo riscrizioni a livello di sorgente supportate da evidenze solo quando è richiesta un'evidenza più profonda.
3. Evidenze di Valutazione: Presentano evidenze di valutazione standardizzate e aware della distribuzione su un benchmark Ascend 950 derivato da NPUKernelBench, riportando guadagni aggregati, il miglior round e tracce cross-layer rappresentative.

4. Risultati Sperimentali

Il sistema è stato implementato su Triton per NPU Ascend ed è stato valutato su 37 operatori convertiti con successo da un benchmark standardizzato Ascend 950 (derivato da NPUKernelBench). Ogni operatore è stato testato su 5 casi (185 casi totali).

• Guadagni di Prestazioni:
  • Media Geometrica di Speedup: 4.35× (dal kernel Triton iniziale a quello ottimizzato).
  • Mediana di Speedup: 2.73×.
  • Distribuzione: 22/37 voci hanno superato un aumento di velocità di 2×; 13/37 hanno superato 5×. I risultati mostrano una distribuzione a coda lunga, con guadagni sostanziali che coesistono con voci vicine alla baseline.
• Contabilità del Processo:
  • L'indice mediano del miglior round era 10, indicando che i risultati ottimali richiedevano spesso molteplici iterazioni.
  • Profondità dell'Evidenza: Sebbene 33 dei "migliori" round selezionati siano stati attribuiti al Pattern Triage, il sistema ha dimostrato l'utilità degli strati più profondi: 32 traiettorie hanno raggiunto il Profiling Diagnosis, 3 hanno raggiunto l'IR Attribution e 4 hanno raggiunto il Compiler-Source Escalation.
  • Caso di Studio: Una traccia rappresentativa (Hyena FFT-size padding) ha mostrato come il profiling abbia identificato un collo di bottiglia scalare, l'attribuzione IR abbia identificato specifici costi di sincronizzazione nel lowering e l'analisi del compilatore abbia spiegato il vincolo, portando a un robusto aumento di velocità di 2.08×.

5. Significato e Claim

Il paper sostiene che il grounding del compilatore debba essere selettivo piuttosto che sempre attivo. Trattando l'ottimizzazione come una diagnosi progressiva, il sistema rende pratica l'analisi profonda del compilatore, invocandola solo quando un'evidenza più leggera non riesce a produrre una riscrittura sufficientemente fondata.

Gli autori sottolineano che i loro risultati dimostrano un margine sostanziale su un sottoinsieme di operatori, ma non rivendicano un miglioramento uniforme in tutti i casi. Dichiarano esplicitamente che la valutazione attuale non fornisce un'isolazione causale (ad esempio, confrontando l'escalation progressiva contro strategie di solo profiling o di profondità costante sotto un toolchain congelata). Il lavoro si posiziona come un "argomento di sistema" supportato da benchmark trasparenti e da evidenze di processo, riservando la decomposizione causale della gerarchia a futuri esperimenti controllati.

Il lavoro evidenzia che la relazione tra codice sorgente e prestazioni di runtime è mediata da trasformazioni invisibili del compilatore, e che colmare questo divario con un approccio gerarchico guidato dalle evidenze è essenziale per ottimizzare i kernel su acceleratori emergenti e opachi come le NPU.