🤖 AI Summary
This work addresses the decoding bandwidth bottleneck in long-context generation caused by frequent reads of the key-value (KV) cache, which becomes especially severe when the cache is offloaded to GPU external memory. The authors propose a retrieval-completion attention mechanism that, without altering the model architecture or KV cache format, computes exact attention over the top-K tokens most relevant to the query along with fixed anchor tokens at the sequence boundaries. Contributions from intermediate regions are approximated using a fixed-size feature map constructed during the prefill phase, and a unified normalization recovers the full attention distribution. This approach effectively mitigates the softmax bias introduced by top-K retrieval while requiring no additional KV cache accesses. Under identical KV access budgets, it significantly outperforms pure top-K methods, with particularly notable gains in high-entropy attention heads.
📝 Abstract
Long-context generation is increasingly limited by decode-time key-value (KV) cache traffic, particularly when KV is offloaded beyond GPU memory. Query-aware retrieval (e.g., Top-K selection) reduces this traffic by loading only a subset of KV pairs, but renormalizing the softmax over the subset introduces bias when attention mass is spread over unretrieved tokens. We propose a retrieval-completion attention module that keeps backbone weights and the KV-cache format unchanged. For each query, we compute exact attention over sink/tail anchors and the query-dependent retrieved Top-K tokens, and estimate the remaining mid-region numerator and denominator using a fixed-size feature-map summary computed at prefill time. We add the exact and estimated contributions in the unnormalized domain and apply a single normalization, recovering the missing softmax mass without additional attention-side KV reads. Across long-context benchmarks, the proposed method improves over selection-only Top-K at matched token-equivalent read budgets, with the largest gains in high-entropy heads.