RetrievalAttention: Accelerating Long-Context LLM Inference via Vector Retrieval

发表时间: 2025-12 · arXiv:2409.10516 (NeurIPS 2025)

原文: https://arxiv.org/abs/2409.10516

作者/机构：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 / Microsoft Research, Shanghai Jiao Tong University, Fudan University

速读

一句话结论
RetrievalAttention 提出了一种免训练的注意力感知向量检索方法，通过将大部分 KV 缓存卸载到 CPU 并解决 Query 与 Key 向量的分布外差异问题，使得单张 RTX4090 显卡能够以极低延迟无损运行 128K 长上下文的 8B 参数大模型。

要解决什么问题
大语言模型在处理长上下文时，注意力机制的计算复杂度呈二次方增长，导致推理延迟极高，且保存历史状态的 KV 缓存会消耗海量 GPU 显存（例如 Llama-3-8B 处理 1M tokens 需要约 125GB 显存）。为了缓解这一卡点，现有研究通常利用注意力机制的动态稀疏性，即每个 Query 实际上只与极少量的关键 Key 产生强交互。然而，现有的稀疏化方法要么采用静态启发式规则导致精度严重下降，要么尝试引入近似最近邻搜索（ANNS）来动态检索关键 Key。但直接套用现成的 ANNS 引擎（如 Faiss）会遇到严重的分布外（OOD）问题：在注意力机制中，Query 和 Key 由不同的权重矩阵投影生成，导致 Query 向量与 Key 向量在空间上的马哈拉诺比斯距离极远。传统的向量索引是基于 Key 向量彼此之间的相似度构建的，用分布外的 Query 去检索时，往往需要扫描 30% 到 50% 的全量数据才能保证召回率，这完全抵消了稀疏检索带来的延迟优势。

怎么做的
RetrievalAttention 的核心思路是构建一个专门针对注意力机制优化的向量索引，并采用 CPU-GPU 协同计算的架构。首先，为了绕开 Query 和 Key 的分布外卡点，该方法在构建索引时不再仅依赖 Key 向量之间的距离，而是利用预填充阶段产生的 Query 向量来引导索引构建。具体而言，它先计算出每个 Query 向量对应的精确 K 近邻 Key 向量，建立从 Query 分布到 Key 分布的映射桥梁；随后，将连接到同一个 Query 的 Key 向量相互连接，从而在 Key 向量之间建立起从 Query 视角出发的邻近关系。这种设计使得在生成阶段，新的 Query 只需扫描 1% 到 3% 的 Key 向量就能精准找齐高注意力得分的关键 Token。在系统工程层面，该方法将 KV 缓存分为两部分。少部分具有静态重要性模式的 Token（如初始 Token 和最近的局部窗口 Token）被持久化保留在 GPU 显存中，记为集合 $\mathcal { W }$；绝大部分历史 Token 则被卸载到 CPU 内存中，记为集合 $\Omega$，并用上述算法构建向量索引。在生成每个新 Token 时，GPU 负责计算静态缓存的局部注意力输出 $\mathbf { o } _ { \mathcal { W } }$，同时 CPU 上的向量数据库根据当前 Query 检索出最相关的关键 Token 并计算局部注意力输出 $\mathbf { o } _ { \Omega }$。最后，系统通过重缩放因子 $\gamma _ { 1 }$ 和 $\gamma _ { 2 }$ 将两部分结果精确合并，其合并逻辑定义为：

$$ \mathbf { o } _ { t } = \gamma _ { 1 } \cdot \mathbf { o } _ { \mathcal { W } } + \gamma _ { 2 } \cdot \mathbf { o } _ { \Omega } $$


这种设计既把庞大的存储压力转移给了 CPU，又通过极高召回率的检索保证了注意力计算的精度，同时掩盖了 CPU 与 GPU 之间的数据传输延迟。

