从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₄] = 得到输出
关键点总结:
| Prefill | Decode | |
|---|---|---|
| 输入 | 所有 prompt tokens | 仅 1 个新 token |
| 计算 | 所有 token 的 Q,K,V | 仅新 token 的 Q,K,V |
| K,V 来源 | 当前计算 → 写入 Cache | 新 token 计算 + Cache 读取 |
| Q 的用途 | 每个 q 查询它之前的所有 k | q 查询 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 共享关系变化]