3.4 KV Cache 与 BlockManager

📑 目录

• 1. 为什么 KV 要分块
• 2. Block：最小管理单元
• 3. Sequence：一条请求的状态机
• 4. BlockManager：分配、追加与回收
• 5. 前缀缓存：用哈希复用历史
• 总结
• 自我检验清单
• 参考资料

---

1. 为什么 KV 要分块

Decode 阶段每生成一个 token，都要用到前面所有 token 的 K/V。如果给每条序列预留「最大长度」那么大的一块连续显存，大多数请求用不满，预留出来的部分就浪费了，这就是第 2 章讲过的内部碎片问题。

PagedAttention 的解法借鉴了操作系统的分页：把 KV Cache 切成固定大小的块（Block），序列需要多少就分多少块，块与块在物理上不必相邻。nano-vllm 的 block_manager.py 是这套机制的最小实现，共 120 行。它只管块的账本，不碰任何矩阵运算，所以归在调度层而不是 layers/。

💡 提示：nano-vllm 的块大小默认 kvcache_block_size = 256（每块装 256 个 token），比 vLLM 常见的 16 大得多。这个值不是随意选的：flash-attn 读分页 KV 时要求块大小是 256 的倍数（见 3.2），256 已经是能取的最小值。代价是最后一块没填满时浪费更多，前缀缓存的匹配粒度也更粗（见第 5 节）。

2. Block：最小管理单元

每个 Block 对应 KV Cache 里一段物理显存，但 Block 对象本身只是一条元数据，记录这块被几条序列引用、装了哪些 token：

class Block:
    def __init__(self, block_id):
        self.block_id = block_id       # 物理块编号
        self.ref_count = 0             # 引用计数：几条序列在共享这块
        self.hash = -1                 # 块内容的哈希（-1 表示还没填满、没登记）
        self.token_ids = []            # 块内 token（用于校验前缀是否真的相同）

    def update(self, hash, token_ids): # 块被填满、算好哈希时登记
        self.hash = hash
        self.token_ids = token_ids

    def reset(self):                   # 从空闲队列取出、分配给新序列时
        self.ref_count = 1
        self.hash = -1
        self.token_ids = []

三个字段的用途：

• 🔑 ref_count（引用计数）：多条序列共享同一个前缀块时，这块的 ref_count 大于 1，减到 0 才能回收。被共享的块都是已经填满的块，内容不会再变，所以 nano-vllm 不需要写时复制（copy-on-write）：每条序列新生成的 token 总是写进自己独占的最后一块。
• 🔑 hash：块内容的指纹。只有被填满的块才算哈希并参与前缀缓存；没填满的块 hash = -1，不能被别的序列复用，因为它还会继续追加 token。
• token_ids：存一份块内 token，哈希命中后再逐个比对，防止哈希碰撞导致复用错块。

3. Sequence：一条请求的状态机

sequence.py 里的 Sequence 跟踪单条请求从进来到结束的全部状态。它用一个枚举表示生命周期：

class SequenceStatus(Enum):
    WAITING = auto()     # 在等待队列里，还没开始算
    RUNNING = auto()     # 正在被调度执行
    FINISHED = auto()    # 已结束（命中 EOS 或达到 max_tokens）

Sequence 的构造函数如下：

class Sequence:
    block_size = 256                 # 类变量：所有序列共享同一块大小
    counter = count()                # 类变量：自增 ID 生成器

    def __init__(self, token_ids, sampling_params=SamplingParams()):
        self.seq_id = next(Sequence.counter)
        self.status = SequenceStatus.WAITING
        self.token_ids = copy(token_ids)   # prompt + 已生成的 token（拷贝一份，不改调用方的列表）
        self.last_token = token_ids[-1]    # 最后一个 token，decode 时作为模型输入
        self.num_tokens = len(self.token_ids)
        self.num_prompt_tokens = len(token_ids)
        self.num_cached_tokens = 0         # 已经算完、KV 写进 Cache 的 token 数
        self.num_scheduled_tokens = 0      # 本轮被调度计算的 token 数
        self.is_prefill = True             # 实例属性：是否还在 prefill 阶段
        self.block_table = []              # 逻辑块 → 物理 block_id 的映射
        # 采样超参直接摊平成实例属性（不另存 sampling_params 对象）
        self.temperature = sampling_params.temperature
        self.max_tokens = sampling_params.max_tokens
        self.ignore_eos = sampling_params.ignore_eos

几个关键设计：

