摘要
基于 Transformer 的大型语言模型（LLMs）中的大输入上下文窗口有助于减少幻觉现象，提高输出准确性和个性化程度。然而，随着上下文窗口的扩大，注意力阶段在执行时间中的占比逐渐增加。键值（KV）缓存通过避免重复计算缓解了部分成本，但 KV 缓存本身可能很快超出当今 GPU 高带宽内存（HBM）的容量。在本研究中，我们提出了 LongSight，这是一种用于在大上下文场景下加速注意力计算的算法 - 硬件协同设计框架。LongSight 利用原本为密集检索加速设计的计算使能 CXL 内存设备，实现 KV 缓存存储和检索的卸载。因此，LongSight 有效地将成本相对较低的 LPDDR DRAM 的价值提升至高端 HBM 的水平。我们的实验表明，仅使用单个 GPU 和单个计算使能 CXL 内存扩展器，LongSight 就能为最先进的 Llama 模型高效支持高达 100 万个令牌的上下文长度。
1 引言
预训练大型语言模型（LLMs）需要获取最新且相关的信息，以减少幻觉现象并生成准确、个性化的输出。这些信息通常作为模型输入上下文的一部分提供。事实上，在测试时技术（如思维链、少样本提示、草稿本提示、ReAct 等）、检索增强生成 以及输入数据（如长文档、代码或多轮交互）复杂度不断提升的推动下，LLMs 正越来越多地应用于需要扩展上下文窗口的场景中。
随着上下文长度的增加，推理过程的计算和内存需求也随之增长。特别是，随着上下文窗口的扩大，基于 Transformer 的 LLMs 中的注意力阶段往往会主导执行时间。键值（KV）缓存通过以增加内存压力为代价避免重复计算，在一定程度上缓解了这一问题。然而，对于更大的上下文窗口，KV 缓存的大小可能很快超出当前神经处理单元（NPUs）（如 GPU 或 TPU）上可用的高带宽内存（HBM）容量。
DReX是一种最新的计算使能 CXL 内存扩展器，用于加速密集检索。密集检索正越来越多地被用于实现检索增强生成（RAG），其中可检索项以高维嵌入向量的形式存储在向量数据库中，通过计算查询向量与嵌入向量之间的余弦相似度来识别语义相关项。DReX 在高容量 CXL 内存扩展器的 LPDDR DRAM 芯片内部及附近集成了轻量级加速器。它还引入了符号位过滤机制，无需从 DRAM 中获取完整的嵌入向量即可快速修剪搜索空间，从而显著提高密集检索的性能。
在本研究中，我们提出了 LongSight，这是一种用于在大上下文推理中加速注意力计算的算法 - 硬件协同设计框架。基于 DReX 的基础，LongSight 将其功能扩展到检索增强生成（RAG）之外，重新利用相同的计算使能 CXL 内存扩展器来加速基于 Transformer 的 LLMs 中的注意力机制。因此，LongSight 能够在极长的注意力上下文下实现高性能。
具体而言，LongSight 使神经处理单元（NPU）能够通过 DReX 中 CXL 提供的加载 / 存储接口存储 KV 缓存。LongSight 实现了一种混合密集 - 稀疏注意力算法：神经处理单元（NPU）在其本地高带宽内存（HBM）中保留最近的键值（KV）对滑动窗口，并对该窗口执行密集注意力计算，同时将稀疏注意力计算卸载到 DReX。对于稀疏部分，神经处理单元（NPU）通过 CXL 接口向 DReX 提交包含查询向量的注意力请求。然后，DReX 高效检索与查询向量点积相似度最高的前 k 个键。最后，神经处理单元（NPU）通过对密集键和稀疏键的组合集应用 softmax 函数，完成注意力操作。
我们的实验表明，配备单个 GPU 和单个 DReX 单元的 LongSight，能够为最先进的 Llama-3 1B 和 8B 模型高效支持高达 100 万个令牌的上下文长度。在当前系统中，只有使用 2 个 H100 GPU 才能支持如此长的上下文长度。在单个 GPU 支持的最大上下文长度下，LongSight 为 Llama-3 模型实现了高达 8.1-9.6 倍的吞吐量提升，以及每用户每秒 3.6-11.9 倍的令牌生成速度提升。
我们的主要贡献如下：
- 我们证实了先前的研究发现，即 Transformer 中的注意力主要受一小部分历史令牌的影响，这些令牌的键向量与当前查询向量表现出高的点积相似度。基于这一见解，我们利用密集检索加速的最新进展，实现了大上下文注意力计算。
- 我们提出了一种混合密集 - 稀疏注意力算法，该算法将短期注意力窗口保留在神经处理单元（NPU）的高带宽内存（HBM）中，并将长期注意力实现为键值向量数据库，通过前 k 点积相似度进行访问。
- 我们重新利用了最近提出的计算使能 CXL 内存扩展器 DReX（其最初设计用于密集检索），以加速我们的混合注意力机制。
2 背景
2.1 基于 Transformer 的大型语言模型（LLM）
图 1 展示了最先进的大型语言模型（LLM）架构。大型语言模型（LLM）推理包括预填充阶段和解码阶段。预填充阶段为用户输入提示构建键值（KV）缓存，而解码阶段则利用键值（KV）缓存以自回归方式生成新令牌。两个阶段共享相同的权重矩阵和模型架构，包括令牌嵌入层和多个解码器层。每个解码器层依次执行查询 - 键 - 值（QKV）生成、多头注意力计算、输出投影和前馈网络操作。查询 - 键 - 值（QKV）生成利用该层的权重矩阵，为每个输入令牌并行创建查询（Q）、键（K）和值（V）张量。之后，查询（Q）、键（K）和值（V）张量被划分为多个头，并传递到多头注意力层。在多头注意力层中，每个头计算每个查询（Q）张量与输入序列中所有先前的键（K）张量的注意力分数。注意力分数经过 softmax 函数处理后，再与相应的值（V）张量相乘。使用多个头使模型能够关注不同位置的不同表示子空间的信息。多头注意力使 Transformer 能够捕获多样化的上下文关系 ；例如，一个头可能用于确定词性，而另一个头可能用于确定情感。得到的向量通过输出投影层进行投影，并通过残差连接与输入嵌入相加，随后进入前馈网络。
图 1 的左侧和右侧分别展示了预填充阶段和解码阶段的解码器层。无论处于哪个阶段，查询 - 键 - 值（QKV）生成、输出投影和前馈网络都可以通过在多个用户之间共享权重矩阵来从批处理中受益，从而实现矩阵 - 矩阵乘法。相比之下，由于注意力计算利用的是编码每个输入令牌信息的查询（Q）、键（K）和值（V）张量，因此由于用户提示的差异，键值（KV）数据无法在一个批次中重复使用。输入长度因提示而异：在预填充阶段，输入长度对应于完整的用户提示长度，而解码阶段仅接收单个输入令牌。因此，预填充阶段的注意力计算涉及矩阵 - 矩阵乘法，而解码阶段则需要向量 - 矩阵乘法。
总之，除了解码阶段的注意力计算外，大多数大型语言模型（LLM）操作都是矩阵 - 矩阵乘法，这使得 GPU 能够高效利用计算资源。然而，解码阶段的注意力计算涉及向量 - 矩阵乘法，导致对内存带宽的高需求，并使 GPU 计算资源未得到充分利用。此外，随着输入上下文长度的增加，注意力计算需要更多的键（K）和值（V）张量。先前的研究表明，由于这些原因，解码器阶段的注意力计算可能成为主要的性能瓶颈，显著影响令牌生成吞吐量 。
2.2 分层 GPU 内存与 CXL
近年来，以 GPU 为中心的架构和系统的进步推动了分层 GPU 内存的发展，将 GPU 的字节可寻址内存空间从本地高带宽内存（HBM）扩展到包括主机 DDR 内存甚至 NVMe 固态硬盘 。这些以 GPU 为中心的方法允许 GPU 按需访问内存或存储中的数据，而无需依赖 CPU 来启动或触发此类访问。
分层 GPU 内存的最简单实现是将主机内存映射到 GPU 的地址空间，使 GPU 能够通过 GPU 线程执行的加载 / 存储指令访问主机内存，或启动直接内存访问（DMA）传输。NVIDIA 最近推出的可扩展加速数据访问（SCADA）API，使 GPU 线程能够跨分层内存层次结构对无限大小的数据集执行多粒度和随机访问 。
计算快速链路（CXL）是一项行业标准，旨在提供对分布式内存的低延迟、字节级访问，同时支持传统上通过 PCIe 互连连接且无一致性的设备之间的缓存一致性。CXL 已被超大规模数据中心采用 ，并正被用于构建机架级共享内存系统 。在本研究中，我们重点关注 CXL 通过 PCIe 将基于 DDR 的内存（“3 类” 设备）连接到处理器的能力，使其能够通过标准加载 / 存储指令直接访问。
3 大上下文生成的最新技术
先前的研究表明，随着序列长度的增加，全注意力计算的成本会逐渐主导运行时间。总序列长度包括输入上下文的大小和推理过程中生成的输出令牌数量。这两个维度在现代大型语言模型（LLM）应用中都变得越来越重要。需要更长的输入上下文来提供背景信息，并使大型语言模型（LLM）能够获取新鲜、相关的信息，以进行准确和最新的生成。同时，在需要推理和多步骤规划的任务中，生成更长的输出序列至关重要，特别是在智能体人工智能系统中，如 OpenAI 的 DeepResearch ，该系统利用强化学习，基于先前生成的令牌指导多轮生成和推理。这些系统需要大型语言模型（LLM）能够同时关注大量输入令牌和不断增加的生成输出令牌。
随着总上下文长度的增加，注意力机制在浮点运算次数（FLOPs）和内存容量方面都成为一个显著的瓶颈。此外，与前馈网络或其他受益于批处理的组件不同，注意力计算的成本无法在多个用户之间分摊。随着批处理大小的增加以提高计算效率，这一限制变得越来越突出。因此，这些挑战推动了非二次部分注意力机制（即稀疏注意力）的发展，以及利用内存内和近内存计算架构加速注意力执行的研究。
3.1 基于软件的稀疏注意力
首先，我们讨论基于软件的稀疏注意力方法，这种方法已成为降低注意力计算成本的一种流行手段。Reformer [15] 通过使用局部敏感哈希（LSH）过滤掉不太可能相关的上下文令牌，在软件中实现了稀疏注意力。然后，仅将保留的令牌传递到注意力机制。这种概率过滤降低了后续注意力阶段的计算复杂度。
然而，Reformer 的基于局部敏感哈希（LSH）的过滤引入了具有线性时间复杂度的每令牌开销。此外，Reformer 执行多轮过滤，每轮过滤都需要额外的存储或重新计算哈希桶。因此，在针对密集点积计算进行高度优化的现代硬件上执行时，稀疏性带来的好处可能会被这些开销所抵消。
此外，Reformer 假设查询和键是相同的，这使得无法使用键值头数量少于查询头数量的配置 —— 这是一种常用的最新技术，用于减少内存占用并提高注意力计算效率。
Longformer [2] 采用了不同的稀疏注意力方法，通过结合滑动窗口和有限的全局注意力来实现非精确注意力。滑动窗口机制的一个关键优势是其计算简单性以及与当前硬件的兼容性，能够实现高效执行。
然而，仅靠滑动窗口注意力本身在捕获长距离依赖关系方面存在固有的局限性。为了解决这个问题，Longformer 通过一组少量的全局注意力令牌来增强局部滑动窗口注意力，这些全局注意力令牌可以广泛地关注整个序列。Longformer 中的稀疏注意力掩码是预定义且静态的，允许模型针对特定任务进行微调，以实现可定制的注意力模式。通过根据每个任务的结构定制掩码，这种方法能够相对于基准模型提高稀疏性。
尽管有这些优点，Longformer 设计的一个关键局限性是，注意力掩码必须针对每个任务进行手动配置，这被认为是一个显著的可用性和泛化性挑战 [2]。
深度求索人工智能（DeepSeek AI）[44] 报告称，“块级选择对于在现代 GPU 上实现高效计算至关重要。” 基于这一见解，他们提出了 NSA（可原生训练的稀疏注意力），该方法将块级稀疏性与压缩密集注意力和滑动窗口注意力相结合。虽然块级稀疏性能够在 GPU 硬件上高效实现，但由于其粗粒度，它限制了可实现的整体稀疏性。此外，NSA 需要进行长上下文微调才能有效执行，这带来了额外的计算和财务成本。
总之，现有的基于软件的稀疏注意力实现难以在不进行重大模型或任务特定修改的情况下，同时实现高稀疏性和低过滤开销。
3.2 基于硬件的注意力加速
NeuPIMs [9] 和 AttAcc [29] 探索了混合架构，将传统的神经处理单元（NPUs）（如 GPU）与内存内处理（PIM）硬件相结合，用于大规模大型语言模型（LLM）推理。批处理大型语言模型（LLM）推理工作负载具有显著的异构性；因此，NeuPIMs/AttAcc 在神经处理单元（NPU）上执行计算密集型流水线阶段（如预填充和前馈网络），同时将解码阶段的内存密集型注意力计算卸载到内存内处理（PIM）单元。然而，这些研究的一个关键局限性是它们使用全密集注意力，即使在内存内处理（PIM）硬件上执行，这种注意力计算仍然成本高昂。此外，NeuPIMs 依赖于双行缓冲机制，这带来了重大的实现挑战，需要对 DRAM 电路设计进行大量修改。
CENT [8] 采用了一种系统级、内存中心的方法，将所有 Transformer 操作卸载到近内存和内存内设备。CENT 的一个关键优势是它利用高带宽内存内处理（PIM）单元来加速内存密集型计算。然而，CENT 在内存内为注意力计算实现了 BFloat16 乘加（MAC）单元，这导致了显著的面积和能量开销。此外，CENT 的全系统设计用定制硬件替代了高度优化的 GPU 或神经处理单元（NPUs），以处理计算密集型 Transformer 组件，从而增加了整体系统成本。与 NeuPIMs 一样，CENT 也实现了密集注意力，这使得难以扩展该架构以支持高效的稀疏注意力。
DynaX [42] 采用了不同的方法，利用查询向量中的稀疏性，并对查询和键采用 4 位或 6 位量化，以降低计算近似注意力分数的成本。然后，这些近似分数被用于构建基于块的稀疏注意力掩码。DynaX 通过定制硬件进一步利用过滤块内的结构化稀疏性。然而，其性能最终受到过滤过程中加载键的成本的限制。即使采用量化，至少需要加载键内存占用的 1/4・6/16≈9.4% 来评估注意力分数，这对可实现的加速比构成了限制。
4 LongSight：大上下文注意力的算法 - 系统协同设计
为了克服现有的基于软件和硬件的大上下文注意力方法的局限性，我们提出了 LongSight。LongSight 将基于 Transformer 的大型语言模型（LLM）中的注意力算法与系统和硬件架构进行协同设计，以实现任意长的上下文注意力，同时不影响现有基于 GPU 的机器学习框架的可组合性或可编程性。LongSight 基于以下观察结果引入了一种稀疏注意力算法：Transformer 主要关注先前的令牌，这些令牌对应的键向量与当前查询向量表现出高的点积相似度 [12]。在高层级上，LongSight 将键值（KV）缓存视为一个向量数据库，通过前 k 点积相似度查询进行访问，仅检索语义最相关的键，而不是关注整个键值（KV）缓存历史。
然而，这个向量数据库在几个关键方面与传统向量数据库有显著差异：
- 粒度和规模：它为每个键值头、层和用户维护独立的数据库。例如，在 Llama-3-8B（具有 8 个键值头和 32 个解码器层）中，每个用户对应 256 个独立的数据库。随着用户数量的增加，这个数字呈线性增长。
- 访问速率：数据库的访问速率为每秒令牌数（TPS）× 层数（L）× 注意力头数（H）× 批处理用户数（U），其中每秒令牌数（TPS）是令牌生成速率，层数（L）是解码器层数，注意力头数（H）是注意力头的数量，批处理用户数（U）是批处理中的用户数量。对于在优化系统上运行的 Llama-3-8B，即使是中等批处理大小，每秒令牌数（TPS）也可以达到数百个令牌 / 秒，这意味着每秒需要执行数十万个向量数据库查询。
- 延迟敏感性：由于基于 Transformer 的模型的自回归和顺序特性，这些向量数据库访问位于令牌生成的关键路径上。对于单用户批处理和 32 层，每秒生成 100 个令牌的情况下，每个层的注意力计算延迟预算约为几百微秒。我们将这个延迟预算称为注意力请求的服务水平目标（SLO）。
- 动态更新：与通常静态或变化缓慢的传统向量数据库不同，这个向量数据库需要频繁更新。在预填充阶段，整个数据库被初始化。在解码阶段，为每个用户生成的每个令牌添加一个新的键值对，导致高更新速率。
当前的方法无法满足上述要求。传统向量数据库通常使用穷举最近邻搜索（ENNS）或近似最近邻搜索（ANNS）来检索前 k 个最相似的向量。穷举最近邻搜索（ENNS）速度极慢，违反了注意力请求的服务水平目标（SLO）。实际上，密集注意力相当于对整个键值（KV）缓存执行穷举最近邻搜索（ENNS），使得基于穷举最近邻搜索（ENNS）的 LongSight 操作具有倒退性。
基于聚类和基于图的近似最近邻搜索（ANNS）虽然速度较快，但在索引构建和维护方面会产生显著开销。每次添加新向量时，都必须更新索引 —— 这个过程既昂贵又耗时。这就是为什么先前的研究只能使用近似最近邻搜索（ANNS）支持固定的长上下文 [12]。
为了在大上下文规模下实现准确且高性能的稀疏注意力，LongSight 使神经处理单元（NPU）（如 GPU）能够将用户的键值（KV）缓存存储在单独的计算使能内存设备上，并将大部分注意力计算卸载到该计算使能内存设备。LongSight 基于三个关键思想构建：
思想 1：多阶段内存内和近内存过滤与检索。第一个思想是摆脱依赖预处理索引的传统近似最近邻搜索（ANNS）方法，转而采用由 DReX [34] 开创的分层过滤机制，该机制通过多个阶段提供穷举最近邻搜索（ENNS）的准确性和近似最近邻搜索（ANNS）的速度：
- 阶段 1：内存内过滤。所有量化的键向量都排列在 DRAM 中，以便内存内处理（PIM）单元能够高效访问（图 2b ②’）。内存内处理（PIM）单元快速计算量化查询和键之间的相似度。我们基于键向量维度的全精度数据表示的符号位，对键向量维度进行一位量化（第 5 节）。使用宽松的阈值来过滤掉距离较远的键，同时确保不过滤掉语义相关的键。这极大地简化了存储，并消除了对复杂索引的需求。这使得低开销、动态过滤和高更新速率成为可能。此外，与其他仅支持有限数据类型的内存内处理方法不同，内存内过滤与任何有符号数据类型兼容。
- 阶段 2：近内存前 k 检索。通过阶段 1 过滤的向量被转发到近内存加速器，该加速器执行穷举全精度点积相似度搜索，以识别最相关的键和值（图 2b ③, ④）。然后，得到的前 k 个集合被返回给 GPU，用于最终的注意力计算（图 2b ⑤）。
思想 2：利用 CXL 实现 GPU 的细粒度访问。DReX 内部的这种多阶段内存内和近内存过滤满足了高带宽、低延迟和高精度的要求。为了支持细粒度访问，DReX 使用 CXL 将其内部内存和内存映射 I/O（MMIO）寄存器直接暴露到 GPU 地址空间。现代 GPU API 通过诸如 SCADA 之类的机制支持这一点，允许 GPU 使用标准加载 / 存储指令访问远程内存。
思想 3：混合密集和稀疏注意力。由于 LongSight 需要在生成过程中频繁更新每个头、每层和每个用户的键值（KV）数据库，我们实现了一种混合密集 - 稀疏注意力策略。
//TODO
