Efficient Mixture-of-Experts LLM Inference with Apple Silicon NPUs

1. Il Problema

L'articolo affronta le sfide nell'esecuzione di modelli linguistici su larga scala (LLM) basati su architetture Mixture-of-Experts (MoE) su dispositivi mobili con Apple Silicon, in particolare sfruttando il Neural Engine (ANE), l'unità di elaborazione neurale (NPU) dedicata.

Sebbene i modelli MoE migliorino l'efficienza dell'inferenza attivando solo un sottoinsieme di "esperti" (FFN) per token, la loro natura dinamica si scontra con le limitazioni hardware delle NPU:

• Dinamicità vs. Staticità: Le NPU sono progettate per grafi di calcolo statici e forme tensoriali fisse. Il routing degli esperti nei MoE è imprevedibile e crea forme tensoriali dinamiche, costringendo spesso il sistema a un fallback sulla CPU.
• Operatori Irregolari: Operazioni critiche per i MoE come il routing top-k, lo scatter/gather e l'indicizzazione dinamica non sono supportate nativamente o sono inefficienti sulle NPU.
• Overhead di Dispatch: Lanciare molti piccoli kernel (uno per ogni esperto) genera un elevato overhead di sincronizzazione CPU-NPU e di gestione della coda, riducendo drasticamente l'utilizzo dell'acceleratore.
• Contesa delle Risorse: In scenari di inferenza on-device, la CPU deve gestire il sistema operativo e l'interfaccia utente, mentre la GPU è spesso occupata dal rendering grafico. Spostare il carico computazionale sulla NPU è essenziale per mantenere il sistema reattivo e migliorare l'efficienza energetica, ma le attuali soluzioni non gestiscono bene la dinamicità dei MoE.

2. Metodologia: NPUMoE

Gli autori presentano NPUMoE, un motore di inferenza runtime che accelera l'esecuzione MoE sugli Apple Silicon offloadando le operazioni dense e statiche sulla NPU, mantenendo un percorso di fallback su CPU/GPU per le operazioni dinamiche.

Il sistema si basa su tre tecniche chiave, guidate da una calibrazione offline che stima la capacità e la popolarità degli esperti:

A. Livelli Statici per la Capacità degli Esperti (Static Tiers for Expert Capacity)

• Problema: Il numero di token inviati a ciascun esperto varia dinamicamente.
• Soluzione: Invece di creare grafi dinamici, NPUMoE assegna a ciascun esperto una capacità fissa basata su una stima del suo carico di lavoro atteso.
• Meccanismo: Gli esperti vengono mappati su "tier" (livelli) di capacità statici (es. Tier 1, 2, 3). Gli esperti "popolari" (che ricevono molti token) ottengono capacità maggiori, mentre quelli meno usati ne ricevono di minori.
• Gestione dell'Overflow: Se un esperto riceve più token della sua capacità statica, i token in eccesso vengono "potati" (dropped) basandosi su un punteggio di salienza (L2 norm delle attivazioni) per minimizzare la perdita di accuratezza, invece di usare padding eccessivo che sprecherebbe risorse.

B. Esecuzione Raggruppata degli Esperti (Grouped Expert Execution)

• Problema: Lanciare un grafo di calcolo separato per ogni singolo esperto genera troppi piccoli task che si serializzano sulla coda del dispositivo, aumentando l'overhead di avvio.
• Soluzione: NPUMoE raggruppa più esperti in un singolo grafo di calcolo denso e statico (FFN raggruppato).
• Meccanismo: I token vengono raggruppati per ID esperto e inseriti in un buffer di input unificato. L'NPU esegue un singolo blocco fuso che elabora più esperti contemporaneamente. Questo riduce drasticamente il numero di invocazioni e l'overhead di sincronizzazione, massimizzando l'utilizzo della NPU.

C. Residenza del Grafo di Calcolo Consapevole del Carico (Load-Aware Compute Graph Residency)

• Problema: Spostare tutti gli esperti sulla NPU non è sempre efficiente; piccoli carichi di lavoro "freddi" (cold experts) non giustificano il costo di sincronizzazione CPU-NPU.
• Soluzione: Una politica di placement che decide dinamicamente quali gruppi di esperti risiedere sulla NPU e quali sulla CPU.
• Meccanismo: Basandosi sulla popolarità degli esperti (calibrata offline), i gruppi di esperti "caldi" (ad alta domanda) vengono mantenuti sulla NPU per ammortizzare l'overhead di avvio. Gli esperti "freddi" o gruppi troppo piccoli vengono eseguiti sulla CPU, evitando costose sincronizzazioni per carichi di lavoro bursty.

3. Contributi Chiave

1. Analisi Approfondita: Prima indagine dettagliata sulle sfide e limitazioni dell'esecuzione MoE su NPU mobile (Apple ANE), evidenziando il disallineamento tra dinamicità del modello e rigidità dell'hardware.
2. Tre Nuove Tecniche: Introduzione di Static Tiers, Grouped Execution e Load-Aware Residency per adattare i modelli MoE dinamici alle NPU statiche.
3. NPUMoE: Il primo sistema di inferenza runtime che ottimizza sistematicamente l'efficienza MoE per le NPU mobili, offrendo un motore ibrido CPU/NPU/GPU.
4. Valutazione Completa: Un'analisi end-to-end su dispositivi reali (M2 Max, M2 Ultra) con modelli MoE di punta (Phi-3.5, Qwen3).

4. Risultati Sperimentali

Il sistema è stato valutato su tre modelli MoE (Phi-3.5-MoE, Phi-tiny-MoE, Qwen3-30B-A3B) con carichi di lavoro a contesto lungo (prefill). I risultati mostrano un miglioramento significativo rispetto a baseline come CoreML (solo CPU), CoreML (naïve) e ANEMLL:

• Riduzione della Latenza: Diminuzione della latenza di prefill (TTFT) da 1.32x a 5.55x.
• Efficienza Energetica: Miglioramento dell'efficienza energetica da 1.81x a 7.37x (meno energia per token).
• Utilizzo CPU: Riduzione del consumo di cicli CPU da 1.78x a 5.54x, liberando risorse per altre attività del sistema.
• Accuratezza: Perdita di accuratezza trascurabile (< 1.1%), dimostrando che la potatura dei token e il padding controllato non compromettono significativamente le prestazioni del modello.

5. Significato e Impatto

Questo lavoro dimostra che è possibile eseguire modelli LLM MoE complessi ed efficienti su dispositivi mobili sfruttando le NPU, nonostante la loro natura statica.

• Sostenibilità: Permette di eseguire modelli più grandi e precisi su dispositivi consumer senza sovraccaricare la CPU o la GPU, migliorando l'autonomia della batteria.
• Privacy: Abilita l'inferenza on-device completa per modelli MoE, garantendo la privacy dei dati utente.
• Co-progettazione Sistema-Hardware: Dimostra che un'attenta progettazione del sistema (software) può colmare il divario tra le esigenze dinamiche dei modelli AI moderni e le restrizioni hardware delle NPU mobili, aprendo la strada a un'adozione più diffusa di modelli MoE su edge devices.

In sintesi, NPUMoE trasforma le NPU da acceleratori per modelli densi statici a motori efficienti per modelli MoE dinamici, risolvendo i colli di bottiglia di sincronizzazione e gestione della memoria attraverso un'architettura ibrida intelligente.