Block Manager：让 KV cache 变成可管理的资源

上一篇文章里，我们讨论了为什么 vLLM 要把 KV cache 做成分页式管理。

核心问题是：在线 LLM serving 里的 KV cache 不是一个固定张量，而是一种动态增长、动态释放、会被大量请求同时占用的显存资源。

如果每个请求都提前分配一整段连续 cache，会浪费大量空间。

如果按需申请连续空间，又容易出现碎片化。

PagedAttention 的思路是把 KV cache 切成固定大小的 block。Sequence 逻辑上看到的是连续 token 序列，但物理上，它的 KV cache 可以分散在不同 block 里。Sequence 只需要维护一张 block table，记录自己的逻辑 block 对应哪个物理 block。

这一篇，我们先不实现 PagedAttention 的 attention 计算，而是先实现它的前置模块：Block Manager。

如果说 Scheduler 是推理引擎的调度中心，那么 Block Manager 就是 KV cache 的资源管理中心。

Scheduler 决定哪些请求可以运行。

Block Manager 决定这些请求有没有足够的 KV cache 空间可以运行。

ModelRunner 负责执行模型。

PagedAttention 负责根据 block table 读取正确的 K 和 V。

这几个模块要配合起来，推理引擎才真正具备高并发能力。

从连续 cache 到 block table

在最简单的 KV cache 版本里，一个 Sequence 可能直接持有自己的 past_key_values。

这个对象通常按层保存 K 和 V，并且对上层来说像是一个完整的 cache。

这种方式很好理解，但问题是系统很难控制底层显存布局。每个请求的 cache 如何增长，如何释放，如何复用，基本交给框架内部处理。

而我们现在要实现的是推理引擎自己的 KV cache 管理。

所以需要把思路倒过来。

Sequence 不再直接拥有一整份 cache。

全局有一个 KV cache pool，它被切成很多固定大小的 physical blocks。

每个 physical block 可以存一小段 token 的 K 和 V，比如 16 个 token。

Sequence 自己只保存一张 block table。

Sequence A:
  logical block 0 -> physical block 3
  logical block 1 -> physical block 8
  logical block 2 -> physical block 1

这张表表达的是：Sequence A 的逻辑上下文被切成了三个逻辑 block，它们在物理 cache pool 中分别位于 3、8、1 号 block。

假设 block size 是 16，那么 logical block 0 存 token 0 到 token 15，logical block 1 存 token 16 到 token 31，logical block 2 存 token 32 到 token 47。

如果 Sequence 当前长度只有 40，那么最后一个 block 只用了前 8 个位置。

这就是 block table 的意义。

它把“一个请求的连续上下文”和“显存中的离散物理 block”连接起来。

没有 block table，PagedAttention 就不知道该从哪里读取历史 K 和 V。

没有 Block Manager，block table 就不知道该如何创建、扩展和释放。

Block Manager 要维护哪些东西

Block Manager 的第一职责，是管理所有 physical blocks 的生命周期。

我们可以把整个 KV cache pool 想象成一组编号的 block。

physical block 0
physical block 1
physical block 2
...
physical block N - 1

一开始，所有 block 都是空闲的。

Block Manager 维护一个 free block list。

free blocks: [0, 1, 2, 3, 4, 5, ...]

当某个 Sequence 需要 cache 空间时，Block Manager 从 free list 里取出 block，加入这个 Sequence 的 block table。

当 Sequence 结束时，Block Manager 把它占用的 block 归还给 free list。

所以最小的 Block Manager 至少要支持三个动作：

allocate_blocks(num_blocks)
free_blocks(block_ids)
get_num_free_blocks()

但对于推理引擎来说，这还不够。

因为 prefill 和 decode 对 block 的需求不同。

prefill 是一次性为 prompt 分配多个 token 的空间。

decode 是每轮追加一个 token，可能只需要写入当前最后一个 block，也可能因为最后一个 block 已满而申请新 block。

所以 Block Manager 还要理解 Sequence 的长度和 block table。

它至少要提供这些能力：

allocate_for_prefill(sequence, num_prompt_tokens)
append_slot_for_decode(sequence)
free_sequence(sequence)
can_allocate(num_tokens)

这里的 slot 是一个很重要的概念。

decode 每生成一个 token，都要把这个 token 的 K 和 V 写进某个位置。

这个位置可以理解成：

physical block id + block 内 offset