效果如何
实验在单张 RTX4090（24GB）和 A100 GPU 上进行，测试了 Llama-3-8B、Yi-6B 和 Yi-9B 模型。对比基线包括全注意力机制、vLLM（代表标准 KV 缓存路线）、StreamingLLM 和 SnapKV（代表静态或启发式丢弃路线）、InfLLM 和 Quest（代表分块检索路线）、InfiniGen（代表推测预取路线），以及 Flat 和 IVF（分别代表传统精确扫描和聚类向量检索路线）。在 $\infty$-Bench 和 RULER 长文本基准测试中，RetrievalAttention 展现出了与全注意力机制几乎一致的任务精度，而 StreamingLLM 等静态方法在长文本下精度出现断崖式下跌。在 128K 上下文设置下，单张 RTX4090 运行 Llama-3-8B 时，RetrievalAttention 的单 Token 生成延迟仅为 0.188 秒，相比精确扫描的 Flat 方法提速 4.9 倍，相比传统向量检索的 IVF 方法提速 1.98 倍。此外，微观测试表明其向量搜索环节的耗时仅占总延迟的 34%，大幅降低了 CPU 检索的内存带宽竞争。作者也指出该方法的局限性：目前的设计专注于加速解码生成阶段，依赖于预填充阶段的 Query 向量来构建索引，因此假设长提示词的预填充是提前完成的（例如利用现有的上下文缓存技术）。

主要贡献

随着基于Transformer的大型语言模型（LLMs）在处理长上下文方面的能力不断提升，其推理效率面临着严峻挑战。由于注意力计算的二次时间复杂度，将LLMs扩展到更长上下文时，会导致极慢的推理速度，并且为了缓存键值（KV）向量会消耗巨大的GPU内存。

为了解决这一核心问题，本文提出了RetrievalAttention，这是一种无需训练（training-free）的方法，旨在同时加速注意力计算并降低GPU内存消耗。该方法的核心创新点和研究目标如下：
1. 利用动态稀疏性与CPU向量检索：基于注意力机制的动态稀疏性，RetrievalAttention提出在CPU内存中为KV向量构建近似最近邻搜索（ANNS）索引，并在生成阶段通过向量搜索检索最相关的向量。
2. 解决分布外（OOD）挑战：作者观察到，现成的ANNS索引在注意力机制的检索任务中通常无效，这是因为查询（Query）向量和键（Key）向量之间存在分布外（Out-of-Distribution, OOD）问题。RetrievalAttention通过设计一种能够适应查询向量分布的“注意力感知（attention-aware）”向量搜索算法，成功解决了这一OOD挑战。
3. 极低的资源消耗与近乎无损的精度：评估表明，RetrievalAttention仅需访问 $1\% - 3\%$ 的数据即可实现接近全注意力的精度。这大幅降低了长上下文LLMs的推理成本和GPU内存占用。特别地，RetrievalAttention仅需单张NVIDIA RTX4090（24GB）显卡即可为8B参数模型提供128K token的上下文服务，且生成一个token仅耗时0.188秒。

Figure 1: RetrievalAttention 实现了与全注意力相似的任务精度，但展现出极低的解码延迟。
Figure 1: RetrievalAttention 实现了与全注意力相似的任务精度，但展现出极低的解码延迟。

表1展示了在单张A100 GPU上，Llama-3-8B模型在不同上下文长度下的解码延迟和KV缓存所需内存（无KV缓存优化的情况）：

提示词长度 128K 256K 512K 1M
总延迟 (s)
FFN (s)
Attention (s)		 32.8
7.6
25.2		 111
15
96		 465
31
434		 1,765
70
1,695
KV缓存内存 (GB) 15.6 31.2 62.5 125

背景知识/关键Observation/设计原则

LLM与注意力操作：在第 $t$ 个token的生成过程中，注意力操作计算查询向量 $\mathbf{q}_t \in \mathbb{R}^{1 \times d}$（$d$ 为隐藏层维度）与所有前置token的键向量 $\mathbf{k}_i \in \mathbb{R}^{1 \times d}$（$i \leq t$）之间的点积。该乘积通过 $d^{-\frac{1}{2}}$ 缩放，并通过Softmax函数归一化以产生注意力分数 $a_{t,i}$。这些分数随后对值向量 $\mathbf{v}_i$ 进行加权，得到输出 $\mathbf{o}_t$。公式如下：

$$ z_i = \frac{\mathbf{q}_t \cdot \mathbf{k}_i^T}{\sqrt{d}}, \quad a_{t,i} = \frac{e^{z_i}}{\sum_{j=1..t} e^{z_j}}, \quad \mathbf{o}_t = \sum_{i=1..t} a_{t,i} \cdot \mathbf{v}_i $$


