ShadowKV: Cache KV em Sombras para Inferência de LLM de Longo Contexto de Alto Rendimento

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

Resumo

Com a ampla implantação de modelos de linguagem grandes de contexto longo (LLMs), houve uma crescente demanda por suporte eficiente de inferência de alto rendimento. No entanto, à medida que o cache chave-valor (KV) se expande com o comprimento da sequência, o aumento da pegada de memória e a necessidade de acessá-lo para cada geração de token resultam em baixo rendimento ao servir LLMs de contexto longo. Embora vários métodos dinâmicos de atenção esparsa tenham sido propostos para acelerar a inferência mantendo a qualidade da geração, eles falham em reduzir suficientemente o consumo de memória da GPU ou introduzem latência significativa de decodificação ao transferir o cache KV para a CPU. Apresentamos o ShadowKV, um sistema de inferência de alto rendimento para LLMs de contexto longo que armazena o cache de chaves de baixa ordem e transfere o cache de valores para reduzir a pegada de memória para tamanhos de lote maiores e sequências mais longas. Para minimizar a latência de decodificação, o ShadowKV emprega uma estratégia precisa de seleção KV que reconstrói pares KV esparsos mínimos sob demanda. Ao avaliar o ShadowKV em uma ampla gama de benchmarks, incluindo RULER, LongBench e Needle In A Haystack, e modelos como Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K e Qwen2-7B-128K, demonstramos que ele pode suportar até 6 vezes maiores tamanhos de lote e aumentar o rendimento em até 3,04 vezes em uma GPU A100 sem sacrificar a precisão, superando até mesmo o desempenho alcançável com tamanho de lote infinito sob a suposição de memória de GPU infinita. O código está disponível em 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.