Recall Before You Rank: Similarity-Guided Top-K Reuse for Efficient Long-Context Attention
Wenshuai Yao, Wenyong Zhou, Hanyong Shao, Yizhe Chen, Zhiyuan Ning, Yuannuo Feng, Ru Huang, Kechao Tang
School of Integrated Circuits, Peking University, Beijing, China · Department of Electrical and Computer Engineering, The University of Hong Kong, Hong Kong SAR, China · School of Integrated Circuit Science and Engineering, Beihang University, Beijing, China · arXiv preprint arXiv:2607.27692

论文概览

论文标题： Recall Before You Rank: Similarity-Guided Top-K Reuse for Efficient Long-Context Attention
方法名称： ReTopK
论文形式： arXiv 预印本，arXiv:2607.27692
发布时间： 2026 年 7 月 30 日

ReTopK 关注动态 Top-K 稀疏注意力中的索引发现成本。即使最终只对 K 个 KV 条目执行 Softmax 和 Value 聚合，Exact Top-K 通常仍需要让当前 Query 扫描完整 Key Cache，并在所有分数上执行全局 Top-K。传统方法稀疏化了注意力聚合，却没有真正稀疏化索引发现。

论文提出一种 training-free 的“先召回、再排序”机制：为每个 Attention Head 保存少量历史 Query 及其 Top-K 支持集合；新 Query 到来时，先检索最相似的历史 Query，合并它们曾访问的索引与最近窗口，再仅对这个紧凑候选集计算当前 Query 的精确 QK 分数并重排。

其核心思想可以概括为：

Top-K 索引不是一次性的中间结果，而是一种具有时间局部性和表示局部性的运行时决策，可以被缓存、召回和校正。

研究背景与问题

Exact Top-K 的隐藏成本

设 Decode Step t 的 Query 为 q_t，当前 KV Cache 中共有 L_t 个 Key。Exact Top-K 首先计算完整 QK 分数：

e_{t,i}=\frac{q_t^T k_i}{\sqrt{d}},\quad i=1,\ldots,L_t

随后从全部 L_t 个分数中选出最大的 K 个索引：

S_t=\operatorname{TopK}_{i\in[1,L_t]}(e_{t,i},K)

尽管最终 Softmax 和 Value 聚合只处理 K 个位置，索引发现仍包含两个随上下文增长的步骤：完整 Key Cache 的 QK 扫描，以及在完整长度上执行全局 Top-K。论文的延迟分解显示，随着上下文从 4K 增长至 128K，这两部分在 Exact Top-K 注意力中的占比持续上升。

ReTopK 因此试图回答：能否在不扫描完整历史 Key 的情况下，为当前 Query 恢复足够准确的 Top-K 候选集合？

相似 Query 具有相似支持集合

论文用下面的指标衡量两个 Query 的 Top-K 重叠率：

\operatorname{Overlap}(t,j)=\frac{|S_t\cap S_j|}{K}

在 Qwen2.5-7B 上，Query 余弦相似度与 Exact Top-K 支持重叠率表现出明显相关性。论文报告，对全部历史 Query 对以及当前 Query 与最相似历史 Query 两种统计方式，相关系数分别约为 0.70 和 0.73。

这意味着当前 Query 不必只复用最近一步的决策，也可以回看一个短历史窗口，寻找表示更相似的状态。

多个相似 Query 的并集具有高覆盖率

设与当前 Query 最相似的 R 个历史 Query 为 j_1...j_R，其支持并集为：

U_t(R)=\bigcup_{r=1}^{R}S_{j_r}

论文报告的覆盖率如下：

单个历史支持不够可靠，但多个相似状态可以共同形成高召回候选集。

索引召回率不等于注意力质量

被漏掉的索引可能只承载很小的归一化注意力权重。论文在一个代表性 Head 中观察到，ReTopK 的 Exact Top-K 索引召回率为 78.9%，但仍保留 99.45% 的注意力质量，注意力分布余弦相似度达到 99.95%。

在全部层和 Query Head 上，Reuse Path 的平均支持召回率仅为 56.9%，保留注意力质量和 Head 输出余弦相似度仍分别达到 91.6% 和 97.2%。因此，ReTopK 不必恢复 Exact Top-K 的每个索引，只需要覆盖承载主要注意力权重的位置。

核心方法

1. Per-Head Query-Support Cache

ReTopK 为每个 Query Head 独立维护一个容量为 C 的 FIFO 缓存：

\mathcal{B}_t=\{(\bar q_c,\hat S_c)\}_{c=1}^{|\mathcal{B}_t|}