PagedAttention 读取历史 KV 时需要 block table。

模型写入新 KV 时需要 slot mapping。

Block Manager 要能告诉 ModelRunner：当前这个新 token 应该写到哪个 physical block 的哪个 offset。

block size 的取舍

在实现 Block Manager 之前，我们必须先决定 block size。

block size 表示一个 physical block 可以存多少个 token 的 KV cache。

比如：

block_size = 16

这意味着每个 block 可以容纳 16 个 token 在所有层上的 K 和 V。

注意这里的 block 不是只存某一层，也不是只存某一个 head。工程实现上，底层张量布局可以有不同设计，但从 Block Manager 的抽象看，一个 block 对应的是某个 Sequence 的一段 token 位置。

block size 不是随便选的。

如果 block size 太大，内部碎片会增加。

假设 block size 是 128，一个请求只用了 129 个 token，就需要 2 个 block。第二个 block 只用了 1 个位置，剩下 127 个位置浪费了。

如果 block size 太小，block table 会变长，管理开销会上升。attention kernel 读取历史 cache 时，需要跨更多 block，寻址成本也会变高。

所以 block size 是空间效率和访问效率之间的取舍。

小 block 更省空间，但管理和访问更碎。

大 block 管理简单，访问局部性更好，但内部浪费更高。

在教学版 mini vLLM 里，可以先选择一个常见且容易理解的值，比如 16。这个值不是绝对最优，但足够帮助我们观察分页式 KV cache 的行为。

真正的工程系统会结合模型结构、kernel 实现、硬件特性和 workload 分布来选择 block size。

这也是系统设计中很典型的一类问题：一个参数并没有脱离场景的全局最优，只有在具体目标下的合适选择。

prefill 时如何分配 block

prefill 是 Block Manager 第一次为一个 Sequence 分配 KV cache。

假设 prompt 长度是 40，block size 是 16。

那么需要的 block 数量是：

ceil(40 / 16) = 3

也就是 3 个 block。

逻辑上：

logical block 0: token 0 到 token 15
logical block 1: token 16 到 token 31
logical block 2: token 32 到 token 39

最后一个 block 只使用了 8 个位置。

Block Manager 会从 free list 中拿出 3 个 physical blocks，然后写入 Sequence 的 block table。

free blocks before:
  [3, 8, 1, 5, 7, ...]

allocate 3 blocks:
  [3, 8, 1]

Sequence block table:
  [3, 8, 1]

free blocks after:
  [5, 7, ...]

从 Sequence 的视角看，它拥有了 3 个逻辑 block。

从系统的视角看，它只是引用了 3 个物理 block。

这里有一个关键点：prefill 分配发生在模型执行之前，而不是执行之后。

因为模型执行时，需要有地方写入 prompt tokens 的 K 和 V。

也就是说，在调度一个 waiting 请求做 prefill 之前，系统必须确认它有足够的 free blocks。

这会影响 Scheduler。

以前 Scheduler 只要看 token budget 就可以决定某个 prompt 能不能进来。

现在 Scheduler 还必须问 Block Manager：

block_manager.can_allocate(num_prompt_tokens)

如果 block 不够，即使 token budget 够，这个请求也不能进入本轮 prefill。

这就是 KV cache 管理对调度器的反向约束。

decode 时如何追加 slot

decode 和 prefill 不一样。

prefill 需要为完整 prompt 一次性分配多个位置。

decode 每轮只为每个 running Sequence 追加一个 token。

假设某个 Sequence 当前长度是 40，block size 是 16。

它已经有 3 个 block。

block 0: token 0 到 15
block 1: token 16 到 31
block 2: token 32 到 39

下一轮 decode 要写 token 40。

token 40 属于哪个 logical block？

logical_block_id = 40 // 16 = 2
offset = 40 % 16 = 8

也就是说，它仍然写在 logical block 2 里，offset 是 8。

因为 block 2 还没有满，所以不需要新分配 block。

再往后写到 token 47，也都还在 block 2 里。

当要写 token 48 时：

logical_block_id = 48 // 16 = 3
offset = 48 % 16 = 0

这时 logical block 3 还不存在，就需要从 free list 里申请一个新的 physical block，并把它追加到 Sequence 的 block table。

所以 decode 追加 slot 的逻辑是：