• 🔑 block_table 相当于页表：一个 Python list，第 i 个元素是本序列第 i 个逻辑块对应的物理块号。有了它，分散在显存各处的物理块就能还原成一条逻辑上连续的序列。
• 🔑 is_prefill 是实例属性而非 property：调度器在第一次把它排进 decode batch 时置为 False（见 3.5 的 decode 分支）；被抢占时重置为 True，下次重新 prefill。
• num_scheduled_tokens 与 num_cached_tokens：支持 chunked prefill，一条长 prompt 可能分几步才算完。num_scheduled_tokens 记本轮算几个，num_cached_tokens 记累计已写进 Cache 几个；两者之和等于 num_tokens 时 prefill 才算做完。命中前缀缓存时，allocate 会直接把 num_cached_tokens 设成命中块数 × 256。
• append_token 把新 token 追加到 token_ids，同时更新 last_token 和 num_tokens。

常用的派生属性都做成了 property，比如：

@property
def num_blocks(self):                        # 需要几个块（向上取整）
    return (len(self) + self.block_size - 1) // self.block_size

@property
def last_block_num_tokens(self):             # 最后一块装了几个 token
    return len(self) - (self.num_blocks - 1) * self.block_size

📌 关键点：num_blocks 和 last_block_num_tokens 是 BlockManager 判断要不要开新块的依据。序列长度刚好跨过块边界时，就需要再分配一个物理块。

3.1 跨进程时只传必要字段

TP>1 时，主进程每一步都要把这批 Sequence pickle 后通过共享内存发给其他 rank（见 3.6）。一条序列的 token_ids 可能有几千个元素，每步全量发送很浪费。Sequence 自定义了 __getstate__，只打包 worker 需要的字段：

def __getstate__(self):
    # prefill 需要整段 token；decode 只需要最后一个 token
    last_state = self.last_token if not self.is_prefill else self.token_ids
    return (self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens,
            self.num_scheduled_tokens, self.block_table, last_state)

__setstate__ 按 last_state 的类型恢复：是 list 就还原 token_ids，否则 token_ids 置空、只保留 last_token。decode 阶段 worker 收到的每条序列只有几个整数和一个 block_table，这也是 prepare_decode 读 seq.last_token 而不读 seq.token_ids[-1] 的原因。采样参数没有被打包，因为采样只在 rank 0 做。

4. BlockManager：分配、追加与回收

BlockManager 持有全部物理块、一个空闲队列和一张哈希表，对外提供分配、追加、回收三组操作。初始化：

class BlockManager:
    def __init__(self, num_blocks, block_size):
        self.block_size = block_size
        self.blocks = [Block(i) for i in range(num_blocks)]
        self.hash_to_block_id = {}                        # 前缀缓存：哈希 → 块号
        self.free_block_ids = deque(range(num_blocks))   # 空闲块队列（先进先出）
        self.used_block_ids = set()

① 分配（prefill 前一次性分够块）

调度器先调 can_allocate(seq)：它沿哈希链逐块查前缀缓存，统计命中了几块，再判断空闲块够不够装下剩下的部分。放不下返回 -1，否则返回命中块数。确认放得下之后，调度器再调 allocate(seq, num_cached_blocks) 把物理块挂到 seq.block_table 上。两个函数的匹配细节放在第 5 节。

新块由 _allocate_block 从空闲队列队头取出：

def _allocate_block(self):
    block_id = self.free_block_ids.popleft()
    block = self.blocks[block_id]
    assert block.ref_count == 0
    # 这块以前登记过哈希，且表里还指向它：要被覆盖了，先把登记删掉
    if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id:
        del self.hash_to_block_id[block.hash]
    block.reset()                  # ref_count = 1，清空 hash 和 token_ids
    self.used_block_ids.add(block_id)
    return block_id

② 追加（decode 时每步可能新增一块）

def can_append(self, seq):
    # 只有「最后一个 token 正好落在新块的第一格」时才需要一个空闲块
    return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

def may_append(self, seq):
    if len(seq) % self.block_size == 1:
        seq.block_table.append(self._allocate_block())

decode 调度时，上一步采样出的 token 已经追加进序列，这一步要计算它的 K/V。len(seq) % block_size == 1 为真，说明这个 token 是新块的第一个 token，需要一个空闲块（布尔值 True 参与比较时等于 1）；否则当前块还有空位，0 个空闲块就够。may_append 在真正跨块时从空闲队列取块追加到 block_table。

③ 回收（序列结束或被抢占）

def deallocate(self, seq):
    for block_id in reversed(seq.block_table):   # 从最后一块往前放
        block = self.blocks[block_id]
        block.ref_count -= 1
        if block.ref_count == 0:                 # 没人用了才放回空闲队列队尾
            self._deallocate_block(block_id)
    seq.num_cached_tokens = 0
    seq.block_table.clear()

