ShadowKV: KV-cache in schaduwen voor high-throughput Long-Context LLM-inferentie

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

October 28, 2024

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

cs.AI

Samenvatting

Met de wijdverbreide inzet van lang-contextuele grote taalmodellen (LLMs) is er een groeiende vraag naar efficiënte ondersteuning van high-throughput inferentie. Echter, naarmate de key-value (KV) cache groeit met de sequentielengte, leiden de toenemende geheugenfootprint en de noodzaak om er toegang toe te hebben voor elke token-generatie beide tot een lage throughput bij het bedienen van lang-contextuele LLMs. Hoewel verschillende dynamische schaarse aandachtsmethoden zijn voorgesteld om de inferentie te versnellen terwijl de generatiekwaliteit behouden blijft, slagen ze er niet in om de GPU-geheugenconsumptie voldoende te verminderen of introduceren ze aanzienlijke decoderingslatentie door de KV-cache naar de CPU te verplaatsen. We presenteren ShadowKV, een high-throughput lang-contextueel LLM-inferentiesysteem dat de low-rank key cache opslaat en de value cache uitbesteedt om de geheugenfootprint te verkleinen voor grotere batchgroottes en langere sequenties. Om de decoderingslatentie te minimaliseren, maakt ShadowKV gebruik van een nauwkeurige KV-selectiestrategie die minimaal schaarse KV-paren on-the-fly reconstrueert. Door ShadowKV te evalueren op een breed scala aan benchmarks, waaronder RULER, LongBench en Needle In A Haystack, en modellen zoals Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K en Qwen2-7B-128K, tonen we aan dat het tot 6 keer grotere batchgroottes kan ondersteunen en de throughput tot 3,04 keer kan verhogen op een A100 GPU zonder nauwkeurigheid op te offeren, zelfs de prestaties overtreffend die haalbaar zijn met oneindige batchgrootte onder de aanname van oneindig GPU-geheugen. De code is beschikbaar op 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.