Continuous Batching与PagedAttention：LLM推理引擎的核心技术

大语言模型的推理效率直接影响着 AI 应用的商业可行性和用户体验。在众多优化技术中，Continuous Batching（连续批处理） 和 PagedAttention（分页注意力） 堪称现代 LLM 推理引擎的两大核心支柱。本文将深入解析这两项技术的工作原理、实现细节和性能影响。

一、为什么需要这些技术？

1.1 传统批处理的痛点

在传统的推理部署中，批处理是静态的：

传统静态批处理：
┌──────┬──────┬──────┬──────┐
│ 请求A│ 请求B│ 请求C│ 请求D│ ← 必须等所有请求到齐才开始
├──────┴──────┴──────┴──────┤
│     批次处理 (batch=4)     │ ← 所有请求必须同时完成
├──────┬──────┬──────┬──────┤
│ 结果A│ 结果B│ 结果C│ 结果D│
└──────┴──────┴──────┴──────┘
      ↓ 等待所有请求完成 → 浪费 GPU 资源

问题在于：

• 碎片化等待：最早到达的请求必须等待最晚到达的请求
• 利用率低：GPU 在等待期间处于空闲状态
• 延迟波动大：长请求拖慢所有短请求

1.2 KV Cache 的内存困境

LLM 推理的内存需求计算公式：

KV Cache 大小 = 2 × num_layers × num_heads × head_dim × seq_len × num_tokens × precision_bytes

以 LLaMA-70B 为例，每个 token 需要约 1.5MB KV Cache。在 80GB 的 A100 上，传统预分配方式导致：

• 最大序列长度受限
• 内存利用率不足 30%
• 无法同时服务多个长上下文请求

二、Continuous Batching 深度解析

2.1 核心思想

Continuous Batching（也称为 Dynamic Batching 或 In-flight Batching）由 Hugging Face TGI 和 vLLM 率先实践。它打破了”批次必须同时开始和结束”的限制：

Continuous Batching：
时间 →
┌──────┐
│ 请求A│ ← 请求A 到达，立即开始
├──────┤
│ 请求A│ 请求B│ ← 请求B 到达，动态加入
├──────┼──────┤
│ 请求A│ 请求B│ 请求C│ ← 请求C 到达，动态加入
├──────┼──────┤     │
│ 请求A│ 请求B│ 请求C│
├──────┼──────┤     │
│      │ 请求B│ 请求C│ ← 请求A 完成，退出批次
│      ├──────┤     │
│      │      │ 请求C│ ← 请求B 完成，退出批次
│      │      ├──────┤
│      │      │      │ ← 请求C 完成，批次清空
└──────┴──────┴──────┘

2.2 实现机制

# Continuous Batching 简化实现示意
class ContinuousBatchingScheduler:
    def __init__(self, max_batch_size=64):
        self.max_batch_size = max_batch_size
        self.running_seqs = []     # 正在运行的序列
        self.waiting_seqs = []     # 等待中的序列
        self.completed = []        # 已完成序列
    
    def add_request(self, seq):
        """添加新请求到等待队列"""
        self.waiting_seqs.append(seq)
    
    def schedule(self):
        """调度逻辑：每个解码步骤调用"""
        # 1. 移除已完成的序列
        self.running_seqs = [s for s in self.running_seqs if not s.finished]
        self.completed.extend(s for s in self.running_seqs if s.finished)
        
        # 2. 从等待队列拉取新序列，填满批次
        slots_available = self.max_batch_size - len(self.running_seqs)
        if slots_available > 0 and self.waiting_seqs:
            to_add = self.waiting_seqs[:slots_available]
            self.waiting_seqs = self.waiting_seqs[slots_available:]
            self.running_seqs.extend(to_add)
        
        return self.running_seqs
    
    def step(self):
        """执行一个解码步骤"""
        batch = self.schedule()
        if not batch:
            return
        
        # 构造批次输入（各序列当前最后一个 token）
        input_ids = [s.last_token_id for s in batch]
        
        # 执行前向传播（整个批次共享一次模型计算）
        logits = model.forward(input_ids)
        
        # 采样更新每个序列
        for seq, logit in zip(batch, logits):
            next_token = sample(logit)
            seq.add_token(next_token)

2.3 调度策略对比

2.4 性能收益

三、PagedAttention 深度解析

3.1 问题起源

传统 KV Cache 分配存在严重的内存碎片问题：

# 传统 KV Cache：预分配连续内存
# 假设 batch_size=4, max_seq_len=4096
# 预分配: 4 × 4096 × KV_Size_per_token = 大量浪费

# 实际使用情况：
# ┌────────────┬────────────┬────────────┬────────────┐
# │ 请求A      │ 请求B      │ 空置      │ 请求C      │
# │ (200 tokens)│ (50 tokens)│ (浪费 4096)│ (300 tokens)│
# └────────────┴────────────┴────────────┴────────────┘

问题总结：

• 内部碎片：每个请求预分配最大长度，实际只用了 5-20%
• 外部碎片：请求完成后，释放的内存无法被新请求利用
• 隔离性差：一个请求的内存波动会影响其他请求

3.2 核心设计

PagedAttention 借鉴操作系统的分页内存管理，将 KV Cache 划分为固定大小的块（block/pages），按需分配：