LLM推理包含两个阶段：预填充（prefill）阶段和解码（decoding）阶段。预填充阶段仅发生一次，时间复杂度为 $O(n^2)$。解码阶段通过缓存过去的KV状态，可将复杂度降至 $O(n)$。

长上下文服务的昂贵代价：由于注意力操作的二次时间复杂度，处理长序列输入成本极高。当提示词长度达到100万token时，生成每个token需要1,765秒，其中超过96%的延迟花在注意力操作上。虽然KV缓存可以降低解码延迟，但100万token需要125GB的GPU内存，远超RTX4090（24GB）或A100（80GB）的容量。

动态与稀疏注意力：注意力计算在LLMs中表现出显著的稀疏性。尽管上下文很长，但只有一小部分具有最高注意力分数的token（即关键token）对输出有显著贡献。通过计算前 $k$ 个关键token的注意力分数累加和（恢复率），可以量化这种稀疏性。如图2蓝线所示，动态准确选择前1000个关键token，大多数注意力头能恢复超过90%的全注意力分数（平均89%）。然而，随着新token的生成，关键键向量会根据当前查询向量动态变化。如图2橙线所示，如果静态使用生成第一个token时确定的前1000个关键token来生成后续token，平均恢复率会骤降至71%。这证明了为每个查询动态选择重要token的必要性。

Figure 2: Llama-3-8B模型在100,000个token的KV检索测试中各层和头的动态稀疏性。蓝线显示动态选择前1000个关键token的平均恢复率为89%；橙线显示静态使用初始前1000个关键token会导致恢复率降至71%。
Figure 2: Llama-3-8B模型在100,000个token的KV检索测试中各层和头的动态稀疏性。蓝线显示动态选择前1000个关键token的平均恢复率为89%；橙线显示静态使用初始前1000个关键token会导致恢复率降至71%。

现成向量搜索的挑战：为了在次线性时间复杂度内准确识别关键token，使用内积作为相似度函数的ANNS索引与注意力机制的目标高度契合。然而，直接应用现成的向量索引（如Faiss中的IVF或HNSW）效果很差，根本原因在于查询向量（$Q$）和键向量（$K$）之间存在分布外（OOD）问题。如图3a所示，对于 $Q$ 到 $K$ 的搜索，基于聚类的IVF需要扫描大约30-50%的数据才能达到0.95以上的召回率，而基于图的HNSW陷入了局部最优。通过马哈拉诺比斯距离（Mahalanobis distance）量化发现（图3b），$Q$ 向量与 $K$ 向量的距离显著偏远（OOD），而 $K$ 向量彼此之间非常接近。传统的仅基于键向量之间接近度构建的索引并不符合注意力机制的要求。相比之下，$K$ 到 $K$ 的搜索或同分布的搜索则能轻松以极低的扫描率（1-5%）实现高召回。因此，高效的向量搜索索引必须在设计上考虑注意力的OOD特性。

Figure 3: (a) 从Yi-9B和Llama-3-8B（提示词长度128,000）中导出的查询向量(Q)和键向量(K)。现成的ANNS索引在Q到K搜索中表现不佳，但在K到K搜索中表现良好。(b) 查询向量与键向量距离较远，而键向量彼此之间很接近。
Figure 3: (a) 从Yi-9B和Llama-3-8B（提示词长度128,000）中导出的查询向量(Q)和键向量(K)。现成的ANNS索引在Q到K搜索中表现不佳，但在K到K搜索中表现良好。(b) 查询向量与键向量距离较远，而键向量彼此之间很接近。

方法细节

近似注意力（Approximated Attention）
基于标准的注意力公式，RetrievalAttention通过选择性地利用与高注意力分数（即 $a_{t,i}$）相关的KV向量来近似完整的注意力输出 $\mathbf{o}_t$。具体而言，首先定义 $\mathcal{T}_{t, \epsilon}$ 为注意力分数超过阈值 $\epsilon$ 的token索引子集。因此，仅考虑位于 $\mathcal{T}_{t, \epsilon}$ 中的token的稀疏注意力机制可以定义如下：

