跳转至

使用 PagedAttention 为 LLMs 提供服务的高效内存管理

Abstract

介绍

导言

  • 运行大模型很昂贵,需要大量如 GPU 的硬件加速器资源,成本可能很高。因此我们想让 LLM serving system 提高吞吐,从而降低每个请求的平均成本。
  • LLMs 的核心是自回归(autoregressive)的 Transformer 模型,每次生成会基于当前的输入和之前的输出 token 序列。这里我们会缓存 Key 和 Value tensor,这就是 KV cache。

    • 在评估 LLM 时由于模型权重是恒定的,并且激活仅占用 GPU 内存的一小部分,因此 KV 缓存的管理方式对于确定最大批处理大小至关重要。当管理效率低下时,KV 缓存会显著限制批处理大小(即一次处理多少个请求),从而限制 LLM 的吞吐量。

  • 论文观察到现有的 LLM serving 系统无法有效管理 KV cache。

    • 这主要是因为我们在分配 KV cache 时,像传统的 tensor 一样分配连续空间。但是 KV cache 有独特特征:随着模型生成更多的 token,KV cache 会动态增长和缩小,而且他的生存期是未知的。
    • 这会导致两个问题:内部和外部碎片;无法利用内存共享的机会。具体分析见后。
  • 因此论文提出了 PagedAttention,以像在操作系统的虚拟内存中一样,以更灵活的方式管理 KV 缓存:可以将 blocks 视为 pages,将 tokens 视为 bytes,将 requests 视为 proecesses,贡献如下:
    • 确定了服务 LLMs 中内存分配方面的挑战,并量化了它们对服务性能的影响。
    • 提出了 PagedAttention
    • 设计并实现了 vLLM,一个基于 PagedAttention 的分布式 LLMs 服务引擎。
    • 评估了 vLLM 的性能,其性能大大优于之前的 SOTA,比如 FasterTransformer 和 Orca。
  • 源代码:https://github.com/vllm-project/vllm

