Sparse Attention across Multiple-context KV Cache

TL;DR

This paper introduces SamKV, a novel method for sparsifying attention in multiple-context KV caches, significantly reducing memory and computation in retrieval-augmented generation without sacrificing accuracy.

Contribution

SamKV is the first approach to apply attention sparsification to multiple-context KV caches, effectively compressing sequence length and improving efficiency in RAG scenarios.

Findings

Reduces sequence length to 15% of original

Maintains accuracy comparable to full-recompute methods

Boosts throughput in multi-context retrieval scenarios

Abstract

Large language models face significant cost challenges in long-sequence inference. To address this, reusing historical Key-Value (KV) Cache for improved inference efficiency has become a mainstream approach. Recent advances further enhance throughput by sparse attention mechanisms to select the most relevant KV Cache, thereby reducing sequence length. However, such techniques are limited to single-context scenarios, where historical KV Cache is computed sequentially with causal-attention dependencies. In retrieval-augmented generation (RAG) scenarios, where retrieved documents as context are unknown beforehand, each document's KV Cache is computed and stored independently (termed multiple-context KV Cache), lacking cross-attention between contexts. This renders existing methods ineffective. Although prior work partially recomputes multiple-context KV Cache to mitigate accuracy loss from missing cross-attention, it requires retaining all KV Cache throughout, failing to reduce memory overhead. This paper presents SamKV, the first exploration of attention sparsification for multiple-context KV Cache. Specifically, SamKV takes into account the complementary information of other contexts when sparsifying one context, and then locally recomputes the sparsified information. Experiments demonstrate that our method compresses sequence length to 15% without accuracy degradation compared with full-recompuation baselines, significantly boosting throughput in multi-context RAG scenarios.