$$ \mathbf{o}_t = \sum_{i \in \mathcal{T}_{t,\epsilon}} a_{t,i} \cdot \mathbf{v}_i + \sum_{i \notin \mathcal{T}_{t,\epsilon}} a_{t,i} \cdot \mathbf{v}_i \approx \sum_{i \in \mathcal{T}_{t,\epsilon}} \widetilde{a}_{t,i} \cdot \mathbf{v}_i \quad \mathrm{where} \quad \widetilde{a}_{t,i} = \frac{e^{z_i}}{\sum_{j \in \mathcal{T}_{t,\epsilon}} e^{z_j}} $$


基于上述近似原理，RetrievalAttention被设计为仅考虑那些持久化在GPU缓存中以及通过向量索引动态检索出来的重要键值向量（即 $\mathcal{T}_{t, \epsilon}$）。

注意力感知向量搜索（Attention-Aware Vector Search）
对于每一对键和值向量，系统首先决定将其保留在CPU还是GPU内存中。被卸载到CPU内存的KV向量将由 $\mathbf{k}_i \in \mathcal{R}^d$ 进行索引，并由 $\mathbf{q}_t$ 进行查询，以找到最相关的向量。为了在token生成期间加速向量搜索，RetrievalAttention摒弃了仅考虑键向量之间接近度来构建索引的传统方法。相反，它利用预填充阶段现有的查询向量来指导键向量的索引构建，从而有效缓解分布差距。在索引构建期间，RetrievalAttention显式地建立从查询向量到其最近键向量（即精确的 $k$ 最近邻，KNN）的连接。这些KNN结果可以通过GPU高效计算，形成从查询向量分布到键向量分布的映射。使用这种结构，解码查询向量可以先搜索其最近的查询向量，然后通过分布映射获取最相关的键向量。

然而，由于需要额外存储和访问查询向量，这种连接结构在内存开销和搜索效率上仍存在缺陷。为了解决这个问题，该方法引入了最先进的跨模态ANNS索引RoarGraph 【8，Roargraph: A projected bipartite graph for efficient cross-modal approximate nearest neighbor search+2024+VLDB+https://doi.org/10.14778/3681954.3681959 】 中的投影技术，以消除所有的查询向量。具体操作为：通过连接那些与相同查询向量相连的键向量，将KNN连接直接投影到键向量中，从而高效地简化了搜索过程。这一过程将从查询向量视角看来相近的键向量连接起来，使得未来的查询向量能够进行高效的索引遍历。评估表明，通过有效建模查询和键向量之间的邻近关系，向量数据库仅需扫描 $1\% - 3\%$ 的键向量即可达到高召回率，与IVF索引 【39，Video google: A text retrieval approach to object matching in videos+2003+ICCV】 相比，索引搜索延迟显著降低了 $74\%$。

Figure 4: (a) RetrievalAttention 将大部分KV token卸载到CPU中的向量数据库，在解码阶段检索这些数据库以找到与查询最相关的KV token。(b) 在索引构建期间，我们将每个查询链接到其精确的top-k最近键向量(KNN)。
Figure 4: (a) RetrievalAttention 将大部分KV token卸载到CPU中的向量数据库，在解码阶段检索这些数据库以找到与查询最相关的KV token。(b) 在索引构建期间，我们将每个查询链接到其精确的top-k最近键向量(KNN)。

CPU-GPU 协同执行（CPU-GPU Co-execution）
为了充分利用GPU并行性并加速注意力计算，RetrievalAttention将注意力计算分解为两个不相交的KV缓存向量集：在GPU上的可预测向量集，以及在CPU上的动态向量集，最后将部分注意力输出合并在一起。系统利用在预填充阶段观察到的模式来预测在token生成期间持续激活的KV向量。与StreamingLLM 【44，Efficient streaming language models with attention sinks+2024+ICLR】 类似，当前实现使用固定的初始token和上下文的最后滑动窗口作为静态模式，并将它们持久化在GPU缓存中。在预填充阶段，系统在物理上将GPU内存中的静态token与剩余token分离，剩余token被卸载到CPU内存并由ANNS建立索引。为了最小化缓慢的PCIe接口上的数据传输，RetrievalAttention受到FlashAttention 【14，FlashAttention: Fast and memory-efficient exact attention with IO-awareness+2022+NeurIPS】 的启发，在CPU和GPU组件上独立计算注意力结果，然后将它们合并。

