vLLM 推理引擎：PagedAttention 与 KVCache 管理

zs宝啊

385人浏览 · 2026-04-28 05:00:02

zs宝啊 · 2026-04-28 05:00:02 发布

vLLM 推理引擎：PagedAttention 与 KVCache 管理

📖 引言

随着大语言模型（LLM）的参数规模从数十亿增长到数千亿，模型推理的内存开销和计算延迟成为了生产环境中的核心瓶颈。传统的推理框架如 HuggingFace Transformers 在处理批量请求时，往往受限于 GPU 显存的碎片化问题，导致显存利用率不足 40%，吞吐量严重受限。

vLLM（version LLM）作为 UC Berkeley 推出的高性能推理引擎，通过创新的 PagedAttention 机制和受操作系统虚拟内存启发的 KVCache 管理策略，在保持模型精度的同时，将推理吞吐量提升了 3-24 倍。本文将深入 vLLM 0.4.0+ 源码，剖析其核心设计原理与工程实现。

核心问题：传统推理引擎为何显存利用率低？PagedAttention 如何解决 KVCache 的内存碎片问题？vLLM 的连续批处理（Continuous Batching）机制又是如何工作的？

🔑 核心概念

KVCache：推理加速的关键

在大模型推理过程中，生成每个 token 都需要计算它与之前所有生成 token 的注意力分数。直接计算会导致复杂度呈二次方增长（O(n²)），KVCache 通过缓存每个位置的 Key 和 Value 向量，将复杂度降至 O(n)。

传统 KVCache 的问题：

预分配浪费：为每个请求预分配最大序列长度的内存（如 2048 tokens），但实际生成长度往往远小于此
内存碎片化：不同请求的序列长度差异大，导致显存碎片严重
静态批处理：同批次内所有请求必须等待最慢的请求完成，才能开始下一批次

PagedAttention：虚拟内存的灵感

vLLM 借鉴操作系统虚拟内存的分页机制，将 KVCache 划分为固定大小的块（Blocks），每个块包含多个 token 的 KV 对。这种设计带来了三大优势：

动态内存分配：按需分配块，避免预分配浪费
内存共享：多个请求可共享相同的 KVCache 块（如系统提示词）
高效内存管理：类似 CPU 的页表管理，实现紧凑的显存布局

连续批处理（Continuous Batching）

vLLM 引入了迭代级调度（Iteration-level Scheduling），在每个生成步骤后动态重组批次，允许新请求加入、完成的请求退出。相比传统静态批处理，显著提高了 GPU 利用率。

🧠 源码深度解析

1. 核心数据结构

vLLM 的 KVCache 管理围绕 Block 和 BlockManager 展开，以下是关键数据结构（源码路径：vllm/vllm/vllm/attention/backends/interface.py，v0.4.0）：

# 物理块：实际存储 KV 数据的显存块
class PhysicalTokenBlock:
    def __init__(self, device: str, block_index: int, block_size: int):
        self.device = device           # GPU 设备 ID
        self.block_index = block_index # 物理块索引（全局唯一）
        self.block_size = block_size   # 每个块的 token 数量（默认 16）
        self.ref_count = 0             # 引用计数，用于内存共享优化

# 逻辑块：请求视角的虚拟块，映射到物理块
class Block:
    def __init__(self, block_id: int, block_size: int):
        self.block_id = block_id       # 逻辑块 ID
        self.block_size = block_size   # 块大小（token 数）
        self.token_ids = []            # 块内存储的 token IDs
        self.physical_block_id = None  # 映射的物理块 ID（由 BlockManager 分配）

# 块表：维护逻辑块到物理块的映射关系
class BlockTable:
    def __init__(self):
        self.mapping = {}  # {logical_block_id: physical_block_id}
    
    def add_block(self, logical_id: int, physical_id: int):
        """添加一个逻辑块到物理块的映射"""
        self.mapping[logical_id] = physical_id
    
    def get_physical_block(self, logical_id: int) -> int:
        """查询逻辑块对应的物理块"""
        return self.mapping.get(logical_id, -1)

2. BlockManager：显存分配器

BlockManager 负责物理块的分配、回收和映射管理（源码路径：vllm/vllm/vllm/attention/backends/paged_attn/block_manager.py）：

