在介绍 Page Attention 之前我们先来介绍一下 LLM 中的 KV-Cache,KV-Cache 是 LLM 中的一个重要组件,它的作用是缓存 Key 和 Value,以减少计算量。
Attention 机制的计算公式如下:
$$ \begin{aligned} & q_i=W_q x_i, k_i=W_k x_i, v_i=W_v x_i \\ & a_{i j}=\frac{\exp \left(q_i^{\top} k_j / \sqrt{d}\right)}{\sum_{t=1}^i \exp \left(q_i^{\top} k_t / \sqrt{d}\right)}, o_i=\sum_{j=1}^i a_{i j} v_j \end{aligned} $$
其中 qi, ki, vi 分别是 Query, Key, Value,Wq, Wk, Wv 是权重矩阵,xi 是输入向量,d 是维度。aij 是第 i 个 Query 和第 j 个 Key 的 Attention 分数,oi 是第 i 个 Query 的输出。
在 Decode 阶段,每次的输出都需要计算一次 Attention,这样会导致计算量过大,KV-Cache 的作用就是缓存 Key 和 Value,以减少计算量。
KV-Cache 是一种用空间换时间的策略,通过缓存 Key 和 Value,可以减少计算量。
但是对于 LLM 推理的 Decode 阶段,这个阶段我们事先不知道什么时候生成会停止,也不知道总共会生成多少个 Token,但 KVCache 需要保证在显存中的连续性。因此,我们需要在显存里预留出一片足够大的连续空间。实际场景中训练框架一般会提前申请一份远大于我们实际请求需要的连续显存空间。
下图展示了两个请求:请求 A 的最大可能序列长度为 2048, 请求 B 的最大为 512。现有系统中的这种预分配内存块的方式存在三种主要的内存浪费来源:
在处理请求时,外部内存碎片(即由于分配方式产生的小空隙)是不会被用到的,这一点在处理请求之前就已经知道。而内部内存碎片(因为为最大序列长度预留了过多空间)也是浪费的,但只有在请求结束后,我们才会发现这部分空间其实没被用上。也就是说,这两部分内存都白白占用了空间。
虽然有些预留的内存最终会被使用,但如果预留的时间太长,特别是当预留的空间很大时,这些内存本可以用于其他请求,结果却被占用了,导致系统效率降低。这就是为什么预留过多内存会造成性能瓶颈。

picture 4
Page Attention 的目标就是解决这个问题,它可以动态地分配显存,以减少内存浪费。
与传统的注意力算法不同,PagedAttention1 允许将连续的键(key)和值(value)向量存储在不连续的内存空间中。具体来说,PagedAttention 将每个序列的 KV 缓存分成若干 KV 块。每个块包含一定数量的 token 的键和值向量,我们称其为 KV 块大小(B)。定义键块为 Kj = (k(j − 1)B + 1, …, kjB),值块为 Vj = (v(j − 1)B + 1, …, vjB)。注意力计算公式可以转化为以下按块计算的形式:
$$ A_{i j}=\frac{\exp \left(q_i^{\top} K_j / \sqrt{d}\right)}{\sum_{t=1}^{[i / B\rceil} \exp \left(q_i^{\top} K_t 1 / \sqrt{d}\right)}, o_i=\sum_{j=1}^{\lceil i / B\rceil} V_j A_{i j}^{\top}, $$