259 lines
6.7 KiB
Markdown
259 lines
6.7 KiB
Markdown
# KV Cache与Continuous Batching
|
||
|
||
## KV Cache原理
|
||
|
||
### 自回归生成的挑战
|
||
|
||
在自回归（Autoregressive）语言模型生成中，每个token的预测需要依赖之前所有token：
|
||
|
||
```
|
||
生成第t个token: P(token_t | token_1, token_2, ..., token_{t-1})
|
||
```
|
||
|
||
朴素实现的问题：每步生成都需要重新计算所有历史token的注意力，计算量随序列长度平方增长。
|
||
|
||
### KV Cache解决方案
|
||
|
||
KV Cache通过缓存已计算的Key和Value向量，避免重复计算：
|
||
|
||
```
|
||
Layer i 的 KV Cache:
|
||
- K_cache[i]: (batch, heads, seq_len, head_dim)
|
||
- V_cache[i]: (batch, heads, seq_len, head_dim)
|
||
```
|
||
|
||
**前向传播修改**：
|
||
|
||
```python
|
||
def forward_with_kvcache(layer, x, past_key_value=None):
|
||
# x: (batch, seq_len, hidden)
|
||
|
||
if past_key_value is None:
|
||
# 首次生成，初始化缓存
|
||
k, v = compute_kv(x) # 完整计算
|
||
past_kv = (k, v)
|
||
else:
|
||
# 后续生成，只计算新token
|
||
k_new, v_new = compute_kv(x[:, -1:]) # 只处理最后一个token
|
||
past_kv = update_kvcache(past_key_value, k_new, v_new)
|
||
k, v = past_kv # 完整K, V用于注意力计算
|
||
|
||
# 注意力计算
|
||
attn_output = attention(q, k, v)
|
||
|
||
return attn_output, past_kv
|
||
```
|
||
|
||
### 计算量对比
|
||
|
||
| 方法 | 注意力计算复杂度 | 序列长度S处计算量 |
|
||
|------|----------------|-----------------|
|
||
| 朴素（无Cache） | O(S²) | 每步重新计算历史 |
|
||
| KV Cache | O(S) per step | 只计算新token |
|
||
| 节省比例 | - | S/(S + 1) → 接近1（长序列） |
|
||
|
||
## KV Cache内存占用
|
||
|
||
### 内存计算公式
|
||
|
||
```
|
||
KV Cache显存 = 2 × batch_size × num_layers × num_heads × seq_len × head_dim × bytes_per_param
|
||
|
||
FP16时 bytes_per_param = 2
|
||
```
|
||
|
||
### 示例计算
|
||
|
||
| 模型配置 | 参数 | 批量=1, 序列长度=2048 |
|
||
|---------|------|----------------------|
|
||
| LLaMA-7B (32层, 32头, 128维) | 2×32×32×2048×128×2 | ~512MB |
|
||
| LLaMA-13B (40层, 40头, 128维) | 2×40×40×2048×128×2 | ~1GB |
|
||
| LLaMA-65B (80层, 80头, 128维) | 2×80×80×2048×128×2 | ~4GB |
|
||
|
||
批量大于1时，显存占用线性增加，成为推理瓶颈。
|
||
|
||
### 内存优化技术
|
||
|
||
1. **序列并行**：将KV cache分片到多个GPU
|
||
2. **页式管理**：vLLM的Paged Attention（见专题）
|
||
3. **MQA/GQA**：减少KV头数量（见下文）
|
||
|
||
## Multi-Query Attention (MQA) 与 Grouped-Query Attention (GQA)
|
||
|
||
### 标准Multi-Head Attention (MHA)
|
||
|
||
每个注意力头都有独立的K、V矩阵：
|
||
|
||
```
|
||
K_heads: (batch, heads, seq, head_dim)
|
||
V_heads: (batch, heads, seq, head_dim)
|
||
计算量: heads × seq²
|
||
```
|
||
|
||
### Multi-Query Attention (MQA)
|
||
|
||
所有注意力头共享同一份K、V：
|
||
|
||
```
|
||
K: (batch, 1, seq, head_dim) # 只有一个头
|
||
V: (batch, 1, seq, head_dim)
|
||
Query: (batch, heads, seq, head_dim)
|
||
```
|
||
|
||
**优点**：KV Cache显存降至 1/heads
|
||
**缺点**：可能降低模型表达能力
|
||
|
||
### Grouped-Query Attention (GQA)
|
||
|
||
MHA和MQA的折中方案：
|
||
|
||
```
|
||
num_kv_heads = num_query_heads / ratio (ratio通常为4-8)
|
||
|
||
LLaMA 2使用: 32 query heads, 8 kv heads
|
||
```
|
||
|
||
### 显存对比
|
||
|
||
| 注意力类型 | KV Cache大小（相对MHA） |
|
||
|-----------|----------------------|
|
||
| MHA (32 heads) | 1x |
|
||
| GQA (8 kv heads) | 0.25x |
|
||
| MQA (1 kv head) | 1/32 x |
|
||
|
||
## Continuous Batching（动态批处理）
|
||
|
||
### Static Batching（静态批处理）的问题
|
||
|
||
传统方法将多个请求组成固定大小的batch一起处理：
|
||
|
||
```
|
||
问题1: 不同请求序列长度差异大
|
||
问题2: 短序列需等待长序列完成
|
||
问题3: GPU资源利用率低
|
||
```
|
||
|
||
```
|
||
Static Batch Example:
|
||
Batch = [req1(len=100), req2(len=500), req3(len=200)]
|
||
所有请求需等待最长的req2完成 → 大量空闲等待
|
||
```
|
||
|
||
### Continuous Batching原理
|
||
|
||
Continuous Batching在iteration级别进行调度，新请求可随时加入，完成的请求立即退出：
|
||
|
||
```python
|
||
def continuous_batching_scheduler():
|
||
running_requests = []
|
||
|
||
while True:
|
||
# 1. 检查完成的请求
|
||
finished = [r for r in running_requests if r.is_done()]
|
||
for r in finished:
|
||
yield r.output
|
||
free_slot(r)
|
||
|
||
# 2. 尝试加入新请求
|
||
if has_pending_requests() and has_free_slot():
|
||
new_req = get_next_request()
|
||
running_requests.append(new_req)
|
||
|
||
# 3. 执行一步迭代
|
||
for req in running_requests:
|
||
req.step() # 生成一个token
|
||
|
||
# 4. 同步（如果有依赖）
|
||
synchronize_if_needed()
|
||
```
|
||
|
||
### 对比分析
|
||
|
||
| 特性 | Static Batching | Continuous Batching |
|
||
|------|----------------|---------------------|
|
||
| 调度时机 | batch开始时 | 每个iteration |
|
||
| 序列结束处理 | 等待batch内全部完成 | 立即释放 |
|
||
| GPU利用率 | 较低（有空闲） | 较高 |
|
||
| 延迟 | 取决于最长请求 | 取决于平均长度 |
|
||
| 吞吐 | 受限于单个长序列 | 充分利用GPU |
|
||
|
||
### 实现框架
|
||
|
||
```python
|
||
# vLLM Continuous Batching示例
|
||
from vllm import LLM, SamplingParams
|
||
|
||
llm = LLM(model="meta-llama/Llama-2-7b-hf")
|
||
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
|
||
|
||
# 提交多个请求
|
||
prompts = ["Hello", "Hi", "How are you"]
|
||
outputs = llm.generate(prompts, sampling_params)
|
||
|
||
# vLLM自动处理Continuous Batching
|
||
for output in outputs:
|
||
print(output.outputs[0].text)
|
||
```
|
||
|
||
### Iteration-level调度详解
|
||
|
||
```
|
||
时间步 t:
|
||
Batch = [req1, req2, req3, req4]
|
||
生成: [tok1, tok2, tok3, tok4]
|
||
|
||
时间步 t+1:
|
||
req3 完成 → 移除
|
||
req5 新加入
|
||
Batch = [req1, req2, req4, req5]
|
||
生成: [tok1, tok2, tok4, tok5]
|
||
```
|
||
|
||
### 调度策略
|
||
|
||
1. **First-Come-First-Served (FCFS)**：按请求到达顺序处理
|
||
2. **Priority Scheduling**：高优先级请求优先
|
||
3. **Preemption**：当显存不足时，暂停低优先级请求
|
||
|
||
### 与Prefix Caching的结合
|
||
|
||
Continuous Batching可与Prefix Caching配合：
|
||
|
||
```python
|
||
# 共享系统提示前缀
|
||
system_prompt = "You are a helpful assistant."
|
||
|
||
# 多用户共享相同前缀 → KV Cache可复用
|
||
req1 = system_prompt + "User1 question"
|
||
req2 = system_prompt + "User2 question"
|
||
|
||
# Prefix KV Cache只需计算一次
|
||
```
|
||
|
||
## 系统设计考量
|
||
|
||
### 显存管理策略
|
||
|
||
| 策略 | 描述 | 适用场景 |
|
||
|------|------|---------|
|
||
| 静态分配 | 预分配固定KV Cache大小 | 简单实现 |
|
||
| 动态分配 | 按需分配/释放 | 通用场景 |
|
||
| Paged Attention | 分页管理（见专题） | 高吞吐 |
|
||
|
||
### Batch大小选择
|
||
|
||
```
|
||
最大batch_size = GPU显存 / (单请求KV Cache + 模型参数)
|
||
|
||
对于7B模型，FP16，约20GB显存用于模型
|
||
剩余约20GB → 最大batch_size受KV Cache限制
|
||
```
|
||
|
||
### 延迟-吞吐权衡
|
||
|
||
```
|
||
高延迟 + 高吞吐: 大量请求排队，大batch处理
|
||
低延迟 + 低吞吐: 小batch及时处理，短等待队列
|
||
```
|
||
|
||
Continuous Batching通过动态调整实现两者平衡。 |