Distance-KV improves memory and speed for long text processing in AI
Distance-KV: Exploiting Relative Distance for Efficient Long-Context Inference
Machine LearningComputation and Language
Summary
Processing very long texts with AI models takes a lot of memory and time. The authors found that how far apart words are in a text strongly affects how the AI remembers related information. They created Distance-KV, a way to smartly keep only the most useful memory parts based on these distances, without needing to check importance every time. This method cuts memory use and speeds up the AI while keeping good performance on several long text tasks.
What this means in practice
- •For ai infrastructure teams: Reduce memory usage and speed up inference for large language model deployments handling very long documents.
- •For cloud service providers: Optimize backend hardware resources by compacting model internal memory during long-context text generation.
Authors
Xianpeng Shang, Canbin Huang, Jiang Li, Tian Lan, Qianyi Cai, Xiaojun Quan, Xiangdong Su
Abstract
The memory usage and decoding latency of LLM inference grow rapidly with context length. To reduce these costs, key-value (KV) cache compression methods selectively retain cached states based on token importance or differences in attention patterns across heads. However, we discover that retrieval capability varies substantially with relative distance, even within the same attention head. To exploit this structure, we introduce Distance-KV, which learns a static KV retention pattern over the joint space of layers, attention heads, and relative distances. The pattern is learned offline with the language model frozen and reused across inputs to prune and compact the KV cache without online importance scoring. Across three backbone models and four long-context benchmarks, Distance-KV consistently achieves the best overall performance among competing KV cache compression methods, exceeding the strongest compression baseline by up to 9.3 points on RULER at 128K. On Llama-3.1-8B-Instruct at 128K, Distance-KV reduces KV cache memory by 65.4% and achieves a $1.66\times$ decoding speedup relative to Dense. Together, these results identify relative distance as an important structural dimension for understanding how LLMs retrieve information over long contexts and for designing more efficient inference methods.