Optimización de Inferencia de LLMs en GPUs: Cuantificación y Compresión para Modelos de Contexto Largo

La inferencia eficiente de Modelos de Lenguaje Grandes (LLMs), especialmente aquellos con arquitecturas de Mixture-of-Experts (MoE) y capacidades de contexto extendido, presenta desafíos significativos en entornos de producción a escala. El problema fundamental radica en la gestión de la memoria de la GPU, que se convierte en el cuello de botella principal, limitando tanto la cantidad de contexto que un modelo puede procesar como el número de solicitudes concurrentes que un hardware puede manejar. Tradicionalmente, la precisión de los datos (BF16 para el KV cache y FP8 para los pesos) ha sido un estándar, pero para modelos de miles de millones de parámetros y contextos de millones de tokens, esta precisión consume rápidamente la memoria disponible.

Este artículo aborda este problema mediante la aplicación de técnicas de cuantificación y compresión. La cuantificación del KV cache a FP8 y la compresión de los pesos del modelo a INT4 son estrategias directas para reducir la huella de memoria, permitiendo una mayor concurrencia y, en algunos casos, mejorando el throughput al reducir la presión sobre el ancho de banda de memoria. La necesidad de estas optimizaciones se ha vuelto crítica con la proliferación de LLMs más grandes y la demanda de latencias bajas en aplicaciones interactivas, empujando los límites de la infraestructura de inferencia existente.

Arquitectura del Sistema

La arquitectura de inferencia descrita se basa en un diseño disociado que separa las fases de 'prefill' (procesamiento inicial del contexto) y 'decode' (generación secuencial de tokens). Esta separación es crucial porque las fases tienen diferentes cuellos de botella: 'prefill' es típicamente compute-bound, mientras que 'decode' es memory-bound. Esta distinción permite aplicar optimizaciones específicas a cada fase sin comprometer el rendimiento general.

Para el KV cache, se implementa la cuantificación de BF16 a FP8 (e4m3). Esto reduce el tamaño del cache a la mitad, duplicando la capacidad de contexto en memoria. Aunque la cuantificación introduce una pequeña sobrecarga computacional por token en el kernel de atención (conversión de FP8 a BF16), el beneficio principal es la capacidad de mantener más solicitudes concurrentes residentes en la GPU, lo que resulta en un mayor throughput global en la fase de 'decode'. Para la fase de 'prefill', donde el rendimiento por token es crítico y la memoria no es el factor limitante, el KV cache se mantiene en BF16.

Para los pesos del modelo, se aplica la compresión de FP8 a INT4. Esta técnica reduce el tamaño de los pesos en aproximadamente un 40%, liberando una cantidad significativa de memoria de GPU. En la fase de 'decode', donde el ancho de banda de memoria es el cuello de botella, la reducción del volumen de datos a mover resulta en una mejora sustancial del throughput. Sin embargo, en la fase de 'prefill', la necesidad de expandir los pesos INT4 de nuevo a una precisión mayor para las operaciones de multiplicación introduce una sobrecarga computacional, haciendo que esta fase sea más lenta. Por lo tanto, los pesos INT4 se utilizan para 'decode', mientras que los pesos FP8 originales se mantienen para 'prefill'.

Finalmente, para mitigar los riesgos asociados con la alta concurrencia y el uso compartido de memoria, se introduce un sistema de verificación de integridad del KV cache. Cada página física del cache recibe una etiqueta que cambia con la reasignación, y el servidor rastrea qué páginas y etiquetas espera usar cada solicitud. Antes de las operaciones de lectura, estas asignaciones se verifican en un chequeo por lotes separado, no fusionado con el kernel de atención para evitar condiciones de carrera. Si hay una discrepancia, la solicitud afectada se aborta. Esta verificación tiene un costo marginal (<1% en throughput y latencia p95) y es configurable por despliegue.

Fundamentos Teóricos

La cuantificación de modelos y datos para reducir la huella de memoria y mejorar la eficiencia computacional es un campo de investigación activo en el aprendizaje automático, con raíces en la compresión de datos y la aritmética de precisión limitada. Conceptos como la cuantificación de pesos y activaciones se han explorado en papers como 'Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference' de Jacob et al. (2018), que sentó las bases para la inferencia con enteros de baja precisión. La idea de cuantificar el KV cache se alinea con la investigación más reciente sobre la optimización de la memoria en arquitecturas Transformer, donde el costo del KV cache es un factor dominante. La separación de las fases de prefill y decode y la aplicación de optimizaciones específicas para cada una reflejan principios de diseño de sistemas distribuidos donde se identifican y optimizan cuellos de botella específicos, similar a la optimización de pipelines de procesamiento de datos. La verificación de integridad del KV cache se relaciona con los principios de seguridad y robustez en sistemas concurrentes, donde la validación de estados compartidos es crucial para prevenir errores lógicos o corrupción de datos, un problema clásico en la computación paralela y distribuida.