vLLM 核心工程深度剖析：PagedAttention 与 KV Cache 调度如何重塑 LLM 推理基础设施

引言

大规模语言模型（LLM）的推理部署正面临一个核心矛盾：模型参数规模以每年 10 倍速度增长（从 GPT-3 的 175B 到 GPT-4 的 1.8T MoE），而 GPU HBM 带宽和容量的提升却远不及此。在自回归推理中，每个 token 的生成都需要访问之前所有 token 的 Key-Value Cache（KV Cache），这导致推理吞吐量的瓶颈从计算密度转向了内存访问效率。vLLM 提出的 PagedAttention 机制，借鉴了操作系统中虚拟内存和分页管理的思想，革命性地解决了 KV Cache 的内存碎片化和预分配浪费问题，将 LLM 推理服务的吞吐量提升了 2-4 倍。本文将从工程实现层面深入剖析 PagedAttention 的核心数据结构、内存分配器设计、Copy-on-Write 优化、以及前缀缓存（Prefix Caching）策略，并给出生产级部署的调优指南。

1. KV Cache 的内存挑战

在标准的自回归 Transformer 推理中，每个已生成 token 对应的 KV 向量需要在解码阶段被反复访问。对于一个有 L 层、隐藏维度为 d 的模型，单个 token 的 KV Cache 大小为 2 × L × d × sizeof(dtype) bytes。

以 Llama-2-70B 为例（L=80, d=8192, FP16）：每个 token 的 KV Cache 约为 2.5 MB。一个 2048 token 的序列就需要 5 GB 的 GPU 显存仅用于 KV Cache。而在一个 batch 中同时处理数十个请求时，KV Cache 的内存管理效率直接决定了系统的吞吐能力。

1.1 传统实现的三大弊病

2. PagedAttention：操作系统思维的跨界应用

PagedAttention 的核心洞察是：KV Cache 的访问模式与操作系统中进程的内存页面访问高度相似——顺序访问为主、需要快速地址映射、支持共享与写时复制。vLLM 将这一思想完整移植到 GPU 显存管理中。

2.1 核心抽象三层模型

┌─────────────────────────────────────────────────────┐
│  Logical Block Table (每个请求一个)                    │
│  [0]→3  [1]→7  [2]→1  [3]→5  (logical→physical)     │
├─────────────────────────────────────────────────────┤
│  Physical Block Pool (GPU 显存池)                     │
│  [0] [1] [2] [3] [4] [5] [6] [7] ...              │
│   │              │                                   │
│   └── Req A ─────┘                                   │
├─────────────────────────────────────────────────────┤
│  BlockPool Manager (CPU 端)                          │
│  free_list: [0, 4, 6, ...]                          │
│  ref_count: [0→1, 3→2 (共享), ...]                   │
└─────────────────────────────────────────────────────┘

逻辑块（Logical Block）：每个请求维护一个逻辑块号到物理块号的映射表。逻辑块是连续的（从 0 开始递增），但物理块可以任意分散在显存池中。

物理块（Physical Block）：固定大小的内存块（通常 16 个 token 的 KV Cache），是 GPU 显存分配的最小单元。块大小在 vLLM 启动时通过 --block-size 参数指定。

块表（Block Table）：每个请求一个，存储逻辑块到物理块的映射。在 CUDA kernel 执行时传入，kernel 通过间接寻址访问非连续的物理内存。

2.2 显存分配器实现

class BlockPool:
    def __init__(self, num_blocks: int, block_size: int, gpu_device: str):
        self.block_size = block_size          # 每个块存储的 token 数 (默认16)
        self.free_list = list(range(num_blocks))  # 空闲物理块号
        self.ref_count = [0] * num_blocks    # 引用计数
        self.block_tables: dict[int, List[int]] = {}  # req_id -> physical blocks

    def allocate(self, req_id: int, num_blocks: int) -> List[int]:
        """为请求分配物理块"""
        if len(self.free_list) < num_blocks:
            raise OutOfMemoryError(f"GPU KV Cache exhausted")
        blocks = [self.free_list.pop() for _ in range(num_blocks)]
        self.ref_count[b] += 1 for b in blocks
        self.block_tables[req_id] = blocks
        return blocks

    def append_slot(self, req_id: int) -> Optional[int]:
        """追加一个 token 的空间：需要新块时分配"""
        blocks = self.block_tables[req_id]
        # 检查最后一个块是否有空槽
        if self._last_block_has_space(blocks):
            return None  # 无需分配新块
        # 需要新块
        if not self.free_list:
            return None  # OOM
        new_block = self.free_list.pop()
        blocks.append(new_block)
        self.ref_count[new_block] += 1
        return new_block

    def free(self, req_id: int):
        """释放请求的所有物理块"""
        for b in self.block_tables[req_id]:
            self.ref_count[b] -= 1
            if self.ref_count[b] == 0:
                self.free_list.append(b)
        del self.block_tables[req_id]

