Cómo la escasez de KV logra una aceleración de 1.5x para vLLM

Zachary, Experto en Algoritmos en Novita AI

Enfoque de Investigación: Aceleración de Inferencia

Introducción

Durante el último año, desde la aparición de H2O, ha habido una proliferación de artículos de investigación sobre la escasez de KV. Sin embargo, un desafío significativo en las aplicaciones del mundo real es la gran brecha entre la investigación académica y la implementación práctica. Por ejemplo, frameworks como vLLM adoptan PagedAttention, una técnica de memoria basada en paginación que es incompatible con la mayoría de los algoritmos de escasez, o donde estos algoritmos funcionan peor que PagedAttention. Estos problemas impiden que los algoritmos de escasez se utilicen de manera efectiva en entornos de producción.

Al revisar artículos académicos recientes del último año en el campo de la escasez de KV, realizamos modificaciones al framework vLLM basadas en la escasez de KV. Estas modificaciones se alinearon con las características de optimización inherentes del framework vLLM, como Continuous Batching, FlashAttention y PagedAttention. Finalmente, al comparar el framework modificado con los frameworks de inferencia SOTA, utilizando los modelos, parámetros y hardware más comúnmente desplegados en producción, logramos una aceleración de 1.5x en el rendimiento de inferencia.

Antes de discutir la escasez de KV, es esencial abordar la característica de Massive Activations en los LLMs. En los LLMs, un pequeño número de valores de activación son significativamente más activos que otros, a veces superando a los menos activos por un factor de 100,000 o más. En otras palabras, un subconjunto pequeño de tokens juega un papel crítico. Esta observación forma la base para aplicar técnicas de escasez de KV, donde retener solo los tokens importantes puede mejorar el rendimiento de inferencia.

Los datos de análisis específicos de Llama2 y Llama3 respaldan aún más esta observación, como se muestra en las Fig. 1 y Fig. 2.

Fig. 1. Magnitudes de Activación (eje 𝐳) en LLaMA2-7B. Los ejes 𝐱 e 𝐲 son las dimensiones de secuencia y características. Para este modelo específico, observamos que las activaciones con magnitudes masivas aparecen en dos dimensiones de características fijas (1415, 2533), y dos tipos de tokens: el token inicial y el primer punto (.) o token de nueva línea (\ ).

A diferencia de Llama2, Llama3 no presenta la característica de Massive Activations en todas las capas. En las capas inferiores, las activaciones están más uniformemente distribuidas, mientras que las capas medias muestran patrones de atención localizados. Solo en las capas superiores el fenómeno de Massive Activations se vuelve prominente. Esta observación explica por qué modelos como Llama3 adoptan técnicas de escasez jerárquica, que adaptan la escasez a diferentes capas según sus patrones de activación específicos.

Fig. 2. Patrones de atención de generación aumentada por recuperación a través de capas en LlaMa (Touvron et al., 2023a;b) revelan que en las capas inferiores, el modelo exhibe un modo de atención de amplio espectro, distribuyendo las puntuaciones de atención de manera uniforme en todo el contenido. En las capas medias, la atención se vuelve más localizada dentro de cada documento, lo que indica una agregación de información refinada (formas triangulares rojas punteadas en las capas 6 y 10). Esto culmina en las capas superiores, donde la “atención masiva” se enfoca en unos pocos tokens clave (barras de atención concentradas después de la capa 18), extrayendo información esencial de manera eficiente para las respuestas.

Modelo del Sistema

Después de introducir la característica de Massive Activations en los LLMs, ahora tenemos una comprensión básica de los principios detrás de la escasez de KV. A continuación, discutiremos cómo se implementa la escasez de KV. En el proceso de inferencia de modelos grandes, se encuentran comúnmente tres cuellos de botella: capacidad de memoria de GPU, potencia computacional y ancho de banda de E/S. Para evitar cálculos redundantes, KV a menudo se almacena en caché durante la inferencia, lo que lleva a un consumo significativo de memoria de GPU. Mientras tanto, la carga computacional y el rendimiento de E/S se convierten en cuellos de botella críticos en las etapas de Prefill y Decode, respectivamente.

Nuestra modificación del framework vLLM se basa en la característica de Massive Activations de los LLMs, implementando escasez por capas (es decir, diferentes niveles de escasez para diferentes capas). Este enfoque reduce significativamente la sobrecarga asociada con el almacenamiento en caché de KV, acelerando así el proceso de inferencia de los modelos LLM.

Como se muestra en la Fig. 3, durante el proceso de inferencia, aplicamos escasez de KV a diferentes capas del modelo. Al emplear una estrategia de poda, eliminamos pares de KV con puntuaciones más bajas mientras retenemos aquellos con puntuaciones más altas y mayor proximidad. Este enfoque reduce el uso de memoria, así como la sobrecarga computacional y de E/S, lo que finalmente conduce a una inferencia acelerada.

