MaxKernel: generazione agentica di kernel per TPU

Abstract degli autori

Progettare e scrivere kernel personalizzati ad alte prestazioni per gli acceleratori è un compito complesso che richiede una profonda conoscenza dell’hardware. I modelli linguistici di grandi dimensioni (LLM) possono essere impiegati insieme al feedback del compilatore in tempo reale per creare sistemi agentici per la generazione di kernel. In questo lavoro presentiamo MaxKernel, un sistema multiagente che implementa tre approcci distinti allo sviluppo di kernel per TPU: (1) un agente Human-in-the-Loop (HITL) per una progettazione collaborativa, passo dopo passo; (2) un agente Autonomous (Auto) che esegue un ciclo di ottimizzazione interamente automatizzato, guidato da metriche e tracce; e (3) una Graph-Based Autonomous Search che estende l’agente Auto per esplorare globalmente lo spazio delle soluzioni progettuali. Tutti e tre gli approcci si avvalgono di un insieme condiviso di sottoagenti specializzati nella pianificazione, nell’implementazione, nel debug autonomo, nei test e nella profilazione dell’hardware. Valutiamo MaxKernel su JaxBench, una suite completa di 50 attività diversificate relative ai kernel per TPU, oltre che su carichi di lavoro complessi e reali tratti da modelli open source all’avanguardia. Dimostriamo che MaxKernel genera con costanza implementazioni altamente ottimizzate, eguagliando le baseline ottimizzate manualmente da esperti e offrendo prestazioni di rilievo nell’intero benchmark. Il nostro agente è open source ed è disponibile all’indirizzo https://github.com/AI-Hypercomputer/accelerator-agents/tree/main/MaxKernel.