Search papers, labs, and topics across Lattice.
This paper introduces a retrieval-completion attention mechanism to reduce KV-cache traffic during long-context generation by selectively retrieving and attending to the top-K relevant tokens. The method estimates the contribution of unretrieved tokens using a fixed-size feature map summary computed during prefill, avoiding bias introduced by softmax renormalization in selection-only methods. Experiments on long-context benchmarks demonstrate that this approach outperforms selection-only Top-K retrieval, especially in high-entropy attention heads, while maintaining backbone weights and KV-cache format.
Mitigating the bias introduced by Top-K retrieval in long-context transformers can be achieved by estimating the contribution of unretrieved tokens with a fixed-size feature map, significantly improving performance without increasing KV-cache reads.
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.
