Reading Notes: “Efficient Memory Management for Large Language Model Serving with PagedAttention”

date
Mar 8, 2025
slug
paged-attn
status
Published
tags
MLSys
summary
type
Post

Motivation

LLM Serving 要达到高吞吐量需要尽可能将多个请求打包到一个批次，然而，目前的 serving 系统由于没有办法很好地管理 KV Cache, 导致显存利用不佳，限制了批次的大小，进而限制了吞吐量。
本文提出了一种新的 attention 算法 Paged Attention, 借鉴了操作系统中分页虚拟内存的思想来减少 KV Cache 中的显存浪费，增加更灵活的显存共享方式，以解决目前系统中内存管理的缺陷。

Background

LLM Service & Autoregressive Generation

在训练完成后，LLM 通常封装成服务供外界调用，一个请求包括一个 prompt token 的列表 ，LLM 根据这些 prompt 以自回归的方式逐步生成输出 .
在每一步的生成过程中，都需要用到已经存在的 token 的 K 和 V 矩阵，因此这两个矩阵通常被缓存在 HBM 中，称为 KV Cache. 这意味着即使同一个 token，如果出现在不同的序列中，或者是同一个序列中的不同位置，也会在 KV Cache 中不同。
给定了 prompt 之后 LLM Service 生成过程可以分为两个阶段：
  • Prefill: 接受整个 prompt token 列表，生成他们的 key 和 value vectors，并计算下一个 token 的条件概率分布。由于 prompt 是事先给定的，这个过程中可以用矩阵乘法来并行化计算，以高效利用 GPU。
  • Decoding: 顺序生成输出 token，在第 个时间步，只有 被重新计算并保存，因此在这个过程中通常只做矩阵-向量乘法，导致 GPU 利用率不高，使得这个过程成为一个 memory bound 的步骤，并且带来更高的延迟。

Batching Techniques for LLMs

在 LLM Serving 中对请求做 batching 可以提高 GPU 的利用率，然而，对请求做 batching 并不是一个 trivial 的操作：
  • 请求可能在不同时候到来，简单的 batching 方法可能导致后来的请求要等前一个批次全部完成才能开始。
  • 请求中包含的 prompt 通常用不同长度，简单的 batching 通常 pad 到最长长度，导致 GPU 内存的浪费。
为了解决上述的问题，有研究者提出了细粒度的 batching 方法，比如 cellular batching 和 iteration level scheduling. 相比于传统的 batching 方法，这些方法在 iteration 的级别做 batching，在每个 iteration 之后，完成的请求会被移除，新的请求可以加入。通过这种方式，一个新到达的请求最多只需要等一个 iteration，而不需要等前一个批次完成。此外，通过特殊的 GPU kernel, 这些技术也消除了需要 padding 的限制，减少 GPU 内存的浪费。

Memory Chanllenges in LLM Serving

尽管通过细粒度的 batching 方法可以做更灵活的 batching，能够 batching 的请求数量仍然被 GPU HBM 的容量限制，主要是被需要分配给 KV Cache 的空间限制。因此，整个系统的吞吐量是 memory bound 的，要解决这个问题需要克服以下的问题：
  • KV Cache 的管理：除了模型参数之外，KV Cache 是显存开销的最大来源，如果不能高效利用和管理 KV Cache 将导致成为瓶颈。
  • 对复杂解码策略的特殊支持：在某些场景下，KV Cache 可以通过共享一部分来提高显存利用，比如做 beam search 的时候 prompt 部分的 KV Cache 可以共享，但是输出部分的不可以，此外，不同的解码策略可能导致可以共享或者不可以共享的部分不同。
  • 考虑未知输入输出长度下的调度：对于输入，需要能够支持一系列长度的 prompt length，对于输出，由于在动态增长，所以 KV Cache 也需要能够动态增长；此外，某些情况下可能显存会消耗完，这时候系统需要能够做出调度决策，比如要删除或者 swap out 某些部分。
Fragmentation in current systems
Fragmentation in current systems
由于目前的深度学习框架中许多操作要求 tensor 需要存储在连续的内存空间下，因此目前的 serving 系统大多也都采用连续的存储空间来保存 KV Cache 的内容，由于输出长度无法预测，因此通常会静态分配能够容纳最大长度的空间。这样的策略造成了内存中的 internal fragmentation，以及由于 chunk allocation 不可避免地带来了 external fragmentation.

