Bodega Inference Engine: Otimizando a Inferência de LLM para a Memória Unificada do Apple Silicon

Bodega é um mecanismo de inferência projetado especificamente para a arquitetura de memória unificada do Apple Silicon, construído ao longo de 2,5 anos com otimizações próximas à camada Metal no MLX. Ele aborda as limitações fundamentais de throughput que os desenvolvedores enfrentam ao executar LLMs no hardware Mac.

Por que o Apple Silicon Requer Otimização Diferente

O Apple Silicon usa memória unificada onde CPU, GPU e neural engine compartilham um pool físico único por meio de um barramento on-chip. Isso difere fundamentalmente de GPUs discretas como as da NVIDIA, que possuem pools separados de VRAM e RAM do sistema conectados por PCIe. A largura de banda da memória varia de ~400 GB/s no M1 Max a ~800 GB/s no M3 Ultra (com penalidade cross-die reduzindo o throughput real para 1,6-1,8x o desempenho de single-die).

Implicações arquiteturais principais:

• Decodificação é limitada por largura de banda de memória - cada token requer carregar pesos do modelo do barramento compartilhado
• Prefill é limitado por computação - dominado por TFLOPS da GPU para multiplicação matriz-matriz
• O barramento de memória é compartilhado com tudo - cache KV, pesos do modelo, sistema operacional e aplicativos competem pela mesma largura de banda de 400-800 GB/s

Essa arquitetura torna portes diretos das implementações de batching do vLLM ou llama.cpp ineficazes no MLX, pois foram projetados para arquiteturas de memória diferentes.

O que o Bodega Constrói

O desenvolvedor estudou os componentes internos principais do vLLM, incluindo batch contínuo, decodificação especulativa, prefill em blocos e cache de prefixo, então reconstruiu cada componente para o MLX e o modelo de memória unificada da Apple.

A percepção central para batch contínuo: gerar um único token para uma única sequência carrega todos os pesos do modelo para uma multiplicação matriz-vetor, o que é ineficiente em hardware com largura de banda de 400+ GB/s. A solução executa múltiplas sequências simultaneamente usando pesos × matriz de vetores em vez de pesos × vetor único.

O gerenciamento do cache KV foi redesenhado para memória unificada, onde evacuar blocos de cache tem implicações de custo diferentes comparado a sistemas VRAM isolados.

Implicações Práticas

O desenvolvedor relata testes em múltiplas configurações do Apple Silicon, incluindo dois M3 Ultra (256GB e 512GB), um M4 Max 128GB e um M1 Max 64GB. O limite comum identificado é o throughput de usuário único com uma requisição por vez e GPU ficando majoritariamente ociosa.

O repositório inclui benchmarks que podem ser verificados com um script curl simples para configuração.