L'accelerazione GPU per PyTorch su Apple Silicon

La piattaforma Apple Silicon ha consolidato la sua posizione come scelta popolare per l'esecuzione locale di Large Language Models (LLM). Fino a poco tempo fa, gli utenti di ExecuTorch su macOS si trovavano spesso limitati a backend basati esclusivamente sulla CPU, come XNNPACK, o al backend AOTI Metal, che non sempre garantivano le performance ottimali per carichi di lavoro intensivi.

In questo scenario, il team di PyTorch ha rilasciato il delegate MLX, una novità che introduce l'Inference completamente ottimizzata e accelerata dalla GPU sui Mac dotati di Apple Silicon. Questa innovazione sfrutta il framework MLX di Apple, progettato specificamente per l'apprendimento automatico sui chip della casa di Cupertino, aprendo nuove possibilità per gli sviluppatori e gli architetti di sistemi che cercano soluzioni di Deployment on-premise o edge per i loro modelli di AI.

Il delegate MLX: architettura e integrazione nel workflow

Il delegate MLX si configura come un nuovo backend per ExecuTorch, con il compito di compilare ed eseguire i modelli PyTorch direttamente sulle GPU di Apple Silicon. Il processo è progettato per essere trasparente per l'utente: dopo aver esportato il modello attraverso la pipeline standard di ExecuTorch, il delegate si occupa della partizione del grafo computazionale, della sua serializzazione in un formato ottimizzato e dell'invio delle operazioni ai kernel GPU Metal di MLX durante l'esecuzione.

Dal punto di vista del workflow, l'approccio rimane coerente con quello di altri backend ExecuTorch. Gli sviluppatori esportano il modello con torch.export, lo abbassano con to_edge_transform_and_lower utilizzando MLXPartitioner e infine eseguono il file .pte risultante con il runtime di ExecuTorch. Il delegate supporta attualmente circa 90 operazioni ATen, coprendo l'intero spettro necessario per l'Inference dei transformer, inclusi matmul quantizzati, multi-head attention, rotary position embeddings, mixture-of-experts routing e operazioni di stato-spazio ricorrenti.

Un aspetto cruciale per l'efficienza e la flessibilità è il supporto per diverse opzioni di Quantization e tipi di dati. Il delegate gestisce BF16, FP16 e FP32 per pesi e attivazioni, oltre alla Quantization affine a 2, 4 e 8 bit tramite l'API quantize_ di TorchAO. Quest'ultima adotta lo stesso schema di Quantization dei backend XNNPACK e Vulkan, consentendo a una singola definizione di modello quantizzato di essere compatibile con più backend. È inoltre supportata la Quantization NVFP4, che utilizza il tipo di dati FP4 di NVIDIA, e gli embeddings quantizzati condivisi per i modelli che riutilizzano i pesi tra il layer di embedding e l'head del modello di linguaggio.

Vantaggi chiave per il Deployment locale

L'introduzione del delegate MLX porta con sé vantaggi significativi, in particolare per chi valuta soluzioni di AI on-premise o edge. Il primo e più evidente è il miglioramento delle performance: il delegate MLX raggiunge un throughput da 3 a 6 volte superiore nei carichi di lavoro di AI generativa rispetto ai delegate ExecuTorch preesistenti su macOS. Questo incremento è fondamentale per applicazioni che richiedono risposte rapide, come le chat basate su LLM o la trascrizione audio in tempo reale.

Un altro punto di forza è l'integrazione nativa con lo stack di export di PyTorch 2. Il delegate si connette direttamente a torch.export per la cattura del grafo e a TorchAO per la Quantization, gli stessi strumenti utilizzati da tutti gli altri backend ExecuTorch. Ciò significa che i nuovi modelli o le tecniche di Quantization che approdano in PyTorch diventano immediatamente disponibili per il delegate MLX senza richiedere lavoro aggiuntivo, garantendo una maggiore agilità nello sviluppo e nel Deployment. Infine, ExecuTorch offre un'unica API di runtime che funziona su tutti i backend. Un'applicazione sviluppata con il runtime C++ o Python di ExecuTorch può eseguire modelli esportati per MLX, XNNPACK, CoreML, Vulkan o CUDA senza modifiche al codice dell'applicazione. Questa portabilità è un fattore abilitante per architetture ibride e per la gestione del TCO, riducendo la complessità e i costi di manutenzione in ambienti eterogenei.

Modelli supportati e prospettive future

Il delegate MLX ha dimostrato la sua efficacia su un'ampia gamma di architetture di modelli. Per quanto riguarda i Large Language Models, supporta transformer densi come Llama 3.2 1B, Qwen 3 (0.6B, 1.7B, 4B), Phi-4 mini (3.8B) e Gemma 3 (1B, 4B), inclusi quelli con sliding window attention. È inoltre compatibile con i modelli Sparse Mixture-of-Experts, come Qwen 3.5 35B-A3B, grazie a operazioni di gather personalizzate che instradano efficientemente i token agli esperti corretti sulla GPU.

Nel campo dello speech-to-text, il delegate abilita modelli per la trascrizione offline, tra cui OpenAI Whisper (dalla versione tiny alla large-v3-turbo), NVIDIA Parakeet TDT (0.6B) con timestamp a livello di parola e Mistral Voxtral (3B). Per i casi d'uso in tempo reale, supporta Mistral Voxtral Realtime (4B), gestendo l'input da microfono live, le KV caches a buffer circolare e lo sliding window attention. Oltre a questi modelli di punta, oltre 30 modelli aggiuntivi sono stati validati attraverso le suite di test del backend, coprendo transformer densi, architetture encoder-decoder e modelli di visione. È importante notare che il delegate MLX è attualmente in fase sperimentale e in sviluppo attivo, il che implica che le API e le funzionalità supportate potrebbero evolvere. Per chi valuta soluzioni di AI on-premise, l'emergere di strumenti come il delegate MLX sottolinea l'importanza di analizzare i trade-off tra performance, flessibilità e controllo dei dati, aspetti che AI-RADAR esplora in dettaglio nei suoi framework analitici su /llm-onpremise.