点击下方卡片，关注 「AI视界引擎」 公众号
自注意力是大型语言模型（LLM）的一个关键组成部分，但同时也是长序列推理延迟的一个重要来源。在多租户LLM服务场景中，可以通过使用多个LLM请求在前缀中共享系统提示的概率来优化自注意力的计算和内存操作成本。在本文中，作者介绍了ChunkAttention，这是一个能感知前缀的自注意力模块，它可以检测多个请求中匹配的提示前缀，并在运行时在内存中共享它们的Key-Value张量，以提高KV缓存的内存利用率。这是通过将单一的Key-Value张量分解为较小的块，并将它们结构化为辅助前缀树来实现的。因此，在基于前缀树的KV缓存之上，作者设计了一个高效的自注意力核，其中实现了一个两阶段分区算法，在存在共享系统提示的情况下，改进自注意力计算中的数据局部性。实验表明，与最先进的实现相比，ChunkAttention可以将自注意力核的速度提高3.2-4.8倍，系统提示的长度从1024到4096不等。
1 Introduction
在过去的几年里，大型语言模型（LLM）发展了各种能力，从上下文学习到思维链推理，并在与自然语言处理相关的广泛任务中取得了显著的成功。具有代表性的模型包括GPT、LLaMA、PaLM和Gemini系列。在ChatGPT和GPT商店的成功之后，基于LLM的应用开始激增，优化LLM推理成本的需求成为了一个新的研究兴趣领域。
自注意力模块作为LLM（大型语言模型）中的关键组成部分，在推理过程中表现不佳（表1），因为它对上下文标记（KV缓存）的Key-Value张量执行密集的内存操作，并且受内存限制。内存复杂度与上下文长度成线性增长。随着对更多上下文标记的需求成为一种趋势（GPT-4需要32K），性能变得更糟。KV缓存还限制了批处理大小和系统吞吐量。例如，在GPT-3(17SB)中，使用FP16，每个标记的KV缓存需要4.5MB的内存。拥有8*A100（80G）的推理服务器内存仅能容纳70000个标记，或者说35个包含2K上下文标记的序列。
另一方面，在设计基于大型语言模型（LLM）的应用时，系统提示作为一种常见做法，会导致Key-Value（KV）缓存中的冗余（Anthropic, 2023）。通常，由于高昂的训练和推理成本，大型语言模型（LLM）会预先训练并在多租户架构中部署，以便多个应用共享。系统提示对于LLM获取每个应用的领域知识并生成更好结果至关重要。由于多个请求共享相同的系统提示，因此在提示前缀（SS2.1）方面存在重大重叠。
一个重要的问题是，作者是否可以利用系统提示的共享特性来使自注意力模块更快、更节省内存。据作者所知，唯一相关的工作是由Kwon等人（2023年）提出的方案，其中服务提供商为应用开发者预定义的一组系统提示的Key-Value张量预留内存。该方案的局限性在于：
- 预定义的系统提示是静态的，在大规模部署中频繁刷新时不灵活，因为应用开发者和服务提供商都参与了操作循环；
- 在系统提示较长和低命中率的情 况下存在内存浪费；
- 在存在共享系统提示的情况下，尚未有工作对自注意力核进行优化。
为了填补这一空白，作者提出了ChunkAttention，一个具有前缀感知KV缓存（PAKV）和两阶段划分（TPP）特点的新颖自注意力模块。ChunkAttention中的KV缓存是一个由分块上下文标记和Key-Value张量构建的前缀树。因此，KV缓存能够前缀感知，并在运行时动态检测和移除冗余，无需人工干预。KV缓存只存储当前解码中的序列的Key-Value张量，并且无内存浪费。
此外，前缀树结构为ChunkAttention提供了重新设计一个高度优化的自注意力核的上下文，该核具有两阶段划分：首先是块阶段，然后是序列阶段。具有匹配提示前缀的序列的 Query 张量被一起批处理，以与Key-Value张量执行注意力。
本文的主要贡献如下：
- 作者发现系统提示可以是长的（SS2.1），这为优化自注意力提供了机会；
- 作者提出使用前缀树来实现KV缓存，这在去除冗余方面是即插即用的、可扩展且健壮的；
- 作者实现了一个两阶段分区算法，以加速对前缀感知的KV缓存上的自注意力核；
- 作者证明了在各种系统配置下，自注意力可以从共享的系统提示中获得可行性并实证量化了收益。作者的实验表明，随着共享系统提示长度的增长，ChunkAttention可以显著加快速度，并且在没有共享系统提示的情况下，与现有高度优化的实现相比，性能没有下降。
2 Preliminaries
Shared System Prompt
在设计基于大型语言模型（LLM）的应用程序时，一种范式一直是引入系统提示。它为LLM提供指导、少量样本示例以及外部知识作为上下文，以使LLM生成更好的结果。最终传递给LLM的提示是由系统提示和特定任务输入的拼接而成。系统提示在多个请求之间共享，可能非常长。这可以在各种基于LLM的应用程序中观察到，从在线聊天机器人到离线实验。
Toolformer或使用外部工具成为LLM获取最新信息或进行精确数学计算的重要技能。它通过插件在类似ChatGPT的在线聊天机器人应用程序中实现。GPT系列模型通过函数调用提供了等同的能力。在幕后，可用的函数规范被默默注入到系统提示中。实验表明，激活6个插件后，共享系统提示的标记长度可以达到1766（附录A）。
另一个共享系统提示的来源是对LLM进行的以研究为重点的离线实验。在这些情景中，研究行人经常创建大量具有相同指令、示例或外部知识的模板请求，并迅速地向LLM发出。示例工作包括：
- Chameleon在ScienceQA和TabMWP数据集上复用策略规划和工具调用提示进行复合推理；
- CREATOR使用一个思维链（CoT）提示模板从TabMWP和MATH数据集中构建问题集；
- PDFTriage将PDF文档元数据注入提示，并在文档上执行多个问答（QA）任务；
- ToolQA进一步发布了一个QA数据集，并复用系统提示对LLM的问答进行评估。
表2显示了系统提示共享的标记数量统计。
LLM Inferencing
大型语言模型的典型推理过程包括两个阶段：预填充和解码。在接收到序列后，服务器开始进行预填充。在预填充过程中，它将所有的个提示符token 输入到大型语言模型中，计算注意力Key-Value张量，并将它们缓存起来以加速后续计算。然后，服务器进行解码。解码是自回归的，输入到大型语言模型的token是来自前一次解码迭代的完成token（或输出token）。这个过程一直持续到生成结束序列token或最大完成token为止。
当服务器同时解码（批大小）个序列时，尽管它们处于不同的迭代中，服务器仍然可以在迭代粒度上进行批处理，并一起预测所有序列的下一个标记，而不是分别预测，这被称为基于迭代的批处理。
具体来说，基于迭代的批处理将多个序列（每个序列一个标记）的最后一个输入标记拼接成一个单一输入，并在自注意力之前计算QKV投影，自注意力之后的输出投影和多层感知器。中间的自注意力没有共享权重，并且需要为每个序列独立计算。在解码过程中，新的序列可以加入，已完成的序列可以离开，显著增加了形成大批次的可能性。基于迭代的批处理已经被vLLM和文本生成推理服务器实现。本文中的ChunKAttention假设基于迭代的批处理被启用，以便为其 Kernel 高效地形成批处理。
3 Our Approach
Prefix Aware KV Cache (PAKV)
传统上，KV缓存存储在大小为 的密集张量中，其中 是批处理大小， 是头数， 是序列长度， 是头维度大小。
当多个序列共享相同的前缀标记时，Key-Value张量是相同的，因此在内存中可以共享。例如，一个特定的LLM推理服务器首先接收到序列，然后接收到序列。对于的KV缓存只能在内存中有一个物理副本。
鉴于这一特性，作者认为Key-Value（KV）缓存应当具备前缀感知能力，即把所有待解码序列的Key-Value缓存组织成前缀树。具体来说，作者会沿着序列长度维度，在内存中连续地将单一的Key-Value张量切割成片段。图1展示了存储在前缀树中的Key-Value缓存的结构。每个节点定义了一个块，存储三个基本元素：
- 由序列共享的个上下文 Token 的片段，以支持前缀树的操作；
- 大小为的键张量的切片，对应个 Token ；
- 相应的值张量的切片。
前缀树中的每一条路径定义了一个序列。在服务器中同时可能存在多个树（即森林）。例如，应用开发者设计不同的系统提示时，就可能产生这种结构。
在推理过程中存在三种可能的情况：一是新序列加入，二是完成序列移出，三是所有序列一起解码一个标记。每种情况都可以转化为前缀树操作。当一个新序列加入时，搜索并更新前缀树以插入一条新路径。当一个完成序列移出时，更新前缀树以删除其路径。在每次解码迭代中，作者将新标记附加到叶块中，或者当叶块已满时生长一个新块。
给定一个固定的块大小，内存管理是高效的。在 ChunKAttention 中，默认采用基于池的内存分配器。它同时跟踪已使用和空闲的块列表。当一个新块被请求时，分配器从空闲列表中返回一个块，或者从操作系统（OS）中分配新的内存。一旦一个序列完成，未使用的块会被返回给分配器，但分配器不会将内存释放给OS，从而防止了不必要的内存分配。一些用于对齐的内存空间未被使用。考虑到序列长度为，内存损失由限定。
通过共享公共前缀，可以同时处理的序列数量大约增加了倍。共享比例由共享 Token 的百分比定义，而是完成 Token 的数量。在内存受限的推理场景中，这有助于增加批处理大小，从而提高吞吐量。
父子关系定义了每个块覆盖的序列子集。根节点覆盖所有序列，而叶节点只覆盖一个。前缀树的一个关键属性是，前缀树中每个块覆盖的序列在序列索引维度上是连续的。因此，在自注意力中的 Query 张量切片在 Kernel 计算期间特别高效，这将在下一节中详细讨论。
Two-phase Partition (TPP)
在本节中，作者深入探讨了在独特的、对前缀敏感的KV缓存存储之上实现的自注意力核。
在预填充过程中，作者执行前缀查找，以避免对匹配提示前缀的KV投影和位置嵌入的重复计算。对于不匹配的后缀标记，仍然计算KV投影和位置嵌入，并将Key-Value张量分块并插入到前缀树中。然后，作者在整个Key-Value张量上应用现有的高度优化自注意力核，例如 FlashAttention。
在迭代解码过程中，自注意力被划分为块优先和序列优先两个阶段。这两个阶段关注 Query 张量、KV缓存块的不同切片，并使用不同的并行策略。该过程如图2所示。由于头维总是被划分的，这里省略不提，在作者的讨论中是隐含的。
块优先阶段。在块优先阶段，作者只处理被多个序列共享的块。由于GPU的流式多处理器数量（A100为108）超过了 Head 的数量（Llama 7B为32），而且按 Head 划分会低效利用硬件资源，作者对Key-Value进行了额外的划分。分块已经提供了便利。作者采用了在线softmax算法，以避免分区之间的同步要求。
计算是通过遍历前缀树中的共享块来执行的，执行部分注意力核 partial_attn 并将部分注意力结果保存到内存中，如算法1所示。序列的数量（批大小）用表示。 是由在最新解码迭代中所有个序列的最后一个标记连接形成的 Query 。
Further Optimizations
前缀树结构保存在CPU内存中。为了在GPU上运行两阶段分割 Kernel ，作者必须从前缀树生成一定的上下文，包括块，以及其覆盖序列的起始索引和结束索引，并将上下文(,,)从CPU复制到GPU内存。例如，在图2中，作者需要生成并复制，，，和。
ChunkAttention通过以下两种方式来管理开销：
- 隐藏延迟。CPU上的上下文生成步骤可以与GPU上自注意力之前的其他 Kernel 重叠。
- 延迟上下文复制。前缀树不会在每个解码迭代中改变。
作者可以将上下文缓存在GPU内存中，并且仅在树结构改变时触发内存复制。触发条件包括每次迭代块满、新序列加入，以及完成序列离开。这样的开销是摊销的。
在块优先阶段为部分注意力分配的临时内存可以通过在执行完_partial_attn_后立即执行_attn_reduce_来消除，直接将部分注意力结果合并到最终结果中。由于在前缀树中具有父子关系的多个共享块写入的同一片，_attn_reduce_需要被序列化。在GPU设备上，原子操作很重，作者不使用这种方法。然而，在CPU设备上，序列化的开销微不足道，可以通过自旋锁实现归约。
4 Experiments
评估分别在自注意力微核 Level 和端到端的GPT风格模型 Level 进行。微核 Level 的评估仅捕捉在自注意力CUDA Kernel 中花费的时间。PAKV和TPP的副作用，例如前缀树操作，在端到端的评估中被捕捉。作者使用NVIDIA A100 GPU（80G）和CUDA 11.8运行所有实验。
Microkernel Evaluation
**Baseline **。作者选择了四种自注意力实现作为 Baseline ：通过公式 的简单PyTorch实现，在xformers中实现的内存高效自注意力，集成在PyTorch中的FlashAttention，以及在vLLM中的PagedAttention。
由于Naive、xformers和FlashAttn都是建立在单一KV张量之上的，它们不能通过部分共享提示前缀的KV缓存来做到前缀感知。PagedAttn也没有实现PAKV。然而，其分页设计使作者能够手动创建一个固定的页表，将虚拟上不共享的内存映射到相同的物理内存。它模拟了KV缓存共享的场景，并帮助作者观察了PagedAttn的CUDA Kernel 性能，这部分被表示为PagedAttn*。没有任何 Kernel 支持TPP算法。
工作负载。序列以批量模式处理，批量大小为 。同一批次内的所有序列同时开始和结束。每个序列都预填充了 个提示标记，并且前 个标记是公共前缀。任务是通过迭代解码接下来的 个完成标记。作者测量解码延迟 和吞吐量（以每秒标记数或 TPS， 来定义）。对于所有实验，头维度 为 128，头的数量 为 32，块大小 为 64。所有张量均采用 FP16 格式。
结果。作者进行了实验，通过改变以下系统超参数来观察PAKV和TPP带来的性能提升：提示和共享 Token 数量、完成 Token 数量以及批处理大小。
表3展示了在给定不同提示和共享 Token 数量时，自注意力实现的延迟。ChunkAttn和PagedAttn的表现优于Naive、xformers、FlashAttn和PagedAttn，后者对共享 Token 数量不敏感。Naive的延迟分别是ChunkAttn和PagedAttn的6.6倍和2.1倍（=4096）。通过比较PagedAttn*和PagedAttn，作者观察到了物理共享KV缓存内存带来的性能提升。
尽管PagedAttn没有实现PAKV或TPP，但硬件缓存使其延迟比PagedAttn最多降低52%（=4096）：反复访问相同的物理内存块能显著提高性能。通过比较PagedAttn和ChunkAttn，可以进一步看到TPP的好处。ChunkAttn的性能比PagedAttn高出2.8-3.2倍，的范围从1024到4096。当没有 Token 共享时（=0，表3中的ChunkAttn与PagedAttn比较），TPP不会导致性能退化。因此，应始终启用TPP。
在解码过程中，序列开始出现分歧，如图3所示，ChunkAttn的性能增益逐渐减少。在给定2048个共享 Token 的情况下，当达到512时，ChunkAttn相比于PagedAttn实现了3.6的 Token 速率提升，而当达到2048时，加速比下降到2.3。然而，这仍然是一个显著的改进。
由于PagedAttn得益于物理上共享的KV缓存内存，因此ChunkAttn相对于PagedAttn的改进较低，这里只有TPP产生影响。但是，在给定=2048的情况下，当分别达到512和2048时，ChunkAttn仍然比PagedAttn*快2.0(145K对比73K)和1.5(70K对比46K)。
图4重点关注批处理大小的变化。对于所有实现方法，除了ChunkAttn和PagedAttn*之外，由于内存限制，当批处理大小达到16时，吞吐量达到峰值。给定为2048，由于更好的数据局部性和提高的算术强度，ChunkAttn的吞吐量在批处理大小从16增长到96的范围内，持续从155k增长到224k toks/s。
End-to-end Evaluation
ChunkLlama建立在Huggingface Llama和vLLM优化 Kernel （层归一化和旋转嵌入）的基础上，并遵循Apache-2.0许可，但注意力模块被ChunkAttn所替代。作者在FP16精度下对所有实验使用Open Llama2 7B模型进行运行。
基准。作者选择了两个被广泛使用并经过优化的大型语言模型服务工具包，它们具有经过验证的生产用途：最先进的vLLM 0.2.7 和 Huggingface的文本生成推理（TGI）1.3.4 [12]。
工作负载。请求按照参数为 的泊松到达过程[10]随机到达服务器，这里的 表示每秒平均请求数（RPS）。实际的批处理大小由每个系统在解码过程中动态调整，作者将其最大值均等地配置为32。应用开发者没有为服务提供商提供关于共享提示前缀的信息以预先配置。作者按照vLLM的方法测量归一化延迟（毫秒/标记或1/TPS），即每个请求的端到端延迟 （包括排队时间）的平均值除以完成标记数 ，以及KV缓存使用的峰值内存字节。
结果。如 Figure 5 所示，ChunkLlama 的推理速度最快。在共享 1024 和 2048 个前缀 Token 时，与 vLLM 相比，ChunkLlama 能实现 1.6（2.9 对 1.8）和 2.3（2.3 对 1.0）的更高吞吐量，同时保持小于 40 ms/ Token 的正常化延迟。
表4 比较了作者的 ChunkLlama 与 vLLM 的延迟和 KV 缓存内存使用情况。在没有共享前缀 Token 的情况下，ChunkLlama 没有观察到性能退化。使用长共享前缀，KV 缓存内存使用减少了 70%-90%。由于 ChunkLlama 可以解码得更快，峰值批量大小也减少了 20%-40%。
6 Conclusion
在本文中，作者提出了ChunkAttention，一个新颖的自注意力模块，以有效管理KV缓存并加速LLMs推理的自注意力核。作者成功采用前缀树创建了一个前缀感知的KV缓存。它解决了在运行时检测和移除冗余KV缓存的挑战。作者在各种配置和不同层次上评估了ChunkAttention，证明了其可行性，并且副作用可以得到控制。
实验表明，在不使用共享系统提示的情况下，ChunkAttention核可以达到与SOTA PagedAttention核相当的处理量，并且在使用1024到4096个 Token 的共享系统提示的A100（80G）上，通过应用前缀感知KV缓存和两阶段划分，其性能可以超出3.2-4.8倍。
7 Limitations
系统提示的位置。 为了在内存中共享Key-Value张量，必须在序列的开始处出现共享的系统提示。尽管这在许多作品和系统中是最常见的做法，但这并非强制性的。Liu等人（2023年）指出，当改变相关信息的位置时，语言模型的性能会显著下降，这表明模型在访问和使用长输入上下文中的信息时存在困难。
特别是，当模型必须使用长输入上下文中间的信息时，性能通常最低。因此，当应用开发者在评估后出于性能考虑而没有将系统提示放在开头，或者由于无意中的错误而没有这样做时，整个序列的KV缓存将有所不同，在这种情况下，PAKV无法节省内存，尽管它们有大量共同的标记。
微调。 除了使用系统提示外，微调是另一种注入领域知识到大型语言模型（LLM）的有前景的方法。由于高昂的训练和部署成本，LLM通常是预先训练并集中托管的，以便多个应用共享。对于每个应用来说，微调模型和部署私有实例并不具有成本效益。然而，随着硬件和软件环境的演进，微调可能会变得更加实用和流行。在这种情况下，作者不再需要为每个应用程序设计冗长的系统提示，系统提示的共享机会也减少了。到目前为止，作者还没有看到在这方面比使用系统提示更具有前景和成本效益的微调及托管解决方案。
模型和硬件兼容性。 为了达到最佳性能，ChunkAttention采用了低级CUDA编程实现两阶段划分 Kernel ，而不是利用cuDNN或PyTorch中的高级原语。作者针对最常见的LLM配置（例如，128个头维度大小）和硬件（例如，NVIDIA A100，GeForce RTX 4090和Intel Xeon CPU）调整其性能。对于其他配置和硬件，作者需要逐个案例调整和验证性能，这增加了显著的开发成本。作者相信需要社区的努力来泛化两阶段划分算法，并使其兼容更多模型配置和硬件。
参考
[1].ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition
点击上方卡片，关注 「AI视界引擎」 公众号