3. CUDA Kernel 的分页注意力实现

PagedAttention 的核心创新不仅在于内存分配，更在于对应的 CUDA kernel 能否高效地在分页 KV Cache 上执行注意力计算。vLLM 的 paged_attention_v1/v2 kernel 实现了所谓的"间接内存访问"。

3.1 Kernel 签名与分块策略

// paged_attention_v1 kernel (简化)
template<typename scalar_t, int HEAD_SIZE, int BLOCK_SIZE, int NUM_THREADS>
__global__ void paged_attention_v1_kernel(
    scalar_t* __restrict__ out,           // [num_seqs, num_heads, head_size]
    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
    const scalar_t* __restrict__ k_cache, // [num_blocks, block_size, num_kv_heads, head_size]
    const scalar_t* __restrict__ v_cache, // [num_blocks, block_size, num_kv_heads, head_size]
    const int* __restrict__ block_tables, // [num_seqs, max_num_blocks_per_seq]
    const int* __restrict__ context_lens, // [num_seqs]
    int max_num_blocks_per_seq,
    float scaling, const float* alibi_slopes) {
    // 1. 通过 block_table 将逻辑块索引转换为物理块偏移
    // 2. 对每个 KV head，循环遍历所有块
    // 3. 在线 softmax 计算（online softmax：两遍扫描，数值稳定）
    // 4. V 加权求和
}

3.2 分页与连续的性能对比

间接内存访问是否会损失性能？vLLM 团队的 benchmark 显示：

• 计算受限场景（大 batch、长序列）：分页注意力与连续注意力性能差距 <5%，因为瓶颈在矩阵乘法而非内存访问
• 内存受限场景（小 batch、短序列）：分页注意力可能慢 10-20%，但此时整体推理延迟本身就很小
• 综合收益：虽然单次注意力计算略慢，但显存利用率提升带来的 batch 增大效应使 端到端吞吐量提升 2-4 倍

4. Copy-on-Write 与序列并行

vLLM 的另一个核心优化是 KV Cache 的 Copy-on-Write（CoW）机制，主要用于 Beam Search 和并行采样场景。

4.1 Beam Search 中的 CoW

在 Beam Search 推理中，多个候选序列共享相同的 prefix tokens。当某个 beam 扩展一个新 token 时：

1. 新 beam 与父 beam 共享所有 prefix 对应的物理块（引用计数 +1）
2. 仅在最后一个物理块中写入新的 KV 条目
3. 如果最后一个块已满，分配新块（而非复制所有已有块）
4. 使用 cudaMemcpyAsync 实现块级 CoW，带宽占用仅为实际修改部分的 1/block_size

# vLLM 的 CoW 实现逻辑（简化）
class BlockAllocator:
    def fork(self, parent_req_id: str, child_req_id: str):
        """子请求继承父请求的块，共享物理内存"""
        parent_blocks = self.block_tables[parent_req_id]
        self.block_tables[child_req_id] = parent_blocks.copy()
        for b in parent_blocks:
            self.ref_count[b] += 1

    def swap_out(self, req_id: str) -> Tuple[List[int], torch.Tensor]:
        """将 KV Cache 从 GPU swap 到 CPU 显存（用于抢占调度）"""
        blocks = self.block_tables[req_id]
        cpu_tensor = torch.empty(...)
        for i, block in enumerate(blocks):
            copy_block_to_cpu(gpu_blocks[block], cpu_tensor[i])
        self.free(req_id)
        return blocks, cpu_tensor  # CPU 上的副本

5. 前缀缓存（Automatic Prefix Caching）

vLLM 0.4+ 引入了自动前缀缓存（APC），进一步消除具有相同 system prompt 或对话历史的请求之间的 KV Cache 冗余存储。

5.1 哈希树索引

