Motor de Inferencia Bodega: Optimizando la Inferencia de LLM para la Memoria Unificada de Apple Silicon

Bodega es un motor de inferencia diseñado específicamente para la arquitectura de memoria unificada de Apple Silicon, desarrollado durante más de 2.5 años con optimizaciones cercanas a la capa Metal en MLX. Aborda las limitaciones fundamentales de rendimiento que enfrentan los desarrolladores al ejecutar LLMs en hardware Mac.

Por qué Apple Silicon requiere una optimización diferente

Apple Silicon utiliza memoria unificada donde la CPU, GPU y el motor neuronal comparten un solo grupo físico a través de un único bus en el chip. Esto difiere fundamentalmente de las GPU discretas como las de NVIDIA, que tienen grupos separados de VRAM y RAM del sistema conectados por PCIe. El ancho de banda de memoria varía desde ~400 GB/s en M1 Max hasta ~800 GB/s en M3 Ultra (con penalización entre chips que reduce el rendimiento real a 1.6-1.8 veces el rendimiento de un solo chip).

Implicaciones arquitectónicas clave:

• La decodificación está limitada por el ancho de banda de memoria: cada token requiere cargar los pesos del modelo desde el bus compartido
• El prellenado está limitado por el cómputo: dominado por los TFLOPS de la GPU para la multiplicación matriz-matriz
• El bus de memoria se comparte con todo: la caché KV, los pesos del modelo, el sistema operativo y las aplicaciones compiten por el mismo ancho de banda de 400-800 GB/s

Esta arquitectura hace que las adaptaciones directas de las implementaciones de procesamiento por lotes de vLLM o llama.cpp sean ineficaces en MLX, ya que fueron diseñadas para diferentes arquitecturas de memoria.

Lo que construye Bodega

El desarrollador estudió los componentes internos centrales de vLLM, incluido el procesamiento por lotes continuo, la decodificación especulativa, el prellenado por fragmentos y el almacenamiento en caché de prefijos, luego reconstruyó cada componente para MLX y el modelo de memoria unificada de Apple.

La idea central para el procesamiento por lotes continuo: generar un solo token para una sola secuencia carga todos los pesos del modelo para una multiplicación matriz-vector, lo cual es ineficiente en hardware con ancho de banda de 400+ GB/s. La solución ejecuta múltiples secuencias simultáneamente usando pesos × matriz de vectores en lugar de pesos × vector único.

La gestión de la caché KV fue rediseñada para memoria unificada, donde evadir bloques de caché tiene implicaciones de costo diferentes en comparación con sistemas VRAM aislados.

Implicaciones prácticas

El desarrollador informa haber probado en múltiples configuraciones de Apple Silicon, incluyendo dos M3 Ultra (256GB y 512GB), un M4 Max 128GB y un M1 Max 64GB. El límite común identificado es el rendimiento de usuario único con una solicitud a la vez y la GPU mayormente inactiva.

El repositorio incluye puntos de referencia que pueden verificarse con un simple script curl para la configuración.