La inferencia eficiente de Modelos de Lenguaje Grandes (LLMs) en GPUs a escala de hyperscaler presenta un desafío fundamental en la gestión de recursos de memoria y ancho de banda. Modelos como Kimi y GLM, con arquitecturas Mixture-of-Experts y contextos largos, demandan una cantidad significativa de memoria para sus pesos y, crucialmente, para el KV cache (Key-Value cache) que almacena las representaciones de atención de tokens previamente procesados. Este KV cache, que crece linealmente con la longitud del contexto, a menudo se convierte en el cuello de botella de memoria antes que los propios pesos del modelo, limitando la concurrencia y, por ende, el throughput total del sistema.
La necesidad de servir a millones de usuarios con latencias bajas y costos controlados impulsa la búsqueda de optimizaciones que permitan maximizar la utilización de cada GPU. Esto implica no solo reducir la huella de memoria de los componentes del modelo, sino también gestionar de forma robusta los recursos compartidos en entornos de alta concurrencia. La solución a este problema fundamental de la computación distribuida y el rendimiento de hardware especializado radica en la aplicación de técnicas de cuantificación y compresión, junto con mecanismos de seguridad para la integridad de datos, que permitan un uso más denso y seguro de la memoria de la GPU.
Arquitectura del Sistema
La arquitectura de inferencia descrita se basa en un diseño disociado (disaggregated) que separa las fases de 'prefill' y 'decode' de la inferencia de LLMs, permitiendo aplicar optimizaciones específicas a cada fase. El sistema utiliza SGLang como framework de serving de inferencia, lo que facilita la implementación de técnicas avanzadas de gestión de memoria y scheduling.
Para el KV cache, se implementa la cuantificación de BF16 a FP8 (e4m3). Esta técnica reduce el tamaño del cache a la mitad, duplicando la capacidad de tokens en memoria. Aunque la cuantificación introduce una pequeña sobrecarga computacional por token debido a la conversión en el kernel de atención, el beneficio principal es el aumento significativo en la cantidad de solicitudes concurrentes que pueden residir en la GPU, lo que resulta en un mayor throughput global. La fase de 'prefill', al ser compute-bound, mantiene el KV cache en BF16 para maximizar su throughput ligeramente superior, mientras que la fase de 'decode' utiliza FP8 para optimizar la concurrencia.
Para los pesos del modelo, se aplica la compresión de FP8 a INT4. Esta compresión reduce el tamaño de los pesos en aproximadamente un 40%, liberando memoria adicional en la GPU. En la fase de 'decode', la reducción del tamaño de los pesos mejora la velocidad al disminuir la cantidad de datos que deben ser transmitidos desde la memoria de la GPU, ya que la velocidad de 'decode' está limitada por el ancho de banda de memoria. Similar al KV cache, la fase de 'prefill' utiliza pesos FP8, ya que la expansión de INT4 a FP8 antes de las operaciones de multiplicación introduce una penalización de rendimiento en esta fase compute-bound. Finalmente, para proteger la integridad del KV cache compartido en un entorno de alta concurrencia, se implementa un mecanismo de verificación de integridad. Cada página física del cache recibe una etiqueta que cambia con la reasignación, y el servidor registra las páginas y etiquetas esperadas por cada solicitud. Antes de las operaciones de 'decode', estas asignaciones se verifican, abortando las solicitudes con inconsistencias para prevenir el retorno de datos incorrectos. Este chequeo se realiza como una verificación por lotes separada para evitar condiciones de carrera en el kernel de atención, con un costo de rendimiento inferior al 1%.
Flujo de Inferencia Optimizado con Cuantificación
- 1 Solicitud de Inferencia Usuario envía una solicitud a Workers AI para un LLM.
- 2 Fase de Prefill (BF16) Procesamiento inicial del contexto, compute-bound. KV cache en BF16.
- 3 Fase de Decode (FP8/INT4) Generación de tokens, memory-bandwidth-bound. KV cache en FP8, pesos en INT4.
- 4 Verificación KV Cache Chequeo de integridad de páginas del KV cache antes de lectura.
- 5 Generación de Token Modelo genera el siguiente token usando KV cache y pesos.
- 6 Actualización KV Cache Nuevo token y sus K/V se añaden al KV cache.
- 7 Respuesta al Usuario Tokens generados se envían de vuelta al usuario.
| Capa | Tecnología | Justificación |
|---|---|---|
| compute | NVIDIA H200 GPUs | Hardware principal para la ejecución de la inferencia de LLMs. Despliegue disociado (disaggregated) para prefill y decode. |
| data-processing | SGLang | Framework de serving de inferencia open-source para LLMs, utilizado para gestionar y optimizar la ejecución de modelos. Colaboración con el equipo de SGLang para upstream de parches y características. |
| storage | KV Cache (FP8) | Almacenamiento de claves y valores de atención para tokens procesados, cuantificado a 8-bit floating point (e4m3) para reducir el consumo de memoria. vs BF16 KV Cache Usado en fase de decode; BF16 en fase de prefill. |
| storage | Model Weights (INT4) | Pesos del modelo comprimidos a 4-bit integers para reducir el footprint de memoria y mejorar el ancho de banda. vs FP8 Model Weights Usado en fase de decode; FP8 en fase de prefill. |
Trade-offs
Ganancias
- ▲▲ Capacidad de tokens en KV cache
- ▲▲ Concurrencia de solicitudes en GPU
- ▲ Throughput total del sistema
- ▲ Costo por token
- ▲ Latencia de decode (con INT4 weights)
Costes
- △ Throughput por token (KV cache FP8 vs BF16)
- △ Throughput de prefill (INT4 weights vs FP8)
- △ Costo de verificación de integridad (throughput y p95 latency)
Fundamentos Teóricos
El problema de la gestión eficiente de memoria en sistemas de computación de alto rendimiento, especialmente con estructuras de datos grandes y dinámicas como el KV cache en LLMs, tiene raíces profundas en la investigación de sistemas operativos y bases de datos. Conceptos como la paginación de memoria virtual, la gestión de caché y la compresión de datos han sido estudiados extensivamente. La cuantificación de modelos y datos, en particular, se ha explorado en el campo del aprendizaje automático y el procesamiento de señales durante décadas, buscando reducir el footprint de memoria y la latencia computacional a expensas de una posible pérdida de precisión.
La idea de cuantificar pesos y activaciones en redes neuronales para reducir el consumo de memoria y acelerar la inferencia se remonta a trabajos como 'Quantized Neural Networks: Training Neural Networks with Low-Precision Weights and Activations' de Courbariaux et al. (2016). Más recientemente, la aplicación de formatos de punto flotante de menor precisión como FP8 ha sido objeto de investigación activa, con papers como 'FP8 versus INT8: Empirical Study of Low-Precision Quantization for LLM Inference' de Dettmers et al. (2022) explorando sus implicaciones en la precisión y el rendimiento. La separación de fases 'prefill' y 'decode' y la gestión de paged attention se basan en principios de scheduling y gestión de recursos que optimizan el uso de hardware, análogos a cómo los sistemas operativos gestionan la memoria y el tiempo de CPU para múltiples procesos concurrentes.