Paged Attention

T4mako4/17/26About 15 min

Paged Attention

LLM 推理过程通常分为两个阶段：prefill 和 decode。通常会使用 KV cache 加速推理

Prefill（预填充阶段）：在这个阶段，将整段 prompt 给模型做 forward，若采用 KV cache，在这个阶段我们会把 prompt 经过 $W_k$ 、 $W_v$ 后得到的 $X_k$ 、 $X_v$ 保存在 cache_k 和 cache_v 中。这样计算后面 token 的 attention 时，就不需要对前面的 token 重复计算 $X_k$ 、 $X_v$ 了
Decode（生成阶段）：在这个阶段，我们根据 prompt 的 prefill 结果，一个 token 一个 token 地生成 response，如果采用了 KV cache，就把对应 response token 的 KV 值存入 cache 中，能加速计算

因此 KV cache 面临的问题：

随着 prompt 数量变多和序列变长，KV cache 也变大，对 GPU 显存造成压力
由于输出的序列长度无法预先知道，所以我们很难提前为 KV cache 量身定制存储空间

为 KV cache 分配存储空间的常规方式

客户端会发送一个或多个 request 到服务端 server，每个 request 包含一个或多个 prompt

当 server 收到一条 request 时，server 会为 request 中的 prompts 分配 gpu 显存，其中包括 KV cache 的分配。由于推理所生成的序列长度无法预知，所以大部分框架会按照 (batch_size, max_seq_len) 的固定尺寸，在 gpu 显存上预先为一条请求开辟一块 连续的存储空间。然而，这样的分配方法很容易引起 gpu 显存利用不足的问题，进而影响模型推理时的吞吐量。）（内存碎片过多）

PagedAttention 原理

PagedAttention 的设计灵感来自操作系统中虚拟内存的分页管理技术

操作系统的虚拟内存

程序运行时，会将代码、数据等内容存放在物理内存上，原始的做法是程序直接对物理内存进行操作，决定使用它的哪些存储地址**。**但如果需要运行多个进程时，由于直接操作了物理内存地址，所以我在为自己的进程分配物理内存时，还要考虑别的进程是如何分配物理内存的。这样不同进程间的耦合性太高了，给开发带来难度。

因此，需要一种办法让各个进程间的开发能够相互独立，一种直觉的做法是**：**

给每个进程分配一个虚拟内存。每个进程在开发和运行时，可以假设这个虚拟内存上只有自己在跑，这样它就能大胆操作。
**虚拟内存负责统一规划代码、数据等如何在物理内存上最终落盘。**这个过程对所有进程来说都是透明的，进程无需操心

分段管理

在分段式内存管理中，虚拟内存会尽量为每个进程在物理内存上找到一块连续的存储空间，让进程加载自己的全部代码、数据等内容。

从上述例子可以看出，当一些进程结束后，尽管有足够的空间运行另外的进程，但是内存空间不连续。如果通过磁盘交换空间又会卡顿

此时自然会想到：不一定要给所有进程都预分配一个固定的存储块（段），可以让进程运行到哪里，或者我想用哪个具体功能时，再加载这部分相关的内容去内存，以此让整个内存分配更加动态

分页管理

虚拟内存的分业管理技术总结起来就是：

将物理内存划分为固定大小的块，称每一块为页（page）。从物理内存中模拟出来的虚拟内存也按相同的方式做划分
对于 1 个进程，我们不需要静态加载它的全部代码、数据等内容。我们想用哪部分，或者它当前跑到哪部分，我们就动态加载这部分到虚拟内存上，然后由虚拟内存帮我们做物理内存的映射。
对于 1 个进程，虽然它在物理内存上的存储不连续（可能分布在不同的page中），但它在自己的虚拟内存上是连续的。通过模拟连续内存的方式，既解决了物理内存上的碎片问题，也方便了进程的开发和运行。

PagedAttention

单个请求

假设现在你向模型 server 发送一条请求，prompt 为 Four score and seven years ago our，你希望模型能做续写。PagedAttention 的运作流程如下图：