背景

  • 基于 Transformer 的 LLMs
  • LLM 服务与自回归生成

    • 一个请求(request)给到 LLM 服务时,会提供一个 input prompt tokens 的列表 \((x_1,\ldots,x_n)\),随后 LLM 服务生成一个 output token 的列表 \((x_{n+1}, \ldots, x_{n+T})\),我们把这两个列表拼接在在一起称为一个序列 sequence
    • LLM 服务的计算可以分为两个阶段:
      • prompt phase:将用户的整个 prompt \((x_1,\ldots,x_n)\) 作为输入,计算第一个 token 的概率(即 \(P(x_{n+1}|x_1,\ldots,x_n)\) 得到第一个输出的 token。同时还生成了 n 个 key vecotor \(k_1,\ldots,k_n\) 和 value vector \(v_1,\ldots,v_n\)。整个过程可以通过矩阵乘法来并行计算。
      • auto-regressive phase:串行生成剩余的 token。本轮迭代只需要生成 \(k_{n+t}, v_{n+t}\) 即可,其他向量已经被缓存。每次把 \(x_{n+t}\) 作为输入,通过 kv 向量 \(k_1,\ldots,k_{n+t}\)\(v_1,\ldots,v_{n+t}\) 计算下一个 token 的概率(即 \(P(x_{n+t+1}|x_1,\ldots,x_{n+t})\)。这个过程因为有依赖关系,只能用向量-矩阵乘法,是 memory-bound。
  • LLMs 的批处理技巧

    • 请求共享相同的模型全在,移动权重的开销可以在批处理请求中均摊,当批处理大小足够大时这个开销是可以被计算开销超过的。
    • 但是批处理请求存在这样的问题:
      • 不同的请求到来时间不同,需要等待。
      • 不同请求长度不同,需要填充,这样会浪费 GPU 的计算和内存。
    • 为此有细粒度的批处理机制,这些技术在迭代级别运行。每次迭代后,都会从批处理中删除已完成的请求,并添加新的请求。因此,可以在等待单次迭代后处理新请求,而不是等待整个批处理完成。此外,使用特殊的 GPU 内核,这些技术无需填充输入和输出。

算法

LLM Serving 中的内存挑战

可以批处理在一起的请求数量仍然受到 GPU 内存容量的限制,即 memory-bound。要克服这个问题,我们需要解决内存管理中的如下挑战:

  • KV cache 大
    • 随着请求数量变多,KV cache 随之快速增长。e.g. 13B 的 OPT 模型里,单个 token 的会消耗 \(2\times 5120\times 40\times 2=800KB\) 空间(分别代表 key+value,隐藏层维度,层数,fp16 的字节),同时 OPT 模型最大的序列长度是 2048,那么单个请求的 KV cache 就会占用 1.6GB 的空间。那么这样即使 GPU 内存全用来放 KV cache,也只能处理几十个请求。
    • 低效的内存管理会进一步减小批处理大小。
    • GPU 的计算增长速度快于内存容量的增长速度。
  • 复杂的 decoding 算法
    • 解码算法可以有很多种,每种算法对内存管理都有不同的影响和要求。
    • KV 缓存共享的程度取决于所采用的特定解码算法。
  • 未知的输入输出长度导致的调度

现有系统的内存管理方式

假设有两个请求 A 和 B,A 的最大长度为 2047,B 为 512,那么我们的分配结果如下:

这里会有 reserved slots,用来放接下来要输出的 token;内部碎片,因为预分配的空间大于实际需要的空间;外部碎片,因为两个请求的分配导致中间一段内存不可用。其中内外部碎片都不会再被使用,是 pure memory waste;reserved slots 最终会被用掉,但是保留这段空间可能会持续很长一段时间,而他们本来可以用来服务其他请求,却一直被保留。

方法

  • PagedAttention

    PagedAttention 将每个序列的 KV cache 划分为 KV blocks,每个 block 包含固定数量 token 的 KV cache。于是,原有的 attention 计算将被替换成 block-wise 计算:

    \[ \begin{aligned} & a_{ij}=\dfrac{\text{exp}(q_i^\top k_j\sqrt{d})}{\sum_{t=1}^i \text{exp}(q_i^\top k_t/\sqrt d)}, o_i = \sum\limits_{j=1}^i a_{ij}v_j\\ \Rightarrow & A_{ij}=\dfrac{\text{exp}(q_i^\top K_j\sqrt{d})}{\sum_{t=1}^{\lceil i/B\rceil} \text{exp}(q_i^\top K_t/\sqrt d)}, o_i = \sum\limits_{j=1}^{\lceil i/B\rceil} V_{j}A_{ij}^\top \end{aligned} \]

    这里 \(B\) 表示块大小,\(K_j=(k_{(j-1)B+1},\ldots,k_{jB})\)\(V_j=(v_{(j-1)B+1},\ldots,v_{jB})\),而 \(A_{ij}\) 表示第 j 个 KV block 的 attention score 行向量。

  • KV Cache Manager

    • 一个请求的 KV cache 表示为一系列逻辑 KV 块,在生成新 token 时从左到右填充。最后一个 KV 块可能会有没有填充满的位置,用来存放输出的 token。
    • 在 GPU worker 上,一个 block engine 分配连续的 GPU DRAM 并切分为物理 KV 块(像 CPU 的 RAM 一样)。
    • 块管理器还维护一个 block tables,用来记录逻辑 KV 块和物理 KV 块之间的映射关系。
  • Decoding with PagedAttention and vLLM

    Example

    对于请求 1 有 7 个 tokens,因此我们把前两个逻辑块(0,1)分配他,这里 vLLM 映射了两个物理块(7,1)。在 prefill 这一步,vLLM 会生成 prompt 的 KV cache,把前 4 个 token 的 KV cache 放进逻辑块 0,剩下的放进逻辑块 1。随后在第一个自回归解码中,会生成第一个 token,放进逻辑块 1 的最后一个插槽。进行第二次解码时,vLLM 会新分配一个物理块并存储映射。

  • 在其他解码场景的应用:

    • Parallel sampling:对于一个 prompt,LLM 给出多个输出,便于用户选择其中一个。

      • 我们可以让输入 prompt 的 KV cache 在多个进程间共享,这样就只需要保存一份内容即可。
      • 同时我们需要记录每个物理块的引用计数,当引用计数为 0 时,释放这个物理块。
      • vLLM 实现了 copy-on-write 写时复制的方法,当需要修改某个共享的物理块时(如最后一个没有填满 slots 的块,我们要放输出的 token),会先复制这个物理块,修改引用技术,然后写在新的块上。

      Example

    • Beam search:每次迭代只保留 top-k 个候选序列。beam search 不仅共享 prompt 的 KV cache,还会共享不同候选序列之间的其他块(即公共前缀的 KV cache blocks)。

      Example

    • Shared prefix:通常 LLM 都会为用户提供 system prompt,用以对任务的描述。

      Example

      序列 A 和序列 B 有同样的 system prompt,因此可以提前计算 system prompt 的 KV 值并缓存下来作为共享部分。

  • Scheduling and Preemption

    • 在 vLLM 里,我们采用 first-come-first-serve (FCFS) 的调度策略来处理请求,这样可以保证公平避免 starvation。
    • LLM 服务有一个独特的挑战:输入和输出的长度都是未知的,这意味着可能用光 GPU 的物理内存。因此我们需要回答:踢出哪个块?如果被踢出的块还有用我们应该如何恢复?
      • vLLM 采用的策略是,对于一个 sequence 的所有块,采用 all-or-nothing 策略,要么全部赶出去要么全部保留。进一步地,单请求多序列的情况也可以被组织安排为一个序列组。
      • 对于第二个问题,我们可以使用 Swapping / Recomputation 的做法。交换和重新计算的性能取决于 CPU RAM 和 GPU 内存之间的带宽以及 GPU 的计算能力。

        • Swapping: 当一个块被踢出时,我们会把这个块的内容写到 CPU 的 DRAM 上。

          当 vLLM 用光所有自由的物理块时,它会选择驱逐一个序列的集合并把他们的 KV cache 转移到 CPU。一旦它抢占一个序列,vLLM 会停止接受新的请求直到所有被抢占的序列完成。当有请求完成时它的所有块都被释放,这时可以把被抢占的序列搬回来继续执行。

        • Recomputation: 当被抢占的序列被重新调度时,我们会重新计算这个序列的 KV cache。

  • Distributed Execution

    • 不同 GPU worker 共享一个 KV cache 管理器。
    • 每轮迭代 scheduler 首先为批处理中的每个请求准备带有输入令牌 ID 的消息,以及每个请求的块表。接下来,调度程序将此控制消息广播给 GPU 工作线程。然后,GPU 工作线程开始使用输入令牌 ID 执行模型。在注意力层中,GPU 工作人员根据控制消息中的块表读取 KV 缓存。在执行过程中,GPU 工作线程将中间结果与 all-reduce 同步,而无需调度程序的协调。最后,GPU 工作线程将此迭代的采样令牌发送回调度程序。

实现

vLLM 是一个端到端的服务系统,具有 FastAPI 前端和基于 GPU 的推理引擎。

  • 内核级优化:融合 reshape 和块写入,融合块读取和 attention,融合块复制。
  • 支持多种解码算法,提供了 fork, append, free 三个方法,用来创建新序列、将新的 token 添加到已有序列、释放序列。

最后更新: 2024年7月2日 18:04:56
创建日期: 2024年7月2日 18:04:56

评论