ANE Inference Engine: Como usei ANE + GPU em paralelo no Apple Silicon

Rodar modelos de linguagem grandes (LLMs) localmente é um problema de largura de banda de memória, não de poder de computação. O gargalo é o quanto de dados você consegue mover por segundo entre a memória e os chips de processamento.

No speculative decoding, a ideia é usar um modelo menor (draft model) pra sugerir os próximos tokens, e o modelo principal só precisa verificar, não gerar do zero. Isso é muito mais rápido — quando o draft acerta.

O problema: toda implementação convencional roda o draft model e o main model sequencialmente, ou no mesmo chip. O Apple Silicon tem uma vantagem enorme que ninguém estava usando: três chips completamente separados, cada um com seu próprio canal de memória.

O Apple Silicon (M1/M2/M3/M4) tem três unidades de processamento com canais DCS (Die-to-Die Coherency Streaming) completamente independentes:

• •CPU — 8 GB/s de bandwidth de memória
• •GPU (Metal) — 55 GB/s de bandwidth de memória
• •ANE (Apple Neural Engine) — 22 GB/s teórico, 25.7 GB/s real medido no M4

Isso significa que rodar ANE e GPU simultaneamente NÃO causa contenção de memória. Os dois acessam a RAM por caminhos físicos diferentes. Medido empiricamente: rodar a ANE enquanto o GPU está inferindo causa apenas 11% de degradação no throughput do GPU — e a ANE fica ativa apenas 12% do ciclo total.

Essa é a base do projeto: enquanto o GPU verifica os tokens que o draft sugeriu, a ANE já está gerando os candidatos para o próximo ciclo.

A primeira implementação ingênua deu 0.61x — mais lenta que GPU alone. Quatro otimizações mudaram tudo:

1. Merged Bonus Eval (economia de ~109ms por ciclo)

O speculative decoding padrão faz 2 chamadas separadas ao GPU: uma pra verificar os draft tokens, outra pra calcular o bonus token. Fundir essas duas chamadas em uma só eliminou 109ms por ciclo. Parece pouco, mas em tokens por segundo faz diferença absurda.

2. ANE/GPU Pipelined Overlap (70%+ de prefetch hit rate)

Enquanto o GPU verifica os tokens do ciclo N, a ANE já está gerando os candidatos para o ciclo N+1. Isso é overlap real de computação — os dois chips trabalhando ao mesmo tempo em coisas diferentes. Com warm-up adequado, 70% dos ciclos aproveitam o prefetch da ANE.

3. INT8 Quantization na ANE

O modelo draft em CoreML com INT8 reduziu a latência de inferência ANE de 26ms para 18ms. Menor precisão, mas o draft model é pequeno (0.5B params) e a taxa de aceitação se manteve em 74% — muito acima do limiar de break-even.

4. Redução de Sequence Length 128→64

O CoreML não tem KV cache. Cada chamada é um recompute completo da atenção. Reduzir o contexto de entrada de 128 para 64 tokens cortou a latência de 18ms para 14ms. N=1 speculative depth é ótimo por causa disso — mais profundidade aumenta a latência da ANE sem ganho proporcional.

Tentei três abordagens que pareciam promissoras e falharam:

Distillation de modelos menores (43M e 80M parâmetros)

O vocabulário do Qwen2.5 tem 151k tokens. Só a camada de embedding + lm_head consome 90% dos parâmetros em um modelo de 43M. O resultado: acceptance rate de apenas 9-10%, completamente inviável para speculative decoding. O modelo tiny não tem capacidade de modelar a distribuição de tokens do modelo grande.

PARD tokens (Parallel Draft Tokens)

Uma técnica experimental onde tokens especiais são gerados em paralelo. No meu setup: 0.3% de acceptance rate. Completamente inútil com a arquitetura atual.

Layer removal sem fine-tuning

Remover camadas do modelo base pra criar um draft mais rápido sem retreinar. Sem fine-tuning pós-remoção, a qualidade cai tanto que o acceptance rate vai pra menos de 20%. O overhead de chamadas ao GPU não compensa.

Uma das descobertas mais interessantes: a eficiência do setup varia muito com o tamanho do main model:

• •1.5B main model → 0.83x (mais lento que baseline)
• •7B main model → 0.43x (absurdamente mais lento)
• •14B main model → 1.14x (sweet spot atual)
• •32B main model → 1.7x projetado (ainda não testado)

Por que modelos maiores se beneficiam mais? Porque o overhead fixo da ANE (setup, pipeline, memory copy) representa uma fração menor do tempo total quando o GPU está verificando um modelo de 14B ou 32B. Com modelos pequenos, o overhead domina o tempo de ciclo.

O 32B projetado em 1.7x seria transformador — um modelo de qualidade GPT-4 nível rodando localmente a velocidade de um modelo menor.