nano-vllm 源码解析(三):BlockManager 与 PagedAttention

nano-vllm 源码解析（三）：BlockManager 与 PagedAttention

如果说调度器是 “大脑”，那么 BlockManager 就是 “海马体”（记忆管理中枢）。

这一章解决的是 LLM 推理中最昂贵的资源——显存的管理问题。这是 nano-vllm 能比朴素实现快几倍甚至几十倍的核心黑科技，也是 vLLM 项目成名的绝技：PagedAttention。

1. 为什么需要“页式”管理？

在没有 BlockManager 之前，显存管理有两个巨大的痛点：

❌ 痛点 1：预分配浪费

• • 传统做法：为了防止显存不够，系统会给每个请求预分配 最大长度（比如 4096）的显存。
• • 现实：用户可能只说了“你好”两个字。
• • 后果：就像为了住一晚，强行包下整个总统套房。99% 的空间被锁死，别人进不来。

❌ 痛点 2：碎片化

• • 传统做法：请求结束了，释放显存。
• • 现实：长长短短的请求来了又走，显存里留下了无数个“小空洞”。
• • 后果：虽然总空闲显存有 1GB，但都是碎的，连不成片，存不下一个大的 KV Cache。

✅ 解决方案：分页 (Paging)

• • 思路：模仿操作系统的虚拟内存。
• • 做法：把 KV Cache 切成 固定大小的小块 (Block)（比如一块存 16 个 Token）。
• • 好处：
1. 1. 按需分配：写满一页纸，再拿下一页。
2. 2. 零碎片：物理上不连续的块，逻辑上可以是连续的。

2. 显存的最小单位：Block

Block 是一个显存格子的“身份证”。

📝 代码显微镜：Block 类

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

class Block:
    def __init__(self, block_id):
        self.block_id = block_id  # 物理门牌号 (例如: 1024号柜子)

        # 🌟【核心机制 1：引用计数】
        # 有几个人正在用这个柜子？
        # 0 = 空闲
        # 1 = 独占
        # >1 = 共享 (比如多个人都在问同样的问题前缀)
        self.ref_count = 0

        # 🌟【核心机制 2：指纹】
        # 柜子里存的数据的 Hash 值。用于 Prefix Cache 查重。
        # -1 表示“还没写满”或者“脏了”，不能作为指纹。
        self.hash = -1

        # 这里的 token_ids 只是为了校验防撞车，不参与模型计算
        self.token_ids = []   

    def update(self, hash, token_ids):
        # 只有写满一个块时，才会调用这个，打上指纹
        self.hash = hash
        self.token_ids = token_ids

    def reset(self):
        # 归还柜子时，清空信息
        self.ref_count = 1  # 刚分配出去时，肯定有1人在用
        self.hash = -1
        self.token_ids = []

3. 核心黑科技：链式哈希 (Chain Hashing)

怎么判断两个请求能不能共享显存？看它们的 内容（Token IDs） 是否一样。
但直接比对 Token 列表太慢了，我们用 哈希（Hash）。

⚠️ 为什么简单哈希不行？（上下文陷阱）

Attention 是 因果 (Causal) 的。一个词的意思，取决于它前面所有的词。

• • 情况 A：[我, 爱] -> [苹果]
• • 情况 B：[我, 恨] -> [苹果]

虽然最后一个块里装的都是 [苹果]，但因为前缀不同，算出来的 KV Cache 是完全不同的！绝对不能复用！

🔗 解决方案：区块链式哈希

我们在计算当前块的哈希时，把 上一个块的哈希 也加进去算。

1
2
3
4
5
6
7
8
9
10
11
12
13
14

@classmethod
def compute_hash(cls, token_ids, prefix=-1):
    # 初始化一个哈希计算器 (这里用的是 xxhash 算法，因为它的计算速度极快)
    h = xxhash.xxh64()

    # 🌟 关键：把前缀的指纹也喂进去
    if prefix != -1:
        h.update(prefix.to_bytes(8, "little"))

    # 再喂当前的内容
    h.update(np.array(token_ids).tobytes())

    return h.intdigest()

这样，情况 A 的“苹果”和情况 B 的“苹果”，哈希值就会截然不同。

4. 管理员：BlockManager

它是整个仓库的调度员。它手里维护着四本账簿，每一本都有特定的用途。

📝 代码显微镜：BlockManager 初始化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28

