← 返回投肯智能知识库

vLLM PagedAttention 原理与源码深度解析:从KV Cache分页到连续批处理

高级 📖 60 分钟 🕐 更新:2026-06-02 👁️ 标签:vLLM / PagedAttention / KV Cache / 推理优化 / 源码

📑 目录

一、背景:LLM 推理的内存瓶颈

在生产环境中部署大模型推理服务,最大的挑战之一是 GPU 内存的精细化管理。对于一个 7B 参数的模型,单次推理可能需要 14GB 显存(FP16),其中 KV Cache 占用的显存往往 超过模型权重本身

传统推理框架(HuggingFace Transformers)的 KV Cache 管理存在两个核心问题:

  1. 内存碎片化:每个序列预分配最大长度的连续显存块,浪费严重
  2. 无法动态批处理:一个长序列会阻塞整个 batch(队头阻塞)

vLLM 通过借鉴操作系统的 虚拟内存分页机制,完美解决了这两个问题,吞吐量比 Transformers 提升 14-24 倍

二、KV Cache 原理回顾

2.1 自回归生成的 KV Cache

Transformer 解码时,每生成一个新 token,都需要 attention 之前所有的 K/V 向量。这些向量被缓存下来,避免重复计算:

# 单个序列的 KV Cache 大小计算
# 假设:32层, 32头, 头维度128, 序列长度2048, FP16

per_token_kv = 2 * 32 * 32 * 128 * 2  # 2字节(FP16)
# = 524,288 bytes = 512 KB per token

seq_kv_total = 512 * 1024 * 2048
# = 1,073,741,824 bytes = 1 GB per sequence!

# 100 并发 × 1GB = 100GB 显存
# 仅仅 KV Cache 就要 100GB!

2.2 传统分配的痛点

传统方案:pre-allocate max_len * per_token_kv 连续显存

三、PagedAttention 核心思想

3.1 借鉴操作系统的虚拟内存

PagedAttention 把 KV Cache 分成 固定大小的 block(类似 OS 的 page),通过 块表(block table) 维护逻辑到物理的映射:

序列:  [Hello][, ][world][!][   ][   ][   ]
逻辑:  block0  block1  block2  block3  ...

块表:  [块3] → [块7] → [块1] → [块2]
         ↓
物理显存:
        块0: [未使用]
        块1: [world][!]    ← 已分配
        块2: [   ][   ]    ← 已分配
        块3: [Hello][, ]   ← 已分配
        块7: [world][!]    ← 已分配

3.2 核心数据结构

class Block:
    """物理 block(默认 16 tokens/block)"""
    block_id: int
    ref_count: int  # 引用计数(支持共享)
    token_ids: List[int]  # 实际存储的 token id


class BlockTable:
    """逻辑到物理的映射"""
    seq_id: int
    blocks: List[int]  # 物理 block id 列表


class KVCache:
    """分页 KV Cache 管理器"""
    def __init__(self, num_blocks, block_size=16):
        self.num_blocks = num_blocks
        self.block_size = block_size
        self.free_blocks = deque(range(num_blocks))
        self.block_tables = {}  # seq_id → BlockTable

3.3 写时复制(Copy-on-Write)

Beam Search 时,多个候选序列共享同一个 prefix block。PagedAttention 通过 ref_count 实现 CoW:

def append_token(seq_id, token_id):
    block_table = self.block_tables[seq_id]
    last_block_id = block_table.blocks[-1]
    
    if is_full(last_block_id):
        # 当前 block 已满,分配新 block
        new_block_id = self.free_blocks.popleft()
        block_table.blocks.append(new_block_id)
    else:
        # 复用当前 block
        pass
    
    # 写入 token
    block = self.blocks[last_block_id]
    block.token_ids.append(token_id)

四、Continuous Batching

4.1 传统 Batching 的问题

Static Batching 必须等所有序列都生成完毕才能开始新 batch,GPU 利用率极低:

时间线 →

Static Batch 1:  [seq1████████████████][seq2██████][seq3████████]
                 ↓
Static Batch 2:                                      [seq4████████][seq5██]
                                                     ↓
                                                     GPU闲置等待seq1-seq3

4.2 Continuous Batching 原理

vLLM 在每个 decode step 检查每个序列的状态:

新请求可以在任意 step 加入 batch:

def step(self):
    """每一步调度"""
    # 1. 收集所有活跃序列
    active_seqs = [s for s in self.sequences 
                   if not s.is_finished]
    
    # 2. 准备 batch 输入
    input_ids, block_tables = self.prepare_batch(active_seqs)
    
    # 3. GPU 前向
    logits = self.model(input_ids, block_tables)
    
    # 4. CPU 后处理
    for i, seq in enumerate(active_seqs):
        next_token = self.sample(logits[i])
        seq.append_token(next_token)
        
        # 检查是否结束
        if next_token == EOS or len(seq) >= max_len:
            seq.is_finished = True
            self.free_blocks(seq.block_table)
    
    # 5. 加入新请求
    while self.can_accept_new():
        new_seq = self.waiting_queue.pop()
        self.add_sequence(new_seq)