其中，q̄_c 是归一化后的历史 Query，Ŝ_c 是该 Query 当时最终选中的 Top-K 索引集合。缓存由 Prefill 阶段最后 C 个 Query 初始化，并在每个 Decode Step 后更新。无论当前 Step 走 Reuse Path 还是 Exact Path，新的 Query-Support Pair 都会写入缓存。

默认配置中 C=32。ReTopK 保存的是少量 Query 表示和索引元数据，不保存历史 QK Score、Softmax 权重或 Attention Output，同时完整 KV Cache 仍被保留。

2. 相似 Query 召回

对于当前归一化 Query，ReTopK 计算其与缓存中所有历史 Query 的余弦相似度：

\rho_{t,c}=\bar q_t^T\bar q_c

然后选取最相似的 R 个缓存项，默认 R=4。该阶段只在固定容量 C 的 Query Cache 中检索，而不是扫描长度为 L_t 的全部 Key。

3. 候选集合构造

ReTopK 将召回的历史支持集合取并集，并强制加入最近 W 个 token：

\mathcal{A}_t=\operatorname{Unique}\left(\bigcup_{c\in\mathcal{R}_t}\hat S_c\cup\mathcal{L}_t\right)

默认 W=32。Recent Window 不是普通的局部性增强，而是必要的正确性保护：新追加的 token 尚未出现在任何历史支持集合中，如果不显式加入最近窗口，它们可能长期无法进入候选集。论文消融中，当 W=0 时，PG19 PPL 从 9.178 恶化到 716.0，模型几乎失效。

候选规模满足：

M_t\le RK+W

默认 R=4、K=512、W=32 时，理论上限为 2080；由于历史支持之间存在大量重叠，64K 实验中的平均候选数约为 701。

4. 当前 Query 精确重排

历史支持只负责召回，不直接决定最终结果。ReTopK 对候选 Key 重新计算当前 Query 的精确分数：

\hat e_{t,i}=\frac{q_t^T k_i}{\sqrt d},\quad i\in\mathcal{A}_t

再在候选集合内部执行 Top-K：

\hat S_t=\operatorname{TopK}_{i\in\mathcal{A}_t}(\hat e_{t,i},K)

因此，只要当前真实 Exact Top-K 全部包含在候选集合中，ReTopK 就可以精确恢复 Exact Top-K。即使候选覆盖不完整，最终分数、Softmax 和 Value 聚合也仍然基于当前 Query 和原始 KV，而不是历史近似值。

5. Similarity Fallback 与 Periodic Refresh

历史支持存在两类风险：当前 Query 与缓存状态不相似，以及近似支持反复写回后发生长期漂移。ReTopK 使用两道保护机制：

• 相似度回退： 最大 Query 相似度低于阈值 τ 时执行完整 Exact Top-K，默认 τ=0.85；
• 周期刷新： 每隔 T_r 个 Decode Step 强制执行一次 Exact Top-K，默认 T_r=128。

整体流程如下：

for each decoding step t and query head h:
    q = normalize(current_query[h])
    similarities = cosine(q, cached_queries[h])

    if max(similarities) < tau or t % refresh_interval == 0:
        support = exact_topk(q, all_keys[h], K)
    else:
        history = supports_of_top_r_similar_queries(similarities)
        candidates = unique(history union recent_window(W))
        scores = exact_qk(q, keys[candidates])
        support = topk(scores, K)

    output = sparse_softmax_value(q, KV[support])
    cache[h].push(q, support)

系统或模型设计

复杂度变化

Exact Top-K 的完整 QK 扫描需要 O(L_t d) 工作量，并在 L_t 个分数上执行全局 Top-K。ReTopK 的 Reuse Path 包括：

• O(Cd)：Query Cache 相似度匹配；
• O(RK+W)：构造原始候选索引；
• O(M_t d)：对候选 Key 执行精确 QK；
• 在 M_t 个候选上执行 Top-K。

当 C、R、K 和 W 固定时，Reuse Path 的主要 selector 工作基本不再随上下文长度增长。额外元数据成本为每个 Head O(Cd+CK)，但完整 KV Cache 的容量保持不变。

GPU Kernel

论文实现并优化了三段 GPU 执行路径：

1. Fused Selector：完成 Query Cache 匹配、Top-R Query 选择、历史支持合并、索引排序与去重；
2. Indexed QK Kernel：直接根据候选索引读取 Key 并打分，不显式物化连续 Gathered-Key Tensor；
3. Fused Candidate Top-K + Softmax-SV：完成候选 Top-K、稀疏 Softmax、Value 聚合及 Cache Update。

