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

1. El Problema: La Incompatibilidad entre MoE y NPUs Móviles

Los Modelos de Lenguaje Grande (LLM) basados en Mezcla de Expertos (MoE) han ganado popularidad por su capacidad de escalar la capacidad del modelo mediante activación dispersa (solo un subconjunto de "expertos" se activa para cada token). Sin embargo, ejecutar estos modelos en Unidades de Procesamiento Neuronal (NPU), como el Apple Neural Engine (ANE), presenta desafíos fundamentales:

• Dinamismo vs. Estaticidad: Las NPUs están diseñadas para grafos de cómputo estáticos con formas de tensores fijas (compilados de antemano). La inferencia MoE es inherentemente dinámica: el enrutamiento de tokens a expertos varía según la entrada, creando formas de tensores impredecibles y operaciones irregulares (como top-k, scatter/gather) que las NPUs no soportan nativamente.
• Sobrecarga de Despacho y Sincronización: Ejecutar muchos expertos pequeños individualmente genera una gran cantidad de llamadas al NPU. Esto satura la cola del dispositivo, provoca serialización y genera una sobrecarga de sincronización CPU-NPU que puede consumir más del 60% del tiempo de ejecución total.
• Contención de Recursos: En dispositivos móviles, la CPU debe mantener el sistema operativo y la interfaz de usuario, y la GPU se reserva para renderizado. La inferencia de LLMs a menudo compite por estos recursos, degradando la capacidad de respuesta del sistema.
• Desbalance de Carga: Durante la inferencia, la carga de trabajo de los expertos es altamente desigual (algunos expertos reciben la mayoría de los tokens, otros casi ninguno). Un diseño estático uniforme desperdicia recursos computacionales o provoca pérdida de tokens.

El objetivo del trabajo es habilitar la descarga (offloading) eficiente de la inferencia MoE al NPU, especialmente durante la fase de prefill (que domina la latencia en contextos largos), minimizando la intervención de la CPU/GPU.

2. Metodología: NPUMoE

El sistema propuesto, NPUMoE, es un motor de inferencia en tiempo de ejecución que adapta la ejecución MoE a las restricciones del hardware Apple Silicon. Su principio de diseño es maximizar el cómputo en el NPU mientras utiliza la CPU/GPU como ruta de respaldo para operaciones dinámicas no soportadas.

Se basa en tres técnicas clave derivadas de una calibración offline:

A. Calibración Offline

Antes de la ejecución, el sistema analiza el modelo para estimar:

• La popularidad de los expertos: Qué expertos reciben más tokens en promedio.
• La capacidad esperada: Cuántos tokens maneja cada experto.
  Esto permite definir una política de colocación óptima para la ejecución en tiempo real.

B. Tres Técnicas Principales

1. Niveles Estáticos para la Capacidad de Expertos (Static Tiers for Expert Capacity):

   • Problema: Asignar una capacidad fija uniforme a todos los expertos genera mucho relleno (padding) si la carga es desigual, o pérdida de tokens si la capacidad es insuficiente.
   • Solución: NPUMoE organiza las capacidades en "niveles" (tiers) estáticos (ej. Tier 1, 2, 3). Los expertos "populares" (de alta frecuencia) se asignan a niveles con mayor capacidad, mientras que los expertos "fríos" (poco usados) reciben capacidades menores.
   • Beneficio: Reduce el desperdicio de cómputo por padding y minimiza la pérdida de tokens, adaptando la forma del tensor estático a la distribución real de carga.

2. Ejecución Agrupada de Expertos (Grouped Expert Execution):

   • Problema: Lanzar un grafo de cómputo separado para cada experto genera una sobrecarga de lanzamiento (dispatch overhead) insoportable para el NPU.
   • Solución: Agrupar múltiples expertos en un solo grafo de cómputo denso y estático. Por ejemplo, en lugar de lanzar 8 grafos pequeños, se lanza 1 grafo grande que procesa 8 expertos simultáneamente con una sola llamada.
   • Beneficio: Amortiza la sobrecarga de lanzamiento y sincronización, aumentando la utilización del NPU y reduciendo la serialización en la cola del dispositivo.

