Nano-VLLM全代码解析笔记(2)-block_manager

当前笔记顺序

Engine(当前block_manager.py)->Layers

Engine的运行架构(总结时或者弄不清一些调用顺序时看)：

block_manager.py

作用解析：管理kv_cache Block，实现Block的分配、释放、复用，这里需要理解一个句子能利用到的kv_cache Block，需要是相同链计算出来的，两个Block即使存的token一致，也不一定能复用，这是本文件的核心逻辑，比如下图中的Block（苹果），它来自对不同前缀的计算。

python 复制代码

from collections import deque
#给block生成唯一的hash，用于快速查找和比对
import xxhash
import numpy as np

from nanovllm.engine.sequence import Sequence


class Block:

    def __init__(self, block_id):
        self.block_id = block_id
        #记录该块被多少个sequence使用，只有运行态的sequence才会贡献ref_count
        self.ref_count = 0
        self.hash = -1
        self.token_ids = []

    def update(self, hash: int, token_ids: list[int]):
        self.hash = hash
        self.token_ids = token_ids

    def reset(self):
        #reset设为1而不是0的原因见问题1
        self.ref_count = 1
        self.hash = -1
        self.token_ids = []


class BlockManager:

    def __init__(self, num_blocks: int, block_size: int):
        self.block_size = block_size
        #按递增初始化一个block资源表
        self.blocks: list[Block] = [Block(i) for i in range(num_blocks)]
        #用于实现「前缀缓存」（相同 token 序列复用 block）
        self.hash_to_block_id: dict[int, int] = dict()
        #以队列形式记录空闲的block有哪些
        self.free_block_ids: deque[int] = deque(range(num_blocks))
        #用过的block必不能重复
        self.used_block_ids: set[int] = set()

    @classmethod
    def compute_hash(cls, token_ids: list[int], prefix: int = -1):
        #h不是数字,是一个有状态的对象，初始内部状态完全一样
        h = xxhash.xxh64()
        #如果不是第一个块,就把前一个块的哈希值(64位整数)以小端序转成8字节,塞进当前哈希器中
        #这是构造链式哈希(rolling hash),使得每个块的哈希不仅依赖自己的 token,还依赖整个前缀链。这样做可以避免内容一样位置不同的块使用相同HASH
        if prefix != -1:
            h.update(prefix.to_bytes(8, "little"))
        #此处看can_allocate的逻辑，是把一个块的token进行转换，而不是输入的全部token
        h.update(np.array(token_ids).tobytes())
        #返回int类型的64位摘要
        return h.intdigest()

    #从空闲池取一个 Block 并重置
    def _allocate_block(self) -> int:
        block_id = self.free_block_ids.popleft()
        block = self.blocks[block_id]
        assert block.ref_count == 0
        #此处逻辑见问题2
        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()
        self.used_block_ids.add(block_id)
        return block_id

    #把一个 Block 放回空闲池，只设置引用数为0，不清空token和hash以便再次利用
    def _deallocate_block(self, block_id: int):
        #当释放block时，必须确保没有sequence在使用它
        assert self.blocks[block_id].ref_count == 0
        self.used_block_ids.remove(block_id)
        self.free_block_ids.append(block_id)
    
    #Prefill-问:"这个序列我能缓存多少 Block?"。返回已经缓存的kv_cache block数量，这些块里的token不用重新算，也不用分配新的block，直接让seq.block_table引用即可
    def can_allocate(self, seq: Sequence) -> int:
        #第一个块没有前缀，因此h=-1
        h = -1
        num_cached_blocks = 0
        #num_blocks是实时计算的
        num_new_blocks = seq.num_blocks
        #最后一个 Block 永远是"正在填"的状态（即使空的也会预分配），因此不会缓存命中，需要跳过
        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)
            #必须缓存命中且内容一致，如果不一致就不需要继续，因为kv_cache的利用必须前面的所有token一致，见transformer原理
            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
        #如果需要分配的新block数不能够满足就return
        if len(self.free_block_ids) < num_new_blocks:
            return -1
        return num_cached_blocks

    #Prefill-把上面函数返回的N个缓存的Block拿来,其余新分配
    def allocate(self, seq: Sequence, num_cached_blocks: int):
        assert not seq.block_table
        h = -1
        for i in range(num_cached_blocks):
            token_ids = seq.block(i)
            h = self.compute_hash(token_ids, 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

    #还回序列的所有token，此函数在序列到完成态再执行
    def deallocate(self, seq: Sequence):
        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()

    #Decode-此处是问再decode一个token时,要不要新的Block?
    def can_append(self, seq: Sequence) -> bool:
        #此处逻辑有点难，由于token是一个个生成的，我们需要新的Block时，肯定是当我们token数为n *block_size + 1的时候
        return len(self.free_block_ids) >= (len(seq) % self.block_size == 1)

    #Decode-答:"需要就分配,不需要就不动"
    def may_append(self, seq: Sequence):
        if len(seq) % self.block_size == 1:
            seq.block_table.append(self._allocate_block())

    #把Block的hash写进表里，已经缓存的肯定写过了（毕竟都是靠hash查的缓存），这里是把新分配的Block写入，发生在seq调度后
    def hash_blocks(self, seq: Sequence):
        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

一些问题

python 复制代码

1.reset为什么把self.ref_count = 1而不是0。
reset 仅在「分配 block」时调用（见 _allocate_block 方法），此时该 block 已经被某个 sequence 占用，引用计数初始化为 1 是合理的。

2.if block.hash != -1 and self.hash_to_block_id.get(block.hash) == block_id:
            del self.hash_to_block_id[block.hash]，此处逻辑是否重复？
简而言之，"我要重置自己啦,但如果全局哈希表还指向我,就把那条记录删掉;如果指向别人,就千万别动,免得把别人的缓存给废了。"

此处背景是Block不会被销毁,而是回收复用，即使引用计数归0，也是保留缓存了的token_ids和hash放入
free_block_ids空闲池。如果有新内容要放进这个块，需要清空self.hash_to_block_id[block.hash]，避免下次有人用相同token_ids查询缓存时被误判缓存命中。这是左边的逻辑
右边的逻辑是，为了处理多个block有相同hash的情况，比如：
Block 5空闲中,内容 ["我爱学习"],hash=12345
hash_to_block_id = {12345: 5}

后来 Block 5 被释放,Block 7 也释放了同样的内容,但...

实际上:Block 7 先被分配、并被某个活跃序列使用:
hash_to_block_id = {12345: 7} ← 7 覆盖了 5(因为 7 现在是持有者)

→ 现在轮到分配 block 5 给另一个新序列
→ 检查 hash_to_block_id[12345] == 5 ? ✗ (现在是 7,不是 5)
→ 不删除! ← 关键!因为如果删了,block 7 的缓存就丢了


3.@classmethod和@staticmethod有什么区别。
特性	@classmethod	@staticmethod
第一个参数	必须是类本身（通常命名为 cls（自动传参，无需调用时传入））	无强制参数（和普通函数一样）
访问类 / 实例属性	可通过 cls 访问类属性、创建实例	无法直接访问类 / 实例属性（需显式传参）
用途	类级别的操作（如工厂方法、通用计算）	工具函数（逻辑属于类，但无类 / 实例依赖）


4.h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1是什么意思，compute_hash不是有cls参数吗？为什么他调用的时候没填写。
见问题3

5.block_id = self.hash_to_block_id.get(h, -1)，字典类型他为什么传入了(h,-1)？不应该是单参数吗
dict.get(key, default) 是 Python 字典的内置方法，不是单参数：
第一个参数 key：要查找的键（这里是哈希值 h）；
第二个参数 default：可选，若键不存在时返回的默认值（这里是 - 1）；
代码中 block_id = self.hash_to_block_id.get(h, -1) 的含义：
如果哈希 h 存在于字典中，返回对应的 block_id；
如果不存在，返回 - 1（表示未命中缓存）；
对比 dict[key]：get 不会因键不存在抛异常，更安全。


6.关于「不满的 block 是否会缓存命中」
结论：不可能。代码里有严格的逻辑限制，不满的 block 不会被缓存，自然也不会命中。


7.can_allocate vs allocate (Prefill 阶段)同（can_append vs may_append）
这是经典的 "两阶段提交(two-phase)":

第一阶段:只读不写(can_allocate)、无副作用,可以安全地试探
第二阶段:拿到第一阶段的结果再真正动手(allocate)
为什么要分两步?因为显存可能不够!调度器要先问"能不能",再决定"要不要做"。如果 can_allocate 返回 -1,就不该调用 allocate。

8.整体顺序

can_allocate -> 算出已有的缓存Block -> allocate -> 把block（缓存了的直接写，没缓存的分配再写）写入seq.block_table -> can_append -> 显存够不够 -> may_append -> 可能调用_allocate_block -> 生成结束调用deallocate -> hash_blocks

📚本系列文章(待写完修正)

$1$ Nano-VLLM全代码解析笔记(1)-sequence

$2$ Nano-VLLM全代码解析笔记(2)-block_manager

$3$ Nano-VLLM全代码解析笔记(3)-llm_engine和scheduler

$4$ Nano-VLLM全代码解析笔记(4)-model_runner

$5$ Nano-VLLM全代码解析笔记(5)-laynorm和attention

$6$ Nano-VLLM全代码解析笔记(6)-embed_head和linear

$7$ Nano-VLLM全代码解析笔记(7)-rotary_embedding

$8$ Nano-VLLM全代码解析笔记(8)-qwen3与qwen3_moe

🔗上一篇： $1$ Nano-VLLM全代码解析笔记(1)-sequence

🔗下一篇： $3$ Nano-VLLM全代码解析笔记(3)-llm_engine和scheduler