你有没有想过,为什么大模型生成第一个字需要几秒钟,但后续的字却快得多?为什么同样的模型,处理1000字的上下文比处理100字消耗更多显存?为什么有些70B参数的模型反而比7B模型的KV缓存更小?
答案都指向同一个技术:KV Cache。这个听起来简单的"键值缓存",实际上是大模型推理效率的核心,也是理解大模型部署成本的钥匙。
从一个直观的问题开始
假设你让一个语言模型续写这句话:“人工智能正在改变”。
模型生成"世界"这个词时,需要看前面的所有词——“人工智能”、“正在”、“改变”。当它生成第二个词"的"时,理论上需要重新计算"人工智能"、“正在”、“改变”、“世界"这些词的注意力表示。
但等等——“人工智能"这个词的注意力表示,在生成"世界"时已经算过了,为什么生成"的"的时候还要重新算?
这就像你每次做数学题都要重新背诵乘法口诀——明明上次已经背过了,为什么还要重来?
KV Cache就是这个"记忆”。它把计算过的Key和Value向量存下来,下次用到时直接取用,不需要重新计算。
注意力机制的本质回顾
要理解KV Cache,先要理解Transformer的注意力机制在算什么。
自注意力机制的核心是三个向量:Query(查询)、Key(键)、Value(值)。每个输入token都会产生这三个向量。Query是"我想找什么”,Key是"我是什么",Value是"我的内容"。
graph LR
subgraph "Token处理流程"
A[输入Token] --> B[线性变换]
B --> C[Query Q]
B --> D[Key K]
B --> E[Value V]
C --> F[注意力分数计算]
D --> F
F --> G[加权求和]
E --> G
G --> H[输出表示]
end
注意力分数的计算公式是:
$$\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V$$这里有个关键点:每个位置的Query只需要跟它之前所有位置的Key和Value计算。这就是因果掩码(Causal Mask)的作用——当前位置不能"看见"未来的位置。
自回归生成时,第$t$个token需要与第$1$到第$t-1$个token的Key和Value计算注意力。问题是:每生成一个新token,之前所有token的Key和Value都要重新算吗?
KV Cache的工作原理
不需要。这正是KV Cache解决的问题。
KV Cache的思路很简单:把每个token的Key和Value向量缓存起来。生成新token时,只需要:
- 计算新token的Query、Key、Value
- 把新的Key和Value追加到缓存中
- 用新token的Query与缓存中所有的Key计算注意力
- 用注意力分数对缓存中所有的Value加权求和
sequenceDiagram
participant M as 模型
participant C as KV Cache
participant O as 输出
Note over M,C: Step 1: 处理第一个token
M->>C: 存储K1, V1
C->>M: 返回K1, V1
M->>O: 输出token 2
Note over M,C: Step 2: 处理第二个token
M->>C: 存储K2, V2
C->>M: 返回K1, K2, V1, V2
M->>O: 输出token 3
Note over M,C: Step n: 处理第n个token
M->>C: 存储Kn, Vn
C->>M: 返回K1...Kn, V1...Vn
M->>O: 输出token n+1
这个简单的设计带来了巨大的效率提升。HuggingFace的基准测试显示,启用KV Cache后,推理速度可以提升5倍以上。
两阶段推理:Prefill与Decode
KV Cache的工作方式将LLM推理分成两个截然不同的阶段:
Prefill阶段:模型一次性处理所有输入token。这阶段是计算密集型的,GPU利用率高,因为可以并行计算所有token的表示。同时,这阶段会建立KV Cache的初始内容。
Decode阶段:模型逐个生成输出token。每生成一个token,就要从KV Cache读取所有之前的Key和Value,计算注意力,然后把新token的Key和Value追加到缓存中。
graph TB
subgraph "Prefill阶段 - 计算密集"
P1[输入: 所有prompt tokens] --> P2[并行计算]
P2 --> P3[建立完整KV Cache]
P3 --> P4[输出: 第一个生成token]
end
subgraph "Decode阶段 - 内存带宽密集"
D1[输入: 单个新token] --> D2[读取完整KV Cache]
D2 --> D3[计算注意力]
D3 --> D4[追加新K,V到缓存]
D4 --> D5[输出: 下一个token]
end
P4 --> D1
这两个阶段的性能特征完全不同。Prefill是计算受限(compute-bound),Decode是内存带宽受限(memory-bound)。在Decode阶段,GPU大部分时间在等待从显存读取KV Cache数据,而不是在计算。
这就是为什么第一个token生成很慢(Prefill阶段需要处理整个输入),而后续token生成较快(Decode阶段每次只处理一个新token)。但当上下文很长时,Decode阶段也会变慢——因为要读取的KV Cache变大了。
KV Cache到底占用多少内存?
这是部署大模型时必须回答的问题。让我们推导精确的计算公式。
存储结构分析
对于单个注意力层,单个token,需要存储:
- Key向量:形状为
[num_kv_heads, head_dim] - Value向量:形状为
[num_kv_heads, head_dim]
对于整个模型,所有层,序列长度为$T$:
- Key缓存:
[num_layers, batch_size, num_kv_heads, T, head_dim] - Value缓存:
[num_layers, batch_size, num_kv_heads, T, head_dim]
graph TD
subgraph "KV Cache存储结构"
A["KV Cache Tensor<br/>形状: [L, B, H_kv, T, D_h]"]
A --> B["L: Transformer层数"]
A --> C["B: Batch Size"]
A --> D["H_kv: KV头数"]
A --> E["T: 序列长度"]
A --> F["D_h: 头维度"]
end
subgraph "内存计算"
G["每个元素: 2 bytes (FP16)"]
H["总元素数 = L × B × H_kv × T × D_h"]
I["总内存 = 2 × 总元素数 × 2 bytes"]
end
A --> G
G --> H
H --> I
内存计算公式
KV Cache的总内存为:
$$M_{kv} = 2 \times L \times B \times T \times H_{kv} \times D_h \times \text{bytes}$$其中:
- $2$:Key和Value两个矩阵
- $L$:Transformer层数
- $B$:batch size(并发请求数)
- $T$:序列长度(token数)
- $H_{kv}$:Key-Value头数
- $D_h$:每个头的维度
- $\text{bytes}$:每个元素的字节数(FP16为2字节,FP32为4字节)
对于标准的多头注意力(MHA),$H_{kv}$等于查询头数$H_q$,且$H_q \times D_h = d_{model}$(隐藏层维度)。此时公式可以简化为:
$$M_{kv} = 2 \times L \times B \times T \times d_{model} \times \text{bytes}$$实际计算示例
让我们用几个热门模型来计算:
| 模型 | 层数$L$ | 隐藏维度$d_{model}$ | KV头数$H_{kv}$ | 4K上下文单请求KV Cache |
|---|---|---|---|---|
| GPT-2 Small | 12 | 768 | 12 | ~75 MB |
| LLaMA 7B | 32 | 4096 | 32 | ~2.0 GB |
| LLaMA 2 70B | 80 | 8192 | 8 | ~1.0 GB |
| LLaMA 3 70B | 80 | 8192 | 8 | ~1.0 GB |
等等,70B模型的KV Cache反而比7B模型小?这正是Grouped Query Attention的功劳,后面会详细解释。
对于Llama-3-70B,8K上下文,32个并发请求:
$$M_{kv} = 2 \times 80 \times 32 \times 8192 \times 8 \times 128 \times 2 = 81.9 \text{ GB}$$这已经超过了一张A100 80GB的显存——仅用于KV Cache,还没算模型权重。
上下文长度的线性代价
KV Cache内存与序列长度成正比。这意味着:
graph LR
subgraph "KV Cache随上下文长度增长"
A["2K上下文<br/>~2GB"] --> B["4K上下文<br/>~4GB"]
B --> C["8K上下文<br/>~8GB"]
C --> D["32K上下文<br/>~32GB"]
D --> E["128K上下文<br/>~128GB"]
end
style A fill:#90EE90
style B fill:#FFD700
style C fill:#FFA500
style D fill:#FF6347
style E fill:#FF0000,color:#fff
- 2K上下文 → 2GB KV Cache
- 8K上下文 → 8GB KV Cache
- 32K上下文 → 32GB KV Cache
- 128K上下文 → 128GB KV Cache
这个线性增长关系没有捷径可绕。每个token的Key和Value必须存储,否则无法正确计算注意力。
更关键的是,当上下文超过某个阈值后,KV Cache会比模型权重占用更多内存。这个交叉点可以通过以下公式计算:
$$T_{crossover} = \frac{M_{weights}}{2 \times L \times d_{model} \times \text{bytes}}$$对于LLaMA 7B(FP16权重约14GB),这个交叉点大约在32K token。超过这个长度,KV Cache成为主要的内存消费者。
PagedAttention:从操作系统借来的智慧
传统的KV Cache实现有一个致命问题:内存碎片化。
连续分配的困境
早期系统(如HuggingFace Transformers)为每个请求预分配一块连续内存,大小按最大序列长度计算。这导致三种内存浪费:
- 内部碎片:请求的实际长度往往远小于最大长度,预分配的空间大量闲置
- 外部碎片:不同请求的预分配大小不同,buddy allocator无法高效利用
- 预留浪费:生成过程中预留的"未来token空间"在当前时刻无法被其他请求使用
vLLM论文中的测量数据显示,传统系统只有20-38%的KV Cache内存实际存储了有用的token状态,其余60-80%都是浪费。
分页内存管理
2023年,UC Berkeley的vLLM团队提出了PagedAttention,核心思想来自操作系统的虚拟内存管理:
- 把KV Cache切分成固定大小的块(blocks)
- 每个块可以存储若干个token的Key和Value
- 块在物理内存中不需要连续存储
- 通过**块表(block table)**映射逻辑块到物理块
graph TB
subgraph "逻辑视图(每个请求看到的)"
L1[逻辑块0]
L2[逻辑块1]
L3[逻辑块2]
end
subgraph "物理存储(实际GPU内存)"
P1[物理块7]
P2[物理块1]
P3[物理块3]
P4[物理块9]
end
subgraph "块表"
BT["块表映射:<br/>逻辑0 → 物理7<br/>逻辑1 → 物理1<br/>逻辑2 → 物理3"]
end
L1 --> BT
L2 --> BT
L3 --> BT
BT --> P1
BT --> P2
BT --> P3
这个设计带来了几个重要优势:
近乎零的内存浪费:只有最后一个块可能有少量未填满的空间,整体碎片率低于4%。相比传统系统的60-80%浪费,这是质的飞跃。
按需分配:生成过程中,只有当需要存储新token时才分配新块。不需要预先为"可能生成的token"预留空间。
内存共享:这是PagedAttention最精妙的设计。当多个请求共享相同的prompt前缀(比如系统提示词),它们的KV Cache可以共享相同的物理块。
Copy-on-Write机制
当共享的物理块需要被修改时(比如不同请求生成了不同的token),PagedAttention采用写时复制(Copy-on-Write)机制:
stateDiagram-v2
[*] --> 检测共享块
检测共享块 --> 引用计数>1?: 检查引用计数
引用计数>1? --> 分配新物理块: 是
引用计数>1? --> 直接写入: 否
分配新物理块 --> 复制内容
复制内容 --> 更新块表
更新块表 --> 减少原块引用计数
减少原块引用计数 --> 写入完成
直接写入 --> 写入完成
写入完成 --> [*]
这个机制使得beam search、parallel sampling等解码算法的内存效率大幅提升。vLLM的测试显示,beam search可以节省37-66%的内存。
性能影响
PagedAttention对性能的影响是双面的:
开销:非连续内存访问增加了attention kernel的复杂度。vLLM的attention kernel比FasterTransformer慢20-26%。
收益:由于内存利用率大幅提升,可以支持更大的batch size,整体吞吐量提升2-4倍。
这是一笔划算的交易:单次计算慢一点,但能同时处理更多请求。
GQA:用架构设计减少缓存
如果说PagedAttention是在系统层面优化内存管理,那么Grouped Query Attention(GQA)则从模型架构层面减少KV Cache的大小。
从MHA到MQA到GQA
标准的多头注意力(MHA):每个查询头都有自己对应的Key头和Value头。如果模型有32个查询头,就有32个Key头和32个Value头。
Multi-Query Attention(MQA):所有查询头共享同一组Key和Value。32个查询头,只有1个Key头和1个Value头。KV Cache减少32倍,但模型质量下降明显。
Grouped Query Attention(GQA):折中方案。将查询头分成若干组,每组共享一组Key和Value。比如64个查询头分成8组,每组8个查询头共享一组KV。KV Cache减少8倍,模型质量下降很小。
graph LR
subgraph "MHA (32头) - 每头独立KV"
M1[Q1→K1,V1]
M2[Q2→K2,V2]
M3["..."]
M4[Q32→K32,V32]
M5["KV Cache: 32x"]
style M5 fill:#FF6347
end
subgraph "MQA (32头) - 全共享"
Q1[Q1]
Q2[Q2]
Q3["..."]
Q4[Q32]
K1[共享的K,V]
Q1 --> K1
Q2 --> K1
Q3 --> K1
Q4 --> K1
M6["KV Cache: 1x"]
style M6 fill:#90EE90
end
subgraph "GQA (32头, 8组) - 分组共享"
G1["组1: Q1-Q8 → K1,V1"]
G2["组2: Q9-Q16 → K2,V2"]
G3["..."]
G4["组8: Q25-Q32 → K8,V8"]
G5["KV Cache: 8x"]
style G5 fill:#FFD700
end
GQA的内存节省
GQA的KV Cache节省比例可以直接计算:
$$\text{节省比例} = \frac{H_q}{H_{kv}}$$其中$H_q$是查询头数,$H_{kv}$是Key-Value头数。
LLaMA 2 70B和LLaMA 3 70B使用64个查询头,8个KV头。这意味着:
- 传统MHA:KV Cache大小 = $2 \times 80 \times 8192 \times 64 \times 128 \times 2$
- GQA配置:KV Cache大小 = $2 \times 80 \times 8192 \times 8 \times 128 \times 2$
- 节省:$64/8 = 8$倍
这就是为什么70B模型的KV Cache反而比7B模型小的原因。
主流模型的GQA配置
| 模型 | 查询头数 | KV头数 | KV Cache节省 |
|---|---|---|---|
| LLaMA 7B | 32 | 32 | 1x (无节省) |
| LLaMA 2 70B | 64 | 8 | 8x |
| LLaMA 3 70B | 64 | 8 | 8x |
| Mistral 7B | 32 | 8 | 4x |
| Mixtral 8x7B | 32 | 8 | 4x |
| DeepSeek-V2 | 128 | 2 | 64x |
GQA已经成为大模型的标准配置。它在不显著损害模型质量的前提下,大幅降低了推理的内存需求。
KV Cache量化:进一步压缩
即使使用了GQA,长上下文场景下KV Cache仍然是内存大户。量化是另一个重要的优化方向。
精度与内存的权衡
| 数据类型 | 每元素字节数 | 相对FP16节省 |
|---|---|---|
| FP32 | 4 | -100% (更大) |
| FP16/BF16 | 2 | 基准 |
| FP8 (E4M3/E5M2) | 1 | 50% |
| INT8 | 1 | 50% |
| INT4 | 0.5 | 75% |
graph LR
subgraph "量化后的KV Cache大小对比 (以LLaMA 7B 4K上下文为例)"
A["FP32: ~8GB"] --> B["FP16: ~4GB"]
B --> C["FP8: ~2GB"]
C --> D["INT4: ~1GB"]
end
style A fill:#FF0000,color:#fff
style B fill:#FFA500
style C fill:#FFD700
style D fill:#90EE90
FP8量化是目前最受欢迎的方案。它将KV Cache从FP16的2字节压缩到1字节,内存减半,同时对模型质量的影响很小。
INT8和INT4量化需要额外的缩放因子(scale factor),但可以进一步压缩。实测表明,INT8 KV Cache几乎无损,INT4在大多数场景下也可以接受。
主流框架支持
vLLM、TensorRT-LLM、llama.cpp等主流推理框架都已经支持KV Cache量化。以vLLM为例:
# FP8 KV Cache量化
--kv-cache-dtype fp8
# INT8 KV Cache量化
--kv-cache-dtype int8
量化的一个额外好处是减少内存带宽压力。在Decode阶段,瓶颈往往是内存带宽而不是计算。更小的KV Cache意味着更少的数据传输,即使计算量不变,整体速度也会提升。
Prefill与Decode:两个世界的性能特征
理解KV Cache的另一个关键角度是Prefill和Decode阶段的根本差异。
性能瓶颈分析
可以用算术强度(Arithmetic Intensity)来量化:
$$I = \frac{\text{FLOPs}}{\text{Bytes transferred}}$$对于A100 GPU,峰值算力312 TFLOPS(FP16),峰值带宽2 TB/s。计算-带宽比约为156 FLOPs/Byte。
如果操作的计算强度低于156,就是内存受限;高于156,就是计算受限。
graph LR
subgraph "A100 GPU性能边界"
A["计算-带宽比: 156 FLOPs/Byte"]
B["计算强度 < 156: 内存受限"]
C["计算强度 > 156: 计算受限"]
end
D["Prefill阶段"] --> E["计算强度高<br/>计算受限"]
F["Decode阶段"] --> G["计算强度低(~0.5)<br/>严重内存受限"]
style E fill:#90EE90
style G fill:#FF6347
Decode阶段的attention计算强度约为:
$$I_{decode} \approx \frac{2 \times d_{model}}{4 \times d_{model}} = 0.5 \text{ FLOPs/Byte}$$这远远低于156的阈值,证实了Decode阶段是严重的内存带宽受限。
分离式推理
理解了两个阶段的差异后,一个自然的想法是:能不能把它们分开处理?
DistServe和Mooncake等系统提出了分离式推理架构:
graph TB
subgraph "分离式推理架构"
Client[客户端请求] --> Router[请求路由]
Router --> Prefill[Prefill节点<br/>高算力GPU]
Router --> Decode[Decode节点<br/>大显存带宽]
Prefill -->|传输KV Cache| Decode
Decode -->|返回结果| Client
end
subgraph "Prefill节点特点"
P1[计算密集型]
P2[需要高FLOPs]
P3[处理完后可复用]
end
subgraph "Decode节点特点"
D1[内存带宽密集]
D2[需要大显存]
D3[保持KV Cache]
end
- Prefill节点:专注于处理prompt,需要高算力
- Decode节点:专注于生成token,需要大显存带宽
两个节点之间通过高速网络传输KV Cache。这种架构可以:
- 针对不同阶段选择不同的硬件
- Prefill节点处理完后可以立即开始处理下一个请求
- Decode节点保持KV Cache,支持多轮对话
Mooncake的测试显示,分离式架构可以提升525%的吞吐量。
实践指南:容量规划
部署大模型时,如何估算显存需求?
总内存公式
$$M_{total} = M_{weights} + M_{kv} + M_{activation} + M_{overhead}$$- $M_{weights}$:模型权重,FP16下约等于参数量×2字节
- $M_{kv}$:KV Cache,使用前面的公式计算
- $M_{activation}$:激活内存,约等于权重的10-20%
- $M_{overhead}$:框架开销,约1-2GB
快速估算示例
以LLaMA 3 8B为例,假设8K上下文,batch size = 1:
模型权重:8B参数 × 2字节 = 16 GB
KV Cache:2 × 32 × 1 × 8192 × 32 × 128 × 2 = 4.3 GB (使用GQA后更小)
激活内存:约 1.6 GB
框架开销:约 1 GB
总计:约 23 GB
一张24GB的RTX 4090刚好够用。
如果batch size = 8:
KV Cache:4.3 GB × 8 = 34.4 GB
总计:16 + 34.4 + 1.6 + 1 = 53 GB
需要A100 80GB才能容纳。
graph TB
subgraph "显存分配决策树"
A[选择模型] --> B{上下文长度}
B -->|短 2K| C[KV Cache小<br/>权重主导]
B -->|中 8K| D[KV Cache中等<br/>需平衡]
B -->|长 32K+| E[KV Cache大<br/>KV主导]
C --> F{Batch Size}
D --> F
E --> F
F -->|单请求| G[可用消费级GPU]
F -->|多并发| H[需要数据中心GPU]
G --> I[计算总内存需求]
H --> I
end
关键决策点
- 上下文长度:每翻倍,KV Cache翻倍
- Batch size:每翻倍,KV Cache翻倍
- 是否使用GQA:根据模型架构决定
- 量化策略:FP8可以减半KV Cache
KV Cache优化的演进历程
从2017年Transformer诞生至今,KV Cache的优化经历了几个重要阶段:
timeline
title KV Cache优化技术演进
2017 : Transformer诞生
: 标准MHA架构
: KV Cache概念出现
2019 : MQA提出
: 所有头共享KV
: 质量损失明显
2020 : GQA提出
: 分组共享KV
: 质量损失小
2023 : PagedAttention
: vLLM发布
: 解决内存碎片
: 吞吐量提升2-4x
2024 : KV Cache量化
: FP8成为主流
: 进一步减少内存
2025 : 分离式推理
: Prefill/Decode分离
: 大幅提升吞吐量
总结
KV Cache是大模型推理的核心技术,它把重复的注意力计算转化为内存查找,实现了空间换时间的优化。理解KV Cache的本质、内存消耗和优化技术,是部署大模型的必备知识。
几个关键认知:
KV Cache的内存增长是线性的,与上下文长度和batch size成正比。长上下文场景下,KV Cache可能超过模型权重成为主要的内存消费者。
PagedAttention解决了内存碎片化问题,将KV Cache利用率从20-38%提升到96%以上。这是系统层面的优化。
GQA从架构层面减少KV Cache,现代大模型普遍采用8个KV头配合64个查询头的配置,实现8倍的内存节省。
Prefill和Decode有根本不同的性能特征。Prefill是计算受限,Decode是内存带宽受限。分离式推理是面向未来的架构方向。
当你下次思考"这个模型能跑在什么显卡上"时,记得把KV Cache算进去——它往往比模型权重更需要你的关注。