Fig. 3. El impacto de la escasez de KV en el modelo

Resultados de Aceleración de Inferencia

La principal preocupación para la mayoría de los investigadores es el impacto real de la escasez de KV en el framework vLLM. Antes de profundizar en los detalles, primero presentamos los resultados de la evaluación de rendimiento.

1.Rendimiento de inferencia

En aplicaciones reales de LLM, una longitud de entrada/salida de 4000/500 es la configuración más comúnmente observada. Además, la GPU RTX 4090 es una de las GPU más utilizadas en la industria. Basándonos en estas condiciones, realizamos pruebas de comparación de rendimiento por lotes. El grupo de control es vLLM 0.6.1.p2, mientras que el grupo experimental es Novita AI Sparse 0.5.1 (una versión modificada con escasez de KV de vLLM 0.5.1). Se llevaron a cabo múltiples rondas de pruebas de comparación de rendimiento entre ambos, con métricas clave que incluyen TTFT (tiempo hasta el primer token, crítico para la experiencia del usuario) y Throughput (la velocidad real de inferencia).

Los resultados finales de las pruebas se muestran en la Tabla 1. Al aplicar la escasez de KV, pudimos aumentar el rendimiento de vLLM aproximadamente 1.58x, manteniendo un TTFT utilizable (con P50 por debajo de 1 segundo). Además, en escenarios de lotes más grandes, vLLM 0.6.1.p2 alcanzó su límite de rendimiento en un nivel de concurrencia de 10, mientras que Novita AI Sparse 0.5.1 pudo mantener un rendimiento TTFT estable incluso en un nivel de concurrencia de 20. Esto demuestra la robustez de la escasez de KV para garantizar la estabilidad del rendimiento en entornos de producción.

Tabla 1: Comparación de rendimiento entre vLLM 0.6.1.p2 y Novita AI Sparse 0.5.1

2.Precisión del modelo

Dado que la escasez de KV es un algoritmo de compresión con pérdida, también es importante evaluar el rendimiento del modelo. Realizamos evaluaciones de rendimiento basadas en Llama3–8B, incluyendo evaluaciones con MMLU y otros benchmarks. Los resultados indican que la pérdida de precisión se mantiene dentro de aproximadamente el 3%.

Como se muestra en la Tabla 2, probamos el rendimiento del modelo utilizando MMLU y benchmarks de humanidades:

Tabla 2: Resultados de evaluación de rendimiento de Llama3–8B utilizando MMLU y benchmarks de humanidades

En segundo lugar, realizamos evaluaciones de tareas de QA en escenarios de texto largo comúnmente encontrados en aplicaciones prácticas, con longitudes de entrada que van desde 7k hasta 30k. La relación de compresión se mantuvo en más de 10 veces, como se muestra en la Tabla 3. Los resultados indicaron una pérdida de precisión general de aproximadamente el 10%.

Tabla 3: Resultados de evaluación de tareas de QA de texto largo con longitudes de entrada de 7k a 30k

Técnicas Clave

1.Escasez por capas y paralelismo de tensores

El framework vLLM emplea una estrategia de programación a nivel de iteración, comúnmente conocida como Continuous Batching. Este enfoque de programación se caracteriza por no esperar a que cada secuencia en un lote complete la generación; en su lugar, facilita la programación a nivel de iteración, lo que resulta en una mayor utilización de GPU en comparación con el batching estático. Como se mencionó al principio de este artículo, modelos como Llama3–8B poseen inherentemente características como la escasez jerárquica. En consecuencia, nuestras modificaciones basadas en Continuous Batching han encontrado una serie de desafíos, como se muestra en la Fig. 4.

Fig. 4. Completando siete secuencias usando Continuous Batching. A la izquierda, ves el lote después de una sola iteración, mientras que a la derecha, el lote después de varias iteraciones. Tan pronto como una secuencia emite un token de fin de secuencia, introducimos una nueva secuencia en su lugar (por ejemplo, las secuencias S5, S6 y S7). Este método aumenta la utilización de GPU al evitar que la GPU tenga que esperar a que todas las secuencias terminen antes de iniciar una nueva.

El principal desafío que abordamos fue la gestión de la memoria en diversas capas. Dado que vLLM utiliza un modelo unificado de gestión de memoria paginada y la limitación de la estrategia de cola a solo soportar Full Cache o Sliding Window, realizamos modificaciones en la estructura central de vLLM para habilitar el soporte de ambas opciones.

