Level 22 | vLLM PagedAttention | vLLM 分页注意力

Mission 1 概念练习

Prefill：先向上取整，再做一次完整容量检查

🎯 先猜一猜

prompt_len=5、block_size=4，空闲池只有 1 个块。

allocate_for_prefill 应怎样处理？

需要 2 块，分配前直接抛 RuntimeError("OOM") 只分 1 块并忽略最后一个 token 先拿走 1 块再返回成功

先猜一个答案，再对照下面的讲解。可以随时修改选择，猜错会得到解释。

先补的知识

physical_kv_cache 已在当前练习的 __init__ 中给出，shape 是 [num_blocks, block_size, head_dim]；不要重复实现参考解析里的初始化 TODO。
req.seq_len 是 prompt token 数，req.block_table 是 List[int]，按逻辑顺序记录物理 block_id。
free_blocks 是可用物理块编号列表；pop(0) 会返回队首编号并从列表移除。
最后一块即使只使用一个 token，也必须占用完整物理块，因此需要整数向上取整。

图解原理

把物理块当成固定容量的箱子。Prompt 要一次性装完：先算需要几个箱子，再确认库存足够，最后才逐个拿走。先检查完整容量能避免分到一半才 OOM，留下 block_table 和 free_blocks 不一致的半更新状态。

需求req.seq_len 个 token 容量每块 block_size 个位置取整(seq_len + size - 1) // size 分配free_blocks → req.block_table

长度 6、block_size 4

token 0..3：第一块token 4..5：第二块

6 // 4 只得到 1，会漏掉尾部两个 token；向上取整得到 2。

块表只记录地址

逻辑块 0 → 物理 ID 7逻辑块 1 → 物理 ID 2

物理 ID 可以不连续，列表顺序负责表达逻辑顺序。

needed_blocksPython int；必须覆盖全部 req.seq_lenfree_blocksList[int]；每次 pop 后长度减一req.block_tableList[int]；按逻辑块顺序 append不足分支抛 RuntimeError("OOM")可见用例长度 6、块大小 4 → 2 块，10 块池剩 8 块

语法热身：给订单分配固定容量的货箱

item_count = 11
box_capacity = 5
boxes_needed = (item_count + box_capacity - 1) // box_capacity
available = [20, 7, 31, 9]
assigned = []

if len(available) < boxes_needed:
    raise RuntimeError('FULL')
for _ in range(boxes_needed):
    assigned.append(available.pop(0))

从语法例子迁移到当前练习

item_count：对应 req.seq_len
box_capacity：对应 self.block_size
boxes_needed：对应 TODO 1 的 needed_blocks
available：对应 self.free_blocks
assigned：对应 req.block_table；异常文本要按 Notebook 写成 OOM

巩固一下

为什么应在 for 循环 pop 之前检查 len(free_blocks)？

为了让 block_id 变成浮点数避免失败时已经拿走部分块，造成半更新状态为了把 seq_len 清零

学完这一段，试着做

用一个小动作确认自己理解了；最后再进入官方题目。

完成当前练习 TODO 1–2：整数向上取整、完整容量检查、循环 pop(0) 并 append。运行 Prefill 与 OOM 用例：长度 6 应分 2 块；只有 1 块却要容纳 5 个 token 时必须抛含 OOM 的 RuntimeError。若块数少一，检查是否用了向下取整；若 OOM 后 free_blocks 已减少，检查是否把容量检查写进了分配循环。

Mission 2 概念练习

Decode：先看已经加一的新长度，只在新块第一个位置补块

🎯 先猜一猜

block_size=4，请求当前长度 8；allocate_for_decode 开头先把长度加到 9。

这一步是否需要新 Block？

需要，因为 9 % 4 == 1 不需要，因为 8 % 4 == 0 需要一次分配 4 个 Block

先猜一个答案，再对照下面的讲解。可以随时修改选择，猜错会得到解释。

先补的知识

函数开头已经执行 req.seq_len += 1；TODO 看到的是新 token 加入后的长度。
block_size=4 时，长度 1–4 使用第一块，5–8 使用第二块，9–12 使用第三块。
当前 Notebook 的判断契约是 (req.seq_len % self.block_size) == 1，表示刚进入新块的第一个槽位。
Decode 每次只新增一个 token，因此跨界时最多追加一个物理块；空闲列表为空要抛 RuntimeError("OOM")。

图解原理

不要孤立背取模公式，先画位置：长度从 4 加到 5 时，新 token 落在第二块第一个位置，所以补块；从 5 加到 6 时仍在第二块，不补；从 8 加到 9 时再补第三块。

1req.seq_len += 1Notebook 已经把本轮新 token 计入长度

2seq_len % block_size检查新 token 在当前逻辑块中的位置

3余数 == 1表示刚进入新块的第一个位置

4pop + append只有跨界时才申请一个新物理块

加一后的长度	余数（block_size=4）	动作
7	3	继续使用现有块
8	0	现有块刚好填满，不提前分配
9	1	进入新块，追加一个 block_id

常见偏一位错误是判断余数等于 0。这里长度已经先加一；余数 0 表示新 token 填满旧块，余数 1 才表示它进入下一块。

语法热身：给连续座位按排扩容

seats_per_row = 4
current_people = 8
current_people += 1

needs_row = (current_people % seats_per_row) == 1
free_rows = [12, 13]
assigned_rows = [4, 9]
if needs_row:
    if not free_rows:
        raise RuntimeError('FULL')
    assigned_rows.append(free_rows.pop(0))

