静态批处理的瓶颈

传统 LLM 推理服务采用静态批处理：每个批次固定大小，所有请求必须等最长的那个请求完成后才能返回。

静态批处理示例（batch=4）：
  Request A: [tok][tok][tok][tok][EOS]            # 生成 5 个 token
  Request B: [tok][tok][tok][tok][tok][tok][EOS]  # 生成 7 个 token
  Request C: [tok][tok][tok][EOS]                 # 生成 4 个 token
  Request D: [tok][tok][tok][tok][tok][tok][tok][tok][EOS]  # 生成 9 个 token

批处理完成时间：等待 D 完成（9 个 token）
实际计算：A 只需要 5 步，但等了 9 步
GPU 利用率：≈ 25/36 = 69%（因为要等待最慢的请求）

问题：

  1. 首 token 延迟：新请求必须等当前批次完成才能进入
  2. GPU 闲置：短请求完成后，GPU 资源被浪费
  3. 吞吐低下：无法动态调度，无法"边生成边接收"

连续批处理原理

连续批处理（Continuous Batching, Orca, 2023）的核心思想是：动态管理批次，随时加入新请求，随时移除已完成的请求。

连续批处理示例：
  T=0: 启动 [A, B, C]（立即开始）
  T=1: 加入 D，批处理变为 [A, B, C, D]
  T=2: 移除 C，加入 E，批处理变为 [A, B, D, E]
  T=3: 移除 A, B，加入 F, G，批处理变为 [D, E, F, G]
  ...

核心机制

class ContinuousBatcher:
    def __init__(self, model, max_batch_size=32, max_seq_len=4096):
        self.model = model
        self.max_batch = max_batch_size
        self.max_seq = max_seq_len
        
        # 请求队列
        self.waiting_queue = []  # 等待中的请求
        self.running_batch = []  # 正在运行的请求
        
        # KV Cache 管理器（PagedAttention）
        self.cache_manager = PagedCacheManager(
            num_blocks=10000,
            block_size=16,
            num_layers=model.config.num_layers,
            num_heads=model.config.num_heads,
            head_dim=model.config.head_dim
        )
    
    def step(self):
        """执行一个推理步骤"""
        # 1. 尝试从等待队列加入新请求
        self._try_add_requests()
        
        if not self.running_batch:
            return None
        
        # 2. 执行前向传播
        outputs = self._forward_step()
        
        # 3. 处理输出，检查是否完成
        self._process_outputs(outputs)
        
        return outputs
    
    def _try_add_requests(self):
        """尝试加入新请求"""
        while self.waiting_queue and len(self.running_batch) < self.max_batch:
            request = self.waiting_queue.pop(0)
            
            # 尝试分配 KV Cache
            if self.cache_manager.can_allocate(len(request.tokens)):
                request.cache_blocks = self.cache_manager.allocate(
                    len(request.tokens) + request.max_tokens
                )
                request.generated_tokens = []
                self.running_batch.append(request)
            else:
                # 显存不足，放回队列
                self.waiting_queue.insert(0, request)
                break
    
    def _forward_step(self):
        """执行一次前向传播"""
        # 收集所有请求的输入
        # 每个请求只需要处理最后一个生成的 token
        input_ids = []
        positions = []
        
        for req in self.running_batch:
            if req.generated_tokens:
                # 已有生成，只需要处理最后一个 token
                input_ids.append(req.generated_tokens[-1])
                positions.append(len(req.prompt_tokens) + len(req.generated_tokens) - 1)
            else:
                # 首次生成，处理整个 prompt
                input_ids.extend(req.prompt_tokens)
                positions.extend(range(len(req.prompt_tokens)))
        
        # 构建批处理输入
        input_tensor = torch.tensor(input_ids, device='cuda').unsqueeze(0)
        position_tensor = torch.tensor(positions, device='cuda')
        
        # 执行模型
        outputs = self.model(input_tensor, position_tensor, self.cache_manager)
        
        return outputs
    
    def _process_outputs(self, outputs):
        """处理输出，检查完成状态"""
        completed = []
        
        for i, req in enumerate(self.running_batch):
            if req.generated_tokens:
                # 解码步骤：取第 i 个输出
                next_token = outputs[i]
            else:
                # 预填充步骤：取最后一个输出（对应 prompt 末尾）
                prompt_len = len(req.prompt_tokens)
                next_token = outputs[prompt_len - 1]
            
            req.generated_tokens.append(next_token)
            
            # 检查是否完成
            if (next_token == self.model.eos_token or 
                len(req.generated_tokens) >= req.max_tokens):
                completed.append(req)
        
        # 移除完成的请求
        for req in completed:
            self.running_batch.remove(req)
            # 释放 KV Cache
            self.cache_manager.free(req.cache_blocks)
            # 调用回调
            req.callback(req.generated_tokens)
    
    def submit(self, prompt, max_tokens, callback):
        """提交新请求"""
        request = Request(
            tokens=self.model.tokenize(prompt),
            max_tokens=max_tokens,
            callback=callback
        )
        self.waiting_queue.append(request)
    
    def run(self):
        """主循环"""
        while self.waiting_queue or self.running_batch:
            self.step()

