ShadowKV: KV-Cache im Schatten für Hochdurchsatz Long-Context LLM-Inferenz

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

October 28, 2024

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

cs.AI

Zusammenfassung

Mit der weit verbreiteten Bereitstellung von langen Kontexten und großen Sprachmodellen (LLMs) besteht ein wachsender Bedarf an effizienter Unterstützung für Hochdurchsatz-Inferenz. Allerdings führen die zunehmende Speicherbelegung des Schlüssel-Wert (KV)-Caches mit der Sequenzlänge sowie die Notwendigkeit, bei jeder Token-Generierung darauf zuzugreifen, bei der Bedienung von LLMs mit langem Kontext zu einer geringen Durchsatzrate. Obwohl verschiedene dynamische spärliche Aufmerksamkeitsmethoden vorgeschlagen wurden, um die Inferenz zu beschleunigen und gleichzeitig die Generierungsqualität zu erhalten, scheitern sie entweder daran, den GPU-Speicherverbrauch ausreichend zu reduzieren oder führen durch die Auslagerung des KV-Caches auf die CPU zu einer erheblichen Decodierungsverzögerung. Wir stellen ShadowKV vor, ein Hochdurchsatz-Inferenzsystem für lange Kontexte von LLMs, das den Cache für Schlüssel mit niedriger Rangordnung speichert und den Cache für Werte auslagert, um den Speicherbedarf für größere Batch-Größen und längere Sequenzen zu reduzieren. Um die Decodierungsverzögerung zu minimieren, verwendet ShadowKV eine genaue KV-Auswahlstrategie, die minimale spärliche KV-Paare dynamisch rekonstruiert. Durch die Evaluierung von ShadowKV anhand einer breiten Palette von Benchmarks, einschließlich RULER, LongBench und Needle In A Haystack, sowie Modellen wie Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K und Qwen2-7B-128K, zeigen wir, dass es bis zu 6-mal größere Batch-Größen unterstützen und den Durchsatz auf einer A100 GPU um bis zu 3,04-mal steigern kann, ohne die Genauigkeit zu beeinträchtigen, und sogar die mit unendlicher Batch-Größe erreichbare Leistung übertrifft, unter der Annahme von unendlichem GPU-Speicher. Der Code ist verfügbar unter 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.