从语法例子迁移到 TODO 3

current_people += 1：对应已经写好的 req.seq_len += 1
seats_per_row：对应 self.block_size
needs_row：对应 is_new_block_needed
free_rows：对应 self.free_blocks；为空时异常文本改为 OOM
assigned_rows：对应 req.block_table，只在跨界时 append 一个 ID

巩固一下

长度 6 连续调用三次 decode 后，为什么第 9 个 token 才让块表从 2 块变成 3 块？

因为 block_id 必须是 9 因为每三次调用固定扩容长度 7、8 仍在第二块，长度 9 才进入第三块

学完这一段，试着做

用一个小动作确认自己理解了；最后再进入官方题目。

完成 is_new_block_needed 和跨界分配分支；不要重复 req.seq_len += 1，也不要每次 decode 都分块。按可见测试从长度 6 手算到 9：第一次后仍 2 块，第三次后变 3 块，空闲池从 8 变 7。若长度 8 时提前扩容，检查是否判断余数 0；若长度 9 不扩容，检查是否在加一前计算了条件。

Mission 3 概念练习

按块表恢复逻辑序列：先索引每块，再沿 token 维拼接并截尾

🎯 先猜一猜

block_size=4、head_dim=8，块表有 2 个 ID，请求 seq_len=5。

cat 前、cat 后、截断后的 shape 依次是什么？

每块 [4,8]；拼接 [8,8]；返回 [5,8] 每块 [4,8]；拼接 [2,4,8]；返回 [5,8] 每块 [8]；拼接 [16]；返回 [5]

先猜一个答案，再对照下面的讲解。可以随时修改选择，猜错会得到解释。

先补的知识

physical_kv_cache 的 shape 是 [num_blocks, block_size, head_dim]；用一个 block_id 索引后得到 [block_size, head_dim]。
req.block_table 按逻辑顺序保存物理 ID，因此列表推导式也必须按这个顺序取块。
token 位于每块的第 0 维，所以多个块沿 dim=0 拼接，结果是 [allocated_blocks*block_size, head_dim]。
最后一个块可能没有填满，返回前要切到 req.seq_len，得到精确 [seq_len, head_dim]。

图解原理

物理地址可以跳来跳去，但 block_table 的列表顺序就是逻辑顺序。先按表把每个二维块取出来，再沿 token 方向接成长序列；最后裁掉尾块未使用的槽位，才能得到请求真正拥有的 KV 长度。

块表 [7,2,9]定义逻辑块顺序 逐块索引每块 [block_size, head_dim] cat dim=0拼成逻辑 token 序列 [:seq_len]裁掉尾块空槽

物理池[num_blocks, block_size, head_dim]单个 blocks 元素[block_size, head_dim]cat_blocks[块表长度 × block_size, head_dim]返回 cache[req.seq_len, head_dim]可见边界长度 5、块大小 4 → 两块拼成 8 行，再截成 5 行

不能直接返回 physical_kv_cache[req.block_table]：它保留三维 [num_used_blocks, block_size, head_dim]，而测试要求逻辑连续的二维 [seq_len, head_dim]。

语法热身：按目录顺序拼接分散的数据页

import torch

pages = torch.arange(5 * 3 * 2).reshape(5, 3, 2)
page_table = [3, 0]
ordered_pages = [pages[page_id] for page_id in page_table]
continuous = torch.cat(ordered_pages, dim=0)
used_rows = continuous[:5]

print(continuous.shape)  # [6,2]
print(used_rows.shape)   # [5,2]

从语法例子迁移到 TODO 5

pages：对应 self.physical_kv_cache
page_table：对应 req.block_table，顺序不能排序或去重
ordered_pages：对应 blocks 列表，每项是一个物理缓存块
continuous：对应 cat_blocks，沿 dim=0 拼接 token
used_rows：对应返回值 cat_blocks[:req.seq_len]

巩固一下

为什么 blocks 的读取顺序必须服从 req.block_table？

列表位置表达逻辑块顺序，物理 ID 本身可能不连续因为 head_dim 是 block_id 因为 torch.cat 会自动排序

学完这一段，试着做

用一个小动作确认自己理解了；最后再进入官方题目。

完成列表推导式、torch.cat(dim=0) 和最终切片；确认没有修改 req.block_table 或 free_blocks。运行两组拼装测试：长度 9 应返回 [9,64] 且三段值为 1/2/3；长度 5 应返回 [5,8] 且第二块只取 1 行。若返回三维张量，检查是否漏了 cat；若尾部 shape 太长，检查是否漏了 [:req.seq_len]；若内容顺序错，检查是否排序了 block_id。

vLLM 分页注意力

Prefill：先向上取整，再做一次完整容量检查

🎯 先猜一猜

先补的知识

图解原理

长度 6、block_size 4

块表只记录地址

语法热身：给订单分配固定容量的货箱

从语法例子迁移到当前练习

巩固一下

学完这一段，试着做

Decode：先看已经加一的新长度，只在新块第一个位置补块

🎯 先猜一猜

先补的知识

图解原理

语法热身：给连续座位按排扩容

从语法例子迁移到 TODO 3

巩固一下

学完这一段，试着做

按块表恢复逻辑序列：先索引每块，再沿 token 维拼接并截尾

🎯 先猜一猜

先补的知识

图解原理

语法热身：按目录顺序拼接分散的数据页

从语法例子迁移到 TODO 5

巩固一下

学完这一段，试着做

准备好，去官方题目试一试