PagedAttention：连续批处理的基础

连续批处理需要一个灵活的 KV Cache 管理系统：PagedAttention。

核心思想

将 KV Cache 分成固定大小的块，按需分配：

传统 KV Cache（连续分配）：
  Request A: [Block 0-255]（预分配，但只用了 30）
  Request B: [Block 256-511]（预分配，但只用了 50）
  碎片严重，浪费显存

PagedAttention（分页分配）：
  Request A: [Block 0][Block 1]（按需分配，实际使用）
  Request B: [Block 2][Block 3][Block 4]（按需分配）
  Request C: 等待...（等 A 释放后复用）

块管理实现

class PagedCacheManager:
    def __init__(self, num_blocks, block_size, num_layers, num_heads, head_dim):
        self.block_size = block_size
        
        # 物理块池
        # (num_blocks, num_layers, 2, block_size, num_heads, head_dim)
        # 2 for K and V
        self.cache = torch.zeros(
            num_blocks, num_layers, 2, block_size, num_heads, head_dim,
            device='cuda', dtype=torch.float16
        )
        
        # 空闲块列表
        self.free_blocks = list(range(num_blocks))
        
        # 逻辑块到物理块的映射
        # request_id -> [physical_block_ids]
        self.block_tables = {}
    
    def allocate(self, num_tokens):
        """分配足够的块来容纳 num_tokens"""
        blocks_needed = (num_tokens + self.block_size - 1) // self.block_size
        
        if len(self.free_blocks) < blocks_needed:
            return None  # 无法分配
        
        allocated = []
        for _ in range(blocks_needed):
            block_id = self.free_blocks.pop(0)
            allocated.append(block_id)
        
        return allocated
    
    def free(self, blocks):
        """释放块"""
        self.free_blocks.extend(blocks)
    
    def get_physical_block(self, block_id, layer_idx, is_k):
        """获取物理块的缓存"""
        kv_idx = 0 if is_k else 1
        return self.cache[block_id, layer_idx, kv_idx]
    
    def can_allocate(self, num_tokens):
        """检查是否可以分配"""
        blocks_needed = (num_tokens + self.block_size - 1) // self.block_size
        return len(self.free_blocks) >= blocks_needed

注意力计算

def paged_attention(Q, block_tables, cache_manager, num_tokens_per_seq):
    """
    PagedAttention 计算
    
    Args:
        Q: Query (batch, heads, 1, head_dim) - 只处理最后一个 token
        block_tables: 每个请求的物理块映射
        cache_manager: KV Cache 管理器
        num_tokens_per_seq: 每个请求的实际 token 数
    """
    batch, heads, _, head_dim = Q.shape
    
    outputs = []
    for i in range(batch):
        q_i = Q[i]  # (heads, 1, head_dim)
        blocks = block_tables[i]
        num_tokens = num_tokens_per_seq[i]
        
        # 收集所有历史 token 的 K, V
        all_k = []
        all_v = []
        
        for block_id in blocks:
            block_k = cache_manager.get_physical_block(block_id, 0, True)
            block_v = cache_manager.get_physical_block(block_id, 0, False)
            all_k.append(block_k)
            all_v.append(block_v)
        
        K_i = torch.cat(all_k, dim=0)[:num_tokens]  # (tokens, heads, head_dim)
        V_i = torch.cat(all_v, dim=0)[:num_tokens]
        
        # 标准注意力计算
        K_i = K_i.permute(1, 0, 2)  # (heads, tokens, head_dim)
        V_i = V_i.permute(1, 0, 2)
        
        scores = torch.matmul(q_i, K_i.transpose(-2, -1)) / (head_dim ** 0.5)
        attn_weights = torch.softmax(scores, dim=-1)
        output = torch.matmul(attn_weights, V_i)
        
        outputs.append(output)
    
    return torch.stack(outputs)

调度策略

1. FCFS（先来先服务）