其中：

请求（request）可理解为操作系统中的一个进程
**逻辑内存（logical KV blocks）**可理解为操作系统中的 虚拟内存，每个 block 类比于虚拟内存中的一个 page。每个 block 的大小是固定的，在 vLLM 中默认大小为 16，即可装 16 个 token 的KV 值
**块表（block table）**可理解为操作系统中的虚拟内存到物理内存的 映射表
**物理内存（physical KV blocks）**可理解为操作系统中的物理内存，物理块在 gpu 显存上，每个 block 类比于虚拟内存中的一个 page

图中带圈的序号表示操作步骤

Prefill 阶段
- 划分逻辑块：vLLM 拿到这条 prompt，先按照设定好的 block 大小 B（本例中 B = 4），为 prompt 划分逻辑块（Logical KV blocks）。由于 prompt 中有 7 个 token，所以 vLLM 用 2 个逻辑块（block 0， block 1）来装它们的 KV 值。其中，在逻辑块 1 中目前只装了 "years", "ago", "hour" 这 3 个 token 的 KV 值，有 1 个位置是空余的。这个位置就被称为保留位（reservation）
- 划分物理块：划分好逻辑块后，我们就可以将其映射到物理块中去了。物理块是实际存放 KV 值的地方。我们通过一张 block table 来记录逻辑块和物理块的 映射关系，block table 的主要内容包括：
  - 逻辑块和物理块的映射关系（physical block number）：例如逻辑块 0 对应物理块 7
  - 每个物理块上被填满的槽位（# filled）：例如在 prefill 阶段，对物理块 7，其 4 个槽位都被填满；对物理块 1，其 3 个槽位被填满。
- 正常计算 prompt 的 KV 值，并通过划分好的关系填入物理块中。
Decode阶段「生成第 1 个词」
- 使用 KV cache 计算 attention，生成第 1 个词 fathers。计算时，我们使用的是 逻辑块，即形式上这些 token 都是连续的。与此同时，vLLM 后台会通过 block table 这个映射关系，帮我们从物理块上获取数据做实际计算。通过这种方式，每个 request 都会认为自己在一个连续且充足的存储空间上操作，尽管物理上这些数据的存储并不是连续的。
- 基于新生成的词，更新 逻辑块、物理块 和 block table。对于 block table，vLLM 将它 filled 字段由 3 更新至 4。
- 分配新的逻辑块和物理块。当 fathers 更新进去后，逻辑块已装满。所以 vLLM 将开辟新的逻辑块 2，并同时更新对应的 block table 和物理块。
Decode阶段「生成第 2 个词」与生成第 1 个词类似

多个请求

nano-vllm 的分配函数

def allocate(self, seq: Sequence):
        assert not seq.block_table
        h = -1
        cache_miss = False
        for i in range(seq.num_blocks): # 遍历序列需要的逻辑块数
            token_ids = seq.block(i) # 获取第 i 个逻辑块的 token 序列
            h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1 # 计算当前逻辑块的哈希值（包含前面 block 的 hash）
            block_id = self.hash_to_block_id.get(h, -1) # 检查 hash_to_block_id 中是否存在对应的物理块。
            if block_id == -1 or self.blocks[block_id].token_ids != token_ids: 
                cache_miss = True # 缓存未命中
            if cache_miss:
                block_id = self.free_block_ids[0] # 从空闲块队列中获取第一个空闲块
                block = self._allocate_block(block_id) # 分配新块
            else:
                seq.num_cached_tokens += self.block_size # 增加已缓存 token 数量
                if block_id in self.used_block_ids: # 存在且内容匹配 (token_ids 一致)，则 缓存命中 (Cache Hit) 。
                    block = self.blocks[block_id]
                    block.ref_count += 1 # 增加其 ref_count
                else:
                    block = self._allocate_block(block_id) # 重新分配新块
            if h != -1:
                block.update(h, token_ids)
                self.hash_to_block_id[h] = block_id # 添加映射
            seq.block_table.append(block_id)