根据当前 context length 计算要写的位置
如果这个位置已经有对应 block，直接返回 slot
如果需要新的 logical block，先申请 physical block
把新 block 追加到 block table
返回新 token 的写入 slot

伪代码可以这样表达：

def append_slot_for_decode(sequence):
    token_index = sequence.context_len
    logical_block_id = token_index // block_size
    offset = token_index % block_size

    if logical_block_id == len(sequence.block_table):
        new_block = allocate_one_block()
        sequence.block_table.append(new_block)

    physical_block_id = sequence.block_table[logical_block_id]

    return Slot(
        physical_block_id=physical_block_id,
        offset=offset,
    )

这个 slot 会交给 ModelRunner，用来告诉模型执行结果应该写到哪里。

注意，这里的 context_len 指的是当前已经存在于上下文中的 token 数，也就是 prompt tokens 加 generated tokens。

decode 要写的是下一个位置，所以使用当前 context_len 作为新 token 的 index。

写入完成并采样出下一个 token 后，Sequence 的 context_len 会增加。

这里有一个细节很容易混淆：decode 输入的是 last token，但写入 cache 的位置对应的是这个 last token 在上下文中的位置。实现时一定要清楚当前 token、下一个 token、当前 context length 之间的关系。

在教学版本里，可以先把这个过程写得直观一些，不急着做极致优化。只要保证每一步的位置语义正确，后面接 PagedAttention 时才不会乱。

Block Manager 和 Scheduler 如何协作

引入 Block Manager 后，Scheduler 的逻辑会变得更真实。

因为调度不再只是“这一轮算多少 token”，还包括“这些 token 的 KV cache 有没有地方放”。

对于 prefill，Scheduler 要检查 prompt 需要多少 blocks。

如果 prompt 长度是 100，block size 是 16，那么需要 7 个 block。

ceil(100 / 16) = 7

如果 free blocks 少于 7，这个请求不能进入 prefill。

对于 decode，情况更细。

并不是每个 running Sequence 每一轮都需要新 block。

大部分时候，它只是写入已有 block 的下一个 offset。

只有当当前长度刚好来到新 block 边界时，才需要新 block。

例如 block size 是 16，那么 context length 为 16、32、48 时，下一次写入会开启新 block。

所以 Scheduler 在调度 decode 前，也应该知道这一轮会不会导致某些 Sequence 申请新 block。

如果 free blocks 不够，就不能把所有 running Sequence 都调度进去。

这说明 Block Manager 需要提供一种“预检查”能力。

can_append_slot(sequence)

或者更批量一点：

can_append_slots(sequences)

第一版可以做得简单一些。

在 schedule 阶段，先估算本轮需要的新 block 数。如果 free blocks 足够，就允许调度。否则减少调度数量。

这个机制会让 Scheduler 更像一个真正的资源调度器。

它不只是维护 waiting 和 running 队列，还要根据 KV cache 可用空间做准入控制。

这也是 vLLM 这类系统能够稳定高并发运行的基础。

Block Manager 和 ModelRunner 如何协作

Scheduler 负责决定哪些 Sequence 本轮运行。

Block Manager 负责为这些 Sequence 分配 cache 位置。

ModelRunner 则需要根据这些位置信息构造模型输入。

在 prefill 阶段，ModelRunner 需要知道每个 prompt token 应该写入哪些 cache slot。

对于一个 prompt 长度为 40 的 Sequence，block size 是 16，那么它的 prompt tokens 对应的 slot 是：

token 0  -> block_table[0], offset 0
token 1  -> block_table[0], offset 1
...
token 15 -> block_table[0], offset 15
token 16 -> block_table[1], offset 0
...
token 39 -> block_table[2], offset 7

在 decode 阶段，每个 Sequence 只需要一个写入 slot。

比如：

Sequence A 当前 token 写到 physical block 3, offset 8
Sequence B 当前 token 写到 physical block 9, offset 0
Sequence C 当前 token 写到 physical block 2, offset 15

这些 slot 会被组织成 slot_mapping 之类的结构，供底层执行模块写入 KV cache。

PagedAttention 读取历史 KV 时，还需要 block table 和 context length。

所以 ModelRunner 的输入不再只是 input_ids。

它还会包含：

input tokens
position ids
block tables
slot mapping
context lengths
physical KV cache pool

这也是为什么前面文章里我们一直强调边界：Scheduler 不应该直接构造 tensor。因为真正的模型输入会随着 cache 管理方式变得越来越复杂。