4.3 性能提升

实测数据(A100 80GB, LLaMA-2-7B, 序列长度 512):

框架吞吐量 (req/s)P99 延迟 (ms)显存利用率
HuggingFace Transformers4.2185035%
Text Generation Inference11.592062%
vLLM78.634092%

五、源码逐行解析

5.1 vLLM 工作流

async def serve_request(request):
    # 1. 编码 prompt
    prompt_tokens = tokenizer.encode(request.prompt)
    
    # 2. 创建 Sequence 对象
    seq = Sequence(
        seq_id=next_id(),
        prompt_token_ids=prompt_tokens,
        block_size=16,
    )
    
    # 3. 加入等待队列
    scheduler.waiting_queue.append(seq)
    
    # 4. 调度循环
    while not seq.is_finished:
        # 调度器决定哪些序列进入下一步
        batch = scheduler.schedule()
        
        # 模型前向
        output_tokens = model.step(batch)
        
        # 处理输出
        for seq_id, token in zip(batch.seq_ids, output_tokens):
            sequence = scheduler.get_sequence(seq_id)
            sequence.append_token(token)
            
            if sequence.is_finished:
                # 流式返回结果
                yield format_response(sequence)

5.2 Block Manager 核心逻辑

class BlockManager:
    def __init__(self, num_gpu_blocks, num_cpu_blocks, block_size):
        self.block_size = block_size
        self.gpu_allocator = BlockAllocator(num_gpu_blocks)
        self.cpu_allocator = BlockAllocator(num_cpu_blocks)
    
    def can_allocate(self, seq: Sequence) -> bool:
        """检查是否有足够 block"""
        required_blocks = len(seq.logical_token_blocks)
        return self.gpu_allocator.num_free_blocks >= required_blocks
    
    def allocate(self, seq: Sequence):
        """分配物理 block"""
        block_table = BlockTable()
        
        for logical_block in seq.logical_token_blocks:
            if block_table.is_empty():
                # 第一个 block 可能已存在(prefix sharing)
                block_id = self._get_cached_block_id(logical_block)
            else:
                # 新分配
                block_id = self.gpu_allocator.allocate()
            block_table.append(block_id)
        
        seq.block_table = block_table
    
    def can_append(self, seq: Sequence) -> bool:
        """检查是否可以追加 token"""
        last_block = self.blocks[seq.block_table.blocks[-1]]
        return not last_block.is_full
    
    def append_slot(self, seq: Sequence):
        """追加一个 token slot"""
        last_block_id = seq.block_table.blocks[-1]
        last_block = self.blocks[last_block_id]
        
        if last_block.is_full:
            # 需要新 block
            new_block_id = self.gpu_allocator.allocate()
            seq.block_table.blocks.append(new_block_id)
        
        seq.num_tokens += 1

5.3 注意力计算(核心)

class PagedAttention(nn.Module):
    """PagedAttention 注意力计算"""
    
    def forward(self, q, kv_cache, block_tables, context_lens):
        """
        q: [batch_size, num_heads, head_dim]
        kv_cache: 物理 block 存储的 K/V
        block_tables: 每个序列的 block 映射
        context_lens: 每个序列的实际长度
        """
        # 1. 根据 block_table 收集 K/V
        # 2. 计算注意力分数
        # 3. 应用 Flash Attention 优化
        
        # 使用 block_tables 做 gather
        k_cache, v_cache = self._gather_kv(
            kv_cache, block_tables, context_lens
        )
        
        # 标准 attention 计算
        attn_weights = torch.einsum(
            "bhgd,bhgd->bhg", q, k_cache
        ) / math.sqrt(self.head_dim)
        
        attn_weights = F.softmax(attn_weights, dim=-1)
        output = torch.einsum("bhg,bhgd->bhd", attn_weights, v_cache)
        
        return output

六、性能数据对比

6.1 不同模型规模

在 A100-80GB × 4 上的吞吐量(tokens/秒):

模型TransformersvLLM提升倍数
LLaMA-7B2,10031,40015.0×
LLaMA-13B1,40021,80015.6×
LLaMA-70B3807,60020.0×
Mixtral-8x7B62014,20022.9×
Qwen-72B4208,80021.0×

6.2 长序列场景优势更明显

序列长度 4096 vs 512:

原因:长序列 KV Cache 占主导,分页机制显著降低碎片。

6.3 生产部署建议

七、真实企业案例

7.1 某互联网公司推理服务改造

某头部电商公司原先用 HuggingFace Transformers 部署客服机器人,QPS 30,需要 30 张 A100。改造到 vLLM 后,QPS 提升到 280,资源减少 60%:

