WiSP: A Working-Set View of Mixture-of-Experts Serving on Extremely Low-Resource Hardware

Resumen Técnico: WiSP – Una visión del conjunto de trabajo para el servicio de Mixture-of-Experts en hardware de recursos extremadamente bajos

Planteamiento del problema
Los modelos modernos de Mixture-of-Experts (MoE), como Qwen3-30B-A3B y Mixtral-8×7B, colocan la mayor parte de sus parámetros en capas de expertos, aunque solo una pequeña fracción de estos expertos está activa para cualquier token dado. En GPUs de consumo con VRAM limitada, estos modelos a menudo no pueden caber enteramente en la memoria. La solución estándar, el volcado de capas estático a CPU (por ejemplo, --cpu-offload-gb en vLLM), mantiene bajo el uso de memoria pero transfiere todos los expertos de una capa a través del bus PCIe en cada pasada hacia adelante (forward pass). Este enfoque falla al no explotar la escasez de los modelos MoE, lo que efectivamente anula sus beneficios de eficiencia.

Los autores enmarcan el servicio de MoE de bajos recursos como un problema de gestión del conjunto de trabajo (working-set). En este régimen, la memoria GPU es el almacén rápido limitado, y la memoria del host es el almacén de respaldo. Dos flujos distintos de demanda de memoria compiten por el mismo presupuesto de VRAM:

1. Pesos de los expertos: Un flujo dependiente del enrutamiento donde solo se necesita un subconjunto de expertos por token.
2. KV Cache: Un flujo dependiente del token requerido para la generación autorregresiva.

Los sistemas existentes suelen tratar estos como problemas separados o se dirigen a entornos de centro de datos de alta concurrencia donde la computación puede solaparse con la transferencia de datos. Sin embargo, en el servicio interactivo de baja concurrencia (por ejemplo, asistentes de codificación locales), la computación es mínima por paso, lo que hace que el ancho de banda de PCIe sea el cuello de botella principal, en lugar de la computación o la precisión de la predicción.

Metodología: WiSP (Paginación de Conjunto de Trabajo)
El artículo presenta WiSP, un sistema que trata el conjunto de expertos residentes como un caché sobre el comportamiento de enrutamiento del modelo, gestionado conjuntamente con el KV cache. WiSP consta de dos componentes principales:

1. Paginador de Expertos Consciente del Enrutamiento (Routing-Aware Expert Pager):

   • Mecanismo: WiSP se integra en un motor de servicio sin modificaciones (por ejemplo, vLLM) aprovechando un gancho de indirección de mapa de expertos existente (π). En lugar de cargar capas completas, mantiene un pequeño espacio de trabajo en GPU que contiene solo el "conjunto de trabajo" de expertos que probablemente se reutilizarán.
   • Operación: En cada paso, el paginador busca solo los expertos enrutados faltantes desde la memoria del host (H2D) y expulsa los expertos menos recientemente utilizados (LRU) para hacer espacio.
   • Garantías: Debido a que los kernels y los pesos subyacentes permanecen sin cambios (solo se redirige el mapeo), WiSP produce salidas byte idénticas a las del motor base. No requiere nuevos kernels, cambios de cuantización ni reentrenamiento del modelo.

2. MV-WSA (Asignación de Conjunto de Trabajo de Valor Marginal):

   • El Problema de la Asignación: Dado que los pesos de los expertos y el KV cache comparten el mismo presupuesto de VRAM, aumentar el conjunto de expertos residente reduce la capacidad del KV cache, y viceversa.
   • Estrategia: MV-WSA determina la división óptima ($f$) entre el presupuesto de expertos y el presupuesto de KV. Su objetivo es igualar el beneficio de latencia marginal por byte en ambos grupos.
   • Piso de Admisión: Crucialmente, el asignador impone un tamaño mínimo de KV cache ($\kappa_{min}$) requerido para admitir el nivel de concurrencia actual. Si la división óptima no restringida agotara el KV cache, el sistema limita la división para asegurar que el programador pueda inicializarse y ejecutarse.
   • Modos de Implementación:
     • Configurador Offline: Analiza una traza representativa para establecer una división fija al inicio.
     • Controlador Online: Redimensiona dinámicamente ambos grupos durante el servicio en "barreras de drenaje" (puntos donde ninguna solicitud activa retiene bloques de memoria), moviendo físicamente los bytes entre el espacio de trabajo de expertos y el grupo de KV sin exceder el presupuesto total.

