课程 02 · 显存管理

显存里的那本「页表」

KV cache 为什么必须分页、引用计数在保护什么、前缀哈希怎么做到「算过的不再算」。这一课读完，block_manager.py 的 120 行就没有黑盒了。

约 50 分钟源码：engine/block_manager.py（120 行）前置：01

术语：Block / Block table / ref_count

1先看看不分页会怎样

decode 每生成一个 token，都要读完整条序列的 K/V。这块缓存必须在显存里常驻，而且随生成不断变长。最直接的做法是：给每个请求预留 max_model_len 的连续显存。

三个问题会立刻出现：

内部碎片：请求实际只生成 200 个 token，却占了 4096 的坑。上节课算过，4096 个 token 是 448 MiB——一条短请求就吃掉半张卡的 KV 预算。
外部碎片：请求结束释放后，显存里留下大小不一的空洞，下一个请求要的长度可能刚好塞不进去。
无法共享：两个请求用同一段系统提示词，KV 明明一模一样，却各自存了一份。

PagedAttention 的答案借自操作系统：把 KV 切成固定大小的块，用一张表把「逻辑块」映射到任意物理块[SOSP'23 §3]。

层	操作系统	nano-vllm
分配单位	页（4 KiB）	block（默认 256 token，28 MiB）
地址映射	页表	`seq.block_table`
共享	写时复制 / 共享页	`ref_count` + 前缀哈希
回收	引用归零回收	`ref_count == 0` 回 free 池

2三个数据结构

nanovllm/engine/block_manager.py · L8–33class Block:
    def __init__(self, block_id):
        self.block_id = block_id
        self.ref_count = 0        # 几条序列在用它
        self.hash = -1            # 满块才有哈希
        self.token_ids = []       # 存一份内容，用来防哈希碰撞

class BlockManager:
    def __init__(self, num_blocks, block_size):
        self.blocks = [Block(i) for i in range(num_blocks)]
        self.hash_to_block_id = dict()      # 前缀哈希 → 物理块
        self.free_block_ids = deque(range(num_blocks))
        self.used_block_ids = set()

三个结构各管一件事：

free_block_ids：空闲块池，popleft 分配、append 归还——FIFO 而不是随机，让刚释放的块晚点被复用，减少「刚释放就被覆盖」的概率。
used_block_ids：哪些块正在被引用（即使 ref_count 变成 0，也要先从 used 移除才算真空闲）。
hash_to_block_id：前缀缓存的索引，是「算过的不再算」的全部秘密。

关键洞察

分页解决的是碎片，哈希解决的是重复。两个问题经常被混为一谈——前者靠 block 切分，后者靠 hash_to_block_id + ref_count，缺一不可。

3一条序列怎么拿到块

请求进来时，调度器先问：这条序列的前缀，有多少块已经在缓存里了？ 这就是 can_allocate()[L58–73]：

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):          # 注意：-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

三个细节值得停下来看：

为什么是 num_blocks - 1？ 最后一块可能没填满。半满的块内容还会变，不能登记成前缀，也不能被别的序列复用。
为什么哈希要带 prefix？ 只看本块内容的哈希，会让「A,B」和「C,B」里的 B 撞成同一个键。把前一块的哈希混进来，等于给每个块一个「路径前缀」，链式定位。
为什么还要比对 token_ids？ 64 位哈希有碰撞概率。生产系统可能选择信任哈希，但教学实现用一次列表比较换来正确性——反正只在命中时才比较。

拿到 num_cached_blocks 后，allocate()[L75–92] 分两段做：命中的块只加引用，剩下的从空闲池里取新块。

def allocate(self, seq, num_cached_blocks):
    for i in range(num_cached_blocks):
        block_id = self.hash_to_block_id[h]
        if block_id in self.used_block_ids:
            self.blocks[block_id].ref_count += 1     # 共享：只加计数
        else:
            self.blocks[block_id].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

最后一行是整个 prefix cache 的记账核心：num_cached_tokens 告诉调度器和 prepare_prefill——这些 token 的 KV 已经在显存里了，只需要算后面的。

4引用计数在保护什么

假设两条请求共享同一段 800 token 的系统提示词（3 个满块）。第二个请求命中缓存后，只做了 ref_count += 1，没有复制任何数据。现在第一个请求结束了，调用 deallocate()：

def deallocate(self, seq):
    for block_id in reversed(seq.block_table):
        self.blocks[block_id].ref_count -= 1
        if self.blocks[block_id].ref_count == 0:
            self._deallocate_block(block_id)     # 真的没人用了，回 free 池
    seq.num_cached_tokens = 0
    seq.block_table.clear()