PagedAttention 在不同解码策略下的运作

根据需求不懂，大模型有不同的解码方式：

Parallel Sampling：模型发送一个请求，希望它对 prompt 做续写，并给出三种不同的回答。这个场景叫 parallel sampling。在这个场景中，我们可以将 prompt 复制 3 次后拼接成 1 个 batch 喂给模型，让它做推理。但我们也需注意到，这种方式会产生 prompt 部分 KV cache 的重复存储。
Beam Search：束搜索，这是 LLM 常用的 deocde 策略之一，即在每个 decode 阶段，我不是只产生 1 个 token，而是产生 top k 个 token（这里 k 也被称为束宽）。top k 个 token 必然对应着此刻的 top k 个序列。我把这 top k 个序列喂给模型，假设词表的大小为 |V|，那么在下一时刻，我就要在 k*|V| 个候选者中再选出 top k，以此类推。不难想象每一时刻我把 top k 序列喂给模型时，它们的前置 token 中有大量的 KV cache 是重复的。
Shared prefix：在某些大模型中，所有请求可能都会共享一个前置信息（比如 system message），这些前置信息没有必要重复存储 KV cache

Parallel Sampling

假定我们发给模型 1 个 request，这个 request 中包含 2 个 prompt/sample，记为 Sample A1 和 Sample A2 ，这两个 prompt 完全一致，都为 Four score and seven years ago our，我们希望模型对这两个 prompt 分别做续写任务

Prefill 阶段，vLLM 拿到 Sample A1 和 Sample A2
- 分配逻辑块：对于 A1，vLLM 为其分配逻辑块 block0 和 block1；对于 A2，vLLM 为其分配逻辑块 block0 和 block1。需要注意的是，A1 的逻辑块和 A2 的 逻辑块是独立的（尽管它们都叫 block0 和 block1），你可以将 A1 和 A2 视作操作系统中两个独立运行的进程。
- 分配物理块：对于 A1 和 A2，虽然逻辑块独立，但因为它们的文字完全相同，所以可以在物理内存上共享相同的空间。所以 A1 的逻辑块 block0/1 分别指向物理块 block7/1；A2 的逻辑块 block0/1 分别指向物理块 block7/1。我们设每个 物理块 下映射的 逻辑块数量为 ref count，所以对物理块 block7/1 来说，它们的 ref count 都为 2
decode 阶段
- 将生成的 token 装入逻辑块：对于 A1 和 A2 来说，将其生成的 token 装入各自的逻辑块block1。
- 触发物理块 copy-on-write 机制：由于 fathers/mothers 是两个完全不同的 token，因此对 物理块 block1 触发复制机制，即在物理内存上新开辟一块空间。此时物理块 block1 只和 A2 的逻辑块 block1 映射，将其 ref count 减去 1，物理块 block3 只和 A1 的逻辑块 block1 映射，将其 ref count 设为 1。

总结起来，vLLM 节省 KV cache 显存的核心思想是，对于相同数据对应的 KV cache，能**复用则尽量复用，**无法复用时，再考虑开辟新的物理空间。

Beam Search

假设 beam width = 4 ，虚线位置表示“当前 decoding 时刻”，

在当前阶段我们生成了 top 4 个概率最大的 token，分别装在 block5，block6，block7 和 block8 中

继续使用 beam search 算法做 decoding，继续找出 top 4 个最可能的 next token。经过计算，这top 4 next token，有 2 个来自 beam candidate 1，有 2 个来自 beam candidate 2。因此我们在 block6 中引出 block9 和 block10，用于装其中两个 top 2 next token，对 block7 也是同理。

block 9/10/11/12 中装的 top 4 next token，就成为新的 beam candidates，可以按照和上述一样的方式继续做 beam search 算法。对于 block5 和 block8，它们已经在 beam search 的搜索算法中被淘汰了，后续生成的 token 也不会和它们产生关系，所以可以清除掉这两个逻辑块，并释放它们对应的物理块的内存空间。

