Espresso: inferencia directa en el Neural Engine de Apple sin pasar por CoreML

Fuentes: Espresso: direct transformer inference on Apple's Neural Engine without CoreML

Espresso es un proyecto de código abierto escrito en Swift 6.2 que permite entrenar y ejecutar modelos transformer directamente sobre el Neural Engine (ANE) de los chips Apple Silicon, esquivando por completo la capa CoreML. Para ello, la herramienta accede a APIs privadas del sistema mediante dlopen (_ANEClient y _ANEInMemoryModel), compila programas MIL y los traduce a binarios E5 que el ANE ejecuta de forma nativa. El resultado, según las mediciones del repositorio, es una velocidad de decodificación de 1,08 ms por token (926 tokens/s) frente a los 5,09 ms (196 tokens/s) de CoreML en el mismo modelo de seis capas, lo que supone una mejora de 4,76 veces. Frente a llama.cpp en Metal, la ganancia ronda las 11 veces.

La clave del rendimiento está en tres decisiones técnicas. Primera, el uso de kernels fusionados de tres capas: en lugar de enviar seis dispatch al ANE, un modelo de seis capas transformer se reduce a dos llamadas. Segunda, la lectura y escritura sin copia (zero-copy) mediante buffers IOSurface, con operaciones NEON y argmax por vDSP, lo que evita serializar datos entre CPU, GPU y ANE. Tercera, un bucle de decodificación que compila el programa una sola vez y lo reutiliza en todos los pasos, con caché KV residente en IOSurface y verificación de paridad exacta de dos tokens por paso.

Espresso ofrece además entrenamiento completo sobre el ANE, con pasada hacia delante y hacia atrás, acumulación de gradientes y optimizador Adam, algo poco habitual fuera de los frameworks oficiales de Apple. La arquitectura interna se divide en módulos: ANEInterop (puente con las APIs privadas en ObjC/C), ANETypes (tensores ~Copyable e I/O por superficies), MILGenerator (más de 28 variantes de kernel), CPUOps (RMSNorm, RoPE, embedding, softmax y Adam vía Accelerate/vDSP) y ANERuntime, que gestiona la compilación MIL a E5 y el presupuesto de compilación del sistema.

La plataforma de modelos se articula en torno a un formato portátil .esp y a una caché compilada .espc, con herramientas CLI (espc pack-native, esprun inspect, esprun generate) que empaquetan, inspeccionan y ejecutan modelos. En cuanto a compatibilidad, el proyecto requiere macOS 15 o superior, Swift 6.0+, Apple Silicon con Neural Engine y no tiene dependencias externas. Se ha validado en M1, M2, M3 y M4; los Mac Intel quedan fuera por carecer de ANE. iOS y tvOS no están soportados de serie, ya que los entitlements de las APIs privadas difieren entre plataformas. Por su naturaleza —uso de APIs no públicas—, la librería no es apta para distribuir en la App Store.