Inferencia de LLM MoE en MacBook Pro: Streaming de Expertos desde SSD con Metal

La inferencia de modelos de lenguaje grandes (LLMs) con miles de millones de parámetros presenta un desafío fundamental en la computación: cómo gestionar y acceder eficientemente a modelos que exceden la capacidad de la memoria RAM o VRAM disponible. Los modelos Mixture-of-Experts (MoE) exacerban esto al requerir el acceso selectivo a subconjuntos de pesos (expertos) por cada token, lo que implica un patrón de acceso a datos disperso y de alta latencia si no se gestiona correctamente. Este trabajo aborda el problema de ejecutar un LLM MoE de 397B parámetros en un dispositivo con recursos limitados (48GB de RAM unificada) mediante una estrategia de streaming de expertos desde el almacenamiento persistente (SSD) directamente a la unidad de procesamiento gráfico (GPU).

La relevancia de este enfoque radica en democratizar el acceso a modelos de gran escala, permitiendo su ejecución en hardware de consumo sin la necesidad de costosos clústeres de GPUs. Históricamente, la ejecución de modelos que no caben en memoria ha requerido técnicas como la cuantificación, la poda o el offloading a disco, a menudo con penalizaciones significativas en rendimiento o calidad. Este proyecto demuestra que, con una ingeniería de sistemas cuidadosa y optimizaciones de bajo nivel, es posible lograr un rendimiento competitivo y una calidad de salida de producción, incluso para tareas complejas como la llamada a herramientas (tool calling), que son sensibles a la precisión de la cuantificación.

Arquitectura del Sistema

La arquitectura del sistema se centra en un motor de inferencia escrito en C/Objective-C con kernels de Metal personalizados, eliminando dependencias de frameworks de alto nivel. El modelo Qwen3.5-397B-A17B, con 397 mil millones de parámetros, se almacena en el SSD (209GB en 4-bit). Los pesos de los expertos se cargan bajo demanda desde el SSD utilizando pread() en paralelo con GCD dispatch groups. Solo los K=4 expertos activos por capa se cargan, cada uno de aproximadamente 6.75MB. El sistema confía en el page cache del sistema operativo (aproximadamente 35GB) para gestionar el almacenamiento en caché de los expertos, evitando la complejidad y el overhead de una caché personalizada.

El pipeline de cómputo en la GPU se implementa con shaders de Metal escritos a mano. Estos incluyen kernels optimizados para la multiplicación matriz-vector de 4-bit y 2-bit dequantizada (tiled, SIMD-reduced, con caché de entrada compartida y optimización FMA), activación Fused SwiGLU, normalización RMS, atención GPU por lotes y RoPE fusionado con deinterleave de Q y normalización de K. La combinación de expertos y el residual se realizan en un kernel fusionado en la GPU. Una optimización clave es la ejecución diferida de los pases de expertos (CMD3), donde el cómputo de la GPU se envía sin esperar, permitiendo que la GPU trabaje mientras la CPU prepara la siguiente capa. Para las capas de atención lineal GatedDeltaNet, se utiliza la librería Accelerate BLAS (cblas_sscal, cblas_sgemv, cblas_sger) para la actualización de la matriz de estado, logrando una mejora del 64% sobre el código escalar. La interacción entre GPU y SSD se serializa debido a que el DMA del SSD y el cómputo de la GPU comparten el mismo controlador de memoria unificada, lo que hace que un pipeline serial (GPU → SSD → GPU) sea óptimo para evitar picos de latencia.

```objective-c
// Simplified representation
dispatch_group_t group = dispatch_group_create();
for (int i = 0; i < K_ACTIVE_EXPERTS; ++i) {
    dispatch_group_async(group, dispatch_get_global_queue(QOS_CLASS_USER_INITIATED, 0), ^{
        // Calculate offset and size for expert i
        pread(fd, buffer_for_expert_i, size_i, offset_i);
    });
}
dispatch_group_wait(group, DISPATCH_TIME_FOREVER);
dispatch_release(group);
```

```metal
// Original concept: (nibble * scale + bias) * x
// FMA-optimized:
float scale_x = scale * x;
float bias_x = bias * x;
output = fma(nibble_float, scale_x, bias_x);
```

Fundamentos Teóricos

Este trabajo se conecta con varios principios fundamentales de la computación y la arquitectura de sistemas. El concepto de cargar datos bajo demanda desde almacenamiento secundario para modelos que exceden la memoria principal es una aplicación directa de las técnicas de paginación y memoria virtual, fundamentales en sistemas operativos desde los años 60. La decisión de 'confiar en el OS' para el page cache, en lugar de implementar una caché personalizada, resuena con el principio de que los sistemas operativos modernos están altamente optimizados para estas tareas, como se discute en trabajos clásicos sobre gestión de memoria y E/S. La optimización del kernel de dequantización con FMA (Fused Multiply-Add) es un ejemplo de cómo la comprensión profunda de la arquitectura del procesador y sus unidades funcionales (como las FMA units presentes en CPUs y GPUs modernas) puede llevar a mejoras significativas en el rendimiento, un tema recurrente en la optimización de algoritmos numéricos y procesamiento de señales.

La estrategia de streaming de expertos desde SSD se inspira en el paper 'LLM in a Flash' de Apple, que explora cómo la memoria unificada y el ancho de banda del SSD en Apple Silicon pueden ser explotados para la inferencia de LLMs. La gestión de la concurrencia y la sincronización entre CPU, GPU y operaciones de E/S, utilizando GCD dispatch groups y la ejecución diferida de kernels, se basa en principios de programación concurrente y pipelines de procesamiento, buscando maximizar la utilización de recursos heterogéneos. La elección de algoritmos eficientes para operaciones matriciales, como el uso de BLAS para la atención lineal, subraya la importancia de las librerías numéricas optimizadas, un campo de estudio que se remonta a los trabajos pioneros en álgebra lineal computacional.