class BlockSpaceManager:
    def __init__(self, block_size: int, num_gpu_blocks: int, num_cpu_blocks: int):
        self.block_size = block_size           # 每个块的 token 数量
        self.num_gpu_blocks = num_gpu_blocks   # GPU 可用物理块数
        self.num_cpu_blocks = num_cpu_blocks   # CPU 可用物理块数（用于 offload）
        
        # 显存池：Free 表示空闲，Allocated 表示已分配
        self.gpu_allocator = BlockAllocator(num_gpu_blocks)
        self.cpu_allocator = BlockAllocator(num_cpu_blocks)
        
        # 每个序列的块表：{seq_id: BlockTable}
        self.block_tables = defaultdict(BlockTable)
    
    def allocate(self, seq_id: int, token_ids: List[int]) -> List[int]:
        """
        为序列分配 KVCache 块
        返回：分配的逻辑块 ID 列表
        """
        required_blocks = (len(token_ids) + self.block_size - 1) // self.block_size
        allocated_blocks = []
        
        for _ in range(required_blocks):
            # 优先从 GPU 分配，不足则从 CPU 分配
            if self.gpu_allocator.free_blocks > 0:
                physical_block = self.gpu_allocator.allocate()
            elif self.cpu_allocator.free_blocks > 0:
                physical_block = self.cpu_allocator.allocate()
            else:
                raise OutOfMemoryError("No available blocks for allocation")
            
            # 创建逻辑块并建立映射
            logical_block = Block(len(self.block_tables[seq_id]), self.block_size)
            self.block_tables[seq_id].add_block(logical_block.block_id, physical_block)
            allocated_blocks.append(logical_block.block_id)
        
        return allocated_blocks
    
    def free(self, seq_id: int):
        """释放序列占用的所有块"""
        if seq_id not in self.block_tables:
            return
        
        for logical_id, physical_id in self.block_tables[seq_id].mapping.items():
            # 根据物理块位置归还到对应分配器
            if physical_id < self.num_gpu_blocks:
                self.gpu_allocator.free(physical_id)
            else:
                self.cpu_allocator.free(physical_id - self.num_gpu_blocks)
        
        del self.block_tables[seq_id]

3. PagedAttention 核心算法

PagedAttention 的核心在于计算注意力时，通过块索引高效聚合 KV 数据（源码路径：vllm/vllm/vllm/model_executor/attention.py）：

class PagedAttention(nn.Module):
    def __init__(self, num_heads: int, head_size: int, block_size: int):
        super().__init__()
        self.num_heads = num_heads       # 注意力头数
        self.head_size = head_size       # 每个头的维度
        self.block_size = block_size     # 每个块的 token 数量
    
    def forward(
        self,
        query: torch.Tensor,             # [num_tokens, num_heads, head_size]
        key_cache: torch.Tensor,         # [num_blocks, block_size, num_heads, head_size]
        value_cache: torch.Tensor,       # [num_blocks, block_size, num_heads, head_size]
        block_tables: torch.Tensor,      # [num_seqs, max_blocks_per_seq]
        context_lens: torch.Tensor,      # [num_seqs] 每个序列的实际长度
    ) -> torch.Tensor:
        """
        计算 PagedAttention
        关键：通过 block_tables 索引 key_cache 和 value_cache
        """
        num_tokens, num_heads, head_size = query.shape
        num_seqs = block_tables.shape[0]
        
        # 1. 将 query reshape 为 [num_seqs, num_tokens_per_seq, num_heads, head_size]
        query = query.view(num_seqs, -1, num_heads, head_size)
        
        # 2. 为每个序列收集其对应的 KV 块
        output = []
        for seq_idx in range(num_seqs):
            seq_len = context_lens[seq_idx].item()
            seq_query = query[seq_idx, :seq_len]  # [seq_len, num_heads, head_size]
            seq_block_ids = block_tables[seq_idx]  # [max_blocks_per_seq]
            
            # 计算需要访问的块数
            num_blocks = (seq_len + self.block_size - 1) // self.block_size
            seq_block_ids = seq_block_ids[:num_blocks]
            
            # 3. 从 key_cache/value_cache 中提取该序列的 KV 数据
            seq_keys = key_cache[seq_block_ids]     # [num_blocks, block_size, num_heads, head_size]
            seq_values = value_cache[seq_block_ids] # [num_blocks, block_size, num_heads, head_size]
            
            # 4. 展平为 [seq_len, num_heads, head_size]
            seq_keys = seq_keys.view(-1, num_heads, head_size)[:seq_len]
            seq_values = seq_values.view(-1, num_heads, head_size)[:seq_len]
            
            # 5. 计算注意力分数：Q @ K^T / sqrt(d)
            attn_scores = torch.matmul(
                seq_query.unsqueeze(2),  # [seq_len, 1, num_heads, head_size]
                seq_keys.transpose(1, 2)  # [seq_len, num_heads, head_size, 1]
            ).squeeze(-1) / (head_size ** 0.5)  # [seq_len, seq_len]
            
            # 6. Softmax 归一化
            attn_weights = F.softmax(attn_scores, dim=-1)
            
            # 7. 加权求和 Value
            seq_output = torch.matmul(attn_weights.unsqueeze(2), seq_values).squeeze(2)
            output.append(seq_output)
        
        return torch.cat(output, dim=0)  # [num_tokens, num_heads, head_size]

4. 调度器：连续批处理的核心

调度器负责在每个迭代步骤后动态调整批次（源码路径：vllm/vllm/vllm/engine/scheduler.py）：