3. Residencia de Grafos de Cómputo Consciente de la Carga (Load-Aware Compute Graph Residency):

   • Problema: No todos los grupos de expertos se benefician de ejecutarse en el NPU. Los grupos pequeños o de expertos "fríos" no justifican el costo de sincronización CPU-NPU.
   • Solución: El sistema decide dinámicamente dónde ejecutar los grupos:
     • Grupos "Calientes" (Hot): Expertos de alta demanda se agrupan y residen permanentemente en el NPU para amortizar el costo de lanzamiento.
     • Grupos "Fríos" (Cold): Expertos de baja demanda o grupos muy pequeños se ejecutan directamente en la CPU/GPU para evitar la latencia de sincronización.
   • Beneficio: Equilibra la latencia y la eficiencia energética, evitando mover datos innecesariamente entre CPU y NPU para cargas de trabajo pequeñas.

3. Contribuciones Clave

• Análisis de Desafíos: Identificación profunda de por qué los patrones de ejecución dinámicos y dispersos de MoE chocan con las restricciones estáticas de las NPUs móviles.
• Técnicas Innovadoras: Propuesta de un enfoque tripartito (Niveles de capacidad, Ejecución agrupada, Residencia consciente de la carga) que transforma la inferencia MoE para ser compatible con hardware estático.
• Sistema NPUMoE: Implementación de un motor de inferencia completo que utiliza la API Core ML de Apple y la biblioteca ANEMLL, optimizado específicamente para dispositivos Apple Silicon (M2 Ultra, M2 Max).
• Evaluación Exhaustiva: Demostración de que es posible lograr inferencia MoE eficiente en dispositivos móviles sin sacrificar significativamente la precisión.

4. Resultados Experimentales

El sistema se evaluó en dispositivos Apple M2 (Ultra y Max) utilizando modelos como Phi-3.5-MoE, Phi-tiny-MoE y Qwen3-30B-A3B, comparándolo con tres líneas base (CPU puro, Core ML nativo y ANEMLL).

• Latencia: NPUMoE reduce la latencia de prefill entre 1.32x y 5.55x en comparación con las líneas base.
• Eficiencia Energética: Mejora la eficiencia energética entre 1.81x y 7.37x (menor energía por token).
• Uso de CPU: Reduce el uso de ciclos de CPU entre 1.78x y 5.54x, liberando recursos para otras tareas del sistema.
• Precisión: La degradación de precisión es insignificante (< 1.1%), incluso con tasas de pérdida de tokens (token dropping) del 10-20% en casos extremos, gracias a la poda basada en saliencia.
• Análisis de Desglose: Se demostró que la mayor parte de la ganancia proviene de la optimización del bloque MoE (FFN de expertos), que suele representar más del 86% de la latencia de prefill.

5. Significado e Impacto

Este trabajo es pionero al demostrar que las NPUs móviles pueden ser utilizadas eficazmente para la inferencia de modelos MoE, un tipo de arquitectura que tradicionalmente se consideraba incompatible con hardware estático debido a su naturaleza dinámica.

• Privacidad y Personalización: Permite ejecutar modelos grandes y eficientes directamente en el dispositivo (on-device), manteniendo la privacidad de los datos y la personalización, sin depender de la nube.
• Optimización de Recursos del Sistema: Al descargar la carga de trabajo del NPU, se libera la CPU y la GPU para tareas interactivas (UI, renderizado), mejorando la experiencia general del usuario en dispositivos móviles bajo carga pesada.
• Guía para el Futuro: Proporciona un marco de diseño (co-diseño de hardware/software) que puede aplicarse a otras plataformas NPU, sugiriendo que la combinación de calibración offline y estrategias de agrupamiento dinámico es la clave para ejecutar modelos de IA avanzados en el borde (edge).

En resumen, NPUMoE cierra la brecha entre la flexibilidad algorítmica de los modelos MoE y las restricciones de hardware de las NPUs, logrando un equilibrio óptimo entre velocidad, eficiencia energética y precisión en dispositivos Apple Silicon.