# RetrievalAttention：通过向量检索加速长上下文大模型推理——中文研读笔记

**论文**：Di Liu、Meng Chen、Baotong Lu、Huiqiang Jiang、Zhenhua Han、Qianxi Zhang、Qi Chen、Chengruidong Zhang、Bailu Ding、Kai Zhang、Chen Chen、Fan Yang、Yuqing Yang、Lili Qiu，*RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval*。  
**原文**：[arXiv:2409.10516v3](https://arxiv.org/abs/2409.10516v3)（2024-12-31 修订，19 页）；[本地 PDF](../source/paper.pdf)。本次研读：2026-10-28。以下页码均从 PDF 首页起计。  
**发表**：已发表于 [NeurIPS 2025 主会论文集](https://proceedings.neurips.cc/paper_files/paper/2025/hash/4e36d4049fb0fea195a8267c8dcd0824-Abstract-Conference.html)。  
**代码**：截至 2026-10-28，未查到可核实的 RetrievalAttention 原方法实现。[Microsoft 同名仓库](https://github.com/microsoft/RetrievalAttention)目前公开的是后续系统 RetroInfer，不能直接当作本文实现。

## 论文总览

长上下文解码既要保存不断增长的键值缓存（KV cache），又要在每一步读取相关历史 token。RetrievalAttention 的出发点是：对当前查询向量，真正承载主要注意力权重的历史 token 很少，但它们的位置随查询改变。因此，作者在 CPU 内存保存大部分 KV，并利用预填充阶段已有的查询向量构建**面向注意力查询分布**的近似最近邻索引；解码时从 CPU 动态检索关键 KV，与 GPU 常驻的少量固定位置 KV 分别计算注意力，再按各自的 softmax 归一化量合并。（Sec. 2–3，PDF 第 3–7 页）

核心结果是在论文的单请求、128K 上下文、RTX 4090（24 GB）设置下，Llama-3-8B 的延迟为 **0.188 秒/词**，而精确向量扫描 Flat 为 0.922 秒/词，普通 IVF 索引为 0.373 秒/词；搜索达到 0.95 以上 Recall@100 时，仅扫描约 1–3% 的 key。任务分数接近全注意力，但并非逐任务相等，复杂 KV retrieval 仍需要更高的检索预算。（Tables 2、4–5；Fig. 6，PDF 第 8–10 页）

## 结构地图

| 原文章节 | 在论证中承担的作用 | PDF 定位 |
| --- | --- | --- |
| 1 Introduction | 长上下文的计算和存储压力；提出查询感知向量检索 | 第 1–3 页；Fig. 1、Table 1 |
| 2 Background and Motivation | 注意力公式、动态稀疏性、常规 ANNS 的分布外问题 | 第 3–5 页；Eq. (1)、Figs. 2–3 |
| 3 RetrievalAttention Design | 稀疏近似、查询感知索引、CPU–GPU 并行 | 第 5–7 页；Eq. (2)、Fig. 4 |
| 4 Evaluation | 任务精度、解码延迟、检索质量 | 第 7–10 页；Tables 2–5、Figs. 5–6 |
| 5 Related Works；6 Conclusion | 与缓存压缩、块检索、向量检索方法的区别及总结 | 第 10–11 页 |
| 附录 A–G | 模型和补充任务、A100 与百万 token、合并公式、实现、预算与 70B | 第 15–19 页；Figs. 7–8、Tables 6–11、Eqs. (3)–(5) |

## 一：动机——稀疏但位置不断变化

### 两种互相牵制的成本

有 KV cache 时，解码一步通常要访问长度为 $n$ 的历史 KV，读缓存的代价随 $n$ 增长；把完整缓存放在 GPU 又受显存限制。作者用 Llama-3-8B 举例：1M token 的 KV 缓存约需 **125 GB**；如果不缓存、每步重新处理上下文，单词生成延迟在其 A100 测试中达到 **1,765 秒**。这两个数字描述不同的瓶颈：125 GB 是完整缓存的容量需求，1,765 秒是无缓存重算的延迟，不能合并解读为同一种执行配置。（Table 1，PDF 第 2–3 页）

作者要解决的是**已有长提示词的解码阶段**。Sec. 3 明确假设长提示词的 prefill 已提前完成；在线解码延迟不含完整的索引建立和首轮 prefill 时间。因此，该方法尤其对应重复使用固定上下文或将 prefill 与解码分开的服务形态。（PDF 第 5 页）

### 为什么静态保留一批 token 不够

注意力分数是

$$
z_i=\frac{\mathbf q_t\mathbf k_i^\top}{\sqrt d},\qquad
a_{t,i}=\frac{e^{z_i}}{\sum_{j=1}^{t}e^{z_j}},\qquad
\mathbf o_t=\sum_{i=1}^{t}a_{t,i}\mathbf v_i .
$$

这里 $\mathbf q_t$ 是当前 token 的查询，$\mathbf k_i,\mathbf v_i$ 是历史第 $i$ 个 token 的键和值，$d$ 是该注意力计算所用的向量维度。按 $a_{t,i}$ 从大到小选 token，其注意力权重之和叫作**恢复率**；它量化选中 token 覆盖了多少 softmax 权重质量。（Eq. (1)，PDF 第 3 页）

在 Llama-3-8B 的 100K token KV retrieval 探针中，作者连续生成 20 个 token，每个注意力头保留动态 top-1000，平均恢复率为 **89%**；若固定采用生成第一个 token 时选出的 top-1000，后续平均只有 **71%**。这组探针支持“每步重新选关键位置”的设计，而不是仅用初始固定集合。（Fig. 2，PDF 第 3–4 页）

![图 2：动态与静态 top-1000 的注意力权重恢复率](figures/fig-02-dynamic.png)

*原文 Fig. 2，PDF 第 4 页。横轴按层和头展开；蓝线是每步动态 top-1000，橙线是第一个生成 token 决定的静态 top-1000。虚线标出跨头平均值。*

### 为什么直接套普通向量索引也不够

按内积检索较大的 $\mathbf q_t\mathbf k_i^\top$，与寻找高注意力分数的目标一致。问题在于查询 $Q$ 和键 $K$ 由不同投影产生，普通按 $K$ 之间接近程度组织的索引，未必形成适合 $Q\!\to\!K$ 查询的访问路径。作者在 Yi-9B 和 Llama-3-8B 的 128K 向量上观察到，$Q\!\to\!K$ 的马氏距离分布远离 $K\!\to\!K$；以 IVF 做 $Q\!\to\!K$ 搜索，要达到高于 0.95 的 Recall@100，约需扫描 **30–50%** 的 key，而同分布的 $K\!\to\!K$ 搜索只需约 1–5%。HNSW 的 $Q\!\to\!K$ 曲线也停留在较低召回。（Fig. 3，PDF 第 4–5 页）

![图 3a：常规索引的扫描比例与 Recall@100](figures/fig-03a-ood-recall.png)

*原文 Fig. 3a，PDF 第 4 页。横轴为扫描的 key 占比，纵轴为检索 top-100 与精确 top-100 的重合率。紫、橙曲线对应真正的 $Q\!\to\!K$ 查询。*

![图 3b：查询与键相对键分布的马氏距离](figures/fig-03b-mahalanobis.png)

*原文 Fig. 3b，PDF 第 4 页。蓝色 $Q\!\to\!K$ 距离分布明显远于红色 $K\!\to\!K$；这是本文索引设计的实测动机。*

## 二：方法——用 prefill 查询塑造 key 索引

### 1. 从完整注意力到被选 token 的注意力

设 $\mathcal I_{t,\epsilon}=\{i:a_{t,i}>\epsilon\}$。作者把集合外的低权重项舍去，然后在集合内重新做 softmax：

$$
\mathbf o_t\approx\sum_{i\in\mathcal I_{t,\epsilon}}\widetilde a_{t,i}\mathbf v_i,
\qquad
\widetilde a_{t,i}
=\frac{e^{z_i}}{\sum_{j\in\mathcal I_{t,\epsilon}}e^{z_j}}.
$$

这是 Sec. 3.1 的目标表达式（Eq. (2)，PDF 第 6 页）。$\epsilon$ 描述理想的“重要 token”集合；实验实现则采用固定检索预算，例如默认动态 top-100，较难的 KV retrieval 试 top-2000，另有 GPU 常驻 token。两者不要混成“实现逐个检查 $a_{t,i}>\epsilon$”：后者需要先知道完整注意力。

恢复率解释为什么这种近似可能可行，但它只是权重质量，并非直接给出输出向量误差或任务精度的保证；最终效果要由任务实验检验。

### 2. 注意力感知图索引

建立索引时，作者利用 prefill 中已有的查询向量。对每个采样查询，先用 GPU 算它在 key 集合中的**精确 top-$k$**，连接“查询—其重要 key”。如此得到一个从 $Q$ 的视角观察 $K$ 的邻接关系，而不是仅用 key 与 key 的相似度。直接保存查询节点会增加内存和访问开销，于是借鉴 RoarGraph 的投影：如果两个 key 被相同查询连到，就把这种关系投影成 key 节点之间的边，最后搜索图只保留 key 节点。解码的新查询沿这张图搜索到相关 key。（Sec. 3.2，Fig. 4b，PDF 第 5–6 页）

![图 4b：用 prefill 查询连接最近 key](figures/fig-04b-index.png)

*原文 Fig. 4b，PDF 第 5 页；图中画的是投影前的查询—key 连接。最终的 key 图把共享查询诱导的邻近关系保留下来。*

索引的效果体现在“找得准且少扫描”：在三种模型的 $Q\!\to\!K$ 实测中，RetrievalAttention 扫描约 **1–3%** key 即达到超过 0.95 的 Recall@100。（Fig. 6，PDF 第 10 页）这里的扫描占比和最终参与注意力的 top-100 数量是两个不同指标：前者衡量搜索时读了多少 key，后者是输出的候选规模。

### 3. CPU–GPU 的两路计算

GPU 保留可预期的固定位置 KV：当前实验使用前 **128** 个 token 和最近 **512** 个 token，共 **640** 个。其余 KV 放到 CPU 内存的索引里，解码时按当前查询动态找 $\Omega$。GPU 对常驻集合 $\mathcal W$ 计算注意力；CPU 对检索结果 $\Omega$ 计算另一份注意力，两路可以并行。CPU 直接完成检索结果的注意力，避免把整份 KV 经 PCIe 搬回 GPU。（Sec. 3.3；Fig. 4a，PDF 第 5–7 页）

![图 4a：GPU 常驻 KV 与 CPU 动态检索并行](figures/fig-04a-architecture.png)

*原文 Fig. 4a，PDF 第 5 页。上路计算固定模式的 KV，下路从 CPU 向量索引取关键 KV，两路结果在右侧合并。*

两路的局部 softmax 不能直接把输出相加，因为各自的归一化分母不同。设

$$
m_S=\max_{i\in S}z_i,\quad
Z_S=\sum_{i\in S}e^{z_i-m_S},\quad
\mathbf o_S=\frac{\sum_{i\in S}e^{z_i-m_S}\mathbf v_i}{Z_S},
\qquad S\in\{\mathcal W,\Omega\}.
$$

令 $m=\max(m_{\mathcal W},m_\Omega)$，两路在并集上的权重为

$$
\gamma_S=
\frac{e^{m_S-m}Z_S}
{e^{m_{\mathcal W}-m}Z_{\mathcal W}
 +e^{m_\Omega-m}Z_\Omega},
\qquad
\mathbf o_t^{(\mathcal W\cup\Omega)}
=\gamma_{\mathcal W}\mathbf o_{\mathcal W}
 +\gamma_\Omega\mathbf o_\Omega .
$$

这与附录 B 的 Eqs. (3)–(5) 等价：只要两集合互不重叠且两路局部计算正确，合并结果**精确等于在选中并集上重新做一次 softmax**；相对于完整历史的误差只来自“漏掉了哪些 token”，而非两路合并的计算。（PDF 第 17 页）论文的 $\tilde z$ 未在文字中单独定义，上式用两路最大值显式写出数值稳定的共同基准。

### 4. 实现上的资源安排

预填充时作者让 KV 向 CPU 传输与 GPU 全注意力流水重叠，并按头处理以降低峰值显存。CPU 侧让不同注意力头并行检索。对分组查询注意力（GQA），同组查询头虽共享 KV，但查询分布不同，所以各查询头建独立索引；索引之间通过指针共享 KV 的物理存储。（附录 C，PDF 第 18 页）这些措施有助于解释解码数据流，但论文主要延迟表测的是已经完成 prefill 的解码阶段。

## 三：实验如何检验方法

主测试机是单张 RTX 4090（24 GB）、Intel i9-10900X（10 物理核）和 128 GB DRAM，单请求 batch size 1；模型为 Llama-3-8B-Instruct-262k、Yi-6B-200K、Yi-9B-200K。附录另用 A100 测试，不能把各机器的延迟数字直接横比。任务包括七项 $\infty$-Bench（平均输入超过 100K）、覆盖 4K–128K 的 RULER，以及 Needle-in-a-haystack。（Sec. 4.1，PDF 第 7–8 页）

对照组分三类：完整注意力（无缓存重算，以及 vLLM 的 GPU KV 缓存）；丢弃或选择部分 KV 的 StreamingLLM、SnapKV、InfLLM、Quest、InfiniGen；向量检索的 Flat（精确线性扫描）和 IVF（按 key 聚类的近似检索）。Flat/IVF 与本文共享“保留 640 个固定位置 + 检索候选”的大体用法，较适合拆解索引设计的贡献。主表默认检索 top-100，$\infty$-Bench 的难题 KV retrieval 另列 top-2000。（Sec. 4.1；Table 2，PDF 第 7–8 页）

| 指标 | 读法 |
| --- | --- |
| $\infty$-Bench、RULER 分数 | 任务正确率/得分，越高越好；括号内是相对全注意力的**百分点**差值 |
| Recall@100 | 检索 top-100 与精确 top-100 的重合比例，越高越好 |
| Scanned vectors | 搜索过程访问的 key 占全部 key 比例，越低越省搜索开销 |
| 解码延迟 | 秒/生成 token，越低越快；不包括预先完成的整个 prefill |
| Act. Tokens | 一步注意力实际使用的 token 规模；与扫描比例不同 |

## 四：结果——准确率与解码成本的具体取舍

### 4.2 任务精度

下表摘自 $\infty$-Bench 的七任务平均分。斜杠左右分别是检索 top-100 与 top-2000；全注意力没有这个预算区分。（Table 2，PDF 第 8 页）

| 模型 | 全注意力 | RetrievalAttention top-100 | RetrievalAttention top-2000 |
| --- | ---: | ---: | ---: |
| Llama-3-8B | 50.4 | 48.9 | 49.6 |
| Yi-9B | 52.8 | 50.8 | 52.2 |
| Yi-6B | 45.5 | 45.0 | 45.0 |

平均接近不代表每项接近。最难的 **Retr.KV** 项中，Llama-3-8B 是全注意力 **17.5**、本文 top-100 **9.0**、top-2000 **14.0**；Yi-9B 是 **30.5 / 20.0 / 30.0**。增加预算显著弥补差距，也增加实际参与注意力的 token 数。相对地，较简单的 passKey 和 number retrieval 在多组设置中已近满分。原文对“接近全注意力”的概括应按具体任务和预算理解。（Table 2，PDF 第 8 页）

RULER 在六个上下文长度的平均分，Llama-3-8B 为全注意力 **86.54**、本文 **84.70**；Yi-9B 为 **76.87 / 76.43**；Yi-6B 为 **67.86 / 65.86**。只看 **128K** 一列则分别是 **78.74 / 74.70**、**60.51 / 59.15**、**55.77 / 51.44**。主表的“Effective”列还给出模型本身有效上下文估计；例如 Llama-3-8B 全注意力为 32K，本文为 16K，说明其并非完全复制全注意力的长程行为。（Table 3，PDF 第 9 页）

Needle-in-a-haystack 的二维热图显示，在 Llama-3-8B 的 4K–128K 测试格点上，本文均取到正确目标；附录对 Llama-3-8B-1048K 的 250K–1M 格点也给出全通过的热图。这证明给定这些合成检索位置和提示长度时可以工作；复杂任务的实际差距仍以上述任务表为准。（Figs. 5、8，PDF 第 9、16 页）

![图 5：4K–128K Needle-in-a-haystack](figures/fig-05-needle.png)

*原文 Fig. 5，PDF 第 9 页。横轴是上下文长度，纵轴是目标插入深度；全绿表示该组格点全部通过。*

### 4.3 RTX 4090 的解码延迟

在主机、Llama-3-8B、**128K** 上下文下，论文的逐词延迟如下。（Tables 4–5，PDF 第 9–10 页）

| 方法 | 检索 | 注意力 | 其他 | 总计（秒/词） |
| --- | ---: | ---: | ---: | ---: |
| Flat（精确 KNN） | 0.798 | 0.083 | 0.041 | 0.922 |
| IVF（普通近似索引） | 0.250 | 0.084 | 0.039 | 0.373 |
| RetrievalAttention | **0.064** | 0.081 | 0.043 | **0.188** |

因而相对 Flat 和 IVF 的总延迟分别约降低 **4.9 倍**和 **1.98 倍**（按基线时间除以本文时间）。检索一项从 IVF 的 0.250 秒降到 0.064 秒，约低 **74%**；从 Flat 的 0.798 秒约低 **92%**（原文文字写 91%，取整口径略有差别）。注意力部分三者相近，主要收益来自少扫描且适合查询分布的索引，而不是把 softmax 算得更快。本文检索占总延迟约 **34%**。（Sec. 4.3；Table 5，PDF 第 9–10 页）

Table 4 中更快的 StreamingLLM（0.029 秒/词）和 SnapKV（0.028 秒/词）只处理固定数量 token，不能脱离前述任务精度表称本文“最快”。同表 vLLM 在作者 RTX 4090 配置中标为 OOM；这项硬件约束是 128K 单卡结论的重要条件。值得区分版本的是：[NeurIPS 2025 出版版摘要](https://proceedings.neurips.cc/paper_files/paper/2025/hash/4e36d4049fb0fea195a8267c8dcd0824-Abstract-Conference.html)写 **0.107 秒/词**，而本次所读 **arXiv v3** 的摘要和 Table 4 都写 **0.188 秒/词**；此笔记全部按 v3 原文讨论，没有将两个版本的数字混用。

### 4.4 索引本身是否解释了加速

Fig. 6 将检索质量与扫描开销直接画在一起：在 Yi-6B、Yi-9B 和 Llama-3-8B 的 $Q\!\to\!K$ 查询上，本文索引在约 **1–3%** 扫描比例时达到高于 **0.95** 的 Recall@100；普通 IVF、HNSW 和针对一般分布外检索的 RobustVamana 曲线在相同扫描开销下表现较差。$K\!\to\!K$ 的同分布查询则对多种索引都较容易。因此 Fig. 6 把 Fig. 3 的诊断与最终延迟表连接起来：重要的不是只换一个 CPU 索引，而是让图的邻接结构反映注意力查询看到的 key 关系。（Sec. 4.4，PDF 第 9–10 页）

![图 6：三模型的召回率—扫描比例曲线](figures/fig-06-recall.png)

*原文 Fig. 6，PDF 第 10 页。红色实线是本文的 $Q\!\to\!K$ 检索；小插图放大了低扫描比例且高召回的区域。横轴 0.05 对应全部 key 的 5%。*

### 附录补充：更大长度、机器和模型

- A100（80 GB）+ 220 GB DRAM 的 **128K** Llama-3-8B 测试中，本文为 **0.155 秒/词**、IVF 为 0.345、Flat 为 0.564；完整 GPU KV 的 vLLM 为 **0.033**。这显示有足够显存时，精确 GPU KV 缓存仍可更快；本文主要处理显存不够而必须利用 CPU 的情况。（Table 7，PDF 第 16 页）
- 另一台 A100 + **1.72 TB DRAM** 机器上，Llama-3-8B 从 100K 到 1M，本文延迟由 **0.159** 升至 **0.172 秒/词**，而 IVF 从 0.308 升至 1.889。百万 token 的延迟数字依赖这台大内存服务器，不是单张 RTX 4090 的结果。（Table 8，PDF 第 16 页）
- 附录 F 在 Llama-3-8B、$\infty$-Bench、128K、每头 2000 token 总预算的另一组实验中，按 PyramidKV 分层分配预算使平均分 **49.9→50.1**，Retr.KV **14.5→16.0**，En.QA **8.7→8.5**；说明预算分配有潜力，也有任务取舍。这是附录独立配置，不应与主 Table 2 的均分作直接差值。（Table 10，PDF 第 19 页）
- 70B 补充测试用 **8 张 A100 40 GB**、$\infty$-Bench 的 KV retrieval：本文精度 **23.5**，Flat **24.0**，全注意力 **35.0**；本文延迟 **1.62 秒/词**，Flat **5.68**。相对精确检索基线速度约快 3.5 倍，但对全注意力仍有明显任务精度差距。（Table 11，PDF 第 19 页）

## 五：结论与理解

RetrievalAttention 把长上下文解码中的“找重要历史 token”落实为面向当前查询的向量搜索：先利用 prefill 查询塑造 key 图，再让 CPU 的动态检索与 GPU 的常驻窗口各做一部分注意力，按局部归一化量正确合并。它的最有说服力的证据链是 **Fig. 3 的查询/键分布差异 → Sec. 3.2 的查询感知图 → Fig. 6 的低扫描高召回 → Tables 4–5 的检索延迟下降**。

实际采用时需同时考虑三个论文中可见的条件：其一，主延迟数字针对已完成 prefill 的解码，服务若频繁更换上下文还要计入索引构建成本；其二，KV 与索引主要占 CPU 内存，百万 token 结果使用 1.72 TB DRAM；其三，有限的 top-$k$ 预算在复杂 KV retrieval 和 70B 补充测试中未完全追上全注意力。由此，方法适合显存受限、可用 CPU 内存较多、长上下文可重复使用且允许按任务调节检索预算的场景。（Sec. 3–4；附录 A、G）