Post

从QKV到VLLM的PageAttention

从QKV到VLLM的PageAttention

1. QKV 与 Attention

Transformer 的核心是 Self-Attention,由 Q、K、V 三个角色协作:

  • Q (Query):当前 token 的「提问」—— 我在找什么?

  • K (Key):所有 token 的「标签」—— 我是什么?

  • V (Value):所有 token 的「内容」—— 我能提供什么?

类比图书馆:带着问题 Q 去匹配书架上的标签 K,按相关性加权阅读对应的内容 V

核心公式:

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V

[Q/K/V 图书馆类比]

自回归生成与 KV Cache

KV Cache = 用显存换计算的时空权衡。

Q、K、V 并不是凭空出现的,它们来自同一个输入经过不同权重矩阵的线性投影。先看一个 token 在 Transformer 中的完整流转路径:

```plain text “猫” → token_id: 1234 ↓ Embedding 查表 → 词向量 x (如 4096维) ← 这是第一层的输入 ↓ ┌─ Layer 0 ──────────────────────────────┐ │ LayerNorm(x) │ │ ↓ │ │ Attention(用 x 算出 q,k,v → 输出 a₀) │ │ ↓ │ │ x + a₀ (残差连接) │ │ ↓ │ │ LayerNorm → FFN → 残差连接 → h₀ │ ← h₀ 就是「隐状态」 └────────────────────────────────────────┘ ↓ h₀ 作为下一层的输入 ┌─ Layer 1 ──────────────────────────────┐ │ 用 h₀ 算出该层的 q,k,v → … → h₁ │ └────────────────────────────────────────┘ ↓ h₁ 作为下一层的输入 …重复 32 层… ↓ 最终隐状态h → LayerNorm → 线性层 → logits → 下一个 token 概率

1
2
3
4
5
6
7
8
9
10
**隐状态 h 就是一个 token 在某一层处理完之后的向量表示。** 它从最初的词向量 x 开始,逐层「精炼」——每经过一层 Attention + FFN,就融入更多上下文信息。第 0 层的 h₀ 只有局部语义,到第 31 层的 h₃₁ 已经融合了整个序列的上下文。

在每一层中,该层的 Q、K、V 就是用该层的输入隐状态算出来的:

```plain text
第 i 层:
  q = hᵢ × W_Q⁽ⁱ⁾    ← 同一个 hᵢ
  k = hᵢ × W_K⁽ⁱ⁾    ← 乘该层自己的权重矩阵
  v = hᵢ × W_V⁽ⁱ⁾    ← 得到不同角色的向量

W_Q、W_K、W_V 是三个可学习的权重矩阵(每层各有一套),训练时学会让 Q 擅长「提问」、K 擅长「被匹配」、V 擅长「携带信息」。

推理分两个阶段,QKV 的计算方式不同:

阶段一:Prefill(处理用户输入的 prompt)

所有 prompt token 并行处理,每个 token 都算出自己的 q、k、v:

```plain text 输入: “我 喜欢 猫” (3个token)

h₁,h₂,h₃ = 各token经过前面层的隐状态

Q = [h₁,h₂,h₃] × W_Q → [q₁,q₂,q₃] ← 矩阵乘法,一次性算出 K = [h₁,h₂,h₃] × W_K → [k₁,k₂,k₃] V = [h₁,h₂,h₃] × W_V → [v₁,v₂,v₃]

将 [k₁,k₂,k₃] 和 [v₁,v₂,v₃] 全部存入 KV Cache

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
**阶段二:Decode(逐个生成新 token)**

每步只有 **1 个新 token**,只需要计算它的 q、k、v:

```plain text
新生成的 token "很" → 经过前面层 → 隐状态 h₄

q₄ = h₄ × W_Q   ← 只算这一个 token
k₄ = h₄ × W_K
v₄ = h₄ × W_V

将 k₄,v₄ 追加到 KV Cache → [k₁,k₂,k₃,k₄], [v₁,v₂,v₃,v₄]

Attention 计算:
  0. 从 Cache 读取k&v
  1. q₄ × [k₁,k₂,k₃,k₄]ᵀ   = 得到注意力权重
  2. 加权求和 [v₁,v₂,v₃,v₄] = 得到输出

关键点总结:

 PrefillDecode
输入所有 prompt tokens仅 1 个新 token
计算所有 token 的 Q,K,V仅新 token 的 Q,K,V
K,V 来源当前计算 → 写入 Cache新 token 计算 + Cache 读取
Q 的用途每个 q 查询它之前的所有 kq 查询 Cache 中所有 k
并行度高(多 token 并行)低(单 token)

Q、K、V 都由同一个隐状态 h 乘以不同的权重矩阵得到。Prefill 并行算所有,Decode 只算新的一个,历史 K、V 从 Cache 复用。

模型逐 token 生成时,每一步只需要新 token 的 Q,但需要所有历史 token 的 K 和 V 参与计算:

```plain text Step 1: 处理 “我” → 计算 q₁,k₁,v₁ 缓存 [k₁],[v₁] Step 2: 处理 “喜欢” → 计算 q₂,k₂,v₂ 缓存 [k₁,k₂],[v₁,v₂] Step 3: 处理 “猫” → 计算 q₃,k₃,v₃ 缓存 [k₁,k₂,k₃],[v₁,v₂,v₃]

k₁v₁, k₂v₂ 直接复用,无需重算

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
*   **没有 KV Cache**:每步都要对所有历史 token 重新计算 K、V → 总计算量 **O(n²)**

*   **有 KV Cache**:每步只算新 token 的 K、V,历史的从缓存读 → 总计算量 **O(n)**

`[KV Cache 逐步增长 + 复用示意图]`

***

## 2. 原生 KV Cache 的问题

先算笔账(以 LLaMA-2 7B, FP16 为例):

| 指标                   | 值          |
| -------------------- | ---------- |
| 单 token 的 KV Cache   | **512 KB** |
| 一个序列 (max\_len=2048) | **1 GB**   |
| 模型本身                 | \~14 GB    |

### 2.1 内存预分配浪费

序列长度在生成前未知,必须按 `max_seq_len` 预分配连续显存:

```plain text
预分配 2048 tokens = 1GB
实际只用 256 tokens
分配了但没有用的会浪费; 没有分配的会浪费
██████░░░░░░░░░░░░░|░░░░░░░  → 浪费 87.5%
 used           wasted

2.2 内存碎片化

请求长度不一,随着请求到来和结束,显存出现碎片:

```plain text T0: [Seq_A: 500][Seq_B: 300][Seq_C: 700][Free] T1: [Free: 500 ][Seq_B: 300][Free: 700 ][Free] ← A,C 结束 T2: 需要 800 连续空间 → 失败!总空闲够,但不连续

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
### 2.3 无法复用

多个请求共享相同前缀(System Prompt、Beam Search 共享前缀)时,每份 KV 各存一份,完全浪费。

`[预分配浪费 + 碎片 + 重复存储三合一对比图]`

***

## 3. 操作系统的启示

这三个问题,操作系统在内存管理中**早已解决**——答案是**虚拟内存 + 分页**。

| 问题   | OS 早期(连续分配) | OS 分页方案           |
| ---- | ----------- | ----------------- |
| 外部碎片 | 有           | **消除**(固定大小页)     |
| 内部碎片 | 严重(预分配)     | **≤ 1 页**         |
| 无法共享 | 有           | **Copy-on-Write** |

核心思路:

*   物理内存切成固定大小的**页帧(Page Frame)**

*   进程使用**虚拟地址**,通过**页表**映射到物理页帧

*   物理页帧**不需要连续**

*   **按需分配**:用到时才分配物理页

*   **写时复制(CoW)**:共享页只有在写入时才复制

> **PagedAttention 的核心洞察:KV Cache 的显存管理与 OS 的内存管理是同构问题。**

| 操作系统   | PagedAttention    |
| ------ | ----------------- |
| 物理内存   | GPU 显存            |
| 虚拟地址空间 | 序列的逻辑 KV Cache    |
| 物理页帧   | Physical KV Block |
| 页表     | Block Table       |
| 按需分配   | 动态分配 Block        |
| 写时复制   | Block 级 CoW       |

`[OS 分页 vs PagedAttention 并排对照图]`

***

## 4. PagedAttention 优化

### 4.1 显存分块

GPU 显存预先切成大量固定大小的 **Physical Block**,每个 Block 存储 `block_size` 个 token 的 KV:

```plain text
GPU 显存:
┌────┐┌────┐┌────┐┌────┐┌────┐┌────┐┌────┐┌────┐
│ P0 ││ P1 ││ P2 ││ P3 ││ P4 ││ P5 ││ P6 ││ P7 │  ← 固定大小物理块
└────┘└────┘└────┘└────┘└────┘└────┘└────┘└────┘

4.2 虚拟映射(Block Table)

每个序列维护一张 Block Table,将逻辑块映射到物理块(不要求连续):

```plain text 序列 “The meaning of life is to find joy” (8 tokens, block_size=4)

Block Table: 逻辑块 0 → 物理块 7 (tokens 0-3, 已满) 逻辑块 1 → 物理块 3 (tokens 4-7, 已满)

GPU 显存: ┌────┐┌────┐┌────┐┌────┐┌────┐┌────┐┌────┐┌────┐ │Free││Free││Free││ ★3 ││Free││Free││Free││ ★7 │ └────┘└────┘└────┘└────┘└────┘└────┘└────┘└────┘ ▲ 逻辑块1 ▲ 逻辑块0 物理上不连续,逻辑上连续

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
生成新 token 时:最后一个 Block 有空 slot → 直接写入;已满 → 从 Free List 分配新 Block。

### 4.3 分块复用(Copy-on-Write)

以 Beam Search(束宽=3)为例,三个 beam 共享前缀 "The meaning of life is":

```plain text
共享阶段 (只读):
  Beam 1,2,3 的逻辑块0 → 同一个物理块2 (ref_count=3)

Beam 1 写入 "happiness" 时:
  ref_count > 1 → 触发 CoW
  1. 分配新物理块 9,复制块 2 的数据
  2. Beam 1 的逻辑块 → 物理块 9 (ref=1)
  3. Beam 2,3 继续共享物理块 2 (ref=2)

共享前缀的 Block 始终只存一份,节省显存。

[CoW 触发前后的 Block 共享关系变化]

This post is licensed under CC BY 4.0 by the author.