e2hang/Notes

Fork 0

Files

e2hang 2f5211d8fe New Notes

2026-05-16 17:16:51 +08:00

6.7 KiB

Raw Permalink Blame History

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)

前向传播修改：

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时，显存占用线性增加，成为推理瓶颈。

内存优化技术

序列并行：将KV cache分片到多个GPU
页式管理：vLLM的Paged Attention（见专题）
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级别进行调度，新请求可随时加入，完成的请求立即退出：

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

实现框架

# 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]

调度策略

First-Come-First-Served (FCFS)：按请求到达顺序处理
Priority Scheduling：高优先级请求优先
Preemption：当显存不足时，暂停低优先级请求

与Prefix Caching的结合

Continuous Batching可与Prefix Caching配合：

# 共享系统提示前缀
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通过动态调整实现两者平衡。

6.7 KiB Raw Permalink Blame History Unescape Escape