缓存状态与 Layer Workspace 均预分配并常驻 GPU。论文的 Kernel 优化将 64K 下每 token、28 层的时间从 6.630 ms 降至 2.577 ms，128K 下则从 7.648 ms 降至 3.343 ms。

实验设置

PG19 使用 64 篇固定文档，并在每篇文档的 512-token teacher-forced suffix 上评估 PPL。NIAH 汇总 single-needle、multi-key 和 multi-value 三类任务。LongBench 采用 2WikiMQA、HotpotQA 和 MultiFieldQA，共 550 个样本。

主要结果

注意力级性能

在 Qwen2.5-7B、K=512 下：

ReTopK 的优势随上下文增长而扩大：Exact Top-K 的完整 QK 扫描持续变贵，而 Reuse Path 的候选检索规模基本固定。

质量与速度权衡

论文还报告了 τ=0.90 的质量导向配置。128K NIAH 三任务平均分如下：

“128K 下 PPL 仅增加 0.50%”不能直接等价为所有任务均近似无损。默认性能配置在 128K NIAH 上比 Exact Top-K 低 11.2 个绝对百分点；提高相似度阈值可以恢复质量，但注意力加速会从 3.07× 降至 2.25×。

LongBench 平均分方面，Exact Top-K、ReTopK τ=0.85 和 ReTopK τ=0.90 分别为 52.5、50.1 和 52.7。质量导向配置基本保持了 Exact Top-K 水平。

跨模型泛化

在不针对模型重新调参的情况下，ReTopK 在 Qwen2.5-7B、Llama-3.1-8B 和 Qwen2.5-14B 上实现 82.5%–89.6% 的支持复用比例。相对于 Exact Top-K，PG19 PPL 变化范围为 -0.46% 至 +2.76%，注意力加速为 1.26×–2.66×。

跨模型表格在 128K 使用 K=1024，不能与主实验中 K=512 的 3.07× 结果直接横向比较。

端到端 Decode 性能

128K 时注意力级加速达到 3.07×，但端到端只提升 1.13×，因为非 Attention 计算仍占 Exact Top-K 延迟的 70.1%。当上下文达到 1M–5M，Attention 成为主要瓶颈，端到端收益才显著扩大。

这些结果排除了 Prefill、初始化和编译时间；1M–5M 部分主要用于性能扩展性验证，论文没有同步给出这些长度上的完整任务质量结果。

局限性

1. 不解决 KV Cache 容量问题

ReTopK 保留完整 KV Cache，只减少 QK 扫描和 Top-K 选择成本。它不会降低 HBM 中的完整 KV 容量需求，也不会直接处理 KV Offload、远端存储、写入迁移与生命周期管理。因此它属于索引发现优化，而不是完整的 KV Cache Management 系统。

2. 百万上下文缺少同步质量验证

1M–5M 实验展示了较强的端到端延迟扩展性，但主要用于性能测量。论文没有充分验证百万长度下 Query 相似性、支持集合漂移和检索任务准确率是否保持稳定。

3. Exact Top-K 本身并非无损基线

128K NIAH 上，Full Attention 得分为 93.2，而 K=512 的 Exact Top-K 仅为 64.5。ReTopK 接近 Exact Top-K，不代表接近原始 Dense Attention 质量。评估时必须区分固定 Top-K 稀疏本身造成的损失，以及 ReTopK 对 Selector 的额外近似损失。

4. 缺少真实 Serving 工作负载

论文没有系统评估 Continuous Batching、多请求长度混合、请求级 Cache 隔离、吞吐量、尾延迟以及多租户元数据管理。每个请求、Layer 和 Query Head 都需要独立 Query-Support Cache，高并发时其调度和 Workspace 成本仍需验证。

5. 尚未验证原生稀疏注意力模型

实验模型仍是通过运行时 Exact Top-K 构造动态稀疏注意力。论文没有直接验证 DSA、Lightning Indexer、SFA、NSA 或共享 Indexer Layer Group 等原生稀疏架构，因此该规律能否直接迁移仍是开放问题。

与相关工作的区别

与 KV 淘汰和压缩方法的区别

StreamingLLM、H2O、SnapKV 等方法通过限制或压缩可用 KV 集合，降低容量或计算成本；被淘汰的位置之后可能无法访问。ReTopK 不删除任何 KV，只近似索引发现过程，因此保留全局可访问性，但也不节省 KV 容量。

与 Quest、SparQ、Loki 的区别

这些 Query-aware 方法通常通过页级上界、选定 Query 维度、低秩 Key 空间或其他压缩表示搜索当前相关 KV。ReTopK 缓存的不是 Key 索引结构，而是历史 Query 已经做出的 Top-K 检索决策。