方法细节参考文献引用汇总

  • 【8】Chen et al., 2024a: Roargraph: A projected bipartite graph for efficient cross-modal approximate nearest neighbor search, 2024, VLDB. 引用段落及描述：在“注意力感知向量搜索”中，引用该文献以说明采用了其跨模态ANNS索引中的投影技术，用于消除查询向量，直接在键向量之间建立基于查询视角的连接。
  • 【39】Sivic & Zisserman, 2003: Video google: A text retrieval approach to object matching in videos, 2003, ICCV. 引用段落及描述：在“注意力感知向量搜索”中，引用该文献作为传统IVF聚类索引的代表，用于对比说明本文方法在扫描率和延迟上的显著优势。
  • 【44】Xiao et al., 2024b: Efficient streaming language models with attention sinks, 2024, ICLR. 引用段落及描述：在“CPU-GPU协同执行”中，引用该文献以说明本文采用的保留初始token和最近滑动窗口token在GPU中的静态模式策略与其类似。
  • 【14】Dao et al., 2022: FlashAttention: Fast and memory-efficient exact attention with IO-awareness, 2022, NeurIPS. 引用段落及描述：在“CPU-GPU协同执行”中，引用该文献以说明在CPU和GPU上独立计算部分注意力结果并进行合并的灵感来源。

实验环境

  • 硬件配置：

    • 常规测试环境：单台服务器，配备1张 NVIDIA RTX4090 GPU（24GB显存），Intel i9-10900X CPU（10个物理核，20个逻辑核），128GB DRAM。
    • 高端测试环境（附录）：1张 NVIDIA A100 GPU（80GB显存），AMD EPYC 7V13 CPU（24核），220GB DRAM。极限长文本使用AMD EPYC 7V12 CPU（48核），1.72TB 内存。

  • 模型架构：Llama-3-8B-Instruct-262k、Yi-6B-200K、Yi-9B-200K。所有模型均支持分组查询注意力（GQA）。

  • 软件配置：基线方法包括无KV缓存的全注意力、基于vLLM的KV缓存、StreamingLLM、SnapKV、InfLLM、Quest、InfiniGen，以及基于Faiss库的传统向量搜索方法Flat（精确KNN）和IVF。基准测试在真实的单批次（single-batch）场景下运行。
  • 数据集/Benchmarks：
    • $\infty$-Bench：包含7个任务（检索、代码调试、数学寻找、问答等），平均上下文长度超100K tokens。
    • RULER：包含4大类13个任务，提示词长度从4K到128K不等。
    • Needle-in-a-haystack（大海捞针）：用于测试模型在长文档中准确检索隐藏信息的能力。

实验结果

长上下文任务的精度

  • $\infty$-Bench实验：如表2所示，得益于对重要token的高效动态检索，RetrievalAttention实现了与全注意力相当的精度。而StreamingLLM和SnapKV等静态方法由于无法捕获动态变化的关键token，精度次优。InfLLM和Quest在复杂任务（如KV检索）中由于代表性向量精度低，准确率几乎为零。通过检索更多的token（如top-2000），RetrievalAttention在KV检索中达到了与全注意力几乎相同的精度。值得注意的是，Flat和IVF需要扫描100%和30%的历史键向量才能达到相同精度，而本方法仅需1-3%。
  • RULER实验：表3显示，配备RetrievalAttention的模型在不同上下文长度下（尤其是128K），均达到了与全注意力几乎相同的任务精度，而其他无需训练的基线方法则出现了明显的精度下降。
  • 大海捞针（Needle-in-a-haystack）实验：如图5所示，RetrievalAttention能够有效关注从4K到128K上下文窗口中各个位置的信息，实现了完美的检索效果。
Figure 5: RetrievalAttention在大海捞针测试中的表现。
Figure 5: RetrievalAttention在大海捞针测试中的表现。

解码延迟
* 端到端延迟对比：表4显示，随着上下文增加，全注意力的延迟呈二次方增长；vLLM因显存限制触发OOM。虽然StreamingLLM等静态方法延迟稳定，但精度损失严重。对于128K上下文（Llama-3-8B），RetrievalAttention的单token生成延迟仅为0.188秒，比Flat（0.922秒）和IVF（0.373秒）分别降低了4.9倍和1.98倍。
* 延迟分解：表5分析了128K上下文下的延迟构成。RetrievalAttention的向量搜索仅占总时间的34.0%，而Flat和IVF分别占86.6%和67.0%。与Flat和IVF相比，本方法的索引搜索延迟分别降低了91%和74%，避免了多头在CPU侧并行检索时的内存带宽争用。

