Our new X account is live! Follow @wizwand_team for updates
WorkDL logo mark

TailorKV: A Hybrid Framework for Long-Context Inference via Tailored KV Cache Optimization

About

The Key-Value (KV) cache in generative large language models (LLMs) introduces substantial memory overhead. Existing works mitigate this burden by offloading or compressing the KV cache. However, loading the entire cache incurs significant latency due to PCIe bandwidth bottlenecks in CPU-GPU communication, while aggressive compression causes notable performance degradation. We identify that certain layers in the LLM need to maintain global information and are unsuitable for selective loading. In contrast, other layers primarily focus on a few tokens with dominant activations that potentially incur substantial quantization error. This observation leads to a key insight that loading dominant tokens and quantizing all tokens can complement each other. Building on this insight, we propose a hybrid compression method, TailorKV, which seamlessly integrates quantization and offloading. TailorKV develops an inference framework along with a hardware-friendly implementation that leverages these complementary characteristics. Extensive long-context evaluations exhibit that TailorKV achieves nearly lossless performance under aggressive compression settings, outperforming the state-of-the-art. Particularly, the Llama-3.1-8B with 128k context can be served within a single RTX 3090 GPU, reaching 82 ms per token during decoding.

Dingyu Yao, Bowen Shen, Zheng Lin, Wei Liu, Jian Luan, Bin Wang, Weiping Wang• 2025

Related benchmarks

TaskDatasetResultRank
Long-context UnderstandingLongBench (test)
Avg Score52.9
80
Long-context Language UnderstandingInfiniteBench
En.Sum24.1
63
Long-context UnderstandingInfiniteBench v1 (test)
Dialogue18.5
31
Long-context evaluationRULER 128k
Query Metric (MQ)98
29
Long-context evaluationRULER 64k
VT Score88
29
Long-context UnderstandingLongBench v1 (test)
SD QA49.3
21
Decoding LatencySynthetic Context Sequences (test)
Latency (16k Context)0.041
16
Decoding LatencyLlama-3.1-8B 16k sequence length v1 (inference)
Decoding latency (s)0.045
8
Decoding LatencyLlama-3.1-8B 32k sequence length v1 (inference)
Decoding Latency (s)0.047
7
Decoding LatencyLlama-2-7B 16k sequence length v1 (inference)
Decoding Latency (s)0.041
6
Showing 10 of 14 rows

Other info

Code

Follow for update