class BlockManager:
    def __init__(self, num_blocks: int, block_size: int):
        self.block_size = block_size

        # 1.【物理账本】：Block 对象池
        # 真实的物理柜子（元数据）。
        # 初始化时就建好所有的 Block 对象，避免运行时反复创建销毁对象。
        self.blocks: list[Block] = [Block(i) for i in range(num_blocks)]

        # 2.【索引账本】：Hash -> Block ID (Prefix Cache 核心)
        # 这是一个查找表。
        # 作用：新请求进来，算一下哈希，查查这里有没有？有就是 Cache Hit。
        # 例子：{982734: 5} 表示指纹为 982734 的内容存在 5 号柜子。
        self.hash_to_block_id: dict[int, int] = dict()

        # 3.【空闲账本】：可用柜子的钥匙堆
        # 数据结构：双端队列 (deque)。
        # 作用：要分配新块时，从这里 pop 一个；回收时，append 回去。
        # 初始状态：所有的块 ID (0 到 num_blocks-1) 都在这里。
        self.free_block_ids: deque[int] = deque(range(num_blocks))

        # 4.【占用账本】：正在使用的柜子名单
        # 数据结构：集合 (set)。
        # 作用：快速判断（O(1)复杂度）某个 ID 是不是正在被使用。
        # 关键逻辑：它和 free_block_ids 是【互斥】的。
        # 一个 ID 要么在 free 里，要么在 used 里，绝不可能同时存在。
        self.used_block_ids: set[int] = set()

💡 为什么要多维护一个 used_block_ids？

你可能会问：“如果 free_block_ids 里没了，不就是 used 吗？为什么还要专门存一个 set？”

这在工程上有两个重要意义：

1. 1. O(1) 极速查询：

• • 在 allocate（分配）逻辑中，当我们通过哈希命中了一个 block_id 时，我们必须再确认一眼：这个 block_id 真的是在使用中吗？
• • 如果是 list 或 deque，查询需要遍历，速度慢。
• • 用 set 查询 if block_id in self.used_block_ids 只需要  时间，极其高效。

1. 2. 状态校验：

• • 它构成了系统的安全围栏。在分配和回收时，我们可以断言（Assert）：
• • 分配时：确保 ID 不在used 里。
• • 回收时：确保 ID 一定在used 里。
• • 这能防止出现“双重释放”或“野指针”等严重的内存 Bug。

互斥关系图解：

5. 核心逻辑一：Prefill 阶段分配 (Allocate)

当新请求进来做 Prefill 时，首先要检查显存够不够，然后尝试复用已有的块（Prefix Cache）。

📝 代码显微镜：资源检查与分配

首先是准入检查。在分配之前，必须先看一眼仓库里的库存够不够。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65

def can_allocate(self, seq):
    # 【库存检查】
    # 问：我手里空闲的柜子钥匙数量 (len(free_block_ids))
    # 是否大于等于 这个请求总共需要的柜子数 (seq.num_blocks)？
    # 只有库存充足，才允许进入下一步的 allocate。
    return len(self.free_block_ids) >= seq.num_blocks

def allocate(self, seq):
    # 【安全断言】确保这个请求是新的，之前没分配过显存。
    assert not seq.block_table

    h = -1          # 初始前缀哈希 (第一个块没有前缀，所以是 -1)
    cache_miss = False  # 标记：是否发生过未命中。一旦断过一次，后面全都不认。

    # 【逐块遍历】
    # 比如这个请求需要 3 个块，我们一个一个处理
    for i in range(seq.num_blocks):
        # 取出当前块的 token 数据 (例如 [1, 2, ... 16])
        token_ids = seq.block(i)

        # 1. 计算哈希 (只有满块才配拥有指纹)
        # 如果是最后一个块且没填满，它的哈希强制为 -1 (Open Block)
        h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1

        # 2. 查字典 (Look up)
        # 看看这个指纹 h 以前有没有记录过？
        # 如果没记录，get 返回 -1。
        block_id = self.hash_to_block_id.get(h, -1)

        # 3. 校验：防哈希碰撞 (Double Check)
        # 条件 A: block_id == -1 (字典里没查到)
        # 条件 B: 查到了，但在物理块里存的内容(token_ids)跟现在的对不上 (哈希碰撞)
        if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
            cache_miss = True # 标记：未命中！

        # 4. 决策分支：根据是否命中来决定怎么分配
        if cache_miss:
            # 【分支 A：未命中】
            # 必须拿新柜子。由于之前做了 can_allocate 检查，这里一定有空闲块。
            block_id = self.free_block_ids[0]
            # 调用底层分配函数 (见下文 _allocate_block)
            block = self._allocate_block(block_id)
        else:
            # 【分支 B：命中】
            # 复用旧柜子。不仅省了显存，还省了计算。
            seq.num_cached_tokens += self.block_size  # 记账：省了多少算力

            # 这里有个细节：查到的块可能是在用状态(used)也可能是刚释放但还没被覆盖(free)
            if block_id in self.used_block_ids:
                # 块正在被别人用：直接共享，引用计数 +1
                block = self.blocks[block_id]
                block.ref_count += 1
            else:
                # 块虽然有数据但在 free 列表里（可能是刚被回收）：重新捞回来用
                block = self._allocate_block(block_id)

        # 5. 更新元数据
        # 只有满块(h != -1)才需要登记到哈希字典里
        if h != -1:
            block.update(h, token_ids)
            self.hash_to_block_id[h] = block_id

        # 将最终确定的 block_id 加入请求的页表
        seq.block_table.append(block_id)

🗺️ 逻辑流程图 (Decision Flow)

这张图展示了对每一个 Block 的处理决策链：