PagedAttention原理

为什么会有 KV cache

在 LLM 的 decode 阶段，需要计算当前 token 与历史 token 的注意力。

计算这些 token 注意力，实际是计算最后一个 token 的 $q_t$ 向量与前面 token 的 $K$ 矩阵转置的乘积 $q_t{K}^T$ ，通过 softmax 得到注意力 $p_t$ 然后得到 $p_{t}V$ 。

$$o_t=\mathrm{softmax}(\frac{q_t{K}^T}{\sqrt{d}})V$$

对于之前 $[1,t-1]$ 的计算来说， $k_i,v_i$ 都是第 $h_i\ast W_K,h_i\ast V_k$ 得到，没有必要重复计算，因此可以缓存这一部分的值。

对于本次计算，仅需要计算 $q_t,k_t,v_t$ 以及 softmax 即可。而我们缓存下来的 $K,V$ 就叫做 KV cache

PagedAttention 之前对 KV cache 的处理以及问题

在之前的 Orca 中，将KV cache 视为一个连续的张量处理，但是 KV cache的如下性质注定其工程上不适合用于连续存储：

• 动态增长，且长度不定
• 生命周期未知

如果仅仅将其连续存储，需要根据最大可能长度预留块长，容易产生内部碎片。而回收后，其位置只能存放最大可能长度小于等于该长度的块，容易产生外部碎片。

这与操作系统中进程内存管理面临的问题高度同构。

从虚拟内存的启发

如果考虑虚拟内存：即该存储位置在逻辑上连续，但是物理上是隔开的，每次只申请一个长度为 $B$ 的空间，用完了再申请下一个，则内部碎片最大为 $B$ ,而不会产生外部碎片。

对于KV cache，我们可以做以下操作

• 将 KV 缓存切分为固定大小的块，允许非连续物理存储，在注意力计算时按块读取
• 不再预分配最大长度，而是逐块分配，将内部碎片限制在一个块以内。

LLM 推理自身的特点

虚拟内存的类比解决了单请求的内存碎片问题，但 LLM 推理还有若干自身特有的需求，是经典虚拟内存框架中没有直接对应的：

共享prompt

在操作系统中，对于共享虚拟内存不是一个强制的优化，但是在 LLM 推理中，这是一个性价比很高的选项。常见的共享 prompt 有以下三种

并行采样

例如代码补全、候选答案生成等场景中，一个输入 prompt 可能同时采样多个结果。

这些输出序列的 prompt 部分完全相同，因此 prompt 对应的 KV cache 也完全相同。

如果每个输出序列都单独保存一份 prompt KV cache，会产生明显冗余。

Beam search

不同 beam 候选共享任意前缀块，且共享模式随解码动态变化。

共享前缀

多个不同请求共享同一段预计算好的 KV 缓存

COW 区别

在虚拟内存的眼中，发生了变化，就必须复制一份新的出来，但是由于 LLM 特性，如果这个块还有空位，就直接加在这个块后面，不需要复制。

推理限制

一个序列必须所有块同时在 GPU 上才能执行推理。

一个块必须足够大才能让 GPU 可以并行读取 KV 块（因为warp原因），但是又不能太大导致内部碎片和共享效率低。

PagedAttention 实现

从虚拟内存启发，解决上述问题，就实现了 PagedAttention。

实现一个虚拟内存

设 block size 为 $B$ ，第 $j$ 个 key block 可以表示为：

$$K_j=(k_{(j-1)B+1},\dots ,k_{jB})$$

第 $j$ 个 value block 可以表示为：

$$V_j=(v_{(j-1)B+1},\dots ,v_{jB})$$

原本完整序列上的 attention 计算可以被改写成按 block 进行：

$$A_{ij}=\frac{\exp (\frac{q_i^T{K}_j}{\sqrt{d}})}{\sum _{t=1}^{\lceil \frac{i}{B}\rceil }\exp (\frac{q_i^T{K}_t}{\sqrt{d}})}$$ $$o_i=\sum _{j=1}^{\lceil \frac{i}{B}\rceil }{V}_j{A}_{ij}^T$$

此时 LLM 可以正常推理。

这一张图说明了KV cache在block table里面存储的方式，我们需要记录每个逻辑块对应的物理地址，以及它已经存储元素的个数。

解决共享prompt

如果要共享，在清理这些空间，必须要保证没有请求正在使用它，因此需要加上一个引用计数。

其次，如果共享了一部分prompt，或者有温度 $T\ne 0$ ，那么写入进的 KV cache就不一定相同，这触发了 COW 。

解决COW

PagedAttention 使用引用计数和 copy-on-write 来保证共享 block 的正确性。

每个物理 block 维护一个 reference count ，当一个物理 block被多个逻辑 block 共享时，其引用计数大于 $1$ 。

如果引用计数为 $1$ ，说明不需要被复制，直接在这个块上面写即可。否则，执行以下步骤：

1. 分配一个新的物理 block；
2. 将原共享 block 的内容复制到新 block；
3. 将当前序列的逻辑 block 映射到新物理 block；
4. 将原物理 block 的引用计数减 1；
5. 在新 block 中写入当前序列的新数据。

解决推理限制

all-or-nothing 调度

一个序列如果要参与当前 iteration 的推理，它的所有 KV block 都必须在 GPU 上。

因此 vLLM 采用 all-or-nothing 的抢占策略：

• 要么一个序列的所有 block 都保留在 GPU 上；
• 要么整个序列被抢占；
• 不会只换出一部分 block 后继续执行该序列。

对于 beam search 这类一个请求包含多个序列的场景，相关序列会作为一个 sequence group 一起调度、一起抢占、一起恢复。

抢占与恢复

如果 GPU 内存不足以支持所有活跃序列的话，GPU会抢占某些序列，对于恢复来说，vllm支持两种恢复方式：

swap

直接把这些交换到内存，计算时重新从内存迁移回来，可能出现搬运耗时比重新计算还高。

recomputation

把序列重新 prefill 一遍

block_size 选取

如果 block size 太小，会带来以下问题：

• 每个序列需要更多 block；
• block table 变大；
• attention kernel 中需要处理更多 block；
• indirect lookup 和调度开销增加；
• 单个 block 内的计算量较小，不利于摊销 kernel 启动和内存访问开销。

如果 block size 太大，会带来另一组问题：

• 最后一个 block 未填满时浪费更严重；
• 共享粒度变粗；
• Copy-on-Write 时需要复制更多 KV；
• 对于短序列，显存利用率下降更明显。

假设块长为 $B$ 从期望上来说，我们浪费的内部碎片为 $\frac{B-1}{2}$

从 GPU 效率角度看，block size 又需要足够大，以便：

• 形成较好的 memory coalescing；
• 提高 vectorized load 效率；
• 让每个 thread block / warp 有足够多的工作；
• 减少 block table 查询次数。

实际系统中，常见的 block size 取值包括 8、16、32 等。vLLM 中常见默认值为 16，但最优值通常与模型结构、head dimension、GPU 架构、dtype、batch size 和上下文长度有关，需要结合具体场景调优。

kernel 优化

还没有学这一块，先放了

分布式管理

还没有学这一块，先写完张量并行博客再说