杨艺韬讲堂

Appearance

第5章 KV Cache 管理：寸土寸金的显存 ​

"The cheapest, fastest, and most reliable components are those that aren't there." -- Gordon Bell

本章要点

  • 理解 BlockPool 的数据结构设计：为什么用 __slots__ 和空闲链表
  • 掌握 KVCacheManager 的分配与释放策略
  • 深入引用计数机制：块共享与 Copy-on-Write 的工程实现
  • 理解 LRU 驱逐策略：当显存不足时如何选择牺牲者
  • 认识前缀缓存的基础设施（为第 10 章做铺垫）

5.1 BlockPool：块的资源池 ​

上一章我们知道了物理块的概念。现在的问题是：谁来管理这些块？

答案是 BlockPool（vllm/v1/core/block_pool.py）。它是所有物理块的容器和管理者，职责包括：

  • 维护空闲块列表
  • 分配块给请求
  • 回收已释放的块
  • 管理块的引用计数
  • 维护前缀缓存的哈希索引

KVCacheBlock：极致精简的数据结构 ​

先看 BlockPool 管理的基本单元——KVCacheBlock（vllm/v1/core/kv_cache_utils.py）：

python
class KVCacheBlock:
    __slots__ = (
        "block_id",
        "ref_cnt",
        "block_hash",
        "prev_free_block",
        "next_free_block",
    )

    def __init__(self, block_id: int):
        self.block_id = block_id
        self.ref_cnt = 0
        self.block_hash: Optional[BlockHash] = None
        self.prev_free_block: Optional[KVCacheBlock] = None
        self.next_free_block: Optional[KVCacheBlock] = None

注意 __slots__ 声明——这不是偶然的。vLLM 可能管理数万个块（一张 80 GB A100 上，如果每个块 320 KB，大约有 170,000 个块）。使用 __slots__ 代替普通的 __dict__，每个对象节省约 100 字节的内存和显著的属性访问时间。在管理十万级别对象时，这种微优化累积起来相当可观。

五个字段各有含义：

字段作用
block_id物理块编号（对应 GPU 显存中的位置）
ref_cnt引用计数。0 = 空闲，1 = 独占，>1 = 共享
block_hash哈希值（用于前缀缓存匹配）
prev_free_block空闲双向链表：前驱
next_free_block空闲双向链表：后继

FreeKVCacheBlockQueue：空闲块管理 ​

空闲块用一个双向链表管理——FreeKVCacheBlockQueue。为什么用链表而不是列表（Python list）？

因为需要 O(1) 的插入和删除。块的释放（插入）和分配（删除）是最频繁的操作，每一步推理都会发生。用 Python list 的 append 和 pop(0) 虽然简单，但 pop(0) 是 O(n) 的。双向链表保证了两端操作都是 O(1)。

空闲链表同时充当 LRU 队列的角色：

  • 释放块时，插入到链表尾部（最近释放）
  • 分配块时，从链表头部取出（最久未使用）
  • 驱逐缓存块时，同样从头部驱逐（LRU 策略）

一个数据结构，三种用途——这是工程上的精巧设计。

5.2 KVCacheManager：调度器的左右手 ​

KVCacheManager（vllm/v1/core/kv_cache_manager.py）是调度器与 BlockPool 之间的桥梁。调度器不直接操作 BlockPool，而是通过 KVCacheManager 的高层接口。

分配：为请求分配块 ​

当调度器决定让一个新请求参与计算时，KVCacheManager 需要为它分配足够的块：

python
# 简化逻辑
def allocate_blocks(self, request, num_tokens):
    num_blocks_needed = ceil(num_tokens / block_size)
    current_blocks = len(request.block_ids)
    new_blocks_needed = num_blocks_needed - current_blocks

    if new_blocks_needed <= 0:
        return  # 已有足够的块

    if self.block_pool.free_count < new_blocks_needed:
        return None  # 显存不足，告知调度器

    # 从 BlockPool 分配
    new_block_ids = []
    for _ in range(new_blocks_needed):
        block = self.block_pool.allocate()
        block.ref_cnt = 1
        new_block_ids.append(block.block_id)

    request.block_ids.extend(new_block_ids)
    return new_block_ids

关键点：分配可能失败（显存不足）。此时 KVCacheManager 返回 None，调度器需要决定是等待、还是抢占其他请求来腾出空间。

释放：请求完成后回收 ​

python
def free_blocks(self, request):
    for block_id in request.block_ids:
        block = self.block_pool.blocks[block_id]
        block.ref_cnt -= 1
        if block.ref_cnt == 0:
            if block.block_hash is not None:
                # 有哈希值 → 放入前缀缓存（不立即回收）
                self.block_pool.free_to_cache(block)
            else:
                # 无哈希值 → 直接回收
                self.block_pool.free_block(block)
    request.block_ids.clear()

注意带 block_hash 的块不会立即回收——它们进入前缀缓存，留待后续请求复用。这个机制是第 10 章"前缀缓存"的基础。

引用计数与 Copy-on-Write ​

当多个请求共享同一个块（如并行采样的 Prompt 部分）时，ref_cnt > 1。当某个请求需要修改共享块时（追加新 Token），需要先复制：