Scheduler 只决定谁运行。

Block Manager 决定 cache 写在哪里。

ModelRunner 负责把这些信息转成模型需要的输入格式。

释放 block：请求结束后的资源回收

Block Manager 最重要的动作之一，是释放。

一个 Sequence 结束后，它占用的所有 physical blocks 都应该归还给 free list。

结束原因可能是 EOS、max tokens、stop words，也可能是用户取消。

无论是哪种原因，只要这个 Sequence 不再参与推理，它的 KV cache 就应该释放。

释放逻辑很简单：

def free_sequence(sequence):
    free_blocks(sequence.block_table)
    sequence.block_table.clear()

但释放时有几个细节要注意。

第一，释放必须只发生一次。

如果一个 Sequence 被重复释放，free list 里就可能出现重复 block id。后续两个 Sequence 可能拿到同一个 physical block，导致 cache 覆盖。

所以 Block Manager 最好能维护 block 的状态，至少在调试版本里检查重复释放。

第二，释放后 Sequence 不应该再被调度。

这需要 Scheduler 和 Sequence 状态配合好。只有 finished 的 Sequence 才会释放 block。释放后，它不应该重新回到 running。

第三，取消请求也必须释放。

在线服务里，用户断开连接很常见。如果取消请求不释放 block，系统会出现资源泄漏。请求表面上结束了，但显存中的 KV cache 一直被占着，最终 free blocks 越来越少。

第四，释放顺序要和输出事件协调好。

通常可以先标记 Sequence finished，生成最后一个 RequestOutput，再释放资源。只要后续不再需要读取它的 KV cache，就可以释放。

这类资源生命周期问题在 demo 中不明显，但在服务中非常关键。

一个推理引擎是否稳定，很多时候就取决于这些边界是否处理干净。

正确性不变量：实现 Block Manager 时最应该守住的东西

Block Manager 看起来只是分配和释放 block，但它非常容易写出隐蔽 bug。

因为一旦 block table 错了，模型可能不会立即报错，而是生成质量变差，甚至出现随机、难复现的问题。

所以实现时要明确一些正确性不变量。

第一个不变量：一个 physical block 在同一时间只能属于一个 Sequence。

除非后面实现 prefix cache 或 copy on write，否则不能让两个普通 Sequence 同时写同一个 block。

第二个不变量：free list 里不能包含正在被使用的 block。

这听起来显然，但重复释放、状态错误、异常路径都可能破坏它。

第三个不变量：Sequence 的 block table 顺序必须和逻辑 token 顺序一致。

block table 的第 0 项对应 token 0 到 block_size 减 1，第 1 项对应下一段 token。这个顺序不能乱。

第四个不变量：Sequence 的 context length 必须和已写入 cache 的 token 数一致。

如果 context length 多算了，attention 可能读到未写入的 cache。

如果少算了，模型会丢失部分历史上下文。

第五个不变量：decode 写入 slot 必须对应当前输入 token 的位置。

这个尤其容易出错。因为 decode 的输入是 last token，采样出来的是 next token，cache 写入的则是输入 token 的 K 和 V。实现时要明确每一步谁先发生，谁后发生。

这些不变量可以通过断言来保护。

例如在调试版本里，可以维护一个 block owner 表：

block_owner[physical_block_id] = sequence_id or None

分配时检查 owner 是否为空。

释放时检查 owner 是否确实是当前 Sequence。

这会带来一点开销，但对早期实现非常有帮助。

高性能版本可以再考虑关闭部分检查。

一个请求生命周期中的 block 变化

我们用一个完整例子把流程串起来。

假设 block size 是 4。

用户请求的 prompt 有 6 个 token。

prompt: A B C D E F

prefill 前，Sequence 还没有 block table。

block_table: []
context_len: 0

prefill 需要为 6 个 token 分配 block。

block size 是 4，所以需要 2 个 block。

Block Manager 从 free list 中分配 physical block 10 和 3。

block_table: [10, 3]

逻辑映射是：

token 0 A -> block 10, offset 0
token 1 B -> block 10, offset 1
token 2 C -> block 10, offset 2
token 3 D -> block 10, offset 3
token 4 E -> block 3, offset 0
token 5 F -> block 3, offset 1

prefill 执行后，模型根据 F 位置的 logits 采样出第一个生成 token G。

现在 Sequence 的上下文里已经包含 A 到 G。