调度和抢占

总原则

当采用动态分配显存的办法时，虽然明面上同一时刻能处理更多的 prompt 了，但因为没有为每个 prompt 预留充足的显存空间，如果在某一时刻整个显存被打满了，而此时所有的 prompt 都没做完推理，那该怎么办？

因此需要一个调度策略来安排如何执行这些请求，这个调度原则如下：

先来的请求先服务（Firt-Come-First-Serve，FCFS）
如有抢占的需要，后来的请求先被抢占（preemption）

当一堆请求来到 vLLM 服务器做推理，导致 gpu 显存不足时，暂停这堆请求中最后到达的那些请求的推理，同时将它们相关的 KV cache 从 gpu 上释放掉。等到先来的请求做完了推理，vLLM 调度器认为 gpu 上有足够的空间了，就能恢复那些被中断的请求的执行了

在资源不足的情况下，暂时中断一些任务的执行，这样的举动就被称为抢占

终止和恢复被抢占

对于因 gpu 资源不足而 被抢占 的任务，vLLM 要完成两件事：

暂停它们的执行，同时将与之相关的 KV cache 从 gpu 上释放掉
等 gpu 资源充足时，重新恢复它们的执行

针对上面两件事，vLLM 设计了 **Swapping（交换策略）**和 Recomputation（重计算策略）

Swapping

对于被抢占的请求，vLLM 要将其 KV cahce 从 gpu 释放掉，那么：

该释放哪些 KV cache
- 在 vLLM 中，采取的是 all-or-nothing 策略，即释放被抢占请求的所有 block
要把这些 KV cahce 释放到哪里去
- 对于要释放的 KV block，如果将它们直接丢掉，那未免过于浪费。vLLM 采用的做法是将其从 gpu 上交换（Swap）到 cpu 上。这样等到 gpu 显存充份时，再把这些 block 从 cpu 上重载回来

Recomputation

知道了 Swapping 机制，重计算的过程也很好理解了：对于有些任务（比如 parallel sampling 中并行采样数 n=1 的任务），当它们因为资源不足而被抢占时，可以不做 swap，而是直接释放它们的物理块，把它们重新放入等待处理的队列中，等后续资源充足时再重新从 prefill 阶段开始做推理

总结

当一堆请求来到 vLLM 服务器上时，按照 **First-Come-First-Serve（FCFS）**原则，优先处理那些最早到来的请求。
当 gpu 资源不足时，为了让先来的请求能尽快做完推理，vLLM 会对那些后到来的请求执行“抢占”，即暂时终止它们的执行。
一旦 vLLM 决定执行抢占操作，它会暂停处理新到来的请求。在此期间，它会将被抢占的请求相关的 KV block 全部交换（swap）至 cpu 上。等交换完成后，vLLM 才会继续处理新到来的请求。
部分情况下，对于一些 seq，vllm 会抛弃它的 kv cache，将它重新放入等待队列中，后续重新做 prefill

分布式管理

分布式环境下vLLM的整体架构

在 LLM 推理实操中，某些场景下单卡是完成不了推理的，需要多卡。那么对于多 gpu 这种更普适性的情况，vLLM 是怎么处理的呢？

首先，vLLM 有一个 中央调度器（Scheduler），它负责计算和管理每张卡上 KV cache 从逻辑块到物理块的映射表(block tables)
在做分布式计算时，Schedular 会将映射表广播到各张卡上，每张卡上的 Cache engine 接收到相关信息后，负责管理各卡上的 KV block

上图中给出的例子，是用张量模型并行做分布式推理时的情况，所以图中每个 worker 上写的是 model shard。在张量并行中**，各卡上的 输入数据相同，只是各卡负责计算 不同 head 的 KV cache。所以这种情况下，各卡上的逻辑块-物理块的映射关系其实是相同** 的（用的同一张 block table），只是各卡上物理块中实际存储的数据不同而已。