索引召回率与扫描向量比例
* 图6微观分析了搜索效率。对于 $Q$ 到 $K$ 的搜索，RetrievalAttention仅需扫描极少量的键向量（1-3%）即可达到0.95以上的召回率，而传统索引和针对OOD优化的RobustVamana方法均表现不佳。这证明了本方法有效缓解了查询和键向量之间的OOD问题。

Figure 6: 分别使用查询向量(Q到K)和键向量(K到K)作为查询时，召回率与扫描键向量比例的关系。Q和K从三个长上下文LLM模型中导出。
Figure 6: 分别使用查询向量(Q到K)和键向量(K到K)作为查询时，召回率与扫描键向量比例的关系。Q和K从三个长上下文LLM模型中导出。

结论

本文提出了RetrievalAttention，该方法将大部分KV向量卸载到CPU内存中，并利用向量搜索实现动态稀疏注意力，以最小化推理成本。通过识别查询向量和键向量的不同分布，RetrievalAttention采用注意力感知的方法高效地为模型生成寻找关键token。实验结果表明，在单张RTX4090 GPU上处理128K token上下文时，其解码速度比精确KNN和传统ANNS方法分别提升了4.9倍和1.98倍。RetrievalAttention是首个支持在单张RTX4090（24GB）GPU上运行128K token的8B级别LLMs的系统，不仅延迟成本可接受，且未牺牲模型精度。未来的工作展望包括利用标量量化（如8-bit量化）进一步压缩KV向量，以在不影响精度的前提下减少内存使用。

附录

模型架构差异
附录A.1提供了实验中使用的三个模型的架构对比。所有模型均支持分组查询注意力（GQA），即多个查询头共享一个KV头。如表6所示，Yi-9B模型拥有更多的Transformer层（48层），而Llama-3-8B模型拥有更多的KV头（8个）。

大海捞针测试的额外结果
附录A.2展示了其他基线方法在大海捞针基准测试中的表现。如图7所示，StreamingLLM只能在目标位置处于其静态模式范围内时找到正确答案。InfLLM在较短上下文中保持高性能，但随着长度增加，性能显著下降。

Figure 7: 不同算法和模型在大海捞针测试中的表现。静态模式的大小固定为640（128个初始token + 局部窗口中的512个token）。
Figure 7: 不同算法和模型在大海捞针测试中的表现。静态模式的大小固定为640（128个初始token + 局部窗口中的512个token）。

极端长上下文推理性能
附录A.3测试了RetrievalAttention在极端长上下文中的表现。如图8所示，在使用Llama-3-8B-1048K模型时，当上下文长度从250K扩展到100万时，RetrievalAttention依然通过了所有测试用例，证明了注意力感知索引的鲁棒性。

Figure 8: RetrievalAttention在100万token大海捞针测试中的表现。
Figure 8: RetrievalAttention在100万token大海捞针测试中的表现。

A100上的解码延迟测试
附录A.4测试了方法在高端设备（A100 80GB）上的通用性。表7显示了128K上下文下的单token生成延迟，RetrievalAttention在保持精度的同时，延迟远低于IVF和Flat。表8展示了上下文从100K变化到1M时的延迟变化。当KV缓存启用时，vLLM在上下文超过200K时触发OOM。与Flat和IVF对上下文大小极其敏感不同，RetrievalAttention在上下文从100K增加10倍至1M时，延迟仅有微小的增加（8%）。

结合CPU和GPU侧注意力结果的公式
附录B.1详细推导了合并部分注意力的数学公式。RetrievalAttention将注意力KV向量划分为GPU上的可预测集（$\mathcal{W}$）和CPU上动态检索集（$\Omega$），即 $\mathcal{T}_{t, \epsilon} = \mathcal{W} \cup \Omega$。为了保证近似输出等于在 $\mathcal{T}_{t, \epsilon}$ 上的计算结果，系统计算两个局部的注意力输出 $\mathbf{o}{\mathcal{W}}$ 和 $\mathbf{o}$：