# PagedAttention KV Cache 管理
class PagedKVCache:
    """
    分页式 KV Cache 管理器
    
    设计理念：
    - Block Table：类似虚拟内存的页表
    - Copy-on-Write：支持多个序列共享相同块
    - 延迟释放：利用引用计数精准回收
    """
    
    def __init__(self, block_size=16, total_blocks=100000):
        self.block_size = block_size          # 每个块容纳的 token 数
        self.total_blocks = total_blocks      # 总块数
        self.free_blocks = list(range(total_blocks))  # 空闲块列表
        self.block_tables = {}                # 序列ID → [块ID列表]
        self.ref_counts = {}                  # 块ID → 引用计数
    
    def allocate(self, seq_id, num_tokens):
        """为序列分配足够的 KV Cache 块"""
        num_blocks = (num_tokens + self.block_size - 1) // self.block_size
        blocks = []
        for _ in range(min(num_blocks, len(self.free_blocks))):
            block_id = self.free_blocks.pop()
            self.ref_counts[block_id] = 1
            blocks.append(block_id)
        self.block_tables[seq_id] = blocks
        return blocks
    
    def fork(self, parent_seq_id, child_seq_id):
        """
        Copy-on-Write：子序列共享父序列的块
        用于 beam search 等场景
        """
        if parent_seq_id not in self.block_tables:
            return
        parent_blocks = self.block_tables[parent_seq_id]
        for block_id in parent_blocks:
            self.ref_counts[block_id] += 1
        self.block_tables[child_seq_id] = list(parent_blocks)
    
    def free(self, seq_id):
        """释放序列占用的 KV Cache 块（引用计数管理）"""
        if seq_id not in self.block_tables:
            return
        for block_id in self.block_tables[seq_id]:
            self.ref_counts[block_id] -= 1
            if self.ref_counts[block_id] <= 0:
                self.free_blocks.append(block_id)
                del self.ref_counts[block_id]
        del self.block_tables[seq_id]

3.3 块大小选择

3.4 内存效率对比

传统预分配 vs PagedAttention
（以 1000 个并行请求为例，平均序列长度 500 tokens）

传统预分配：
┌────────────────────────────────────────────┐
│ ████████████████████████████████████████░░│ ← 80% 未使用
│ 已用: 500MB    浪费: 2000MB               │
└────────────────────────────────────────────┘

PagedAttention：
┌────────────────────────────────────────────┐
│ ████████████████████████████████░░░░░░░░░░░│ ← 20% 空闲
│ 已用: 500MB    空闲: 125MB (给新请求用)    │
└────────────────────────────────────────────┘

内存效率提升: 25% → 80% (+220% 有效容量)

四、主流框架的实现

4.1 vLLM

vLLM 同时实现了 PagedAttention + Continuous Batching：

# vLLM 启动示例
vllm serve meta-llama/Llama-2-7b-hf \
    --max-model-len 4096 \
    --max-num-batched-tokens 8192 \
    --gpu-memory-utilization 0.90 \
    --enforce-eager \
    --max-num-seqs 256

关键配置参数：

4.2 Hugging Face TGI

TGI 的 Continuous Batching 实现：

# TGI 启动示例
text-generation-launcher \
    --model-id meta-llama/Llama-2-7b-hf \
    --max-batch-prefill-tokens 4096 \
    --max-total-tokens 8192 \
    --max-input-length 2048 \
    --max-batch-size 64

4.3 TensorRT-LLM

NVIDIA 的框架使用 In-flight Batching 引擎：

# TensorRT-LLM 构建优化引擎
trtllm-build \
    --checkpoint_dir ./llama-checkpoint \
    --output_dir ./engine \
    --max_batch_size 64 \
    --max_input_len 2048 \
    --max_seq_len 4096 \
    --gemm_plugin auto \
    --gpt_attention_plugin auto

4.4 框架特性对比

五、性能调优实战

5.1 关键指标

# 性能监控公式
# Request Throughput = num_requests / total_time
# Token Throughput = total_tokens / total_time  
# TTFT = Time to First Token 首token延迟
# ITL = Inter-Token Latency token间延迟
# TPOT = Time Per Output Token 每输出token时间

5.2 调优策略

# 高吞吐场景：优先批次大小
vllm serve ... \
    --max-num-batched-tokens 16384 \
    --max-num-seqs 512 \
    --gpu-memory-utilization 0.95

# 低延迟场景：限制批次
vllm serve ... \
    --max-num-seqs 32 \
    --gpu-memory-utilization 0.85 \
    --block-size 32

# 平衡场景：默认调优
vllm serve ... \
    --max-num-batched-tokens 8192 \
    --max-num-seqs 128 \
    --gpu-memory-utilization 0.90 \
    --enable-prefix-caching

5.3 瓶颈诊断

六、实践建议

6.1 根据场景选择框架

6.2 部署 checklist

• 确定模型和精度（FP16/INT8/FP8）
• 评估 GPU 显存容量
• 根据并发需求设置 max-num-seqs
• 测试不同 block-size 的影响
• 启用 prefix caching（重复 prompt 场景）
• 设置合理的 gpu-memory-utilization（留 5-10% 余量）
• 配置优雅的重试和限流机制
• 监控 TTFT、ITL、吞吐量三个核心指标

七、未来展望

1. 多模态支持：PagedAttention 扩展到视觉 token 管理
2. 异构内存：CPU + GPU 混合 KV Cache 存储
3. 自适应调度：基于请求特征自动选择调度策略
4. 拆分解码：Attention 和 FFN 分离，独立调度

总结

Continuous Batching 和 PagedAttention 是 LLM 推理引擎从实验走向生产的基石技术。Continuous Batching 通过迭代级动态调度消除了 GPU 等待浪费，PagedAttention 通过分页内存管理将 KV Cache 效率提升 4 倍以上。理解这两项技术的工作原理，是高效部署 LLM 服务的关键前提。