ShadowKV: Caché KV en Sombras para Inferencia de LLM de Largo Contexto y Alto Rendimiento

ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference

October 28, 2024

Autores: Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, Beidi Chen

cs.AI

Resumen

Con la amplia implementación de modelos de lenguaje grandes de contexto largo (LLMs, por sus siglas en inglés), ha habido una creciente demanda de soporte eficiente para inferencia de alto rendimiento. Sin embargo, a medida que la caché clave-valor (KV) se expande con la longitud de la secuencia, el aumento en la huella de memoria y la necesidad de acceder a ella para cada generación de token resultan en una baja velocidad de procesamiento al atender LLMs de largo contexto. Aunque se han propuesto varios métodos de atención dinámica dispersa para acelerar la inferencia manteniendo la calidad de generación, estos métodos no logran reducir suficientemente el consumo de memoria de la GPU o introducen una latencia significativa en la decodificación al transferir la caché KV a la CPU. Presentamos ShadowKV, un sistema de inferencia de alto rendimiento para LLMs de largo contexto que almacena la caché de claves de rango bajo y transfiere la caché de valores para reducir la huella de memoria en tamaños de lote más grandes y secuencias más largas. Para minimizar la latencia de decodificación, ShadowKV emplea una estrategia precisa de selección de KV que reconstruye pares KV dispersos mínimos sobre la marcha. Al evaluar ShadowKV en una amplia gama de pruebas, incluidas RULER, LongBench y Needle In A Haystack, y modelos como Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K y Qwen2-7B-128K, demostramos que puede admitir hasta 6 veces mayores tamaños de lote y aumentar la velocidad de procesamiento hasta 3.04 veces en una GPU A100 sin sacrificar precisión, incluso superando el rendimiento alcanzable con un tamaño de lote infinito bajo la suposición de memoria de GPU infinita. El código está disponible en https://github.com/bytedance/ShadowKV.

English

With the widespread deployment of long-context large language models (LLMs), there has been a growing demand for efficient support of high-throughput inference. However, as the key-value (KV) cache expands with the sequence length, the increasing memory footprint and the need to access it for each token generation both result in low throughput when serving long-context LLMs. While various dynamic sparse attention methods have been proposed to speed up inference while maintaining generation quality, they either fail to sufficiently reduce GPU memory consumption or introduce significant decoding latency by offloading the KV cache to the CPU. We present ShadowKV, a high-throughput long-context LLM inference system that stores the low-rank key cache and offloads the value cache to reduce the memory footprint for larger batch sizes and longer sequences. To minimize decoding latency, ShadowKV employs an accurate KV selection strategy that reconstructs minimal sparse KV pairs on-the-fly. By evaluating ShadowKV on a broad range of benchmarks, including RULER, LongBench, and Needle In A Haystack, and models like Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K, and Qwen2-7B-128K, we demonstrate that it can support up to 6times larger batch sizes and boost throughput by up to 3.04times on an A100 GPU without sacrificing accuracy, even surpassing the performance achievable with infinite batch size under the assumption of infinite GPU memory. The code is available at https://github.com/bytedance/ShadowKV.