# 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通过动态调整实现两者平衡。