与 TokenSelect 的区别

TokenSelect 更强调连续 Query 之间复用单次选择结果。ReTopK 在一个有界历史 Query Cache 中搜索，可召回多个非连续但表示相似的 Query，并对它们的支持并集执行当前 Query 精确重排。

与跨层索引复用工作的区别

IndexCache、Kascade 等方法主要利用 Layer 间 Top-K 索引相似性；ReTopK 主要利用同一 Head 内跨 Decode Step 的 Query 相似性。二者分别覆盖空间尺度与时间或状态尺度，可以进一步组合。

对大模型推理加速和存储优化的启发

1. Top-K 决策应成为一等运行时元数据

多数 KV 管理方法只保存 KV 数据和少量重要性统计。ReTopK 表明，历史 Top-K 索引集合本身具有预测价值。它可以与 Query 表示、Layer、Head、Step、Round 和 Prompt Segment 一起构成访问元数据，用于后续检索、预取和放置。

2. 从最近状态复用转向相似状态复用

连续 Step 的时间局部性只是最简单的信号。当前 Query 可能与更早、甚至上一轮对话中的某个状态更相似。在 Agent 工作负载中，工具调用、结果解析、规划和反思等阶段可能反复进入相似隐状态，因此可以建立跨 Round 的 Query-Support Memory。

3. 作为 HBM–DRAM–SSD 分层预取信号

在分层 KV 系统中，不必将历史支持直接当作最终 Top-K，而可以将其作为 Indexer 完成前的预取候选：

当前 Query 或 Indexer Query
        ↓
搜索相似历史状态
        ↓
召回历史 Top-K 索引
        ↓
高置信候选 → HBM
中置信候选 → DRAM 预取
低置信长尾 → 保留 SSD
        ↓
当前 Indexer 完成后精确确认并补取

这将 ReTopK 从“替代 Selector”转化为“隐藏远端 KV 读取延迟”的预测器。即使候选并不完全准确，只要能提前覆盖大部分最终 Top-K，就可能缩短关键路径。

4. 多个历史支持应进行置信度投票

对于每个候选 KV 索引 i，可以结合 Query 相似度与其在多个历史支持中的出现频率定义访问置信度：

f(i)=\sum_{c\in\mathcal R_t}w_c\cdot\mathbf{1}[i\in S_c]

随后按照置信度决定 HBM、DRAM 和 SSD 层级，而不是将所有候选统一拉入高层缓存。

5. 控制候选与 I/O 放大

在论文默认 K=512 时，候选理论上限为 2080，平均约 701。但原生 DSA 可能使用更大的激活预算。例如 K=2048、R=4 时，候选上限超过 8192。如果这些候选分散在 SSD 上，简单并集会导致随机 I/O、传输和缓存污染放大。

面向存储系统，更合理的策略是优先读取多个支持集合的高置信交集，按出现频率与 Query 相似度排序，分阶段预取，并为 HBM、DRAM 和 SSD 设置不同候选预算。

今夜白的观察

ReTopK 最有价值的地方并不是提出了另一种稀疏 Attention Kernel，而是把 Top-K 索引从“计算后立即丢弃的中间结果”提升为“可复用的决策缓存”。这是一个自然但长期未被充分利用的系统抽象。

从方法完整性看，ReTopK 的设计是成立的：Training-free、候选集上精确重排、不复用陈旧分数，并使用 similarity fallback、recent guard 和 periodic refresh 控制误差。它也实现了融合 GPU Kernel，而不是仅做离线算法模拟。

但论文的收益边界非常清晰。128K 下注意力加速 3.07×，端到端仅 1.13×；真正明显的系统收益主要出现在百万级上下文。同时，它完整保留 KV Cache，因此不会直接缓解 HBM 容量、Host DRAM 成本或 SSD Offload 问题。

对于多时间尺度 KV Cache 管理研究，ReTopK 更适合作为一个预测组件，而不是最终缓存策略。它补充了一条重要尺度：除了 Layer 间、连续 Step 间、Prompt Segment 间和 Round 间，还可以按照 Query 隐状态相似性进行非连续历史召回。进一步的问题是，如何把这些多尺度预测统一为分层 KV 放置和异步预取决策，并同时控制候选覆盖率、I/O 放大、预取时效性和缓存污染。

参考资料

1. 论文 arXiv 页面
2. 论文 PDF
3. 截至 2026 年 8 月 4 日，论文 arXiv 页面及正文未列出可确认属于作者团队的官方代码仓库。