$A^3$: Attention-Aware Accurate KV Cache Fusion for Fast Large Language Model Serving

TL;DR

This work tackles the latency and memory challenges of serving large language models with long contexts by improving KV Cache reuse. It introduces Attention-Aware Accurate KV Cache Fusion ($A^3$), which precomputes KV caches for the system prompt and relevant documents, recovers positional information with RoPE, and selects tokens for recomputation based on user-query attention to align KV updates with the query. The approach yields superior task performance across multiple models and long-context benchmarks while maintaining competitive decoding overhead, and an eviction-based extension further boosts throughput. The results demonstrate practical viability for real-world long-context LLM serving, offering substantial speedups without sacrificing accuracy and enabling scalable retrieval-augmented generation.

Abstract

Large language models (LLMs) have demonstrated strong capabilities in processing long contexts, enabling them to tackle tasks involving long textual inputs such as multi-turn conversations, legal documents, or retrieved documents in Retrieval-Augmented Generation (RAG) systems. However, despite their ability to handle long sequences, the resulting decoding latency and memory overhead remain substantial, posing challenges for real-world deployment. Recent advances in KV Cache reuse have shown potential to mitigate these costs, but still suffer from notable performance degradation. To address this issue, we conduct an in-depth investigation of recomputation-based reuse methods and observe that the recomputed tokens often fail to align with the context segments most relevant to the question. This misalignment hinders proper updates to the critical contextual representations. Therefore, we propose the $\textbf{A}$ttention-$\textbf{A}$ware $\textbf{A}$ccurate KV Cache Fusion algorithm ($A^3$), which precomputes and selectively fuses the KV Cache of text chunks based on their relevance to the question, achieving accurate integration with minimal computational overhead. Extensive experiments on various benchmarks and LLMs demonstrate that $A^3$ achieves the best task performance compared to four baselines while reducing the time-to-first-token (TTFT) by 2$\times$.

Figures (7)

• Figure 1: Blue blocks means the corresponding KVs are recomputed. (Top) Without recomputation. (Middle) Recomputation applied but key information missed. (Bottom) Key information updated based on query-aware attention.
• Figure 2: The mismatch between high-attention tokens and recomputed tokens. (1) The attention heatmap between the question and the document. We observe that both (2) the strategy of selecting tokens with the largest KV differences and (3) the strategy of selecting head and tail tokens of the document result in limited coverage of the high-attention tokens.
• Figure 3: Attention recovery performance under different reuse strategies. The attention map obtained by recomputing question-relevant tokens is more consistent with the original attention map compared to other.
• Figure 4: Results of Needle-in-a-Haystack on LLaMA-3-8B-Instruct. The vertical axis of the figure represents the depth percentage, and the horizontal axis represents the token length. Our method achieves the best retrieval performance.
• Figure 5: Comparison of inference performance among different reuse methods.