03 - KV Cache与批处理：大模型推理的内存管理核心技术

本文是《大模型推理框架深度解析》系列的第三篇，深入剖析PagedAttention与连续批处理如何提升推理效率。

写在前面

在大模型推理中，有两个技术直接影响服务的吞吐量和成本：

1. KV Cache管理：决定了你能同时服务多少用户
2. 批处理策略：决定了GPU利用率有多高

vLLM相比传统框架能实现2-4倍的吞吐提升，核心秘密就在于PagedAttention和连续批处理。本文将揭开这两项技术的原理。

---

一、KV Cache是什么

1.1 为什么需要KV Cache

Transformer的自注意力机制需要计算每个token与之前所有token的关系：

# 简化示意
for i in range(seq_len):
    # 计算第i个token的注意力
    q = W_q @ x[i]           # Query（当前token）
    k = W_k @ x[:i+1]        # Key（所有历史token）
    v = W_v @ x[:i+1]        # Value（所有历史token）
    attention = softmax(q @ k.T / sqrt(d)) @ v

问题：生成第1000个token时，需要重新计算前999个token的Key和Value吗？

答案：不需要。KV Cache就是把这些中间结果存起来，避免重复计算。

1.2 KV Cache的显存占用

以Llama-3.1-70B为例，batch_size=1，seq_len=4096：

每层KV Cache大小 = 2 × num_heads × head_dim × seq_len × sizeof(dtype)
                = 2 × 64 × 128 × 4096 × 2字节
                = 128 MB

80层总KV Cache = 128 MB × 80 = 10.2 GB

如果batch_size=16，则需要163GB显存——远超单卡容量。

---

二、llama.cpp的环形缓冲区

2.1 设计原理

llama.cpp采用预分配的环形缓冲区（Ring Buffer）管理KV Cache：

struct llama_kv_cache {
    // 预分配固定大小
    ggml_tensor* k_cache[80];  // 80层的Key缓存
    ggml_tensor* v_cache[80];  // 80层的Value缓存
    
    uint32_t head = 0;   // 写入位置
    uint32_t size = 0;   // 当前大小
    uint32_t max_size;   // 最大容量（由ctx-size决定）
};

工作流程：

1. 启动时根据--ctx-size预分配全部显存
2. 新token写入head位置
3. head循环移动，覆盖最旧的数据

2.2 优缺点分析

优点：

• 无动态内存分配开销
• 内存连续，访问局部性好
• 实现简单，无外部碎片

缺点：

• 预分配大小固定，无法动态扩展
• 短序列浪费显存（分配了4096但实际只用100）
• 无法实现前缀共享

内存效率：约60-70%

---

三、vLLM的PagedAttention

3.1 核心思想：操作系统分页

PagedAttention借鉴了虚拟内存的思想：

传统方式：每个序列分配连续的显存块
┌─────────────────────────────────────────┐
│ Seq A: [K0,K1,K2,K3,K4,.....,K4095]     │  ← 4096个token连续存储
└─────────────────────────────────────────┘

PagedAttention：将KV Cache划分为固定大小的Block
┌─────┬─────┬─────┬─────┬─────┬─────┐
│ B0  │ B1  │ B2  │ B3  │ B4  │ B5  │  ← 每个Block 16 tokens
│[0-15]│[16-31]│ ... │     │     │     │
└─────┴─────┴─────┴─────┴─────┴─────┘

通过Block Table映射逻辑位置到物理Block
Seq A: [B0 → B1 → B3 → B5]  ← 非连续分配
Seq B: [B0 → B2 → B4]       ← 与Seq A共享B0（前缀相同）

3.2 Block Table机制

class PagedAttention:
    block_size = 16  # 可配置，默认16 tokens
    
    def __init__(self):
        # 全局Block池
        self.gpu_blocks: List[KVBlock] = []
        self.cpu_blocks: List[KVBlock] = []
        
        # 每个序列的Block Table
        self.block_tables: Dict[int, List[int]] = {}
    
    def allocate(self, seq_id: int, num_tokens: int):
        """为序列分配Block"""
        num_blocks = ceil(num_tokens / self.block_size)
        blocks = self._get_free_blocks(num_blocks)
        self.block_tables[seq_id] = blocks
    
    def append_token(self, seq_id: int, token_idx: int):
        """追加新token到序列"""
        block_table = self.block_tables[seq_id]
        block_idx = token_idx // self.block_size
        
        if block_idx >= len(block_table):
            # 需要分配新Block
            new_block = self._get_free_block()
            block_table.append(new_block)

3.3 Copy-on-Write前缀共享

当两个请求有相同前缀时，PagedAttention可以实现零拷贝共享：

请求A: "What is the capital of France?"
请求B: "What is the capital of Germany?"

前缀共享：
┌─────────────────────────────────────────┐
│ Shared Blocks: ["What", "is", "the", "capital", "of"] │
│ Ref Count = 2                           │
└─────────────────────────────────────────┘
        ↓
┌─────────────┐    ┌─────────────┐
│ Seq A: [Shared│ → │ "France"]   │
│ Ref: 1      │    │ New Block   │
└─────────────┘    └─────────────┘
        ↓