Este ajuste otorga a cada capa la flexibilidad de optar por Full Cache o Sliding Window, permitiendo distintos niveles de escasez entre capas, como se representa en el diagrama proporcionado.

El primer desafío a abordar es el problema de gestión de memoria entre diferentes capas. Dado que vLLM utiliza un modelo unificado de gestión de memoria paginada para todas las capas, y la estrategia de cola permite solo Full Cache o Sliding Window, necesitamos ajustar la estructura subyacente de vLLM. Específicamente, esto implica habilitar el soporte tanto para Full Cache como para Sliding Window, de modo que diferentes capas puedan seleccionar independientemente su estrategia de gestión de memoria preferida. Este ajuste finalmente facilita diversos grados de escasez en diferentes capas. El diagrama estructural correspondiente se muestra en la Figura 5.

Fig. 5. Estructura de gestión de memoria por capas en vLLM

En general, los ajustes se pueden clasificar en tres fases distintas:

Fase de inicialización del servicio: Durante el proceso de inicio del servicio, la memoria asignada para las tres capas se inicializa principalmente mediante el cálculo del espacio restante, junto con una inicialización jerárquica de la estructura de metadatos. En contraste, el vLLM original funciona como un Full Cache completo, sin tener en cuenta la lógica de asignación de memoria para diferentes capas.

Fase de preparación de pasos: Durante el proceso de inicialización de los pasos en vLLM, específicamente en las fases que involucran métodos de Block Manager y Model Runner, se emplean diferentes estrategias de gestión de tablas de bloques para diferentes capas para controlar el mecanismo de actualización de KV. Como se ilustra en la figura anterior, utilizamos el mecanismo Full Cache para las primeras tres capas, almacenando todos los KV, mientras que las capas más allá de la tercera emplean la estrategia Sliding Window para reemplazar automáticamente los KV con las puntuaciones más bajas. Este enfoque permite una asignación flexible del espacio de memoria adaptado a las necesidades de cada capa, ahorrando espacio mientras se maximiza la efectividad del algoritmo de escasez.

Fase de cálculo de escasez: Esta fase ocurre típicamente después de la ejecución del modelo, donde se realizan operaciones de escasez basadas en las puntuaciones de escasez calculadas. Una consideración clave en este paso es la compatibilidad con tecnologías como Tensor Parallelism, CUDA graphs y GQA. En implementaciones prácticas de ingeniería, Tensor Parallelism requiere atención particular. Como se ilustra en la figura a continuación, en un escenario de múltiples GPU en una sola máquina, es esencial emplear un mecanismo de partición apropiado para confinar las operaciones de escasez de KV dentro de cada GPU, evitando así la comunicación entre GPU. Afortunadamente, la unidad de Atención emplea un enfoque de partición basado en ColumnLinear, lo que hace que su implementación sea relativamente sencilla.

Fig. 6. Proceso de tres etapas para la gestión de memoria y escasez en vLLM

La discusión anterior se centra principalmente en las modificaciones realizadas a nivel de framework para soportar la escasez jerárquica en vLLM. A continuación, dirigiremos nuestra atención a las optimizaciones implementadas a nivel de CUDA.

2.Modificación de Atención

Las modificaciones a nivel de CUDA se centran principalmente en las mejoras realizadas en la unidad de cálculo de Atención, específicamente en FlashAttention y PagedAttention. Si bien PagedAttention es relativamente simple, enfatizaremos las modificaciones relacionadas con FlashAttention.

Primero, presentamos la fórmula de cálculo para Atención:

Como se muestra en la fórmula, el proceso de cálculo de softmax implica calcular el producto de QK, y es un algoritmo de 2 pasadas (que requiere dos pasadas). Sin embargo, el objetivo final de cálculo O se puede lograr utilizando el algoritmo FlashAttention con una sola pasada (que requiere una pasada). La lógica de implementación de FlashAttention se puede ilustrar mediante la Fig. 7 del artículo de FlashAttention2. En resumen, implica recorrer Q por fragmentos para realizar cálculos de bloque en KV, mientras se actualizan gradualmente los datos de O, P y rowmax hasta que el bucle concluye. Finalmente, dividir O por ℓ permite el cálculo de 1 pasada de FlashAttention.

Fig .7. Lógica de implementación de FlashAttention2

Para nuestra implementación de la escasez de KV, es importante tener en cuenta que durante el proceso de cálculo de FlashAttention, los valores requeridos para softmax — es decir, rowmax/ℓ (el máximo global y los valores acumulativos) — ya están calculados. Sin embargo, FlashAttention no devuelve los resultados de estos dos valores.

Con una comprensión fundamental de FlashAttention, ahora podemos discutir las fórmulas de puntuación para la escasez presentadas en artículos como PYRAMIDKV:

