SpeCache: Speculative Key-Value Caching for Efficient Generation of LLMs

TL;DR

SpeCache tackles the VRAM bottleneck in long-sequence LLM inference by offloading the full KV cache to CPU memory and using a speculative decoding step to prefetch the most relevant KV pairs into VRAM. It maintains a low-bit KV copy in GPU memory and decodes two tokens in parallel, enabling effective latency hiding without retraining. Across LongBench and Needle-in-a-Haystack benchmarks, SpeCache achieves substantial VRAM reductions and throughput gains, maintaining near baseline accuracy even with high KV compression. The method is training-free and complementary to existing KV quantization techniques, making long-context generation and retrieval-augmented tasks more practical.

Abstract

Transformer-based large language models (LLMs) have already achieved remarkable results on long-text tasks, but the limited GPU memory (VRAM) resources struggle to accommodate the linearly growing demand for key-value (KV) cache as the sequence length increases, which has become a bottleneck for the application of LLMs on long sequences. Existing KV cache compression methods include eviction, merging, or quantization of the KV cache to reduce its size. However, compression results in irreversible information forgetting, potentially affecting the accuracy of subsequent decoding. In this paper, we propose SpeCache, which takes full advantage of the large and easily expandable CPU memory to offload the complete KV cache, and dynamically fetches KV pairs back in each decoding step based on their importance measured by low-bit KV cache copy in VRAM. To avoid inference latency caused by CPU-GPU communication, SpeCache speculatively predicts the KV pairs that the next token might attend to, allowing us to prefetch them before the next decoding step which enables parallelization of prefetching and computation. Experiments on LongBench and Needle-in-a-Haystack benchmarks verify that SpeCache effectively reduces VRAM usage while avoiding information forgetting for long sequences without re-training, even with a 10x high KV cache compression ratio.

Figures (5)

• Figure 1: SpeCache use low-bit KV cache and speculative token to "guess" the top-$k$ most relevant 16-bit KV pairs for the next token, and prefetch them before the next decoding step.
• Figure 2: Left: Hit rates of query-dependent top-$k$ attention and greedy cache eviction. Middle: The latency of a single decoding step on the GPU vs the latency of loading the KV cache from the CPU to the GPU. Right: CPU-GPU transfer latency of contiguous and non-contiguous CPU memory. We highlight the full and top-1% KV cache size with context length of 32k. Measured using Mistral-7B-Instruct-v0.2 on NVIDIA A6000 GPU.
• Figure 3: Illustration of SpeCache. In prefilling stage, the KV cache is quantized and offloaded layer by layer. In pre-decoding stage, we use the first output token to compute the first speculative token and prefetch the 16-bit KV pairs needed by the first decoding step. In each step of decoding stage, we simultaneously decode two tokens: the output token and the speculative token. the results of both serve as inputs for the next step. The top-$k$ most relevant 16-bit KV pairs for the speculative token are prefetched before the next step.
• Figure 4: Performance on Needle-in-a-haystack benchmark. We use $g=32$ for quantization, resulting in the KV cache compression ratio of 0.13 and 0.14 for Mistral-7B-Instruct-v0.2 and LLaMA-3-8B-Instruct, respectively.
• Figure 5: Ablation study on $k$. We use 1-bit SpeCache with $g=32$.