⚠️ 注意：回收时先减引用计数、减到 0 才放回空闲队列，被多条序列共享的前缀块不会因为其中一条结束就被回收。另外 _deallocate_block 只把块号移回空闲队列，不清空块的 hash 和 token_ids，哈希表里的登记也还在。这块真正被 _allocate_block 取走覆盖之前，后来的请求仍然可以命中它。

两个顺序上的细节让缓存保留得更久：

• 空闲队列先进先出，_allocate_block 从队头取、回收放到队尾，所以最近释放的块最晚被覆盖。
• deallocate 倒序释放，同一条序列的后段块先进队列、先被覆盖，开头的块排在后面。开头的块最可能是多个请求共享的前缀，这样它们留得最久。

5. 前缀缓存：用哈希复用历史

多条请求常有相同前缀（同一段 System Prompt、同一组 few-shot 例子）。让它们共享这段前缀的 KV，就能省掉重复的 prefill 计算。nano-vllm 用块级哈希实现：给每个填满的块算一个指纹，新请求按指纹去找已经算好的块。

5.1 哈希链

每个块的哈希都把前一块的哈希拌进来，形成一条哈希链：

@classmethod
def compute_hash(cls, token_ids, prefix=-1):
    # 把「前缀哈希 + 本块 token」一起喂给 xxhash
    h = xxhash.xxh64()
    if prefix != -1:
        h.update(prefix.to_bytes(8, "little"))
    h.update(np.array(token_ids).tobytes())
    return h.intdigest()

🔑 为什么要把前缀哈希拌进来：两个块即使内容相同，只要它们前面的历史不同，就不是同一段上下文，不能复用。把前缀哈希链进来，等于让每个块的哈希代表「从序列开头到这块为止的完整前缀」。这和 Merkle 链的思想一致。

5.2 登记：块填满后调用 hash_blocks

块要等填满才算哈希。登记发生在调度器的 postprocess 里（见 3.5）：每步模型跑完，在更新 num_cached_tokens 之前，对这条序列调用一次 hash_blocks：

def hash_blocks(self, seq):
    start = seq.num_cached_tokens // self.block_size                               # 本轮之前已填满的块数
    end = (seq.num_cached_tokens + seq.num_scheduled_tokens) // self.block_size    # 本轮之后已填满的块数
    if start == end: return                                                        # 本轮没有新填满的块
    h = self.blocks[seq.block_table[start - 1]].hash if start > 0 else -1          # 接上链的前一环
    for i in range(start, end):
        block = self.blocks[seq.block_table[i]]
        token_ids = seq.block(i)
        h = self.compute_hash(token_ids, h)
        block.update(h, token_ids)
        self.hash_to_block_id[h] = block.block_id                                  # 直接覆盖旧登记

[start, end) 就是本轮新填满的块。prefill 一步可能填满好几块；decode 每步只算 1 个 token，只有恰好写满一块的那一步才登记 1 块。chunked prefill 的一段停在块中间时，那块不算哈希，等下一段把它填满再登记。链的前一环直接读 block_table[start - 1] 上存好的哈希，不用从头重算。

hash_to_block_id[h] 是直接覆盖的。两块内容相同时（比如两条相同前缀的请求同时 prefill，都没命中），表里只记最后登记的那块。这就是 _allocate_block 删除登记前要先确认 hash_to_block_id.get(block.hash) == block_id 的原因：表里指向的可能已经是别的块了。

5.3 匹配：can_allocate 与 allocate

新请求进来时，调度器先调 can_allocate 沿哈希链逐块查表：

def can_allocate(self, seq):
    h = -1
    num_cached_blocks = 0
    num_new_blocks = seq.num_blocks              # 最坏情况：每块都要占一个空闲块
    for i in range(seq.num_blocks - 1):          # 注意：最后一块不查
        token_ids = seq.block(i)
        h = self.compute_hash(token_ids, h)
        block_id = self.hash_to_block_id.get(h, -1)
        if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
            break                                # 第一次没命中就停
        num_cached_blocks += 1
        if block_id in self.used_block_ids:      # 命中在用的块：不占空闲块
            num_new_blocks -= 1
    if len(self.free_block_ids) < num_new_blocks:
        return -1
    return num_cached_blocks

确认放得下后，allocate 按同样的哈希链把命中的块挂上去，剩下的块从空闲队列新取：

