L’efficienza computazionale per l’AI è diventata una priorità per sviluppatori e ricercatori: la startup giapponese Sakana AI ha introdotto l’AI CUDA Engineer, un framework innovativo progettato per automatizzare la scoperta e l’ottimizzazione dei kernel CUDA, promettendo miglioramenti prestazionali fino a 100 volte rispetto alle operazioni standard di PyTorch.
CUDA, acronimo di Compute Unified Device Architecture, è un’interfaccia di programmazione sviluppata da NVIDIA che consente l’accesso diretto alle GPU per il calcolo parallelo. Sebbene l’ottimizzazione manuale dei kernel CUDA possa portare a significativi incrementi di performance, richiede una profonda conoscenza dell’architettura delle GPU, rendendo il processo complesso e dispendioso in termini di tempo.
L’AI CUDA Engineer di Sakana AI affronta questa sfida utilizzando modelli linguistici di grandi dimensioni (LLM) e tecniche di ottimizzazione evolutiva per automatizzare la conversione del codice PyTorch in kernel CUDA altamente ottimizzati. Questo processo avviene attraverso una pipeline multi-stadio: inizialmente, il sistema traduce le operazioni PyTorch in kernel CUDA funzionanti, spesso migliorando le prestazioni senza necessità di tuning manuale. Successivamente, applica strategie evolutive, come operazioni di ‘crossover’ e l’utilizzo di un ‘archivio di innovazione’, per affinare ulteriormente le performance dei kernel generati.
Secondo Sakana AI, l’AI CUDA Engineer è riuscito a tradurre con successo oltre 230 su 250 operazioni PyTorch valutate, generando più di 30.000 kernel CUDA, di cui oltre 17.000 verificati per correttezza. Circa il 50% di questi kernel supera in performance le implementazioni native di PyTorch, con incrementi di velocità che vanno da 10 a 100 volte. In alcuni casi, i kernel ottimizzati hanno dimostrato di essere fino a cinque volte più veloci rispetto ai kernel CUDA esistenti utilizzati in produzione.
Per promuovere la collaborazione e l’ulteriore sviluppo nel campo, Sakana AI ha reso disponibile il dataset dei kernel generati sotto licenza CC-By-4.0 sulla piattaforma Hugging Face. Questo dataset include implementazioni di riferimento, dati di profilazione e confronti di performance rispetto ai runtime nativi di PyTorch. Inoltre, l’azienda ha lanciato un sito web interattivo che consente agli utenti di esplorare il dataset e visualizzare le classifiche dei kernel ottimizzati, fornendo accesso al codice dei kernel, metriche di performance e dettagli sugli esperimenti di ottimizzazione correlati