Chiedere, non dotare: agenti di coding guidati da domande per l’ottimizzazione profonda dei kernel

Abstract degli autori

Gli agenti basati su LLM per l’ottimizzazione dei kernel GPU si collocano ai due estremi di un compromesso tra struttura e libertà. Gli agenti con procedure predefinite stabiliscono cosa fa il modello a ogni passaggio, tramite una policy appresa, un menu di ottimizzazioni o ruoli fissi: sono efficienti in termini di campioni quando la procedura è applicabile, ma si bloccano quando non lo è. Gli agenti a iniziativa libera permettono a un modello autonomo di modificare una popolazione di kernel, pagando questa libertà in termini di prove: migliaia di valutazioni, modifiche ripetute e candidati errati. Entrambi affrontano obiettivi difficili dotando l’agente di più strumenti: output del profiler analizzati, analisi del disassemblato, recupero di informazioni dalla documentazione e dalle librerie e archivi di insegnamenti. Sosteniamo che a entrambi manchi una direzione, un’ipotesi su dove si trovi il divario residuo, e che una breve domanda posta al momento giusto possa fornirla senza limitare le possibilità di azione successive dell’agente. Presentiamo KernelSeek, un agente guidato da domande che rende più veloce un kernel CUDA esistente senza conoscenze specifiche dell’operatore né un’implementazione più veloce da imitare: un agente di coding generico, un harness minimale che impone un obiettivo verificabile e domande nei momenti decisionali, prima poste da una persona, poi attinte da un repertorio a cui l’agente stesso fa ricorso. Sullo stesso compito, con lo stesso valutatore, hardware e settimana, un agente autonomo ricco di strumenti, costruito secondo lo stile prevalente, si è fermato a 0,50x FlashAttention-3 dopo 133 passaggi guidati dal modello; KernelSeek ha raggiunto 1,13x. Ha poi prodotto kernel forward di FlashAttention senza dipendenze per Hopper e Blackwell: singoli file CUDA senza CUTLASS, CuTe, cuBLAS o cuDNN, corretti in tutti i 182 scenari per ciascuna architettura, con 1,16x per FlashAttention-3 e 1,46x per FlashAttention-4 (media geometrica). Nell’arco di 162 messaggi umani e 11,5 mila chiamate agli strumenti dell’agente, ogni grande miglioramento è seguito a una domanda che ha portato l’agente a un diverso livello del sistema o ha respinto un limite dichiarato, mai a una domanda che fornisse informazioni. Rendiamo disponibili i kernel, l’harness, il repertorio di domande e tutte le tracce.