def schedule_fcfs(waiting_queue, max_batch, max_seq):
    """先来先服务，简单但可能不够高效"""
    return waiting_queue[:max_batch]

2. SJF（短作业优先）

def schedule_sjf(waiting_queue, max_batch, max_seq):
    """短作业优先，优化平均延迟"""
    # 按预测长度排序
    sorted_queue = sorted(
        waiting_queue, 
        key=lambda r: r.predicted_length
    )
    return sorted_queue[:max_batch]

3. 预测调度（Predicted Length）

class PredictiveScheduler:
    def __init__(self, predictor):
        self.predictor = predictor  # 长度预测模型
    
    def schedule(self, waiting_queue, running_batch, max_batch, max_tokens):
        """
        智能调度：
        1. 预测每个请求的生成长度
        2. 组合请求使得总 token 数接近 max_tokens
        """
        # 预测长度
        for req in waiting_queue:
            if not hasattr(req, 'predicted_length'):
                req.predicted_length = self.predictor(req.prompt)
        
        # 组合请求（类似 bin packing）
        selected = []
        total_tokens = sum(
            len(r.prompt) + r.generated_length 
            for r in running_batch
        )
        
        for req in waiting_queue:
            req_tokens = len(req.prompt) + req.predicted_length
            if total_tokens + req_tokens <= max_tokens:
                selected.append(req)
                total_tokens += req_tokens
                if len(selected) >= max_batch:
                    break
        
        return selected

vLLM 架构

vLLM 是目前最成熟的连续批处理推理框架：

┌─────────────────────────────────────────────────┐
│                  vLLM Architecture               │
├─────────────────────────────────────────────────┤
│                                                 │
│  ┌─────────────┐    ┌─────────────────────────┐ │
│  │ API Server  │───▶│   Scheduler            │ │
│  │ (FastAPI)   │    │  - Request Queue       │ │
│  └─────────────┘    │  - Block Manager       │ │
│                     │  - Continuous Batching │ │
│                     └─────────────────────────┘ │
│                              │                  │
│                              ▼                  │
│                     ┌─────────────────────────┐ │
│                     │   Model Executor        │ │
│                     │  - PagedAttention      │ │
│                     │  - CUDA Kernels         │ │
│                     └─────────────────────────┘ │
│                              │                  │
│                              ▼                  │
│                     ┌─────────────────────────┐ │
│                     │   KV Cache Manager      │ │
│                     │  - Block Pool          │ │
│                     │  - Copy-on-Write       │ │
│                     └─────────────────────────┘ │
│                                                 │
└─────────────────────────────────────────────────┘

性能对比

框架吞吐量 (req/s)首 token 延迟 (ms)GPU 利用率
HuggingFace (静态)1285045%
TGI (静态)2542062%
vLLM (连续)8518092%
SGLang (连续)9216595%

生产部署

from vllm import LLM, SamplingParams

# 初始化
llm = LLM(
    model="meta-llama/Llama-3-70B",
    tensor_parallel_size=8,
    gpu_memory_utilization=0.9,
    max_model_len=65536,
)

# 批处理推理
prompts = [
    "Explain quantum computing",
    "Write a poem about AI",
    "What is the capital of France?",
    # ... 可以随时添加
]

sampling_params = SamplingParams(
    temperature=0.7,
    max_tokens=256,
)

# vLLM 自动使用连续批处理
outputs = llm.generate(prompts, sampling_params)

# 流式输出
for output in outputs:
    print(f"Prompt: {output.prompt!r}")
    print(f"Generated: {output.outputs[0].text!r}")

最佳实践

1. 调整 max_num_batched_tokens

# 根据显存调整
llm = LLM(
    model="...",
    max_num_batched_tokens=16384,  # 批处理总 token 上限
    # 更大的值 = 更高吞吐，但需要更多显存
)

2. 预测生成长度

# 为短回复优化
sampling = SamplingParams(
    max_tokens=50,
    stop=["\n\n"],  # 提前停止
)

3. 处理超长请求

# 分块处理超长 prompt
def chunked_generate(llm, prompt, chunk_size=8192):
    if len(prompt) <= chunk_size:
        return llm.generate(prompt)
    
    # 先处理前面的 chunk，再用最后一块生成
    chunks = [prompt[i:i+chunk_size] for i in range(0, len(prompt), chunk_size)]
    # ... 分块推理逻辑

参考文献

  • Yu, G. et al. (2022). Orca: A Distributed Serving System for Transformer-Based Generative Applications
  • Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention
  • vLLM Project (2023). vLLM: Easy, Fast, and Cheap LLM Serving

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。