如果实现中把 G 的 cache 写入放在下一轮 decode，那么 context length 的更新要非常小心。

更常见的理解方式是：prefill 处理的是 prompt tokens，cache 中已经有 A 到 F。prefill 输出 G，G 会作为下一轮 decode 的输入。下一轮 decode 执行时，才会把 G 的 K 和 V 写入 cache，并采样 H。

在这种语义下，prefill 后：

cache contains: A B C D E F
generated tokens: G
next decode input: G

下一轮 decode 输入 G。

此时要为 G 分配写入 slot。

当前 cache length 是 6，所以 G 对应 token index 6。

logical_block_id = 6 // 4 = 1
offset = 6 % 4 = 2

logical block 1 已经存在，对应 physical block 3。

所以 G 写入：

block 3, offset 2

decode 输出 H。

再下一轮 decode 输入 H。

H 对应 token index 7。

block 3, offset 3

decode 输出 I。

下一轮 decode 输入 I。

I 对应 token index 8。

logical_block_id = 8 // 4 = 2
offset = 0

logical block 2 不存在，所以 Block Manager 申请一个新 block，比如 physical block 6。

block_table: [10, 3, 6]

I 写入 block 6, offset 0。

这个例子展示了 block table 如何随着 decode 增长。

注意，generated tokens 和 cache 中已经写入的 tokens 之间可能存在一个“刚采样出来但尚未写入 cache”的 token。具体实现可以有不同方式，但一定要让语义一致。

在教学实现里，建议显式区分：

generated_token_ids: 已经采样出来、需要返回给用户的 tokens
cached_token_count: 已经写入 KV cache 的 tokens
last_token_id: 下一轮 decode 的输入 token

这样可以避免把“生成了 token”和“这个 token 的 KV 已经写入 cache”混为一谈。

这个细节非常重要。

因为在自回归推理里，某个 token 被采样出来后，它会在下一轮 forward 中作为输入。只有经过这一轮 forward，它自己的 K 和 V 才会被写入 cache。

Block Manager 的第一版设计原则

实现第一版 Block Manager 时，不要一上来追求复杂功能。

先做到几个目标就够了。

第一，block 分配和释放正确。

free list 不重复，不丢失，不把正在使用的 block 分配出去。

第二，Sequence 的 block table 正确。

逻辑 token index 能稳定映射到正确的 physical block 和 offset。

第三，prefill 和 decode 的资源检查正确。

调度前能判断是否有足够 block。

第四，请求结束后能释放资源。

EOS、max tokens、cancel 都能走到释放路径。

第五，错误尽早暴露。

在调试版本中，多加断言，多维护 owner 表，多检查 block table 长度和 context length 是否一致。

这一版不需要支持 prefix cache。

也不需要支持 copy on write。

不需要支持 swap。

不需要支持多 GPU。

先把最基本的资源生命周期跑通。

因为只要这个基础不稳，后面 PagedAttention 的任何 bug 都很难排查。

小结

这一篇我们把 KV cache 从一个抽象概念，进一步变成了可以管理的系统资源。

Block Manager 的核心思想是：把全局 KV cache pool 切成固定大小的 physical blocks，然后让每个 Sequence 通过 block table 维护自己的逻辑上下文。

prefill 时，Block Manager 根据 prompt 长度一次性分配多个 block。

decode 时，Block Manager 为当前输入 token 返回写入 slot。如果当前 block 已满，就申请新的 block 并追加到 block table。

请求结束时，Block Manager 把 Sequence 占用的 blocks 归还给 free list。

这套机制带来的关键变化是：KV cache 不再要求每个请求拥有一整段连续显存。Sequence 的上下文逻辑上连续，但物理上可以分散在不同 blocks 中。

这正是 PagedAttention 能工作的前提。

从系统结构上看，Block Manager 也让 Scheduler 变得更真实。调度器不能只看 token budget，还要检查 KV cache block 是否足够。只有计算资源和 cache 资源都满足，请求才能进入本轮执行。

下一篇文章，我们会在这个 Block Manager 的基础上，实现一个教学版 PagedAttention。

我们不会一上来写高性能 CUDA kernel，而是先用更容易读懂的方式，把 logical token index 到 physical KV cache address 的映射过程实现出来。

只有先把这条映射链路讲清楚，后面再谈 Triton、CUDA 或 FlashAttention backend 才有意义。