┌─────────────┐    ┌─────────────┐
│ Seq B: [Shared│ → │ "Germany"]  │
│ Ref: 1 → 0  │    │ New Block   │
└─────────────┘    └─────────────┘

当Seq A需要修改Shared Block时，触发COW复制

3.4 PagedAttention的优势

特性	环形缓冲区	PagedAttention
内存碎片	20-40%外部碎片	<5%内部碎片
动态扩展	不支持	按需分配
前缀共享	不支持	COW零拷贝
长上下文	受限于预分配	理论上无上限
内存效率	60-70%	95%+

---

四、批处理策略对比

4.1 静态批处理（llama.cpp默认）

时间轴 →

请求A: [ Prefill AAAAAAAAAAAAA ] [ Decode A ]
                         ↑
请求B:                          [ Prefill BBBBBBBBB ] [ Decode B ]
                                                  ↑
请求C:                                               [ Prefill CCCCC ] [ Decode C ]

问题：
1. Prefill阶段GPU利用率低（只有一个请求在计算）
2. 新请求必须等待当前batch全部完成
3. 短请求被长请求阻塞

GPU利用率：45-60%

4.2 连续批处理（vLLM Continuous Batching）

时间轴 →

Step 1: [ Prefill AAAAA ]
         GPU: ████████░░░░░░░░░░░░

Step 2: [ Prefill AAAAA ] [ Prefill BBB ]
         GPU: ██████████████░░░░░░

Step 3: [ Decode A ] [ Prefill BBB ] [ Prefill CC ]
         GPU: ██████████████████░░

Step 4: [ Decode A ] [ Decode B ] [ Prefill CC ]
         GPU: ████████████████████

Step 5: [ Decode A ] [ Decode B ] [ Decode C ]
         GPU: ████████████████████

优势：
1. Token级调度，GPU持续满载
2. 请求动态进出，无等待时间
3. 短请求不会被长请求阻塞

GPU利用率：85-95%

4.3 性能对比实测

测试环境：Llama-2-13B on A100，ShareGPT数据集

批处理模式	峰值吞吐	P99 TTFT	GPU利用率
静态批处理	1,200 tok/s	800ms	45-60%
连续批处理	2,800 tok/s	120ms	85-95%

TTFT（Time To First Token）：从请求发起到收到第一个token的时间

---

五、vLLM的关键配置参数

5.1 内存相关

vllm serve model \
    --gpu-memory-utilization 0.85 \  # GPU显存利用率上限
    --max-model-len 32768 \          # 最大上下文长度
    --max-num-batched-tokens 8192 \  # 每batch最大token数
    --max-num-seqs 256               # 最大并发序列数

5.2 调度相关

vllm serve model \
    --scheduling-policy fcfs \       # 调度策略：fcfs或priority
    --preemption-mode swap \         # 抢占模式：swap或recompute
    --swap-space 4                   # 用于抢占的CPU内存(GB)

5.3 参数调优建议

场景	gpu-memory-utilization	max-num-seqs
短文本高并发	0.90	512
长文本低并发	0.80	64
混合负载	0.85	256

---

六、常见误区澄清

误区1：swap-space用于KV Cache溢出到磁盘

真相：--swap-space预留的CPU内存用于请求抢占时的临时存储，而非活跃KV Cache的溢出。

当GPU显存不足时，vLLM会选择性地将部分请求的KV Cache换出到CPU，待GPU空闲时再换入。

误区2：PagedAttention没有内存碎片

真相：PagedAttention消除了外部碎片，但仍有内部碎片。如果序列长度不是block_size的整数倍，最后一个block会有浪费。

例如：block_size=16，序列长度=17，则第二个block只使用了1/16。

误区3：连续批处理适合所有场景

真相：连续批处理在吞吐优先场景下最优，但在延迟敏感场景（如交互式聊天）可能需要限制batch size以保证ITL（Inter-Token Latency）。

---

七、性能调优实战

7.1 监控关键指标

# vLLM暴露的Prometheus指标
vllm:gpu_cache_usage_percents        # GPU KV Cache使用率
vllm:cpu_cache_usage_percents        # CPU KV Cache使用率
vllm:prompt_tokens_total             # 输入token数
vllm:generation_tokens_total         # 输出token数
vllm:time_to_first_token_seconds     # TTFT
vllm:time_per_output_token_seconds   # TPOT/ITL

7.2 调优 checklist

• KV Cache使用率保持在80-90%（过高会OOM，过低浪费显存）
• TTFT P99 < 500ms（交互式场景）
• ITL P99 < 100ms（流式输出场景）
• GPU利用率 > 80%

7.3 常见问题排查

问题1：GPU OOM

# 降低显存利用率上限
--gpu-memory-utilization 0.80

# 或减少并发序列数
--max-num-seqs 128

问题2：TTFT过高

# 限制batch大小，减少排队
--max-num-batched-tokens 4096

# 或启用优先级调度
--scheduling-policy priority

---

参考资源

• PagedAttention论文
• vLLM调度器文档
• Continuous Batching原理

---

文章标签

PagedAttention KV Cache 连续批处理 vLLM 内存优化 大模型推理 GPU利用率