Approach

vLLM System Overview
vLLM System Overview
为了解决上述的挑战，本文提出了 Paged Attention 算法，并在此基础上构建了一个 LLM Serving 系统 vLLM. vLLM 通过一个中心化的调度器来管理和协调多个 GPU，KV Cache Manager 通过分页的方式来管理内存，内存管理的指令通过调度器来分发到具体的 GPU 上。

Paged Attention

notion image
Paged Attention 的核心思想和操作系统中的页式内存管理类似，把 HBM 空间分为固定大小的 block，每个 block 存储固定数量的 token 对应的 k 和 v vectors，在计算 attention 的时候逐块做（需要两个 pass，第一个 pass 计算 score）
通过这种方式，可以允许 KV Cache 不存在连续的物理内存空间。

KV Cache Manager

Block Table illustration
Block Table illustration
KV Cache Manager 类似操作系统中虚拟内存模块的角色，管理逻辑地址到物理地址到映射，在这里 block table 就类似 page table，管理逻辑地址到物理地的映射，以及最后一个 block 还有什么地方没用，以用作未来输出 token 的 KV Cache. 通过这种方式，不需要静态地为未来的 token 预先分配空间，减少了 internal fragmentation，以及分页的方式消除了 external fragmentation.

Decoding with vLLM

Example: Decoding in vLLM
Example: Decoding in vLLM
在每个 iteration, vLLM 首先从候选请求中选择多个请求进行 batching，将 prompt 拼接，然后输入 paged attention kernel 并产生新的 token，在这个过程中产生的新的 token 会被放到前面的 block 的空余部分或者分配新的 block，block table 也会随之更新。
此外，在解码时，对于 prompt 部分，如果解码时会产生多个序列（比如 beam search），那么 prompt 部分的 KV Cache 可以共享（除了最后一个 logical block），而输出部分不行，这可以通过 copy-on-write 机制来实现。

Scheduling and Preemption

对于请求的调度执行，vLLM 采用 FCFS(First Come First Serve) 的策略，先到达的请求会被优先执行完成，晚到达的请求在 preemption 的时候会优先考虑。
在解码时，由于会动态增加空间，因此有可能存在最后空间耗尽的情况，这时有两个调度决策：
  • 需要驱逐哪些 block
  • 如何恢复驱逐的 block
对于第一个问题，由于在解码时一个序列总是被一起访问，因此在驱逐 block 的时候一次性把一整个序列对应的 block 驱逐，此外，对于同一个请求中的多个序列，构成一个 group，同一个 group 也倾向于被一起驱逐。
block 被驱逐后被拷贝到 CPU RAM 中，并且一旦有 block 被驱逐 vLLM 停止接收新请求，因此，CPU RAM 需要使用到交换分区大小会被 GPU HBM 分配给 KV Cache 的空间大小 bound 住。
当请求被重新调度时，只需要重新计算被驱逐时正在计算的步骤即可，由于前序部分的 KV Cache 都已经有了，这个重新计算的延迟开销较小。

Distributed Execution

vLLM 采用 Megatron Style 的张量并行，对于 MLP 进行分割以做分块矩阵乘法，对于 attention，在 attention head 维度分割。在这种情况下，每个 shard 处理的输入 token id 实际上仍然是相同的，因此，所有 shard 共用一个中心化的 KV Cache Manager 以及逻辑地址到物理地址的映射，这样简化了分布式内存管理。
在每个 step 开始时，调度器把输入 token ids 以及内存控制信息（包含 block table）广播给所有 GPU Worker, 每个 GPU Worker 根据得到的 block table 访问数据以及执行模型 forward 逻辑，以及通过 allreduce 在 worker 之间同步，结束时把 sample 的结果发回调度器。整个过程中不需要在 worker 时间进行内存管理上的同步，因为在每个 step 开始时已经和输入一起从调度器获取了。

Experiments and Results

使用 Alpaca 和 ShareGPT 做了合成数据集进行测试，和 SOTA baselines 进行比较：
  • 相同 latency 水平下，吞吐量显著提高（2x-4x）
  • 吞吐量提升以及内存节省绿在长序列，大模型以及复杂解码策略上提升更加明显
 

© Lifan Sun 2023 - 2025