donde [n − α, n] es el rango de los tokens de instrucción. En cada capa l y para cada cabeza h, se seleccionan los mejores k l tokens con las puntuaciones más altas, y se retienen sus respectivos cachés KV. Todos los demás cachés KV se descartan y no se utilizarán en ningún cálculo posterior durante todo el proceso de generación.

Específicamente, calculamos la puntuación de seleccionar el i-ésimo token para su retención en el caché KV como S en cada cabeza de atención h mediante:

En la fórmula anterior, A representa la puntuación softmax, mientras que S denota la suma de las puntuaciones de los a valores softmax más cercanos. Como se indica, softmax en sí mismo es un algoritmo de 2 pasadas, que no puede integrarse directamente en FlashAttention. Esto también implica que la puntuación de escasez no puede implementarse en 1 pasada.

Como se mencionó anteriormente, aunque FlashAttention no genera las puntuaciones softmax, sí calcula valores intermedios como rowmax/ℓ. Por lo tanto, al referenciar Online Softmax: algoritmo de 2 pasadas, podemos utilizar las salidas de rowmax/ℓ para lograr la puntuación específica S para la escasez de KV. En la práctica, solo necesitamos realizar mejoras en FlashAttention ilustradas en la Fig. 7 para cumplir con este requisito.

Fig .7. Modificación de FlashAttention para la puntuación de escasez de KV

• Salida de rowmax/ℓ en FlashAttention: FlashAttention incorpora las salidas de rowmax/ℓ. Dado que estos valores representan el máximo y los totales acumulativos, su tamaño de datos es mínimo, ejerciendo un impacto insignificante en el rendimiento general de FlashAttention. Al reutilizar rowmax/ℓ, podemos evitar la necesidad de una segunda pasada, mejorando así la eficiencia computacional general.
• Adición del paso Online Softmax: Introducimos el paso Online Softmax, que se basa en FlashAttention calculando las puntuaciones softmax usando la fórmula p=e((kq)−row_max)/ℓ (lógica de 1 pasada). Posteriormente, al realizar una suma por filas, podemos calcular la puntuación de escasez S.

En la implementación de Online Softmax, dos detalles críticos requieren atención particular:

• Ajuste del orden del bucle: Dado que S es la suma de puntuaciones basada en la dimensión k, es necesario ajustar el cálculo en Online Softmax de qk a kq, en comparación con el mecanismo de atención. Este ajuste evita comunicación innecesaria entre diferentes bloques al transformar la suma de columnas en una suma de filas, permitiendo que todos los cálculos se completen dentro de los registros.
• Posiciones de inicio y fin de q: Debido a la presencia de codificación causal y la lógica de suma para los a más cercanos, no es necesario recorrer todos los valores de q con k. Al hacer referencia a los indicadores de máscara etiquetados como 0/1 en la Fig .7, podemos determinar si un bloque correspondiente requiere cálculo según las posiciones de inicio y fin de q, lo que conserva significativamente los recursos computacionales. Por ejemplo, si la longitud del prompt de entrada es 5000 y la longitud del a más cercano es 256, el cálculo convencional requeriría operaciones de 50005000. Sin embargo, al registrar las posiciones de inicio y fin, solo se requieren operaciones de 5000256, mejorando así la eficiencia computacional.

PagedAttention realiza un trabajo similar al de FlashAttention, con el objetivo común de calcular la puntuación final S. No se proporciona una explicación detallada de PagedAttention aquí.

Conclusión

Nuestra versión modificada, Novita AI Sparse 0.5.1, basada en vLLM 0.5.1, actualmente soporta principalmente modelos como Llama3–8B y Llama3–70B. Opera en el modo de inferencia utilizando CUDA graphs y se despliega principalmente en GPU de grado de consumo como la RTX 4090. Todavía hay un margen significativo de optimización tanto en aspectos de ingeniería como algorítmicos.

Un aviso importante: Las siguientes características aún no son compatibles en vLLM-0.6.2:

• Chunked-prefill
• Prefix caching
• FlashInfer y otros backends que no sean FlashAttention

Para aquellos interesados en Novita AI Sparse 0.5.1, pueden visitar Novita AI para experimentarlo antes que otros.

Referencias:

[1] H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models

[2] Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference

[3] SnapKV: LLM Knows What You are Looking for Before Generation

[4] PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling

[5] PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference

[6] MiniCache: KV Cache Compression in Depth Dimension for Large Language Models

[7] Layer-Condensed KV Cache for Efficient Inference of Large Language Models

[8] TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding

[9] CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion

[10] KV-Runahead: Scalable Causal LLM Inference by Parallel Key-Value Cache Generation