def allocate(self, seq, num_cached_blocks):
    h = -1
    for i in range(num_cached_blocks):
        h = self.compute_hash(seq.block(i), h)
        block_id = self.hash_to_block_id[h]
        block = self.blocks[block_id]
        if block_id in self.used_block_ids:      # 别的序列正在用：共享
            block.ref_count += 1
        else:                                    # 已释放但还没被覆盖：从空闲队列里摘回来
            block.ref_count = 1
            self.free_block_ids.remove(block_id)
            self.used_block_ids.add(block_id)
        seq.block_table.append(block_id)
    for i in range(num_cached_blocks, seq.num_blocks):
        seq.block_table.append(self._allocate_block())
    seq.num_cached_tokens = num_cached_blocks * self.block_size

整个匹配流程如下，菱形是两处判断，蓝色框是命中后的处理，红色框是从空闲队列新取块：

这段代码有几处细节：

• 🔑 最后一块永远不查：循环只到 num_blocks - 1。即使整条 prompt 都在缓存里，最后一块也要重新算。prefill 需要最后一个 token 的 logits 来采样下一个 token，如果所有 token 都命中，这一轮就没有 token 进模型，也就拿不到 logits。
• 前缀必须连续命中：第一次未命中就 break。中间断了一块，后面的块即使还在表里也用不上，因为它们依赖的前文 KV 已经不在缓存里了。
• 命中块分两种：在用的块直接 ref_count += 1，不消耗空闲块；已释放、还没被覆盖的块要从空闲队列里摘出来重新启用，仍然占一个空闲名额。所以 can_allocate 只在命中在用块时才把 num_new_blocks 减 1。被抢占的序列重新 prefill 时，也可能这样命中自己刚释放的旧块。
• token_ids 二次校验：哈希相等后再逐个比对 token，防止 xxhash 碰撞导致复用错块。

下图用 block_size = 4 演示一条 12 token 的请求：L0 命中另一条请求正在用的 #1，L1 命中空闲队列里的 #3，最后一块不查表，直接从空闲队列队头取 #6（暖色是本请求的块，斜线是其他请求占用的块）：

📌 关键点：前缀缓存的粒度是一整块，加上最后一块不查的规则，block_size = 256 时复用量会被截断。两条请求的公共前缀有 600 个 token，只能复用前 512 个（2 块），其余 88 个要重算；如果一条 512 token 的 prompt 和缓存完全相同，也只能复用第 1 块，第 2 块要重算。块越小，匹配越细，但块表更长、哈希次数更多。

📝 总结

• nano-vllm 把 KV Cache 切成固定大小的物理块（默认 256 token/块，flash-attn 分页 KV 允许的最小值），用 BlockManager 管理块的账本，不碰矩阵运算。
• Block 的 ref_count/hash/token_ids 分别支撑共享回收、前缀缓存与碰撞校验。被共享的块都已填满、不再改动，所以不需要写时复制。
• Sequence 是请求状态机：block_table 充当页表，is_prefill 在首次进入 decode 时置 False、被抢占时重置；__getstate__ 在 decode 阶段只给 worker 发 last_token。
• can_append 用 len(seq) % block_size == 1 判断是否跨块；deallocate 先减引用计数、减到 0 才放回空闲队列队尾，且保留哈希，块被覆盖前仍可命中。
• 前缀缓存靠拌入前缀的块级 xxhash 形成哈希链：hash_blocks 在 postprocess 里登记本轮填满的块，can_allocate 连续匹配、遇到未命中即停，最后一块永远重算。

🎯 自我检验清单

• 能解释 KV 分块相比「预留最大长度连续显存」解决了什么碎片问题
• 能说出 Block 的 ref_count、hash、token_ids 各自的作用，以及 nano-vllm 为什么不需要写时复制
• 能解释 Sequence.block_table 为什么相当于操作系统的页表，以及 __getstate__ 在 decode 阶段省掉了什么
• 能看懂 can_append 里 len(seq) % block_size == 1 为什么能判断「是否需要新块」
• 能说清 deallocate 为什么要先减引用计数再回收，以及释放后的块为什么还能被命中
• 能解释前缀缓存为什么要把前缀哈希链入当前块，以及 hash_blocks 在什么时候登记哪些块
• 能说明 can_allocate 为什么不查最后一块，以及命中在用块和命中空闲块对空闲块需求的影响有何不同
• 能算出 block_size = 256 时，一段 600 token 的公共前缀实际能复用多少 token

📚 参考资料

• nano-vllm GitHub 仓库
• Efficient Memory Management for Large Language Model Serving with PagedAttention（Kwon et al., 2023）
• vLLM 官方文档 - Automatic Prefix Caching
• xxHash 官方仓库