Da CUDA a MLX: L'expertise di K-Search al servizio di Apple Silicon

Una nuova era nel calcolo

Il mondo del calcolo evolve a un ritmo frenetico. Le innovazioni hardware non si limitano a GPU più veloci, ma generano anche una varietà di chip di diversi produttori, ognuno con la propria architettura, spesso progettata per carichi di lavoro specifici legati all'intelligenza artificiale (IA). In parallelo, lo sviluppo software progredisce altrettanto rapidamente, con strumenti di codifica in IA capaci di generare ciò che prima richiedeva mesi in pochi minuti.

In un ambiente in cui l'IA occupa un posto centrale, i kernel GPU diventano essenziali. Questi programmi di basso livello, eseguiti all'interno delle GPU, richiedono un'esperienza considerevole per essere ottimizzati in modo efficace. La migrazione di un kernel da un hardware a un altro rappresenta una sfida ancora maggiore, richiedendo spesso una riscoperta delle ottimizzazioni specifiche per ogni architettura.

L'expertise dei kernel CUDA

L'ecosistema CUDA, sostenuto da decenni di ricerca e ottimizzazioni, ha permesso di sviluppare implementazioni altamente efficienti di funzioni critiche come l'attenzione o i modelli di spazio di stato. Questi risultati rappresentano migliaia di ore di ingegneria. Tuttavia, i nuovi ecosistemi hardware, come Apple Silicon e altri acceleratori IA personalizzati, si sviluppano rapidamente ma mancano di questa profondità di esperienza.

Diventa quindi cruciale chiedersi se questa expertise possa essere trasferita in modo automatico.

K-Search: un passo avanti nell'ottimizzazione dei kernel

K-Search, un framework di ottimizzazione dei kernel scalabile sviluppato all'UC Berkeley, si presenta come una risposta a questa sfida. Questo metodo utilizza l'intelligenza artificiale per ottimizzare i kernel GPU. Estendendo K-Search con un backend per MLX, il framework di apprendimento automatico di Apple, abbiamo creato uno strato di traduzione da CUDA a MLX. Questo approccio consente a K-Search di sfruttare i kernel CUDA esistenti per generarne di nuovi per Apple Silicon, riducendo così la necessità di ricostruire i kernel da zero.

I risultati parlano da soli: il nostro metodo raggiunge prestazioni vicine a quelle di un esperto su Apple Silicon, offrendo un guadagno di velocità di 0,97x rispetto al kernel di attenzione nativo MLX, e fino a 20x in velocità di pre-riempimento rispetto all'implementazione comunitaria mlx-lm sul kernel Mamba SSM.

Perché scegliere MLX?

Dalla fine del 2023, il framework MLX di Apple ha conosciuto un'adozione notevole. Con milioni di MacBook e Mac Studios dotati di Apple Silicon, MLX consente l'inferenza IA locale senza costi legati al cloud. La sua architettura di memoria unificata lo rende particolarmente adatto per modelli di dimensioni intermedie (da 7 a 70 miliardi di parametri sui chip M).

Tuttavia, persiste una lacuna importante: molti kernel critici per le prestazioni, spesso disponibili nell'ecosistema NVIDIA, si rivelano assenti o poco ottimizzati su MLX senza una regolazione specifica per l'hardware. Sebbene MLX esegua correttamente i modelli, spesso lascia da parte un potenziale di prestazione considerevole.

Il funzionamento di K-Search

K-Search si distingue per la sua capacità di iterare su un kernel naïf in base a una specifica hardware. Un modello di IA ragiona sulle ottimizzazioni da applicare, genera kernel candidati, che vengono poi compilati e testati su hardware reale. Le misurazioni ottenute alimentano il processo di ricerca, consentendo un miglioramento continuo fino a quando le prestazioni convergono.

Questo processo è guidato da un documento specifico per un dominio che codifica le regole hardware e le restrizioni matematiche, impedendo la generazione di codice non valido e assicurando che i candidati si compilino ed eseguano in modo efficiente. In pratica, un modello unico (come Gemini 3.5 Pro Preview) svolge sia il ruolo di ragionamento che di generazione di codice, analizzando il kernel prima di proporre ottimizzazioni.

Conclusione

In sintesi, la combinazione dell'expertise accumulata nell'ecosistema CUDA e dell'innovazione portata da K-Search apre nuove prospettive per l'ottimizzazione delle prestazioni su Apple Silicon. Le applicazioni di questo metodo non si limitano a MLX, ma si estendono a qualsiasi ecosistema in cui l'expertise CUDA possa essere trasferita utilmente.

Se desideri saperne di più sull'ottimizzazione delle prestazioni della tua infrastruttura IA o sull'implementazione di strategie simili, non esitare a Contattami.