ShadowKV: Cache KV nelle Ombre per LLM ad Alto Throughput con Lungo Contesto di Inferenza

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

October 28, 2024

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

cs.AI

Abstract

Con la diffusa implementazione di modelli linguistici di lungo contesto a grande scala (LLM), c'è stata una crescente domanda di supporto efficiente per l'inferenza ad alto throughput. Tuttavia, poiché la cache chiave-valore (KV) si espande con la lunghezza della sequenza, l'aumento della memoria e la necessità di accedervi per ogni generazione di token comportano un basso throughput durante il servizio di LLM a lungo contesto. Mentre sono stati proposti vari metodi dinamici di attenzione sparsa per accelerare l'inferenza mantenendo la qualità della generazione, essi non riescono a ridurre in modo sufficiente il consumo di memoria della GPU o introducono un significativo ritardo di decodifica scaricando la cache KV sulla CPU. Presentiamo ShadowKV, un sistema di inferenza LLM ad alto throughput a lungo contesto che memorizza la cache chiave a basso rango e scarica la cache valore per ridurre l'occupazione di memoria per dimensioni di batch più grandi e sequenze più lunghe. Per minimizzare il ritardo di decodifica, ShadowKV utilizza una strategia accurata di selezione KV che ricostruisce coppie KV sparse minime al volo. Valutando ShadowKV su un'ampia gamma di benchmark, tra cui RULER, LongBench e Needle In A Haystack, e modelli come Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K e Qwen2-7B-128K, dimostriamo che può supportare dimensioni di batch fino a 6 volte più grandi e aumentare il throughput fino a 3,04 volte su una GPU A100 senza sacrificare l'accuratezza, superando persino le prestazioni ottenibili con dimensioni di batch infinite nell'ipotesi di memoria GPU infinita. Il codice è disponibile su 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.