如果 ref_count 是 2，减到 1 就停手——第二个请求的 KV 毫发无损。这就是「共享前缀」不需要任何数据拷贝的原因。

易错点

deallocate 会把 num_cached_tokens 清零，但不会删除哈希表里的条目。块的哈希和被 _allocate_block 复用时会清除[L43–51]，所以「哈希指向一个已释放的块」是合法状态：下次命中时按「缓存还在但没人用」认领。

5哈希是何时写进去的

分配时并不写哈希——因为块还没算过，内容还不存在。真正登记发生在 postprocess 里每步调用的 hash_blocks()[L110–120]：

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]]
        h = self.compute_hash(seq.block(i), h)
        block.update(h, seq.block(i))
        self.hash_to_block_id[h] = block.block_id

它只处理「这一轮刚好被填满」的块：start 是已缓存的块数，end 是算完这轮之后的块数。两者相等说明没跨过块边界，直接返回。跨过了，就把新填满的块依次登记，哈希链从上一次的 h 继续。

这就是为什么前缀缓存对相同系统提示词特别有效：所有请求的第一块内容相同、链式哈希相同，第二块、第三块也跟着相同，一路命中到分歧点为止。

6block_size 为什么是 256

这是个纯粹的工程取舍，面试里常被追问：

块越小：内部碎片越少（最多浪费 block_size - 1 个 token 的空间），但块表变长、哈希表条目变多、block_table 这个二维张量变宽，注意力 kernel 的间接寻址开销上升。
块越大：kernel 更高效，但一条 300 token 的请求也要占 256 的整数倍，浪费更明显。

vLLM 默认 block_size=16，nano-vllm 用 256，因为它的注意力走 FlashAttention 的 block_table 参数，而 FlashAttention 的块粒度本来就是 256[config.py L17]。块大小要和底层 kernel 的粒度对齐，否则要么多一层重排，要么浪费。

和面试官聊这个

不要说「256 是随便选的」。正确说法：块大小要在碎片率、块表开销、kernel 对齐三者之间取平衡；nano-vllm 选 256 是为了直接喂给 FlashAttention 的 block table。

7练习 A · 手算一张块表

设 block_size = 4，空闲池里有块 [0,1,2,3,4]，当前无缓存。

序列 A 来了，长度 10（含 prompt 8 + 已生成 2）。请写出它的 block_table、需要的块数、以及哪些块会被 hash_blocks 登记。

展开：答案与推理

num_blocks = ceil(10/4) = 3，块表 [0,1,2]。第 0、1 块是满的（8 个 token），第 2 块只有 2 个 token。

can_allocate 会扫到第 1 块就停（range(num_blocks-1) = range(2)），两块都没命中，返回 0。分配 3 个新块。

这轮 prefill 结束后 num_cached_tokens = 10，end = 10//4 = 2，start = 0，所以登记块 0 和块 1——第 2 块不登记，因为它只有 2 个 token，还没满。

8练习 B · 闭卷测验

9练习 C · 代码寻宝

1. _allocate_block 的 L47–48 在删什么？为什么必须删？

如果这个物理块之前带着哈希（被释放过但哈希还在表里），而 hash_to_block_id 仍然指向它，就要先删掉这条映射。否则哈希表会指向一块马上要被新内容覆盖的显存——前缀缓存就变成了「返回错误数据」。

2. 为什么 free_block_ids 是 deque 且用 popleft，而不是 list.pop()？

FIFO 让块按释放顺序轮转复用，延长每个块的「冷却时间」。若总是复用最近释放的块，刚被别的序列释放的块可能立刻被覆盖，而它的哈希还没来得及被任何人命中。这是缓存系统里常见的抗抖动技巧。

3. deallocate 里为什么用 reversed(seq.block_table)？

释放顺序不影响正确性，但倒序有一个好处：最后一块（部分填充、没有哈希）最先被处理，物理上相邻的块倾向于晚一点进入空闲池，对后续大块分配更友好。属于实现习惯，答「不影响正确性」就算对。

10面试复述卡

「vLLM 的 PagedAttention 到底解决了什么？」

展开：30 秒版本

「它把 KV cache 从『每个请求预留一大段连续显存』改成『按固定大小的 block 分配，用块表做逻辑到物理的映射』。这样内部碎片最多浪费一个块，释放是整块回收、没有外部碎片；更关键的是两个请求可以共享同一批物理块，只要引用计数管住生命周期。在此基础上再做前缀哈希，相同前缀的 KV 只算一次。」

卡住了就问。 特别是「为什么 can_allocate 要返回 -1 而不是 0」这种细节——它决定了调度器会不会跳过这条序列，下一课刚好要用。