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

Título: Inference Eficiente de LLMs Mixture-of-Experts (MoE) com NPUs de Apple Silicon

1. O Problema

Os Grandes Modelos de Linguagem (LLMs) baseados em arquiteturas Mixture-of-Experts (MoE) ganharam popularidade por permitir escalar a capacidade do modelo através da ativação esparsa (apenas um subconjunto de "especialistas" é ativado para cada token). No entanto, executar esses modelos em dispositivos móveis (especificamente em NPUs como o Apple Neural Engine - ANE) apresenta desafios significativos:

• Dinamismo vs. Estaticidade: As NPUs são projetadas para gráficos de computação estáticos e formas de tensor fixas (compilados antecipadamente). O MoE, por natureza, é dinâmico: o roteamento de tokens para especialistas varia a cada entrada, criando formas de tensor imprevisíveis e operadores irregulares (como top-k, scatter, gather).
• Operadores Não Suportados: Muitas operações críticas do MoE (indexação dinâmica, lógica condicional, redução) não são otimizadas ou suportadas diretamente pelo hardware NPU, forçando fallbacks para a CPU.
• Overhead de Disparo e Sincronização: Lançar muitos pequenos kernels (um para cada especialista) na NPU gera um custo de despacho (dispatch) e sincronização CPU-NPU proibitivo. Em casos extremos, o tempo de sincronização pode representar mais de 60% do tempo total de execução.
• Contenção de Recursos: Em dispositivos móveis, a CPU e a GPU já estão sobrecarregadas com tarefas do sistema operacional, renderização e I/O. Executar LLMs inteiramente na CPU/GPU degrada a responsividade do sistema.

O foco deste trabalho é a fase de prefill (processamento do prompt inicial), que domina a latência de inferência em tarefas de longo contexto (80% da latência total e 70% dos ciclos de CPU).

2. Metodologia: NPUMoE

O artigo apresenta o NPUMoE, um motor de inferência em tempo de execução projetado para acelerar modelos MoE em Apple Silicon. A filosofia central é maximizar o offload para a NPU (para operações densas e estáticas) enquanto mantém a CPU/GPU como fallback para operações dinâmicas.

O sistema utiliza três técnicas principais, baseadas em uma calibração offline que estima a popularidade e a carga de trabalho de cada especialista:

1. Níveis Estáticos para Capacidade de Especialistas (Static Tiers for Expert Capacity):

   • Para lidar com a variabilidade do número de tokens roteados para cada especialista, o NPUMoE não tenta criar gráficos dinâmicos. Em vez disso, ele define "níveis" (tiers) de capacidade fixa (ex: Tier 1, 2, 3).
   • Especialistas populares (que recebem muitos tokens) são mapeados para capacidades maiores, enquanto especialistas menos usados recebem capacidades menores.
   • Se o número de tokens exceder a capacidade estática, tokens de baixa relevância (baseados em pontuação de saliência) são descartados (pruning) ou preenchidos (padding), evitando a necessidade de gráficos dinâmicos.

2. Execução Agrupada de Especialistas (Grouped Expert Execution):

   • Para mitigar o overhead de disparo de muitos pequenos gráficos, o NPUMoE agrupa múltiplos especialistas em um único gráfico de computação denso e estático.
   • Os tokens são agrupados por ID de especialista e empacotados em um buffer de entrada contíguo. Isso permite que a NPU execute vários especialistas como um bloco fundido, compartilhando o custo de lançamento e aumentando a utilização do hardware.

3. Residência de Gráfico de Computação Consciente da Carga (Load Aware Compute Graph Residency):

   • Nem todo especialista deve rodar na NPU. O sistema decide dinamicamente onde executar os grupos de especialistas com base na popularidade estimada.
   • Especialistas "Quentes" (Hot): Grupos de especialistas frequentemente ativados são mantidos residentes na NPU para amortizar o custo de sincronização.
   • Especialistas "Frios" (Cold): Especialistas raramente ativados ou grupos muito pequenos são executados na CPU/GPU para evitar o custo de despacho e sincronização da NPU, que seria maior que o benefício computacional.

3. Contribuições Principais

• Análise de Desafios: Uma investigação profunda das limitações de executar modelos MoE dinâmicos em NPUs móveis, identificando gargalos de sincronização e incompatibilidade de formas de tensor.
• Técnicas Inovadoras: Proposta de três técnicas de co-design de sistema: níveis de capacidade estática, execução agrupada e residência consciente da carga.
• Sistema NPUMoE: Implementação de um motor de inferência de primeira geração que otimiza sistematicamente a eficiência do MoE para NPUs, suportando modelos como Phi-3.5-MoE e Qwen3-MoE.
• Avaliação Abrangente: Demonstração de que é possível alcançar ganhos significativos de desempenho e eficiência energética sem sacrificar a precisão do modelo.

4. Resultados Experimentais

Os testes foram realizados em dispositivos Apple M2 (Max e Ultra) com três modelos MoE representativos e quatro cargas de trabalho de longo contexto.

• Latência: Redução de 1.32x a 5.55x na latência de prefill (TTFT) em comparação com baselines (CoreML CPU, CoreML NPU ingênuo e ANEMLL).
• Eficiência Energética: Melhoria de 1.81x a 7.37x na eficiência energética (energia por token).
• Uso de CPU: Redução de 1.78x a 5.54x no uso de ciclos de CPU, liberando recursos do sistema para outras tarefas.
• Precisão: Perda de precisão negligenciável (< 1.1%), demonstrando que o descarte de tokens e o padding não afetam significativamente a qualidade da resposta.
• Comparação com GPU: Embora a GPU seja mais rápida em picos, ela consome muito mais energia e potência. O NPUMoE foca em eficiência sustentada e liberação da CPU/GPU para tarefas interativas.

5. Significado e Impacto

Este trabalho é significativo porque:

1. Desbloqueia o MoE em Dispositivos Móveis: Demonstra que modelos MoE grandes e esparsos podem rodar eficientemente em hardware móvel padrão, algo que antes era considerado inviável devido às restrições de NPUs.
2. Otimização de Recursos do Sistema: Ao transferir a carga computacional pesada para a NPU (dedicada) e manter a CPU livre, o sistema melhora a responsividade geral do dispositivo (UI, I/O, redes).
3. Paradigma de Co-design: Estabelece um novo paradigma de como adaptar modelos de IA dinâmicos a hardware estático através de técnicas de calibração offline e agrupamento inteligente, em vez de depender apenas de hardware mais potente.
4. Sustentabilidade: A melhoria drástica na eficiência energética é crucial para a viabilidade de LLMs em dispositivos alimentados por bateria.

Em resumo, o NPUMoE prova que, através de um design de sistema cuidadoso, é possível superar a incompatibilidade fundamental entre a natureza dinâmica dos modelos MoE e as restrições estáticas das NPUs móveis, tornando a inferência de LLMs de ponta acessível e eficiente em dispositivos pessoais.