class PrefixCache:
    def __init__(self):
        self.hash_to_blocks: Dict[str, List[int]] = {}  # content hash → physical blocks
        self.lru = LRUCache(maxsize=10000)              # LRU 淘汰
    
    def lookup(self, token_ids: List[int]) -> Optional[List[int]]:
        """查找是否有已缓存的前缀"""
        key = hash_token_ids(token_ids)
        if key in self.hash_to_blocks:
            self.lru.touch(key)
            return self.hash_to_blocks[key]
        return None
    
    def cache_prefix(self, token_ids: List[int], blocks: List[int]):
        """缓存已完成请求的前缀"""
        key = hash_token_ids(token_ids)
        self.hash_to_blocks[key] = blocks
        self.lru[key] = True

5.2 缓存命中收益

在实际生产场景中（如多轮对话、共享 system prompt 的 API 调用），前缀缓存可以达到 30-60% 的命中率：

• 首 token 延迟（TTFT）降低 40-70%：命中前缀意味着无需重新计算对应 KV
• 吞吐量提升 20-35%：减少了 GPU 计算量和显存带宽消耗
• 显存效率提升：同一份物理块可以被多个活跃请求共享

6. GPU 显存管理与调度器设计

vLLM 的调度器需要在 GPU 显存约束下做出精细的准入控制决策。其核心流程：

新请求到达
    │
    ▼
┌───────────────────────────────┐
│ 1. 估算新请求所需的 KV 块数     │
│    num_blocks = ceil(num_computed_tokens / block_size) │
│    + ceil((num_prompt_tokens + num_output_tokens) / block_size) │
├───────────────────────────────┤
│ 2. 检查可用物理块              │
│    if free_blocks >= num_blocks: │
│        → 准入，开始预填充       │
│    else:                        │
│        → 抢占（Preemption）     │
│          - SWAP: 将低优先级请求的 KV Cache 移到 CPU │
│          - RECOMPUTE: 释放 KV Cache，后续重算   │
├───────────────────────────────┤
│ 3. 执行预填充（Prefill）        │
│    → 计算所有 prompt token 的 KV │
│    → 填充物理块                 │
│    → 释放计算中间态             │
├───────────────────────────────┤
│ 4. 加入 Decode 队列             │
│    → 按 priority 排队           │
│    → 组装 batch 执行 decode kernel │
└───────────────────────────────┘

6.1 抢占策略

• Recomputation（默认）：直接释放被抢占请求的 KV Cache，恢复时重新计算。适合 GPU↔CPU 带宽受限的场景
• Swapping：将 KV Cache 复制到 CPU 内存，恢复时拷回。适合 CPU 内存充裕、PCIe 带宽充足的场景
• Chunked Prefill：将长 prompt 的预填充拆分为多个 chunk，与其他 decode 操作交错执行，缓解大 prefill 阻塞 decode 的问题

7. 生产部署调优指南

7.1 监控关键指标

• KV Cache 使用率：vllm:gpu_cache_usage_perc — 超过 95% 意味着频繁抢占
• 前缀缓存命中率：vllm:prefix_cache_hit_rate — 低于 20% 可考虑增大哈希树
• 抢占频率：vllm:num_preemptions_total — 高抢占率表明 GPU 显存不足
• 平均首 token 延迟（TTFT）：衡量预填充性能，chunked prefill 可显著降低

8. 与 SGLang RadixAttention 的对比

9. 未来演进方向

• Cross-Attention KV Cache：多模态模型（如 LLaVA）中图像编码 KV 的管理策略
• Speculative Decoding 集成：投机解码与 PagedAttention 的交互优化
• Disaggregated Prefill：预填充与解码分离部署，各自独立显存管理（NVIDIA 正在主导）
• FP4/INT4 量化 KV Cache：减少 KV Cache 显存占用，代价是注意力精度
• Pageable LoRA：将 LoRA 权重也纳入分页管理，支持超大规模模型的多租户服务

10. 总结

vLLM 的 PagedAttention 机制将操作系统级别的虚拟内存管理思想引入 GPU 显存管理，彻底解决了 LLM 推理中 KV Cache 的内存效率问题。其设计哲学——用间接寻址换取灵活性、用局部不连续性换取全局利用率——在工程实践中得到了充分验证。随着 LLM 推理从单模型服务走向多模态、多租户、投机解码的复杂未来，PagedAttention 所奠定的分页管理框架将继续作为 LLM 推理基础设施的核心范式发挥关键作用。对于从事 AI 基础设施建设的工程师而言，深入理解 PagedAttention 的工程细节，是构建高性能、高利用率推理服务的必要条件。