class Scheduler:
    def __init__(self, block_manager: BlockSpaceManager, max_model_len: int):
        self.block_manager = block_manager
        self.max_model_len = max_model_len
        self.running_queue = []    # 正在生成的请求
        self.waiting_queue = []    # 等待资源的请求
    
    def schedule(self) -> SchedulingDecision:
        """
        执行调度决策
        返回：包含批次输入和输出元组的决策对象
        """
        decision = SchedulingDecision()
        
        # 1. 检查已完成请求（生成长度达到上限或生成 EOS token）
        finished_seqs = [seq for seq in self.running_queue if self._is_finished(seq)]
        for seq in finished_seqs:
            decision.finished_seqs.append(seq)
            self.block_manager.free(seq.seq_id)  # 释放显存
            self.running_queue.remove(seq)
        
        # 2. 尝试从等待队列调度新请求
        while self.waiting_queue and self._can_allocate_more():
            new_seq = self.waiting_queue.pop(0)
            try:
                # 尝试为新请求分配 KVCache 块
                self.block_manager.allocate(new_seq.seq_id, new_seq.token_ids)
                self.running_queue.append(new_seq)
                decision.prefill_seqs.append(new_seq)
            except OutOfMemoryError:
                # 显存不足，重新放回队列
                self.waiting_queue.insert(0, new_seq)
                break
        
        # 3. 正在生成的请求继续解码
        decision.decode_seqs = [seq for seq in self.running_queue]
        
        return decision
    
    def _can_allocate_more(self) -> bool:
        """检查是否还有足够的显存块"""
        return (
            self.block_manager.gpu_allocator.free_blocks > 0 or
            self.block_manager.cpu_allocator.free_blocks > 0
        )
    
    def _is_finished(self, seq: Sequence) -> bool:
        """判断序列是否生成完成"""
        return seq.get_last_token() == EOS_TOKEN_ID or seq.get_len() >= self.max_model_len

🎯 实战应用

部署 vLLM 推理服务

以下是一个完整的 vLLM 服务部署示例（基于 v0.4.0）：

# 安装 vLLM（推荐使用 CUDA 11.8+）
pip install vllm==0.4.0

# 启动 OpenAI 兼容的 API 服务
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-7b-hf \      # 模型路径（HuggingFace 格式）
    --tensor-parallel-size 2 \               # 张量并行度（多 GPU）
    --block-size 16 \                        # 每个 KVCache 块的 token 数量
    --gpu-memory-utilization 0.90 \          # GPU 显存利用率（0.9 表示 90%）
    --max-model-len 2048 \                   # 最大序列长度
    --dtype float16 \                        # 数据类型（float16/bfloat16）
    --host 0.0.0.0 \                         # 监听地址
    --port 8000                              # 监听端口

Python 客户端调用示例

from openai import OpenAI

# 初始化客户端（指向 vLLM 服务）
client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="dummy"  # vLLM 不需要真实 API Key
)

# 流式生成
stream = client.chat.completions.create(
    model="meta-llama/Llama-2-7b-hf",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "解释 PagedAttention 的核心原理"}
    ],
    max_tokens=512,
    temperature=0.7,
    stream=True  # 启用流式输出
)

for chunk in stream:
    if chunk.choices[0].delta.content:
        print(chunk.choices[0].delta.content, end="", flush=True)

性能优化技巧

调整块大小（block-size）：
- 小块（8-16）：适合短序列，减少内存浪费，但增加块管理开销
- 大块（32-64）：适合长序列，减少块数量，但可能浪费内存
- 推荐值：16（平衡内存利用率和计算效率）
GPU 显存利用率（gpu-memory-utilization）：
- 保守值：0.85（避免 OOM）
- 激进值：0.95（最大化吞吐量，需监控显存）
张量并行（tensor-parallel-size）：
- 单 GPU：1
- 多 GPU：建议 2/4/8（需网卡支持 NCCL）

📊 对比分析

vLLM vs 传统推理框架

特性	vLLM	HuggingFace Transformers	TGI (Text Generation Inference)
显存管理	PagedAttention（动态分页）	预分配静态 KVCache	PagedAttention（类似 vLLM）
批处理策略	连续批处理（迭代级调度）	静态批处理	连续批处理
显存利用率	80-90%	30-40%	70-80%
吞吐量提升	基准 3-24x	基准 1x	2-5x
部署复杂度	中（需调优块大小）	低（开箱即用）	高（Docker + CUDA）
适用场景	高并发生产环境	研究实验	企业级部署

PagedAttention vs 多查询注意力（MQA/GQA）

注意力机制	KV Cache 存储量	计算复杂度	显存占用	精度影响
MHA（多头注意力）	num_heads × seq_len × 2 × head_dim	O(n²)	高	无
MQA（多查询注意力）	1 × seq_len × 2 × head_dim	O(n²)	低	轻微下降
GQA（分组查询注意力）	num_kv_heads × seq_len × 2 × head_dim	O(n²)	中	轻微下降
PagedAttention	num_heads × seq_len × 2 × head_dim	O(n²)	中（通过分页优化）	无