7.2 某金融科技公司代码补全服务

Code Llama-34B 部署用于内部 IDE 代码补全,关键指标:

八、性能调优实战指南

8.1 Chunked Prefill 优化

长 prompt(10K+ tokens)会阻塞整个 batch,vLLM 用 Chunked Prefill 解决:

# 启用 Chunked Prefill(默认已开启)
vllm serve codellama-34b \
  --enable-chunked-prefill \
  --max-num-batched-tokens 8192 \
  --max-model-len 16384

# 效果:长 prompt 不再阻塞短请求
# 短请求 P99 延迟从 1.2s 降至 280ms

8.2 Speculative Decoding

用小模型 draft + 大模型 verify,加速 2-3 倍:

from vllm import LLM, SamplingParams

# 主模型(700B)
main_model = LLM(model="llama-3-70b")

# 草稿模型(7B)  
draft_model = LLM(model="llama-3-7b")

# Speculative Decoding
output = main_model.generate(
    prompts,
    SamplingParams(temperature=0),
    speculative_model=draft_model,
    num_speculative_tokens=5  # 每次草稿 5 个 token
)

# 实测加速:2.4 倍
# 代价:显存占用增加 20%(需额外加载 draft model)

8.3 Prefix Caching(系统提示词缓存)

如果你的应用有固定的 system prompt(如"你是一个代码助手..."),启用 prefix caching 可以让所有请求共享 prompt 的 KV Cache:

vllm serve qwen-72b \
  --enable-prefix-caching \
  --block-size 16

# 实际效果:
# - 100 用户的 system prompt 相同
# - 只计算 1 次 prompt 的 KV Cache(512 tokens)
# - 节省 99 次重复计算
# - 首 token 延迟从 800ms 降至 50ms

8.4 量化部署

# AWQ 量化(推荐,4-bit 几乎无损)
vllm serve Qwen2-72B-Instruct-AWQ \
  --quantization awq \
  --gpu-memory-utilization 0.95

# GPTQ 量化
vllm serve codellama-34b-gptq \
  --quantization gptq \
  --dtype float16

# FP8 量化(H100 上性能最佳)
vllm serve llama-3-70b \
  --quantization fp8 \
  --kv-cache-dtype fp8

# 显存对比(Qwen-72B):
# FP16: 144GB (需要 2 张 A100-80GB)
# AWQ 4-bit: 42GB (单张 A100-80GB 即可)
# 吞吐量:AWQ 4-bit 仅下降 8%,但能多服务 3 倍用户

九、常见问题排错

9.1 OOM (Out of Memory)

# 错误信息
torch.cuda.OutOfMemoryError: CUDA out of memory.

# 排查步骤:
# 1. 减小 gpu-memory-utilization
vllm serve model --gpu-memory-utilization 0.85

# 2. 启用量化
vllm serve model --quantization awq

# 3. 减少 max-model-len
vllm serve model --max-model-len 8192

# 4. 减少并发
vllm serve model --max-num-seqs 64

# 5. 启用 chunked prefill 减内存峰值
vllm serve model --enable-chunked-prefill

9.2 首 token 延迟高

# 原因:模型没预热,CUDA kernel 首次编译
# 解决:启动时跑预热请求

vllm serve qwen-72b \
  --enable-prefix-caching \
  --warmup 请求(5 个不同长度的)

# 实际数据:首 token 延迟从 1.2s → 80ms

9.3 流式输出断流

# 原因:客户端超时设置太短
# 解决:
# 1. Nginx 反向代理超时
proxy_read_timeout 600s;
proxy_send_timeout 600s;

# 2. 客户端 fetch 取消
const controller = new AbortController();
setTimeout(() => controller.abort(), 300000);  // 5 分钟超时

十、vLLM vs 其他推理框架

框架吞吐量易用性生产成熟度推荐场景
vLLM⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐通用大模型服务
TGI (HuggingFace)⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐快速原型
TensorRT-LLM⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐NVIDIA 极致性能
SGLang⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐⭐结构化生成
DeepSpeed-MII⭐⭐⭐⭐⭐⭐⭐⭐⭐微软生态

十一、未来发展方向

  1. 多模态原生支持:vLLM v0.6+ 已支持图像输入
  2. 分布式推理:跨节点张量并行(Tensor Parallel + Pipeline Parallel)
  3. LoRA 热加载:运行时动态加载不同 LoRA,无需重启服务
  4. 推测解码标准化:更智能的 draft model 选择
  5. 端侧推理:vLLM-mobile 项目(试验阶段)

小刚结论:vLLM 是当前生产环境推理部署的最佳选择,没有之一。哪怕你不理解 PagedAttention,也能享受到它带来的吞吐量红利。但如果想真正发挥 vLLM 性能,理解 KV Cache 分页、Continuous Batching、Chunked Prefill 这些核心概念是必须的。