在上一篇博客中,我们详细探讨了 BlockSpaceManagerNaiveBlockAllocator 的设计与内存管理策略,了解了 vLLM 在生成任务中如何通过分配内存块来支持多任务并发、动态扩展及数据交换等需求。NaiveBlockAllocator 提供了基础的内存分配和回收机制,确保了序列在生成不同阶段(如 prefill 和 decode 阶段)所需的资源。然而,随着任务规模的增大,特别是对于频繁出现相同或相似前缀的请求,简单的内存管理策略在效率上面临瓶颈。

在深度学习生成任务中,特别是长文本生成或多轮对话应用场景中,缓存机制显得尤为重要。每当需要重复生成某一段相同内容时,如果可以将已经生成的部分缓存下来以供复用,就能够显著降低系统的开销,提高任务效率。vLLMPrefixCachingBlockAllocator 就是为了解决这一需求而设计的一种优化器,其通过前缀缓存实现了对已有内容生成结果的重用,避免了无效计算。

在长 prompt 场景或多轮对话生成任务中,经常会遇到相似的请求。例如,对于同一系统 prompt 或在多轮对话场景中需要不断复用的历史上下文,重复计算会大幅增加系统的延迟,降低处理效率。如果能在生成过程中,将相同 prompt 的计算结果以 KV 缓存的形式保存下来,供后续请求复用,不仅可以节省内存和计算资源,还可以减少首 token 的生成时间。

vLLM 中的 PrefixCachingBlockAllocator 主要基于 RadixAttention 算法实现,该算法通过基数树的结构来高效地复用相似前缀内容。在 vLLM 的缓存体系中,PrefixCaching 不仅缓存 prefix 阶段的 KV 数据,还包括生成阶段的 KV 缓存数据。RadixAttention 则通过标识内容的 hash 方式来唯一标识不同的缓存单元,并在必要时动态生成 KV 缓存块。这种设计方式,避免了多轮对话或长 prompt 生成中大量重复计算。

本篇文章将深入分析 PrefixCachingBlockAllocator 的整体架构与内部机制,主要涵盖以下几点:

  1. Prefix Caching 的核心原理:介绍 RadixAttention 算法的应用,及其在 PrefixCaching 中的作用。
  2. 缓存复用的具体实现:讲解 allocate_immutable_block 方法如何实现基于内容的缓存复用。
  3. 状态管理和块分配策略:分析 BlockTrackerEvictor 等核心组件如何协作管理缓存状态,保证高效利用内存。
  4. 多轮对话和生成优化:通过典型应用场景展示 PrefixCaching 的具体优化效果。

通过本篇分析,希望能帮助大家更好地理解 vLLM 的缓存机制,挖掘 PrefixCaching 在生成任务中的优势,并在实践中灵活运用以提升生成效率。

本系列的代码基于 vLLM 的 0.6.3 版本介绍

1. Prefix Caching:RadixAttention 原理及应用

1.1 RadixAttention 的基本原理

RadixAttention 是一种创新的缓存机制,利用基数树(Radix Tree)的数据结构来管理和复用缓存。它的核心思想是通过共享公共前缀的内容,减少在生成任务中重复计算的成本,尤其是在对话生成和长 prompt 场景下尤为有效。

基数树如何用于缓存复用?

基数树是一种压缩的前缀树结构,常用于存储具有相同前缀的字符串。在 RadixAttention 中,基数树的每个节点代表一个内容块,节点之间的边代表着从父节点到子节点的前缀关系。通过这样的树结构,我们可以高效地存储和复用具有相同前缀的内容。例如:

1.2 Radix Tree 与 Prefix Tree 的区别

Radix TreePrefix Tree 在结构上有所不同,尤其是在处理共享前缀的数据场景中,Radix Tree 更为灵活和高效。具体来说:

在多轮对话或长 prompt 场景中,这种动态分片和调整的能力让 Radix Tree 能更好地组织和管理大量重叠的上下文。例如,在对话系统中,如果用户在一轮对话中询问了多种问题,每个问题的开头可能相同,但后续内容不同。Radix Tree 可以在保留共同开头的同时,为每个分支问题创建单独的节点,最大限度地复用共享的前缀内容。

1.3 Prefix + Generated KV Caching