Hallazgos Clave sobre Prefetching
Un hallazgo empírico central del artículo es que el prefetching basado en enrutamiento proporciona poco o ningún beneficio en el decode de un solo flujo (single-stream), de baja concurrencia.

• Cuello de botella: El rendimiento está limitado por el ancho de banda de PCIe, no por la precisión de la predicción. En el decode de batch-1, la computación por paso es demasiado pequeña para ocultar las transferencias de PCIe detrás de ella.
• Resultado: El prefetching especulativo a menudo degrada el rendimiento (reduciendo el throughput aproximadamente un 50% en las pruebas) porque compite por el ancho de banda y expulsa expertos necesarios para el siguiente paso inmediato.
• Reencuadre: La señal de enrutamiento es valiosa no para ocultar la latencia, sino para la asignación de memoria. Ayuda a identificar el tamaño mínimo del conjunto residente "seguro", permitiendo al sistema dimensionar correctamente el caché de expertos por usuario y asignar el VRAM liberado al KV cache.

Resultados
Los autores evalúan WiSP en varias arquitectes de MoE (LLMs, VLMs, Híbridos SSM-MoE y modelos de Difusión) utilizando una NVIDIA H100 (con algunos experimentos emulando presupuestos de memoria más pequeños).

• Ganancias de Throughput: Bajo presupuestos de VRAM emparejados donde el modelo no cabe, WiSP logra hasta un 1.95× de throughput de decode en comparación con el volcado estático a CPU. La ganancia es mayor en presupuestos de memoria intermedios donde WiSP puede cachear una parte significativa del conjunto de expertos reutilizados mientras que el modelo base sigue transmitiendo todo.
• Capacidad de KV: Al reducir la huella de los expertos, WiSP aumenta directamente la capacidad del KV cache. Los experimentos muestran que reducir el tope de expertos puede aumentar la concurrencia máxima en más de un 40% con el mismo presupuesto de memoria total.
• Rendimiento de la Asignación:
  • Offline: El configurador MV-WSA funciona bien tanto en prefill como en decode, manteniéndose a pocos puntos porcentuales de un oráculo por flujo de trabajo en simulación, mientras que las divisiones fijas (por ejemplo, 50/50) son aproximadamente un 20% peores.
  • Online: El controlador de redimensionamiento dual dinámico mejora el throughput en 1.20× sobre la mejor división estática offline al adaptarse a los cambios de carga de trabajo (por ejemplo, pasando de turnos de contexto largo a ráfagas de llamadas a herramientas).
• Generalidad: WiSP sirve con éxito modelos que no caben en una sola GPU (por ejemplo, Jamba-v0.1, MiniMax-M2) y funciona a través de diversas modalidades incluyendo Modelos de Lenguaje-Visión (Kimi-VL) y modelos de Difusión (DiT-MoE, LLaDA-MoE), manteniendo salidas byte idénticas.

Significancia y Reivindicaciones
El artículo afirma que WiSP proporciona un marco unificado para el servicio de MoE de bajos recursos al tratar la residencia de expertos y el cache de KV como un único problema de asignación de conjunto de trabajo.

• Cambio de Paradigma de Diseño: Argumenta que para el servicio de baja concurrencia, el enfoque debe cambiar de "¿cómo prefetch de expertos más rápido?" a "¿cómo asignar VRAM óptimamente entre expertos y KV cache?".
• Impacto Práctico: WiSP permite el servicio de modelos MoE de escala de frontera en hardware de consumo sin sacrificar la corrección de la salida o requerir hardware especializado (como Intel AMX).
• Limitaciones: Los autores señalan explícitamente que WiSP está diseñado para el régimen donde el modelo no cabe en la VRAM; si el modelo cabe, la residencia estática sigue siendo más rápida ya que el paginado añade una indirección innecesaria. Además, se muestra que la optimización de prefetching es ineficaz en el decode de un solo flujo, aunque puede ser relevante para sesiones de agentes de múltiples turnos o de alto batch donde el solapamiento de computación es posible.

El trabajo concluye que, en el decode de un solo flujo limitado por memoria, la señal de enrutamiento es una palanca de memoria (que determina el tamaño del conjunto de trabajo) en lugar de una palanca de velocidad (que oculta la latencia), haciendo que la asignación inteligente sea el principal motor de rendimiento.