$$\begin{aligned} \begin{array}{rl} \mathbf{o}_{\mathcal{W}} &= \mathrm{Attn}(\mathbf{q}_t, \mathbf{K}[\mathcal{W}, :], \mathbf{V}[\mathcal{W}, :]) \\ &= \frac{\sum_{i \in \mathcal{W}} e^{z_i - \tilde{z}_1} \cdot v_i}{\sum_{i \in \mathcal{W}} e^{z_i - \tilde{z}_1}} \\ \mathbf{o}_{\Omega} &= \mathrm{Attn}(\mathbf{q}_t, \mathbf{K}[\Omega, :], \mathbf{V}[\Omega, :]) \\ &= \frac{\sum_{i \in \Omega} e^{z_i - \tilde{z}_2} \cdot v_i}{\sum_{i \in \Omega} e^{z_i - \tilde{z}_2}} \end{array} \end{aligned}$$


最终输出通过重缩放因子进行合并：$\mathbf{o}_t = \gamma_1 \cdot \mathbf{o}_{\mathcal{W}} + \gamma_2 \cdot \mathbf{o}_{\Omega}$。其中 $\tilde{z}_1 = \max_{i \in \mathcal{W}} z_i$ 和 $\tilde{z}2 = \max$，重缩放因子定义为：} z_i$ 是局部最大点积。设全局最大值为 $\tilde{z

$$\begin{aligned} \begin{array}{c} \gamma_1 = \frac{e^{\tilde{z}_1 - \tilde{z}} \cdot \sum_{i \in \mathcal{W}} e^{z_i - \tilde{z}_1}}{\sum_{i \in \mathcal{T}_{t,\epsilon}} e^{z_i - \tilde{z}}} \\ \gamma_2 = \frac{e^{\tilde{z}_2 - \tilde{z}} \cdot \sum_{i \in \Omega} e^{z_i - \tilde{z}_2}}{\sum_{i \in \mathcal{T}_{t,\epsilon}} e^{z_i - \tilde{z}}} \end{array} \end{aligned}$$

整体执行流程算法
附录B.2总结了执行流程的伪代码（Algorithm 1）。在每次生成开始时，预测活跃KV向量，移至GPU并使用FlashAttention计算；并行地在CPU侧向量数据库中检索最相关KV向量并计算；最后合并结果。

# Algorithm 1: RetrievalAttention
# Input: Query vector q_t
# Data: KV Cache in GPU K_W, V_W; CPU-based Vector Database H
# Output: Attention output o_t

# Find the predictable KV vectors
W_prime = PredictActiveTokens(...)
for i in (H union W_prime):
    H.remove(i)
    W.insert(i)  # move to GPU
for i in (not in W_prime and in H):
    W.remove(i)
    H.insert(i)  # move to CPU

# Attention on GPU
o_W = FlashAttention(q_t, K_W, V_W)

# Attention on CPU
# Search in vector database to retrieve most relevant KV vectors
Omega = VectorSearch(q_t)
o_Omega = AttentionCPU(Omega)

# Combine partial attention outputs
o_t = gamma_1 * o_W + gamma_2 * o_Omega

系统实现优化
附录C介绍了几个关键优化：
1. 预填充阶段优化：以流水线方式将KV向量向CPU的移动与GPU上的全注意力计算重叠。为最小化GPU内存峰值，跨多个注意力头顺序执行计算。
2. CPU侧多头并行：利用现代CPU多核能力，为不同注意力头的检索启动多线程并行搜索。对于GQA，由于同一组内不同查询头的分布不同，系统为每个查询头单独构建一个向量索引。
3. 最小化CPU内存使用：同一注意力组内的索引共享一份KV向量副本，索引中仅存储指针。

动态检索预算分配
附录F探讨了根据各层稀疏程度调整检索预算的影响。通过采用PyramidKV的预算分配策略（在较低层分配更多预算，在较高层分配较少），如表10所示，结合PyramidKV的RetrievalAttention在Retr.KV任务中表现更好，平均精度略微超过原始的固定预算版本（49.9% vs 50.1%）。

在更大模型上的性能
附录G在Llama-3-70B-262k模型（使用8张40GB A100）上验证了方法的泛化能力。在最复杂的KV检索任务中，表11显示RetrievalAttention达到了与精确KNN（Flat）几乎相同的任务精度（23.5% vs 24.0%），且优于Quest 80%。由于有效减少了扫描向量数，其解码速度比Flat快3.5倍（1.62s vs 5.68s）。