python
def cow_block(self, block_id, request):
    """Copy-on-Write：复制共享块"""
    block = self.block_pool.blocks[block_id]
    if block.ref_cnt == 1:
        return block_id  # 独占，无需复制

    # 分配新块
    new_block = self.block_pool.allocate()

    # 复制 KV Cache 数据（GPU 显存中的拷贝）
    copy_kv_cache(src=block_id, dst=new_block.block_id)

    # 更新引用计数
    block.ref_cnt -= 1
    new_block.ref_cnt = 1

    return new_block.block_id

GPU 显存中的数据拷贝是通过 CUDA memcpy 完成的，成本不低。但相比完全不共享（3 份 Prompt KV Cache），COW 只在写入时复制一个块（320 KB），远比复制整个序列（可能几 MB）便宜。

5.3 显存预算的计算 ​

系统启动时，vLLM 需要确定能分配多少个 KV Cache 块。这个计算在 Worker 初始化阶段完成：

具体来说，Worker 通过 determine_available_memory() 方法执行一次"试探性"的前向传播：

  1. 加载模型权重
  2. 构造一个最大批次大小的虚拟输入
  3. 执行一次前向传播（不保存 KV Cache）
  4. 测量 GPU 显存的峰值使用量
  5. 总显存 - 峰值使用量 × 安全系数 = 可用于 KV Cache 的显存

这种"先试后算"的方法比静态估算更准确——它考虑了 PyTorch 的内存碎片、CUDA 上下文占用、NCCL 通信缓冲区等难以预测的因素。

gpu_memory_utilization 参数 ​

用户可以通过 --gpu-memory-utilization（默认 0.9）控制 vLLM 使用多少比例的 GPU 显存。设为 0.9 意味着 vLLM 最多使用 90% 的显存，留 10% 给系统和其他进程。

可用显存 = 总显存 × gpu_memory_utilization - 非 KV Cache 开销
KV Cache 块数 = 可用显存 ÷ 单块大小

对于 80 GB A100，模型 26 GB，其他开销 4 GB：

  • 可用显存 = 80 × 0.9 - 26 - 4 = 42 GB
  • 块大小 = 320 KB
  • 块数 = 42 GB ÷ 320 KB ≈ 137,000 个块
  • 可容纳 Token 数 = 137,000 × 16 ≈ 220 万个 Token

这意味着理论上可以同时服务 1,100 个各 2,048 Token 的请求。当然，实际数量会因模型大小、序列长度分布和工作负载特征而异。

5.4 驱逐策略 ​

当所有块都被占用且需要为新请求分配块时，KVCacheManager 可以驱逐前缀缓存中的块。

驱逐遵循 LRU（最近最少使用） 策略：空闲链表头部的块是最久没有被任何请求引用的，优先被驱逐。

驱逐时需要注意**链式块（Chain Blocks）**的处理。前缀缓存中，块通过哈希链关联（第 10 章会详细讲）。驱逐一个块时，依赖它的子块也需要被驱逐，否则哈希链会断裂。

V1 的实现中有一个精巧的优化：同一时间戳释放的块中，链尾的块优先被驱逐。因为链尾块是最具体的（包含最长前缀的最后一段），被复用的概率最低；而链头块（包含公共前缀的前几段）被复用的概率最高，应该尽量保留。

5.5 显存是如何被精确到字节管理的 ​

让我们把本章的所有知识串联起来，看一个完整的场景。

假设系统有 100 个空闲块，收到两个请求：

请求 A: 输入 50 Token, 最终生成 30 Token
请求 B: 输入 80 Token, 最终生成 20 Token

整个过程中：

  • 最大同时占用：A(5块) + B(7块) = 12 块 = 5.9 MB
  • 如果用传统方案（假设 max_len=2048）：A(128块) + B(128块) = 256 块 = 81.9 MB
  • 节省了 93% 的显存

这就是 PagedAttention + 按需分配的威力。

5.6 本章小结 ​

KV Cache Manager 是 PagedAttention 设计理念的工程落地：

  • BlockPool 用 __slots__ 对象和双向链表管理十万级别的物理块，O(1) 分配/释放
  • KVCacheManager 提供高层接口（分配、释放、COW），屏蔽底层细节
  • 引用计数 + COW 实现了高效的块共享，在并行采样中节省大量显存
  • LRU 驱逐 在显存紧张时释放最不可能被复用的缓存块
  • 显存预算 通过试探性前向传播精确计算，配合 gpu_memory_utilization 控制总量

到此为止，我们已经完整理解了 vLLM 引擎核心的四大组件：EngineCore、调度器、PagedAttention、KV Cache 管理。它们构成了 vLLM 的基石。

下一章，我们将走出引擎核心，进入执行层——看看 Worker 和 Executor 如何将调度决策变成真正的 GPU 计算。

源码导航

  • BlockPool：vllm/v1/core/block_pool.py
  • KVCacheBlock：vllm/v1/core/kv_cache_utils.py
  • KVCacheManager：vllm/v1/core/kv_cache_manager.py
  • KV Cache 接口定义：vllm/v1/kv_cache_interface.py
  • Worker 显存探测：vllm/v1/worker/gpu_worker.py（determine_available_memory）

基于 VitePress 构建

Copyright © 2024-present 杨艺韬