大模型推理引擎学习入门-2：带你理解 nano-vllm 的 KV Cache 以及 Paged Attention

茜布
0 阅读1小时+

致谢与声明 本文基于 nano-vllm 项目（MIT License，作者 Xingkai Yu）进行学习和分析。文中引用的代码片段均来自该项目。日志中 TEACHING/... 格式的输出是笔者为了学习自行添加的，非项目原始功能。感谢 nano-vllm 作者提供了这个优秀的教学级推理引擎实现。

前言

本文是「大模型推理引擎学习入门」系列的第二篇。上一篇我们把 nano-vllm 的整体架构走通了——从 llm.generate() 调用到 LLMEngine → Scheduler → ModelRunner 的主循环，搞清楚了 Continuous Batching 怎么调度 prefill 和 decode。

这一篇我们要拆开上次留下的一个黑盒：KV Cache 的显存管理——具体来说，就是 KV Cache 这个优化技巧为什么存在，以及 nano-vllm（借鉴 vLLM）如何用 PagedAttention 这套系统方案来高效管理它。

这两个概念不是同一层面的东西，它们的关系是：KV Cache 是算法层面的优化（把算过的 K/V 缓存下来，避免 decode 阶段重复计算），它把问题从”计算量爆炸”变成了”显存占用爆炸”；PagedAttention 是系统层面的方案（借鉴操作系统的分页机制），专门解决 KV Cache 带来的显存管理难题。先有 KV Cache 的需求，才有 PagedAttention 的解决方案。

上一篇里我们多次提到”KV cache 按块管理”“block_manager 分配块”，但一直没深入看它内部怎么运转。这次我们从第一性原理出发——先搞清楚为什么需要 KV Cache，再看它带来了什么新问题、PagedAttention 是怎么解决的，然后带着这些动机去读 nano-vllm 的实现代码，最后用真实日志把整个流程串起来。

整体方法论

延续上一篇的方法论，同时新增一条：

  1. 黑盒 + 接口：不讲的东西先当黑盒，这次 FlashAttention 的内部计算过程、TP 多进程通信仍然是黑盒
  2. 日志对着代码看：所有流程分析都配合真实运行的 DEBUG 日志
  3. 读函数前先搞清楚输入输出：先知道一个方法的调用者是谁、传了什么参数、期望什么结果
  4. （新增）拆解复杂机制到可执行的小脚本：对于 Paged Attention 这类涉及多个模块协作的机制，先写独立的小脚本模拟核心逻辑（不依赖 nano-vllm），建立直觉后再去读真实代码

上一篇回顾

上一篇的几个关键结论，这篇会直接引用：

  • 推理分两阶段：Prefill（一次性处理整段 prompt）和 Decode（逐 token 生成），两者的计算特征完全不同
  • Continuous Batching：Scheduler 在每个 step 的调度窗口里动态选择做 prefill 还是 decode，序列可以随时加入/退出 batch
  • ModelRunner 的 prepare_prefill/decode：把 Scheduler 选出的序列转换成 GPU 需要的张量格式（input_ids、positions、slot_mapping 等）
  • KV Cache 一次性预分配：引擎启动时根据 GPU 剩余显存计算出 num_kvcache_blocks，一次性分配一大块显存作为 KV Cache 池

导读

这篇文章会带你学什么

  • 从第一性原理理解 KV Cache 为什么存在：没有它的 decode 有多慢，有了它省了多少计算
  • 理解 Paged Attention 的核心思想：为什么要把 KV Cache 切成 block，block_table 和 slot_mapping 是怎么回事
  • 逐段读 nano-vllm 的 BlockManager 实现：allocate、may_append、deallocate、prefix caching 的链式 hash
  • 理解 KV Cache 的初始化流程：warmup 为什么必要、num_blocks 怎么算出来的
  • 用 真实运行日志 追踪两条共享前缀请求的完整生命周期：block 分配 → prefix cache 命中 → decode → block 释放

这篇文章不讲什么

  • FlashAttention 的内部计算（tiling、online softmax 的详细算法）——另开一篇讲
  • store_kvcache 的 Triton kernel 实现——另开一篇和 Triton 入门一起讲
  • Tensor Parallel 下的 KV Cache 切分——TP 专题讲
  • Attention 的数学原理（Q@K^T → softmax → @V）——这篇只讲 KV Cache 和 Paged 管理，不讲 Attention 怎么算

读之前最好有的背景

  • 读过本系列第一篇，了解 nano-vllm 的 LLMEngine → Scheduler → ModelRunner 主循环
  • 知道 Prefill 和 Decode 两个阶段的区别
  • 大致知道 Attention 需要 Q、K、V 三个矩阵（不需要懂计算细节）
  • 了解 GPU 显存的基本概念（上一篇的 memory_allocated vs memory_reserved）

1. 为什么需要 KV Cache

上一篇我们知道了 LLM 的推理分成 Prefill 和 Decode 两个阶段：Prefill 一次性处理整段 prompt，Decode 逐 token 生成。但我们当时没有深究一个问题：Decode 阶段每步只生成 1 个 token，为什么还要反复调用整个模型？能不能省掉一些重复计算？

1.1 Decode 阶段的计算量问题

Attention 的核心操作是 Q @ K^T：当前 token 的 Q 向量要和所有历史 token 的 K 向量做点积，算出”该关注哪些历史 token”。

如果没有任何缓存，decode 的每一步都要：

  1. 对所有历史 token重新做 QKV 投影（F.linear），算出它们的 K 和 V
  2. 用新 token 的 Q 和所有历史的 K 做点积
  3. 按注意力权重加权所有历史的 V

第 1 步的问题最大：假设已经生成了 100 个 token，接下来每步 decode 都要对 100+ 个 token 做 3 次 F.linear（Q、K、V 各一次）。随着序列变长，这个计算量线性增长。生成 L 个 token 的总 QKV 投影次数是 O(L²)。

我们写了一个实验脚本来直观对比（完整脚本见 08a_kvcache_motivation.py）。先看”无 cache”的 decode 长什么样：

def decode_without_cache(hidden_states_history, W_q, W_k, W_v, num_steps=20):
    """没有 KV Cache: 每步重新算所有历史 token 的 K/V"""
    total_qkv_ops = 0
    for step in range(num_steps):
        context_len = len(hidden_states_history)
        all_hidden = torch.stack(hidden_states_history)     # (context_len, hidden_size)
        q_all = F.linear(all_hidden, W_q)  # 实际只需要最后一个 Q，但这里全算了
        k_all = F.linear(all_hidden, W_k)  # 全部重新算！
        v_all = F.linear(all_hidden, W_v)  # 全部重新算！
        total_qkv_ops += context_len * 3
        hidden_states_history.append(torch.randn(hidden_size, dtype=torch.float16))
    return total_qkv_ops

每步 decode 都在对所有历史 token 做 QKV 投影，context_len 有多大就算多少次。

再看”有 cache”的版本：

def decode_with_cache(hidden_states_history, W_q, W_k, W_v, num_steps=20):
    """有 KV Cache: 每步只算 1 个新 token 的 K/V"""
    # Prefill: 一次性算完所有已有 token 的 K/V 并缓存
    all_hidden = torch.stack(hidden_states_history)
    k_cache = F.linear(all_hidden, W_k)
    v_cache = F.linear(all_hidden, W_v)
    total_qkv_ops = len(hidden_states_history) * 3

    for step in range(num_steps):
        new_hidden = torch.randn(1, hidden_size, dtype=torch.float16)
        q_new = F.linear(new_hidden, W_q)   # 只算 1 个 token 的 Q
        k_new = F.linear(new_hidden, W_k)   # 只算 1 个 token 的 K
        v_new = F.linear(new_hidden, W_v)   # 只算 1 个 token 的 V
        k_cache = torch.cat([k_cache, k_new], dim=0)  # 追加到 cache
        v_cache = torch.cat([v_cache, v_new], dim=0)
        total_qkv_ops += 3                  # 每步固定 3 次
    return total_qkv_ops

差别很明显——decode 循环里从 context_len * 3 变成了固定的 3 次。历史的 K/V 已经在 cache 里了，不需要重新算。

初始序列 100 个 token，decode 生成 50 个 token 的对比结果：

无 KV Cache:
  QKV 投影总次数: 18,675 (每步重新算所有历史)
  计算量增长: O(L²)

有 KV Cache:
  QKV 投影总次数: 450 (prefill 300 + decode 150)
  计算量增长: O(L)

  节省: 18,225 次 QKV 投影 (98%)

18,675 次投影里的 98% 都是在重复计算已经算过的历史 token 的 K/V。

1.2 KV Cache 的思路

思路很直接：把每步算出的 K/V 缓存下来，下一步直接读取，不重新计算。

Prefill 阶段:
  输入 100 个 token → 算出 100 个 K/V → 全部存入 cache

Decode 第 1 步:
  输入 1 个新 token → 只算 1 个新的 K/V → 追加到 cache
  Attention: 新 Q × (cache 中的 100+1 个 K)

Decode 第 2 步:
  输入 1 个新 token → 只算 1 个新的 K/V → 追加到 cache
  Attention: 新 Q × (cache 中的 100+2 个 K)

  ...

每步 decode 只需要 3 次 F.linear（1 个 token 的 Q、K、V），不管历史有多长。生成 L 个 token 的总投影次数从 O(L²) 降到 O(L)。

下面这张图直观对比了”有 cache” 和 “无 cache” 两种方式在 decode 阶段的差异：

上半部分：没有 KV Cache 时，每步 decode 都要重算所有历史 token 的 QKV（红色方块），随着序列变长红色方块越来越多，计算量 O(L²) 下半部分：有 KV Cache 时，历史 K/V 缓存为绿色方块，每步只对新 token 计算一次 QKV（1 个红色方块），追加到 cache，计算量 O(L)

1.3 计时对比

理论上 98% 的计算量节省，实际耗时差多少？我们在 H20 上用 Qwen3-0.6B 的参数（hidden_size=1024, num_heads=16, num_kv_heads=8, head_dim=128）做了计时，decode 20 步：

context_len |   无 cache (ms) |   有 cache (ms) |      加速比
------------------------------------------------------------
          50 |          23.41 |           1.29 |    18.1x
         200 |          16.20 |           1.25 |    12.9x
         500 |          21.29 |           1.47 |    14.5x

有 cache 时 decode 耗时基本恒定（~1.3 ms），因为每步只处理 1 个 token。无 cache 时耗时随 context_len 增长，因为每步要处理的 token 数等于整个历史长度。

1.4 KV Cache 的显存代价

KV Cache 用显存换计算量。每个 token 需要缓存 K 和 V 两个向量，在模型的每一层都要缓存。以 Qwen3-0.6B 为例：

每 token 每层: 2(K+V) × 8(kv_heads) × 128(head_dim) × 2(bf16 字节) = 4,096 bytes = 4 KB
每 token 全部 28 层: 4,096 × 28 = 114,688 bytes ≈ 112 KB

看起来 112 KB 不多，但 KV Cache 的总大小 = 每 token 大小 × 序列长度 × 并发请求数，它会随着服务规模快速膨胀：

并发请求数平均序列长度KV Cache 总显存
14,0960.47 GB
84,0963.76 GB
324,09615.0 GB
1284,09660.1 GB

128 条并发请求、每条序列长度 4096 token 时，KV Cache 就要 60 GB（128 × 4096 × 112 KB ≈ 60 GB）——在 H20（96 GB 显存）上已经占了大半，而模型权重本身只有 1.2 GB。KV Cache 才是显存的大头。而且这还只是 0.6B 的小模型，更大的模型（7B、70B）KV Cache 占用更惊人。

这就引出了下一个问题：这么大一块显存，多条请求的长度不同、生命周期不同，怎么管理？

2. KV Cache 带来的新问题 → Paged Attention

KV Cache 解决了计算量的问题，但引入了一个新问题：显存管理。

我们用 nano-vllm 在 H20 上实测，Qwen3-0.6B 引擎初始化后的 GPU 显存占用是 84.6 GiB，其中 KV Cache 预分配占了 83.5 GiB（88%），模型权重只占约 1.2 GiB（1.3%）。KV Cache 才是显存的绝对大头，怎么管理它直接决定了能同时服务多少请求。

2.1 KV Cache 的显存浪费问题

在没有 Paged Attention 之前，传统的 KV Cache 管理方式是：为每条请求预留一段连续的、按最大可能长度分配的显存空间。

这带来了三种浪费，vLLM 论文（Kwon et al., 2023）的 Figure 3 清晰地展示了这个问题：

vLLM Figure 3: KV Cache 的三种显存浪费

图片来源：Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention”, SOSP 2023, Figure 3

三种浪费分别是：

  • Reserved（预留浪费） ：为了应对序列可能增长到最大长度，系统预留了远超实际需要的空间。比如 max_model_len=4096，但请求实际只生成了 200 个 token，剩下 3896 个 token 的空间白占着
  • Internal Fragmentation（内部碎片） ：预留的空间内部有一部分是空的，不会再被使用
  • External Fragmentation（外部碎片） ：多条请求之间的已释放空间不连续，无法分配给新请求

论文的实验数据显示：现有系统浪费了 60%-80% 的 KV Cache 显存，实际有效利用率最低只有 20.4%。

2.2 从操作系统的分页机制到 Paged Attention

先回顾：操作系统是怎么解决内存碎片的

操作系统管理内存也经历过同样的碎片化问题。最早的做法是分段管理——给每个进程在物理内存上分配一块连续的空间。但当进程释放后，留下的空洞可能不够大给新进程用，即使总的空闲空间足够——这就是碎片化。

操作系统的解决方案是分页管理：

  1. 把物理内存划分成固定大小的页（page）
  2. 每个进程有自己的虚拟地址空间，逻辑上是连续的
  3. 通过页表（page table） 记录虚拟页到物理页的映射
  4. 进程不需要在物理内存上连续存储——只要页表映射正确，进程感知到的地址空间就是连续的

这样就消除了碎片化——每个页大小相同，任何空闲页都能分配给任何进程。

PagedAttention：把分页思想用在 KV Cache 上

vLLM 的 PagedAttention 把这套思路直接搬到了 KV Cache 的管理上。对应关系如下：

操作系统PagedAttention
物理内存GPU 上的 KV Cache 显存
页（page）block（固定大小，存 block_size 个 token 的 KV）
进程一条推理请求（sequence）
虚拟地址空间逻辑 block（每条请求看到的连续 KV 空间）
页表（page table）block_table（逻辑 block → 物理 block 的映射）

PagedAttention 带来三个关键改进：

  1. KV Cache 被切分为固定大小的 block：每个 block 存若干个 token 的 K 和 V 向量（vLLM 默认 16，nano-vllm 用 256）。统一了分配粒度，分配和回收都以 block 为单位，有效减少碎片
  2. block 可以存放在非连续的物理显存中：通过 block_table 维护逻辑 block 到物理 block 的映射关系。模型计算时操作逻辑 block（连续的），vLLM 在后台通过 block_table 从物理 block 读取实际数据。每条请求仿佛在一个连续且充足的空间中运行，尽管物理上数据是非连续的
  3. 支持按需分配、回收和共享：不预留最大长度的空间，用到哪分配到哪。多条请求共享相同的 prompt 前缀时，可以共享同一组物理 block（通过引用计数 + copy-on-write 机制）

用论文的例子逐步走一遍

我们用 vLLM 论文中的经典例子来具体看看整个流程。假设 prompt 为 “Four score and seven years ago our”（7 个 token），block_size=4：

vLLM Figure 6: Block table translation

图片来源：Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention”, SOSP 2023, Figure 6

第 ① 步：Prefill 阶段

vLLM 拿到 7 个 token 的 prompt，按 block_size=4 划分逻辑 block：

  • 逻辑 block 0：装前 4 个 token “Four score and seven” 的 KV → 映射到物理 block 7
  • 逻辑 block 1：装后 3 个 token “years ago our” 的 KV → 映射到物理 block 1，还有 1 个空位

block_table 记录两件事：逻辑 block 对应的物理 block 编号，以及每个物理 block 被填了多少个 slot：

Block Table:
  逻辑 block 0 → 物理 block 7, #filled = 4 (满)
  逻辑 block 1 → 物理 block 1, #filled = 3 (还有 1 个空位)

第 ② 步：Decode — 生成第 1 个词 “fathers”

模型基于 KV Cache 计算 Attention，生成 token “fathers”。此时逻辑 block 1 还有 1 个空位（#filled=3），直接写入：

  • 物理 block 1 的第 4 个位置写入 “fathers” 的 KV
  • block_table 更新：#filled 3 → 4

这里有一个关键点：计算时我们操作的是逻辑 block，形式上所有 token 是连续排列的。vLLM 在后台通过 block_table 这个映射关系，从物理 block 上读取实际数据做计算。通过这种方式，每条请求都仿佛在一个连续且充足的存储空间中运行，尽管物理上数据是非连续的。

第 ③ 步：Decode — 生成第 2 个词

逻辑 block 1 已满（4/4），需要新 block：

  • 为逻辑 block 2 分配一个新的物理 block 3
  • 把新生成 token 的 KV 写入物理 block 3 的第 1 个位置
  • block_table 新增一行：逻辑 block 2 → 物理 block 3, #filled = 1

整个过程中，只有在前一个 block 填满后才分配新 block，最大限度减少显存浪费。

多条请求的 KV Cache 共享

当多条请求共享相同的 prompt 前缀时，PagedAttention 可以让它们的 block_table 指向同一组物理 block，不重复存储 KV Cache。下图展示了两条请求如何共享 prompt 的 KV：

vLLM Figure 7: KV cache of two requests

图片来源：Kwon et al., 2023, Figure 7

下面这张动图更直观地展示了 Parallel Sampling（同一个 prompt 生成多个候选回答）场景下的共享机制：

Paged Attention memory sharing

图片来源：vLLM Blog - Easy, Fast, and Cheap LLM Serving with PagedAttention, Apache 2.0 License

两个 sample 共享相同的 prompt，它们的 block_table 指向相同的物理 block，每个物理 block 维护一个 ref_count（引用计数）。当某个 sample 需要写入不同的 token 时，才会触发 copy-on-write——将共享的物理 block 复制一份到新位置，然后再写入。这样在保证数据隔离的同时极大地节省了显存。

这种共享机制适用于多种解码策略：

  • Parallel Sampling：同一 prompt 生成多个回答，共享 prompt 的 KV
  • Beam Search：多个 beam 路径共享公共前缀的 KV，只在路径分叉时触发 copy-on-write
  • Shared Prefix：多条请求共享 system prompt 的 KV

PagedAttention 的效果

论文的实验数据量化了 PagedAttention 的收益：

指标传统 KV Cache 管理PagedAttention
显存有效利用率20.4% ~ 38.2%> 96%
显存浪费60% ~ 80%< 4%（仅最后一个 block）
Parallel Sampling 显存每个 sample 独立分配共享 prompt 的 KV，节省 55%
外部碎片严重无（所有 block 大小相同）

数据来源：Kwon et al., 2023, Table 1 & Section 5

到这里我们已经理解了 Paged Attention 的核心机制。接下来看 nano-vllm 是怎么实现它的——block 怎么分配回收（BlockManager）、token 位置怎么映射到物理地址（slot_mapping）、KV 怎么实际写入显存（store_kvcache kernel）。

3. nano-vllm 的 KV Cache 和 Paged Attention 实现

第 2 章我们用代码从零模拟了 Paged Attention 的完整流程——block 分配、KV 写入、decode 追加、block 回收。现在把目光转向真正的 nano-vllm 代码，看它是怎么把这套机制落地的。

读代码之前先看全局。

3.1 三层架构：从 block_table 到物理 KV 写入

nano-vllm 把 Paged Attention 的职责拆成了三层，每一层只负责一件事，通过 block_id 把三层串起来：

逻辑层 · BlockManager

源码：nanovllm/engine/block_manager.py

它只管一件事：维护”每条序列用了哪些 block”。输入一条 sequence，输出 seq.block_table（一个 list，比如 [3, 7, 12]）。BlockManager 不关心 K/V 的实际内容，甚至不知道 GPU 上有没有 kv_cache tensor 存在——它只是在逻辑层面做 block 的分配和回收。

映射层 · prepare_prefill / prepare_decode

源码：nanovllm/engine/model_runner.py

它把 block_table（逻辑位置）翻译成 slot_mapping（物理行号）。公式很简单：

slot = block_table[i] × block_size + offset

输入是一批 sequence，输出是一个 tensor，告诉下一步的 GPU kernel “这批 token 要写到 kv_cache 的哪几行”。

物理层 · store_kvcache

源码：nanovllm/layers/attention.py（Triton kernel）

它真正执行写入操作。给定 slot_mapping 和刚算出来的 K/V，在 GPU 上并行地把每个 token 的 K/V 写到 kv_cache tensor 的对应位置。这一层是完全物理的，不关心 block 的概念，只看 slot 编号。

block_id 是串起这三层的关键。BlockManager 说”这条序列用 block 3、block 7、block 12”，prepare_* 把它翻译成”写到 kv_cache 的第 768-1023 行（3 × 256 = 768）、第 1792-2047 行（7 × 256 = 1792）…“，最后 store_kvcache kernel 按这些行号把数据实际写进显存。

全局 kv_cache tensor 长什么样

整个引擎只有一个 kv_cache tensor，在引擎启动时一次性分配。我们之前跑 verify_kvcache_size.py 的时候已经从真实日志里看到过这个 tensor 的形状。以 Qwen3-0.6B 在 H20 上为例：

kv_cache shape: (2, 28, 3053, 256, 8, 128)
                 │   │    │     │    │   │
                 │   │    │     │    │   └── head_dim
                 │   │    │     │    └────── num_kv_heads (GQA)
                 │   │    │     └─────────── block_size (每 block 存 256 个 token)
                 │   │    └───────────────── num_blocks (物理 block 总数)
                 │   └────────────────────── num_hidden_layers
                 └────────────────────────── 2: K cache 和 V cache

其中 num_blocks=3053 不是固定配置，是根据 GPU 剩余显存动态算出来的——下一节 §3.2 讲它怎么算的。其他维度都由模型的 config.json 决定：28 层、8 个 KV head、128 head_dim、block_size=256 是 nano-vllm 的配置默认值。

整个 tensor 的大小可以手算出来：

2 × 28 × 3053 × 256 × 8 × 128 × 2 bytes (bf16)
  = 89,640,468,480 bytes
  ≈ 89.6 GB  (或 83.5 GiB)

对照的实测日志：

引擎初始化后 GPU 显存: 84.64 GiB
  其中包括: 模型权重 (~1.2 GiB) + KV Cache 预分配 (83.5 GiB)

数字对上了。这块 83.5 GiB 的显存在引擎启动后就被一次性划走，整个运行期间不会再向 CUDA driver 申请/释放，BlockManager 只是在这块固定大小的池子里做逻辑上的分配。

block_id 到物理位置的映射

以 block_id=3 为例，它在 kv_cache tensor 中对应的切片是：

# block 3 在所有 28 层的 K cache
kv_cache[0, :, 3, :, :, :]   # shape: (28, 256, 8, 128)

# block 3 在所有 28 层的 V cache
kv_cache[1, :, 3, :, :, :]   # shape: (28, 256, 8, 128)

注意 block_id 贯穿所有 28 层——分配一个 block 不只是分配一层的显存，而是在所有层的同一位置都预留了 256 个 token 的 KV 空间。

这也解释了为什么 block_bytes 要把层数乘进去：

block_bytes = 2(K+V) × 28(layers) × 256(tokens) × 8(kv_heads) × 128(head_dim) × 2(bf16)
            = 29,360,128 bytes ≈ 28 MB

每分配一个 block，就是 28 MB 的显存被锁定。3053 个 block × 28 MB = 83.5 GiB，正好是我们看到的 KV Cache 总占用。

下面几节怎么走

三层职责清楚后，我们逐层往下看：

  • §3.2 看 kv_cache tensor 是怎么被创建的——warmup 为什么必要，num_blocks 怎么算出来
  • §3.3 看 BlockManager 的内部设计——数据结构、allocate / may_append / deallocate 三个核心方法
  • §3.4 看 BlockManager 的 prefix caching 机制——链式 hash 和 block 复用
  • §3.5 看 prepare_prefill 怎么把 block_table 翻译成 slot_mapping

3.2 KV Cache 的初始化：显存预算和 warmup

§3.1 最后我们提到 num_blocks=3053 这个数字不是固定配置，是引擎启动时动态算出来的。这一节看它是怎么算的。

先看 nano-vllm 跑一次初始化的日志

直接看 nano-vllm 引擎启动的真实日志

11:21:17 TEACHING/MODEL_WEIGHTS_LOADED   phase=init
           stats={num_files=1, num_keys=311}
11:21:17 TEACHING/MODEL_LOAD_READY        phase=init
           stats={dtype=torch.bfloat16, rank=0}
11:21:17 warmup_model：用假数据跑一轮最大规模 prefill,
           使 cuBLAS 等分配 workspace 并记录 peak;
           否则 allocate_kv_cache 会低估占用、易 OOM。
11:21:19 TEACHING/KV_BUDGET_ESTIMATE      phase=init
           stats={total_gib=95.085, free_gib=93.43, peak_gib=1.599,
                  block_bytes=29360128, num_blocks=3053}
11:21:19 allocate_kv_cache：
           total_mem=95.09GiB  free=93.43GiB  used=1.66GiB
           peak_alloc=1.60GiB  curr_alloc=1.16GiB
           gpu_memory_utilization=0.9
           -> num_kvcache_blocks=3053  block_size=256
              block_bytes≈29360128

整个初始化过程分三步：加载权重 → warmup → 计算并分配 KV cache。

加载权重后模型权重占约 1.2 GiB（Qwen3-0.6B 的 596M 参数 × 2 bytes），这是稳定占用。接下来的问题是：剩下 90+ GiB 显存，多少能分给 KV cache？

为什么不能直接用”当前剩余显存”来算

最朴素的想法：num_blocks = 剩余显存 / block_bytes。但这里有个陷阱—— “当前”剩余显存不等于”未来运行时”的剩余显存。

模型推理的时候，除了权重之外还有两块临时占用：

  1. 激活（activations） ：Attention 的 Q @ K^T 中间矩阵、FFN 的 gate/up 投影结果等，每次 forward 都会临时分配一批，forward 结束释放
  2. cuBLAS / cuDNN workspace：这些库首次调用时会分配内部的 workspace（几百 MB），用来存 GEMM 的中间结果、算法选择的查找表等。这块内存会常驻

如果你在刚加载完权重就去算”剩余显存”，看到的是 ~93 GiB。但真跑一次最大规模的 prefill，激活 + workspace 会占掉 1+ GiB。如果这时候你已经把那 93 GiB 全分给 KV cache 了，激活和 workspace 就没地方放——OOM。

nano-vllm 的解决办法是：先跑一次假数据 prefill（warmup），把这块峰值占用”摸”出来，再按峰值扣除余量。

warmup 做了什么

model_runner.py 的 warmup_model 方法很短：

def warmup_model(self):
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()    # 把历史峰值清零
    max_num_batched_tokens = self.config.max_num_batched_tokens  # 16384
    max_model_len = self.config.max_model_len                    # 4096
    num_seqs = min(max_num_batched_tokens // max_model_len, self.config.max_num_seqs)
    seqs = [Sequence([0] * max_model_len) for _ in range(num_seqs)]
    self.run(seqs, True)                    # 跑一次最大规模 prefill
    torch.cuda.empty_cache()

它构造了一批”最坏情况”的假请求——每条请求长度都是 max_model_len（4096），总 token 数接近 max_num_batched_tokens（16384），然后跑一次 prefill。这次 prefill 会：

  • 触发模型所有层的前向计算，产生完整的激活张量
  • 触发 cuBLAS 分配 workspace
  • 把 PyTorch 的 max_memory_allocated() 记录刷到这次 forward 的峰值

warmup 完成后再看显存，日志里的 peak_alloc=1.60 GiB（warmup 期间达到过的峰值）就是这次 prefill 的峰值占用。

num_blocks 的计算公式

接下来 allocate_kv_cache 用这些数字算出 num_blocks。源码里核心的一行是（model_runner.py）：

free, total = torch.cuda.mem_get_info()
used = total - free
peak = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
current = torch.cuda.memory_stats()["allocated_bytes.all.current"]

config.num_kvcache_blocks = int(
    total * config.gpu_memory_utilization - used - peak + current
) // block_bytes

公式有四个变量，把日志里的数字代进去（单位 GiB）：

total × gpu_memory_utilization     = 95.09 × 0.9     =  85.58
- used  (驱动层面已用)                 =  1.66           → -1.66
- peak  (warmup 峰值)                  =  1.60           → -1.60
+ current (warmup 后稳定占用)          =  1.16           → +1.16
────────────────────────────────────────────────────────
  可用于 KV cache 的字节                                ≈  83.48 GiB

为什么 + current 要加回来？因为 peak 里其实已经包含了当前稳定的那部分（权重 + workspace），如果既从总量里扣 used 又扣 peak，当前占用就被减了两次，所以要 + current 把它加回来避免重复扣除。

最后除以 block_bytes 得到 block 数：

83.48 GiB / 28 MiB ≈ 3053 blocks

和日志里的 num_blocks=3053 完全对得上。

一次性分配大 tensor

算出 num_blocks=3053 后，nano-vllm 就一次性分配整个 kv_cache tensor：

self.kv_cache = torch.empty(
    2,
    hf_config.num_hidden_layers,   # 28
    config.num_kvcache_blocks,     # 3053
    self.block_size,               # 256
    num_kv_heads,                  # 8
    head_dim,                      # 128
)

一次 torch.empty 调用就拿到了 83.5 GiB 的连续显存。整个服务运行期间这块显存不会被释放，也不会再申请新的——BlockManager 只在这个固定大小的池子里做逻辑上的分配和回收。

这么做有三个好处：

  1. 没有碎片化：§1 讲过，反复 malloc/free 会造成碎片化 OOM。一次性分配避免了这个问题
  2. 分配开销归零：每次请求进来不需要 cudaMalloc，BlockManager 只是在 Python 层面操作 list 和 set
  3. 显存预算可控：跑起来之后显存占用不会再变化，不会出现”跑着跑着 OOM”

最后一步：按层绑定到 Attention 模块

分配完 tensor 还有一步关键操作——把 kv_cache 的各层切片绑定到 Attention 模块：

layer_id = 0
for module in self.model.modules():
    if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
        module.k_cache = self.kv_cache[0, layer_id]   # (num_blocks, block_size, kv_heads, head_dim)
        module.v_cache = self.kv_cache[1, layer_id]
        layer_id += 1

这段代码遍历模型的所有子模块，找到带 k_cache / v_cache 属性的（就是 Attention 层），把全局 kv_cache tensor 的对应层切片赋给它。之后每层的 Attention.forward 就直接通过 self.k_cache / self.v_cache 访问自己那一层的 KV 数据。

注意 self.kv_cache[0, layer_id] 返回的是一个 view（不是拷贝），所以虽然我们看到的是 28 个独立的 k_cache，它们在物理显存上其实是同一块大 tensor 的切片。当 store_kvcache kernel 写入第 5 层的 k_cache 时，它写的就是全局 kv_cache[0, 5, :, :, :, :] 这个切片对应的显存。

小结

nano-vllm 的 KV cache 初始化流程可以概括成这样：

1. 加载模型权重                     (~1.2 GiB)
2. warmup: 跑一次最大规模 prefill    → 记录 peak_alloc = 1.60 GiB
3. 算 num_blocks:
     (total × 0.9 - used - peak + current) / block_bytes = 3053
4. 一次性分配 kv_cache tensor        (~83.5 GiB)
5. 按层绑定到每个 Attention.k_cache / v_cache

这里有个细节值得留意：如果你把 gpu_memory_utilization 调低（比如从 0.9 → 0.8），num_blocks 会相应变小，能支持的并发请求数也会减少。这个参数就是给你在”吞吐量”和”显存安全边界”之间做权衡的——越接近 1.0 吞吐量越高，但余量越小，遇到未预料的显存需求就容易 OOM。

3.3 BlockManager 的设计

§3.2 我们看完了 KV cache tensor 是怎么被创建出来的。现在进入 §3.1 三层架构图里的”逻辑层”——BlockManager。

BlockManager 是 nano-vllm 里 PagedAttention 实现的核心。它的代码不长（约 350 行），全部在 nanovllm/engine/block_manager.py 里，但它承载了 §2 我们讲过的全部分页机制：block 分配、回收、引用计数、prefix caching、链式 hash。

按之前定下的源码阅读方法论，我们分五步看：先看核心数据结构，再看上游谁在调用它，最后逐个看三个核心方法（allocate / may_append / deallocate）。Prefix caching 的链式 hash 单独放到 §3.4 讲。

3.3.1 核心数据结构

BlockManager.__init__ 维护四个核心数据结构。以 Qwen3-0.6B 在 H20 上的实际配置（num_blocks=3053, block_size=256）为例，初始状态如下：

class BlockManager:
    def __init__(self, num_blocks: int, block_size: int):
        self.block_size = block_size
        self.blocks: list[Block] = [Block(i) for i in range(num_blocks)]
        self.hash_to_block_id: dict[int, int] = dict()
        self.free_block_ids: deque[int] = deque(range(num_blocks))
        self.used_block_ids: set[int] = set()

四个数据结构的角色：

数据结构类型初始状态 (num_blocks=3053)作用
blockslist[Block]3053 个 Block 对象block 元数据数组，下标即 block_id
free_block_idsdeque[int][0, 1, 2, ..., 3052]空闲 block 队列，分配时 popleft，回收时 append
used_block_idsset[int]{}当前被占用的 block_id 集合
hash_to_block_iddict[int, int]{}链式 hash → block_id，用于 prefix cache 查找

free_block_ids 用 deque 而不是 list 是因为 popleft 和 append 都是 O(1)；used_block_ids 用 set 是为了 O(1) 判断”某个 block 是否在使用”——这个判断在 prefix cache 命中时会用到（§3.4 详讲）。

每个 Block 对象只存元数据，不存 K/V：

class Block:
    def __init__(self, block_id):
        self.block_id = block_id   # 物理块 id, 对应 kv_cache tensor 的第 3 维下标
        self.ref_count = 0          # 引用计数: 多少条序列正在用这块, 0 = free
        self.hash = -1              # 链式 hash 指纹, -1 = 未满块或已回收
        self.token_ids = []         # 该块的 token 内容, 用于精确匹配防 hash 碰撞

实际的 K/V 数据存在 ModelRunner.kv_cache tensor 里——Block.block_id 就是访问那个 tensor 第 3 维的下标。这是 §3.1 里强调的”BlockManager 只管逻辑分配，不动 K/V”在代码层面的体现。

3.3.2 上游调用关系：Scheduler 怎么用 BlockManager

BlockManager 由 Scheduler.__init__ 创建，整个引擎只有一个实例。Scheduler 在调度的不同阶段调用 BlockManager 的不同方法。

下面这张时序图展示了一条请求从 prefill → decode → 结束的完整生命周期：

把图里的调用按顺序梳理一下：

阶段Scheduler 调用BlockManager 做的事
Prefillcan_allocate(seq)检查 free_block_ids 数量 ≥ seq.num_blocks
Prefillallocate(seq)分配 blocks, 填 seq.block_table, 设 seq.num_cached_tokens
Decode (每步)can_append(seq)检查当前 block 还能不能再写 1 个 token
Decode (每步)may_append(seq)block 满了固化 hash, 需要新 block 就分配
序列结束deallocate(seq)逆序释放 block, ref_count 为 0 时回收

can_allocate 和 can_append 是”判断”方法，只返回 bool 不改状态，给 Scheduler 做调度决策用——比如 can_append 返回 False 时，Scheduler 知道当前 free block 不够，会触发 preempt（驱逐其他序列）来腾出空间。

allocate / may_append / deallocate 是真正”动手”的方法。我们逐个看。

3.3.3 allocate：Prefill 阶段的 block 分配

allocate 是 BlockManager 最复杂、也最核心的方法。它做三件事：

  1. 按序列的 token 数算出需要多少个逻辑 block
  2. 对每个逻辑 block 检查 prefix cache 命中情况
  3. 命中就复用已有 block，不命中就从 free_block_ids 取新 block，最后填到 seq.block_table

源码（精简版）：

def allocate(self, seq: Sequence):
    assert not seq.block_table
    h = -1                # 上一块的链式 hash
    cache_miss = False    # 一旦某块未命中, 后续块全部走新分配 (前缀链断裂)

    for i in range(seq.num_blocks):
        token_ids = seq.block(i)
        # 满块才算 hash, 未满块 hash = -1 (不参与 cache)
        h = self.compute_hash(token_ids, h) if len(token_ids) == self.block_size else -1
        block_id = self.hash_to_block_id.get(h, -1)

        # 未命中, 或 hash 碰撞 (token 内容不一致)
        if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
            cache_miss = True

        if cache_miss:
            # 路径 1: ALLOCATE 新块
            block_id = self.free_block_ids[0]
            block = self._allocate_block(block_id)
        else:
            seq.num_cached_tokens += self.block_size
            if block_id in self.used_block_ids:
                # 路径 2: REUSE (block 还在被别的序列用 → ref_count++)
                block = self.blocks[block_id]
                block.ref_count += 1
            else:
                # 路径 3: REACTIVATE (block 已回到 free list, 但数据还在)
                block = self._allocate_block(block_id)

        # 满块要把 hash 注册到 hash_to_block_id, 供后续请求复用
        if h != -1:
            block.update(h, token_ids)
            self.hash_to_block_id[h] = block_id

        seq.block_table.append(block_id)

代码里有几个关键点值得展开：

1. 三种路径：ALLOCATE / REUSE / REACTIVATE

每个逻辑 block 处理后只可能落到三种路径之一：

  • ALLOCATE：prefix cache 没命中（最常见的情况）→ 从 free_block_ids 取新 block
  • REUSE：cache 命中且这个 block 还在 used_block_ids 里 → 直接 ref_count += 1，多条序列共享同一物理块
  • REACTIVATE：cache 命中但 block 已经回到 free_block_ids → 从 free list 重新激活（数据还在没被覆盖）

REACTIVATE 是 nano-vllm 的”惰性回收”策略带来的收益——block 被释放后 hash 条目不会立刻清除，如果新请求恰好匹配，可以直接复用之前算过的 K/V。这个机制 §3.4 详细讲。

2. cache_miss 一旦为 True 就不会变回 False

注意第 9 行：cache_miss = True，但循环里没有把它设回 False 的代码。这是故意的——前缀链一旦断裂，后面的 block 即使内容相同也不能复用。

为什么？因为 §2.2 讲过，链式 hash 的核心是”只有从开头到当前 block 的整个前缀都相同，hash 才会匹配”。如果第 0 个 block 不同，第 1 个 block 即使 token 完全相同，它的链式 hash 也会不同（因为 prefix hash 不一样），自然查不到。

3. 为什么命中后还要比较 token_ids

第 13 行 self.blocks[block_id].token_ids != token_ids 这个判断是防 hash 碰撞——xxhash 是 64 位的，理论上有极小的概率不同输入产生相同 hash，所以查到 block 后还要逐 token 比较，确认内容真的一致才算命中。

用真实日志验证

之前跑过一个实验：两条请求共享一段长 prompt（400+ token），block_size=256，所以前 256 token 的 block 0 是满块。日志（day2_block_experiment.log）：

# seq_id=4 (第一条请求, 406 tokens):
KV_BLOCK_ALLOCATE  block_id=0  logical_block_idx=0  seq_id=4
KV_BLOCK_ALLOCATE  block_id=1  logical_block_idx=1  seq_id=4
allocate seq_id=4 block_table=[0, 1] num_cached_tokens=0 free_blocks=2232

# seq_id=5 (第二条请求, 409 tokens):
KV_BLOCK_REUSE     block_id=0  ref_count=2          seq_id=5
KV_BLOCK_ALLOCATE  block_id=2  logical_block_idx=1  seq_id=5
allocate seq_id=5 block_table=[0, 2] num_cached_tokens=256 free_blocks=2231

逐行对照：

  • seq_id=4 是第一个进来的，hash_to_block_id 还是空的，两个 block 都走 ALLOCATE 路径。block 0 满了之后 hash 被注册
  • seq_id=5 来的时候，第 0 个逻辑块的 hash 命中了 seq_id=4 注册的那个 hash，且 block 0 还在 used 里 → 走 REUSE 路径，ref_count 从 1 变 2
  • seq_id=5 的第 1 个逻辑块只有 153 token（409 - 256 = 153），不满 256，hash=-1 不参与 cache，走 ALLOCATE 路径
  • 最终 seq_id=5 的 block_table=[0, 2]，block 0 是和 seq_id=4 共享的，block 2 是自己独占的
  • num_cached_tokens=256 表示这 256 个 token 的 KV 已经在 cache 里了，prefill 时不需要重新计算

3.3.4 may_append：Decode 阶段的 block 管理

allocate 处理的是 prefill 的”批量分配”。但 decode 是逐 token 生成的，每生成 1 个新 token 都可能涉及 block 边界——比如某个 block 刚好被填满，或者需要分配下一个 block。这就是 may_append 要处理的事情。

源码：

def may_append(self, seq: Sequence):
    block_table = seq.block_table
    last_block = self.blocks[block_table[-1]]

    if len(seq) % self.block_size == 1:
        # 情况 A: 刚跨入新块
        # 上一步刚好把上一个 block 写满了, 这一步的 token 是新 block 的第 1 个
        assert last_block.hash != -1
        block_id = self.free_block_ids[0]
        self._allocate_block(block_id)
        block_table.append(block_id)

    elif len(seq) % self.block_size == 0:
        # 情况 B: 当前块刚好写满
        # 这一步刚好把当前 block 填满了, 固化 hash, 注册到 hash_to_block_id
        assert last_block.hash == -1
        token_ids = seq.block(seq.num_blocks - 1)
        prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1
        h = self.compute_hash(token_ids, prefix)
        last_block.update(h, token_ids)
        self.hash_to_block_id[h] = last_block.block_id

    # 其他情况: 当前 block 没满也没刚满 → 什么都不做

may_append 处理两个边界情况，关键是用 len(seq) % block_size 来判断当前在 block 的什么位置。以 block_size=256 为例：

seq 长度len % 256触发情况操作
2560当前 block 刚好写满固化 block 的 hash, 注册到 hash_to_block_id
2571刚跨入新 block从 free_block_ids 分配新 block, 追加到 block_table
258, 259, …2, 3, …当前 block 还没满什么都不做
5120第 2 个 block 刚好写满固化第 2 个 block 的 hash
5131跨到第 3 个 block分配第 3 个 block

注意这里固化 hash 的时机：block 必须填满才算 hash。这是 prefix cache 设计上的一个权衡——未满的 block 可能还会继续写入新 token，hash 算出来过一会儿就失效，所以干脆等满了再算。

代价是：如果一个序列的最后一个 block 始终没填满（比如生成短文本），它的 hash 永远不会被注册，后续相同前缀的请求就无法复用这个 block 的 KV。

3.3.5 deallocate：序列结束的 block 释放

deallocate 处理两种触发场景：序列正常完成（生成 EOS 或达到 max_tokens），或者被 Scheduler preempt（驱逐）。

源码：

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)
            # _deallocate_block: 从 used_block_ids 移除, 加回 free_block_ids
    seq.num_cached_tokens = 0
    seq.block_table.clear()

deallocate 做的事很简单：逆序遍历 block_table，每个 block 的 ref_count 减 1，如果减到 0 就回收到 free_block_ids。

为什么逆序？因为前面的 block 可能被其他序列共享（比如 prefix cache 的情况），后面的 block 通常是独占的。先释放独占的、再处理共享的，逻辑上更清晰。

继续用之前两条共享 prefix 的请求做例子，看 ref_count 是怎么递减的。两条序列的 block_table 分别是 [0, 1] 和 [0, 2]，所以 block 0 的 ref_count=2（被两条序列共享）。

# seq_id=4 完成, 调用 deallocate(seq_4):
#   逆序处理 block_table=[0, 1]
#   block 1: ref_count 1 → 0  → 回收
KV_BLOCK_RELEASE  block_id=1  free_blocks=2232  seq_id=4
#   block 0: ref_count 2 → 1  → 不回收 (seq_5 还在用)
deallocate seq_id=4 free_blocks=2232

# seq_id=5 完成, 调用 deallocate(seq_5):
#   逆序处理 block_table=[0, 2]
#   block 2: ref_count 1 → 0  → 回收
KV_BLOCK_RELEASE  block_id=2  free_blocks=2233  seq_id=5
#   block 0: ref_count 1 → 0  → 回收!
KV_BLOCK_RELEASE  block_id=0  free_blocks=2234  seq_id=5
deallocate seq_id=5 free_blocks=2234

注意 seq_id=4 完成时只有一条 KV_BLOCK_RELEASE 日志（block 1），因为 block 0 的 ref_count 只是从 2 降到 1，没有真正回收。直到 seq_id=5 完成，block 0 的 ref_count 才从 1 降到 0，触发回收。

free_blocks 的数字也能验证整个过程：从 2231（两条序列都在跑）→ 2232（释放 block 1）→ 2233（释放 block 2）→ 2234（释放 block 0），最后回到初始的 2234。

值得注意的是：deallocate 没有清理 hash_to_block_id。block 0 已经从 used_block_ids 移到 free_block_ids，但它的 hash 还在 hash_to_block_id 里。这就是 §3.4 要讲的”惰性回收”——为下一次相同前缀的请求保留 REACTIVATE 的可能。

3.4 Prefix Caching：链式 hash 和 block 复用

§3.3 里我们多次遇到 “prefix cache 命中”、”REUSE”、”REACTIVATE”、”hash_to_block_id” 这些名词，但一直没仔细讲它们怎么工作。这一节专门讲 prefix caching 的实现。

为什么要做 prefix caching

prefix caching 要解决的是一个很常见的现实问题：多条请求共享相同的前缀。

想想真实的 LLM 服务场景：

  • 同一个 system prompt：一个聊天应用的所有用户请求都以相同的系统提示开头（”你是一个有帮助的 AI 助手，回答要简洁…“，经常几百上千 token）
  • few-shot 提示：很多应用会在 prompt 里放固定的示例，每条请求只有末尾的用户问题不同
  • 同一篇文档的多次提问：RAG 场景下，用户可能基于同一篇长文档问多个不同的问题

这些场景的共同点是：前缀完全相同，后缀不同。如果每条请求都从头算一遍 K/V，前缀部分的计算就是纯浪费——它们的 K/V 和别的请求是完全一样的。

Prefix caching 的思路很直接：检测到”这段前缀和之前某条请求一样”时，直接复用之前算过的 block，不重新算 K/V。§3.3 里的 num_cached_tokens=256 就是这么来的——表示前 256 个 token 的 KV 已经在 cache 里了，prefill 时可以跳过它们。

第一次尝试：对 block 内容算 hash（错误的方式）

判断”前缀相同”最直接的方法是逐 token 比较，但太慢——每次新请求进来都要扫所有历史请求的 token。更合理的做法是给每个 block 算一个 hash 指纹，用 hash 快速查找。

但这里有个陷阱。假设我们只对 block 的 token 内容 算 hash：hash(block) = xxhash(block.token_ids)。考虑这个场景：

请求 A 的 token: [10, 20, 30, 40, 50, 60]    block_size = 3
请求 B 的 token: [77, 88, 99, 40, 50, 60]

Block 分布:
  请求 A: Block 0 = [10, 20, 30]  Block 1 = [40, 50, 60]
  请求 B: Block 0 = [77, 88, 99]  Block 1 = [40, 50, 60]

两条请求的 Block 1 内容完全一样 [40, 50, 60]，所以它们的 hash 也一样。如果 B 查到”Block 1 的 hash 命中了”，直接复用 A 的 Block 1 的 KV——这会产生错误的计算结果。

为什么？回忆一下 Attention 的计算：token 40 的 K/V 不仅取决于它自己，还取决于它之前所有 token 的上下文。A 的 token 40 前面是 [10, 20, 30]，B 的 token 40 前面是 [77, 88, 99]，经过 Transformer 的所有层后，它们的 K/V 向量是完全不同的。

直接复用会让 B 在计算 Attention 时，用的是基于 A 的前缀算出来的 K/V，结果完全不对。

正确的方式：链式 hash

nano-vllm 用的是链式 hash：每个 block 的 hash 把前一个 block 的 hash 也算进去。这样只有”从开头到当前 block 的整个前缀都相同”时，hash 才会匹配。

代码在 block_manager.py 的 compute_hash：

@classmethod
def compute_hash(cls, token_ids: list[int], prefix: int = -1):
    h = xxhash.xxh64()
    if prefix != -1:
        h.update(prefix.to_bytes(8, "little"))   # 上一块的 hash 参与, 形成链
    h.update(np.array(token_ids).tobytes())      # 当前块的 token 内容
    return h.intdigest()

参数 prefix 就是上一个 block 的 hash 值。-1 表示”这是第一个 block，没有前驱”。

用上面的例子走一遍：

请求 A:
  hash_A0 = xxhash(-1, [10, 20, 30])
  hash_A1 = xxhash(hash_A0, [40, 50, 60])    ← 把 hash_A0 算进来

请求 B:
  hash_B0 = xxhash(-1, [77, 88, 99])         ← 和 hash_A0 不同 (token 不同)
  hash_B1 = xxhash(hash_B0, [40, 50, 60])    ← 即使 token 相同, 但 prefix 不同
                                                所以 hash_B1 ≠ hash_A1

虽然 A.Block1 和 B.Block1 的 token 内容完全相同，但它们的链式 hash 不同——因为 hash 的输入里带上了 prefix（上一个 block 的 hash）。这样 B 查 hash_to_block_id 时就不会错误地命中 A 的 Block 1。

只有”从开头到当前 block 的整个前缀完全相同”时，hash 才会匹配。这正好是 prefix cache 需要的语义。

为什么一定要 xxhash，不能用其他 hash

compute_hash 用的是 xxhash.xxh64，这是一个 非加密 的 64 位 hash 函数。和 MD5/SHA256 这类加密 hash 相比，xxhash 快得多（C 实现，一个 64 位 hash 几十纳秒就能算完），但没有密码学强度——理论上可能被恶意构造输入产生碰撞。

prefix caching 不需要密码学强度，只需要速度。每次 allocate 可能要算几十个 block 的 hash，如果用 SHA256 就会变成性能瓶颈。xxhash 是速度和碰撞率的合理权衡。

防碰撞的兜底在 allocate 里（§3.3.3 讲过）：hash 命中后还要比较 token_ids 是否完全一致：

if block_id == -1 or self.blocks[block_id].token_ids != token_ids:
    cache_miss = True   # 即使 hash 命中, token 内容不一致也算 miss

这样即使 xxhash 真的产生了碰撞，也不会导致错误复用。

hash 什么时候被注册

prefix cache 生效的前提是：block 的 hash 被注册到 hash_to_block_id 里。nano-vllm 只在两个时机注册 hash：

  1. allocate 过程中，每当一个满 block 被处理完（不管是 ALLOCATE 还是 REUSE/REACTIVATE），都会把它的 hash 写入 hash_to_block_id（block_manager.py 的 allocate 最后几行）：
if h != -1:   # 只有满块才有 hash
    block.update(h, token_ids)
    self.hash_to_block_id[h] = block_id
  1. may_append 的第二个分支，decode 过程中当前 block 刚好被填满时（len(seq) % block_size == 0），固化它的 hash：
elif len(seq) % self.block_size == 0:
    token_ids = seq.block(seq.num_blocks - 1)
    prefix = self.blocks[block_table[-2]].hash if len(block_table) > 1 else -1
    h = self.compute_hash(token_ids, prefix)
    last_block.update(h, token_ids)
    self.hash_to_block_id[h] = last_block.block_id

共同点是：只有被填满的 block 才会注册 hash。未满的 block hash=-1，不参与 prefix cache。

为什么要这个限制？因为未满的 block 还可能继续写入新 token——你在它还只有 100 个 token 时算 hash，等它写到 256 个 token 时 hash 就过期了。干脆等它填满再算，一次到位，避免反复更新。

代价是：如果一条序列的最后一个 block 始终没填满（比如生成了短文本），它永远不会参与 prefix cache。但这通常不是问题——用户关心的前缀（system prompt、RAG 上下文）一般都很长，都能填满多个 block。

惰性回收：让 REACTIVATE 成为可能

§3.3.5 结尾我埋了一个点：deallocate 没有清理 hash_to_block_id。现在来解释为什么。

考虑这个场景：

  1. 请求 A 来了，占用 block 7 存一段 system prompt 的 KV。满块，hash 被注册到 hash_to_block_id[h] = 7
  2. 请求 A 生成完成，deallocate 被调用，block 7 的 ref_count 从 1 降到 0，被放回 free_block_ids
  3. 这时 block 7 的数据还在——PyTorch tensor 的数据不会因为逻辑上的”回收”就自动清零
  4. 一段时间后，请求 B 来了，它的前缀和 A 相同
  5. allocate 算出的 hash 和当初 A 的一样，查 hash_to_block_id 命中 block_id=7
  6. 但此时 block 7 在 free_block_ids 里，不在 used_block_ids 里

这时候怎么办？allocate 走的就是 REACTIVATE 路径：

if block_id in self.used_block_ids:
    # REUSE: block 还在使用 → 直接 ref_count++
    block.ref_count += 1
else:
    # REACTIVATE: block 已回到 free list, 但数据还在 → 从 free list 重新激活
    block = self._allocate_block(block_id)

REACTIVATE 的好处是：请求 B 完全跳过了前缀部分的 prefill 计算，直接用了几分钟前请求 A 算好的 KV。这对”间歇性相同 system prompt”的场景非常有效——即使请求之间有间隔、block 被释放了，只要它还没被新内容覆盖，就能复用。

这就是为什么 deallocate 不清理 hash_to_block_id——故意保留这些 hash 条目，为未来可能的 REACTIVATE 留一条路。只有当 block 被重新分配给新序列时，_allocate_block 调用 block.reset() 才会把它的 hash 清掉（那时旧 hash 条目还在 hash_to_block_id，但指向的 block 已经存了新内容，下次查找时的 token_ids 比对会 miss，不会错误复用）。

严格说这是一个有限的”惰性”：旧 hash 条目永远不会被显式删除，只会被新 hash 条目替换（如果某个 block 被重新分配并填入新内容，会在 allocate 或 may_append 里重新注册 hash，覆盖旧条目）。长期运行下 hash_to_block_id 会一直增长，但由于 hash 是 64 位整数、value 是 int，单条条目大概几十字节，几十万条也就几 MB 内存，完全可以接受。

小结：prefix cache 的完整工作流

把整个 prefix caching 的机制串起来：

  1. 注册：当一个 block 被填满时（allocate 里命中的满块，或 may_append 里 decode 刚填满），用链式 hash 算出它的指纹，写入 hash_to_block_id

  2. 查找：新请求 allocate 时，对每个满块算链式 hash 查 hash_to_block_id——

    • 命中且 block 仍被使用 → REUSE（ref_count++）

    • 命中但 block 已回到 free list → REACTIVATE（从 free list 重新激活，数据还在）

    • 未命中 → ALLOCATE 新块

  4. 防碰撞：hash 命中后还要逐 token 比对，确认内容一致

  5. 前缀链约束：某个 block 一旦未命中（cache_miss = True），后续 block 即使内容相同也不能复用

  6. 惰性回收：deallocate 不清理 hash_to_block_id，为未来的 REACTIVATE 留机会

这五条加起来构成了一个完整的 prefix cache：快（xxhash + dict 查找）、正确（链式 hash + token 比对防误命中）、高效（惰性回收最大化复用机会）。

3.5 从 block_table 到 slot_mapping

§3.3 和 §3.4 看完了 BlockManager——它的产出是 seq.block_table，一个逻辑块 ID 的列表。但 store_kvcache Triton kernel 不认识 block_table，它只接受物理的 slot_mapping：一个 tensor，告诉 kernel”这批 token 的 K/V 要分别写到 kv_cache 的第几行”。

从 block_table 到 slot_mapping 的转换，就发生在 §3.1 三层架构图里的 映射层——model_runner.py 的 prepare_prefill 方法里。这一节把这个转换的每一步看清楚。

prepare_prefill 的职责

先从全局定位一下 prepare_prefill。上一篇博客讲 ModelRunner.run 的时候，我们知道每一步 forward 的流程大致是：

Scheduler.schedule() → 选出一批 seqs
  ↓
ModelRunner.prepare_prefill(seqs)   ← 把 seqs 转换成 GPU 需要的 tensor
  ↓
ModelRunner.run_model()             ← 执行 forward
  ↓
sampler → token_ids

prepare_prefill 的职责是：接收 Scheduler 选出来的一批 sequence，把它们打平成 GPU 需要的 input tensors。具体来说，它的输出包括：

输出用途
input_ids本步要计算的新 token 序列（对每条 seq 只取未 cached 的部分）
positions每个 token 的绝对位置索引（供 RoPE 使用）
cu_seqlens_qQ 侧的累积序列长度（给 FlashAttention）
cu_seqlens_kK 侧的累积序列长度（给 FlashAttention）
max_seqlen_q / max_seqlen_kFlashAttention 需要的最大长度
slot_mapping每个新 token 的 KV 写到 kv_cache 的哪一行
block_tables当 batch 内存在 prefix cache 时，告诉 FlashAttention 从 paged cache 读 K/V

这一节我们聚焦 slot_mapping（顺带讲 cu_seqlens_q 和 cu_seqlens_k 的差异，因为它们和 prefix cache 紧密相关）。input_ids 和 positions 的构造比较直观，相对次要。

核心公式和遍历规则

slot_mapping 的计算源码（model_runner.py 的 prepare_prefill 内循环）：

# 对 batch 里的每条 seq
for seq_idx, seq in enumerate(seqs):
    ...
    # 关键循环: 只遍历 [num_cached_blocks, num_blocks) 范围内的逻辑块
    seq_slot_mapping = []
    for i in range(seq.num_cached_blocks, seq.num_blocks):
        start = seq.block_table[i] * self.block_size
        if i != seq.num_blocks - 1:
            end = start + self.block_size
        else:
            end = start + seq.last_block_num_tokens   # 末块可能未满
        seq_slot_mapping.extend(list(range(start, end)))
    slot_mapping.extend(seq_slot_mapping)

这段代码不长，但有两个关键设计：

1. 只遍历未 cached 的逻辑块

range(seq.num_cached_blocks, seq.num_blocks) 是整段代码的灵魂。num_cached_blocks 是”前 N 个逻辑块已经在 cache 里了”，对这些 block 完全跳过——它们的 K/V 已经在 kv_cache 里，不需要重新写入。这就是 prefix cache 在 prepare 阶段的体现。

2. 末块特殊处理

最后一个 block 可能没填满。比如 seq 有 430 个 token、block_size=256，那么：

  • Block 0: token 0-255（满，256 个 slot）
  • Block 1: token 256-429（未满，只有 174 个 slot）

所以代码里对末块用了 seq.last_block_num_tokens 而不是 block_size。

核心公式 就是 §2.2 讲过的那个：

slot = block_table[i] × block_size + offset

其中 offset 是 token 在 block 内的位置（0 到 block_size-1）。

用真实例子走一遍

我们之前跑 blog02_e2e_trace.log 的时候，给 prepare_prefill 加了逐 seq 的详细日志。两条共享前缀的请求（seq_id=4 和 seq_id=5）产生了这样的输出：

prepare_prefill seq[0] seq_id=4:
  seqlen=430  num_cached_tokens=0
  seqlen_q=430  seqlen_k=430
  block_table=[0, 1]
  slot_mapping(前10)=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
  slot_mapping_len=430
  cu_seqlens_q=[0, 430]
  cu_seqlens_k=[0, 430]

prepare_prefill seq[1] seq_id=5:
  seqlen=433  num_cached_tokens=256
  seqlen_q=177  seqlen_k=433
  block_table=[0, 2]
  slot_mapping(前10)=[512, 513, 514, 515, 516, 517, 518, 519, 520, 521]
  slot_mapping_len=177
  cu_seqlens_q=[0, 430, 607]
  cu_seqlens_k=[0, 430, 863]

下面这张图把这两条 seq 的 slot_mapping 计算过程可视化了：

逐条对照：

seq_id=4（430 tokens，无 prefix cache）

  • num_cached_tokens = 0 → num_cached_blocks = 0
  • 遍历 range(0, 2)，也就是 Block 0 和 Block 1 都要算
  • Block 0（满）：start = 0 × 256 = 0，end = 0 + 256 = 256 → slots [0, 1, ..., 255]
  • Block 1（未满，174 个 token）：start = 1 × 256 = 256，end = 256 + 174 = 430 → slots [256, 257, ..., 429]
  • 合起来 slot_mapping = [0, 1, ..., 429]，共 430 个

日志里 slot_mapping_len=430 和手算对得上。seqlen_q=seqlen_k=430 也符合——全部 token 都是新的，没有 cached 部分。

seq_id=5（433 tokens，前 256 token 命中 cache）

  • num_cached_tokens = 256 → num_cached_blocks = 1（第 0 个 block 已在 cache）
  • 遍历 range(1, 2)，跳过 Block 0，只处理 Block 1
  • Block 1（未满，177 个 token = 433 - 256）：start = 2 × 256 = 512（注意 block_table[1] = 2），end = 512 + 177 = 689 → slots [512, 513, ..., 688]
  • 合起来 slot_mapping = [512, 513, ..., 688]，共 177 个

日志里 slot_mapping_len=177 和手算对得上。slot_mapping 的第一个值 不是 0 而是 512，因为 Block 0 被整个跳过了，从 Block 2（在 block_table 里索引为 1，物理 block_id 为 2）开始写。

到这里你应该能看到 prefix cache 对 slot_mapping 的影响了：有 cache 时 slot_mapping 会短得多（seq_id=5 只有 177 个 slot 对应 433 个 token，其中 256 个被省掉了）。这直接对应到 prefill 时 store_kvcache kernel 的写入量——只写新计算的那些 token 的 K/V，cached 部分碰都不碰。

cu_seqlens_q 和 cu_seqlens_k：prefix cache 的”指纹”

继续看日志，除了 slot_mapping 还有一对值得研究的 tensor：cu_seqlens_q 和 cu_seqlens_k。

seq_id=4:  seqlen_q=430  seqlen_k=430
seq_id=5:  seqlen_q=177  seqlen_k=433

最终:
  cu_seqlens_q = [0, 430, 607]    (430 + 177)
  cu_seqlens_k = [0, 430, 863]    (430 + 433)

cu_seqlens_* 是 FlashAttention 的变长序列接口（flash_attn_varlen_func）需要的参数，表示 batch 里每条序列在展平后的起止位置。

关键区别在于 seqlen_q 和 seqlen_k 不一定相等：

  • seqlen_q 表示”本步 Q 要计算多少个 token”——只算需要新生成 K/V 的 token，即 seqlen - num_cached_tokens
  • seqlen_k 表示”本步 Attention 的 K/V 覆盖多少个 token”——整段序列的长度，seqlen

对 seq_id=5 来说：

  • seqlen_q = 433 - 256 = 177（Q 只算后 177 个新 token）
  • seqlen_k = 433（K 要覆盖全部 433 个 token，前 256 个从 cache 读）

这个差异是 prefix cache 的数学本质：cached 部分的 token 已经在 cache 里了，不需要再算它们的 Q 和 K；但新 token 的 Attention 需要对所有历史 token 做 Q@K^T，所以 K 侧还是要看到完整的历史。

在 prepare_prefill 末尾有一个判断：

if cu_seqlens_k[-1] > cu_seqlens_q[-1]:
    block_tables = self.prepare_block_tables(seqs)

当 cu_seqlens_k[-1] > cu_seqlens_q[-1]（K 总长 > Q 总长）时，说明 batch 里至少有一条 seq 用了 prefix cache。这时候要额外产出 block_tables——FlashAttention 需要通过它从 paged KV cache 读取 cached 部分的 K/V。

在我们的例子里：

cu_seqlens_q[-1] = 607   ← 本步 Q 总共 607 个 token
cu_seqlens_k[-1] = 863   ← 本步 K 总共 863 个 token
              |    |
           607 < 863  →  存在 prefix cache  →  prefix_block_tables = True

这就是日志最后一行 prefix_block_tables=True 的来源。

把 prepare_prefill 的输出串起来

最后把 prepare_prefill 的几个关键输出梳理一下，看看它们分别被后续的谁用：

输出被谁用用途
input_idsQwen3Model.forward 的 embedding 层查 embedding 得到 hidden states
positionsRoPE计算每个 token 的旋转位置编码
cu_seqlens_q / cu_seqlens_kflash_attn_varlen_func告诉 kernel 每条序列的边界
slot_mappingstore_kvcache_kernel（Triton）告诉 kernel 每个新 token 的 KV 写到哪一行
block_tablesflash_attn_varlen_func有 prefix cache 时，从 paged cache 读 K/V

这些 tensor 一起通过 set_context() 存到一个线程局部的 Context 对象里，Attention.forward 的时候直接 get_context() 拿出来用。这样模型的 forward 函数本身不需要多传参数，也不需要感知 batch 里有没有 prefix cache——所有调度信息都藏在 Context 里。

到这里 §3 整个”nano-vllm 的 KV Cache 和 Paged Attention 实现”就讲完了。我们从三层架构的全局视角开始，看了 kv_cache tensor 的初始化（§3.2）、BlockManager 的数据结构和三个核心方法（§3.3）、prefix caching 的链式 hash 机制（§3.4），最后看了 prepare_prefill 如何把 BlockManager 的产出翻译成 GPU kernel 需要的物理地址（§3.5）。

下一节我们回到真实运行的完整日志，把这些零散的机制串起来，看一条请求从入队到结束的完整生命周期。

4. 用真实日志串起来整个流程

§3 我们逐个看完了 nano-vllm 的每一个零件——KV cache 初始化、BlockManager、prefix caching、slot_mapping。但零件是零件，整机是整机。这一节我们把所有零件串起来，用一条实际跑出来的完整日志，看一次推理从引擎启动到请求完成的全过程。

这个实验我们在前面的章节里已经零零散散引用过几次了，这里系统地走一遍完整的 timeline。

4.1 实验设置

两条共享长前缀的请求发给 nano-vllm，看 BlockManager / prepare_prefill / kv_cache 是怎么协作的。

环境：

  • 模型：Qwen3-0.6B（bfloat16）
  • GPU：NVIDIA H20（95 GiB）
  • tensor_parallel_size=1
  • enforce_eager=True（跳过 CUDA Graph 以便观察真实 forward）
  • gpu_memory_utilization=0.9
  • kvcache_block_size=256
  • max_num_batched_tokens=16384
  • max_model_len=4096

请求：

  • seq_id=4：430 个 token，max_tokens=8
  • seq_id=5：433 个 token，max_tokens=8
  • 两条 prompt 共享一段超过 256 token 的长前缀（system prompt），只有末尾问题不同

运行命令：

CUDA_VISIBLE_DEVICES=4 python example.py \
  --model ~/huggingface/Qwen3-0.6B/ \
  --tp 1 --enforce-eager \
  --log-level DEBUG --trace-teaching-mode --trace-tensor-shapes \
  --max-tokens 8

完整日志保存在 docs/plan/log/blog02_e2e_trace.log。下面分阶段摘取关键行。

4.2 阶段 1：引擎初始化

从进程启动到”引擎就绪”这段，约 4 秒。日志里关键的几行：

11:21:16  LLMEngine init: model=Qwen3-0.6B  tensor_parallel_size=1  enforce_eager=True
          max_num_seqs=512  max_num_batched_tokens=16384  max_model_len=4096
          gpu_memory_utilization=0.9  kvcache_block_size=256

11:21:17  TP_RUNTIME_INIT  phase=init  stats={device=cuda:0, rank=0, world_size=1}

11:21:17  load_model: 自目录 Qwen3-0.6B/ 读取 1 个 .safetensors 文件, 共 311 个权重 key

11:21:17  MODEL_WEIGHTS_LOADED  stats={num_keys=311}
11:21:17  rank0: 模型权重已加载至 GPU  dtype=torch.bfloat16

11:21:17  warmup_model: 用假数据跑一轮最大规模 prefill
11:21:17  prepare_prefill: flat_len=16384 num_seqs=4 max_seqlen_q=4096 max_seqlen_k=4096
11:21:17  run_model: eager (prefill=True enforce_eager=True bs=16384)

11:21:19  KV_BUDGET_ESTIMATE  stats={block_bytes=29360128, num_blocks=3053,
                                     total_gib=95.085, free_gib=93.43, peak_gib=1.599}
11:21:19  allocate_kv_cache: total_mem=95.09GiB  free=93.43GiB  used=1.66GiB
          peak_alloc=1.60GiB  curr_alloc=1.16GiB  gpu_memory_utilization=0.9
          -> num_kvcache_blocks=3053  block_size=256  block_bytes≈29360128

11:21:19  跳过 CUDA Graph (enforce_eager=True), decode 全程 eager
11:21:20  引擎就绪: num_kvcache_blocks=3053

这段日志完美对应 §3.2 讲的初始化流程：

  1. 11:21:17 load_model：从 safetensors 读取 311 个权重 key 到 GPU（约 1.2 GiB）
  2. 11:21:17 warmup_model：构造一批最大规模的假请求（flat_len=16384），跑一次 prefill 摸峰值
  3. 11:21:19 KV_BUDGET_ESTIMATE：num_blocks = int(95.085×0.9 - 1.66 - 1.60 + 1.16) / (29 MiB) ≈ 3053，和 §3.2 手算对得上
  4. 11:21:19 allocate_kv_cache：一次性分配 3053 个 block 的大 tensor（约 83.5 GiB）
  5. 11:21:20 引擎就绪：初始化完成，开始进入调度主循环

注意 warmup 这一步耗时不少（约 2 秒），因为它真的跑了一次 bs=16384 的 prefill——这个开销是必要的，保证了 peak_alloc 是”可能发生的最坏情况”。

4.3 阶段 2：Prefill — Block 分配与 Prefix Cache

接下来两条请求入队。看 add_request 到 step#1 prefill done 这段：

11:21:20  generate: 入队 2 条请求, 循环 step 直至 waiting/running 皆空
11:21:20  add_request  seq_id=4  prompt_len=430  max_tokens=8  → waiting 队尾
11:21:20  add_request  seq_id=5  prompt_len=433  max_tokens=8  → waiting 队尾

11:21:20  KV_BLOCK_ALLOCATE  block_id=0  logical_block_idx=0  seq_id=4
11:21:20  KV_BLOCK_ALLOCATE  block_id=1  logical_block_idx=1  seq_id=4
11:21:20  allocate seq_id=4  block_table=[0, 1]  num_cached_tokens=0  free_blocks=3051

11:21:20  KV_BLOCK_REUSE     block_id=0  ref_count=2  seq_id=5
11:21:20  KV_BLOCK_ALLOCATE  block_id=2  logical_block_idx=1  seq_id=5
11:21:20  allocate seq_id=5  block_table=[0, 2]  num_cached_tokens=256  free_blocks=3050

11:21:20  [调度] engine_step=1 phase=prefill n_seqs=2 new_tokens_in_forward=607
          seq_ids=[4, 5]  cached_prefix_tokens=[0, 256]

11:21:20  prepare_prefill seq[0] seq_id=4: seqlen=430 num_cached_tokens=0
            seqlen_q=430 seqlen_k=430 block_table=[0, 1]
            slot_mapping(前10)=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9] slot_mapping_len=430
            cu_seqlens_q=[0, 430] cu_seqlens_k=[0, 430]

11:21:20  prepare_prefill seq[1] seq_id=5: seqlen=433 num_cached_tokens=256
            seqlen_q=177 seqlen_k=433 block_table=[0, 2]
            slot_mapping(前10)=[512, 513, 514, 515, 516, 517, 518, 519, 520, 521]
            slot_mapping_len=177
            cu_seqlens_q=[0, 430, 607] cu_seqlens_k=[0, 430, 863]

11:21:20  prepare_prefill: flat_len=607 num_seqs=2
          cu_seqlens_q[-1]=607 cu_seqlens_k[-1]=863 prefix_block_tables=True

11:21:21  run_model: eager (prefill=True bs=607)
11:21:21  step#1 prefill done in 851.23ms  num_tokens_stat=865

每一行都能对上 §3 讲过的机制。逐层拆解：

1. BlockManager.allocate 的三种路径都出现了

seq_4 是第一个到的，hash_to_block_id 还是空的，两个逻辑 block 都走 ALLOCATE 路径（§3.3.3 的路径 1）：

  • block_id=0 分配给 seq_4 的第 0 个逻辑块（满 256 个 token，hash 被注册）
  • block_id=1 分配给 seq_4 的第 1 个逻辑块（174 个 token，未满，不注册 hash）
  • 结果：block_table=[0, 1]，num_cached_tokens=0

seq_5 到的时候，它的第 0 个逻辑块的 token 内容和 seq_4 的完全一样（共享前缀），链式 hash 命中 hash_to_block_id 里 seq_4 注册的条目。block 0 还在 used_block_ids 里，走 REUSE 路径（§3.3.3 的路径 2）：

  • block_id=0 的 ref_count 从 1 变成 2
  • 但 seq_5 的第 1 个逻辑块内容和 seq_4 不同（用户问题不同），cache_miss = True，走 ALLOCATE 分配新的 block_id=2
  • 结果：block_table=[0, 2]，num_cached_tokens=256（前 256 个 token 的 KV 已经在 cache 里，不用重算）

free_blocks 从 3053 先降到 3051（seq_4 占了 2 个），再降到 3050（seq_5 又占了 1 个新 block，因为 block 0 是共享的不计入新分配）。

2. prepare_prefill 的 slot_mapping 和 cu_seqlens 对应 §3.5 讲的所有内容

seq_4 没 cached，slot_mapping 从 0 开始（block_table[0]=0 → 0×256+0=0），共 430 个 slot。

seq_5 有 256 个 cached token，slot_mapping 从 512 开始——不是 0！因为第 0 个逻辑块被跳过（已在 cache），从第 1 个逻辑块开始写，而 block_table[1]=2，所以 start = 2×256 = 512。slot_mapping_len=177（433-256=177），只写新 token 的 KV。

cu_seqlens_q[-1]=607 < cu_seqlens_k[-1]=863，Q 总共 607 个新 token（430+177），K 覆盖 863 个 token（430+433）。prefix_block_tables=True 告诉 FlashAttention “这次调用要从 paged cache 读一部分 K/V”。

3. step#1 prefill 耗时 851 ms

这一步发生了：

  • forward 28 层 Qwen3DecoderLayer
  • 每层做 QKV projection → RoPE → store_kvcache Triton kernel 写入 kv_cache → flash_attn_varlen_func 读 kv_cache 做 attention → o_proj → MLP
  • 最后在 rank 0 做 sampling，采出两条序列的第一个 token

num_tokens_stat=865 表示本步统计意义上处理了 865 个 token（607 + 2 × 128 左右，含 overhead）。实际参与 Q 计算的是 607 个——seq_5 前 256 token 的 Q 被彻底省掉了。

4.4 阶段 3：Decode — 逐 token 生成

Prefill 结束后进入 decode 循环。每一步 decode 的结构都一样，我们先看第一次：

11:21:21  [调度] engine_step=2 phase=decode decode_step_index=0 n_seqs=2 seq_ids=[4, 5]
11:21:21  SCHEDULE_DECODE  scheduled=2  running_len=2  waiting_len=0

11:21:21  PREPARE_DECODE_TENSORS  batch_size=2  block_tables_shape=(2, 2)
                                  context_lens_shape=(2,)  slot_mapping_shape=(2,)
11:21:21  decode 详单#1/3: bs=2 input_ids.shape=(2,) context_lens=[431, 434]
          slot_mapping=[430, 689]  block_tables.shape=(2, 2)

11:21:21  run_model: eager (prefill=False bs=2)
11:21:21  step#2 decode done in 65.21ms

关键信息：

  • batch_size=2：两条序列合并成一个 batch 做 forward，这是 continuous batching 的核心

  • context_lens=[431, 434] ：两条序列当前的总长度（seq_4 的 430 + 1 个新 token = 431，seq_5 的 433 + 1 = 434）

  • slot_mapping=[430, 689] ：这一步每条序列只新写 1 个 slot

    • seq_4 的新 slot：block_table[-1]=1 → 1×256+174 = 430（Block 1 的第 174 位，接在前 430 token 之后）

    • seq_5 的新 slot：block_table[-1]=2 → 2×256+177 = 689（Block 2 的第 177 位）

  • block_tables.shape=(2, 2) ：batch 里 2 条序列，每条有 2 个逻辑 block，这个 tensor 交给 flash_attn_with_kvcache 让它从 paged cache 读 K/V

注意 decode 用的是 flash_attn_with_kvcache 而不是 prefill 的 flash_attn_varlen_func——§3.5 里我们列过这两个 API 的区别。decode 时 Q 只有每条序列 1 个 token（总共 bs×1），K/V 完全从 kv_cache 读，所以接口不同。

第一次 decode 耗时 65 ms 比较长，因为首次调用 flash_attn_with_kvcache 需要分配一些内部 buffer 和选择算法。后面稳定下来：

11:21:21  step#3 decode done in 20.49ms
11:21:21  step#4 decode done in 20.56ms
11:21:21  step#5 decode done in 20.69ms
11:21:21  step#6 decode done in 20.40ms
11:21:21  step#7 decode done in 20.48ms

稳态每步约 20 ms，7 步 decode 总共约 140 ms——比 prefill 的 851 ms 少得多。但注意 prefill 一次处理了 607 个 token，decode 每步只处理 2 个 token（batch=2）：

prefill 吞吐: 607 token / 851 ms ≈ 713 token/s
decode 吞吐: 2 token / 20 ms = 100 token/s (per step, batch=2)

这完全符合上一篇博客里讲过的 prefill / decode 性能特性——prefill 是 compute-bound（大矩阵乘），decode 是 memory-bound（每步只算 1 个 token，但要把整个 kv_cache 过一遍） 。

整个 decode 阶段 BlockManager 完全没变——block_table 不扩展，free_blocks 不变。原因：这 7 步总共给每条序列增加 7 个 token，两条序列分别从 430 → 437 和 433 → 440，都没跨过 256 的 block 边界（Block 1 有 174 个 token，再加 7 = 181，离 256 还远）。所以 may_append 每步的判断都落到”当前 block 没满也没刚满”的分支，什么都不做。

4.5 阶段 4：序列完成 — Block 释放

第 8 步 decode 是最后一步，两条序列同时达到 max_tokens=8：

11:21:21  step#8 sampler: token_ids=[13, 13] seq_ids=[4, 5]

11:21:21  KV_BLOCK_RELEASE  block_id=1  free_blocks=3051  seq_id=4
11:21:21  deallocate seq_id=4  free_blocks=3051
11:21:21  STATE_TRANSITION  seq_id=4  RUNNING → FINISHED  num_completion_tokens=8

11:21:21  KV_BLOCK_RELEASE  block_id=2  free_blocks=3052  seq_id=5
11:21:21  KV_BLOCK_RELEASE  block_id=0  free_blocks=3053  seq_id=5
11:21:21  deallocate seq_id=5  free_blocks=3053
11:21:21  STATE_TRANSITION  seq_id=5  RUNNING → FINISHED  num_completion_tokens=8

11:21:21  step#8 decode done in 20.56ms  finished_this_step=[4, 5]
11:21:21  generate: 全部完成, 返回 2 条结果

deallocate 的顺序和 §3.3.5 讲的完全一致：逆序遍历 block_table，ref_count 为 0 才回收。

seq_4 先结束（block_table=[0, 1] 逆序）：

  • block 1：ref_count 从 1 降到 0 → KV_BLOCK_RELEASE，free_blocks 从 3050 → 3051
  • block 0：ref_count 从 2 降到 1 → 不回收（seq_5 还在用），所以日志里没有 block 0 的 RELEASE

seq_5 接着结束（block_table=[0, 2] 逆序）：

  • block 2：ref_count 从 1 降到 0 → KV_BLOCK_RELEASE，free_blocks 3051 → 3052
  • block 0：ref_count 从 1 降到 0 → 这次真的回收，free_blocks 3052 → 3053

整个生命周期结束，free_blocks 从 3050 恢复到 3053（初始值），正好对得上。

注意一个细节：hash_to_block_id 里 block 0 的 hash 条目没被删。这是 §3.4 讲的惰性回收策略——如果下一秒又来一条相同前缀的请求，它还能通过 hash 查到 block 0，直接走 REACTIVATE 路径复用数据（只要这期间 block 0 没被重新分配覆盖）。

4.6 完整生命周期时序图

把整个过程按时间顺序画成时序图，每个关键事件对应日志中的一条记录：

几个值得注意的全局观察：

  1. 总耗时分布：初始化约 4 秒（其中 warmup 占 2 秒），prefill 851 ms，7 步 decode 总共约 160 ms。初始化的 4 秒是一次性成本，生产环境常驻服务会摊销掉。真正服务请求的是 prefill + decode 这 1 秒左右
  2. KV Cache 的”固定”特性：从引擎就绪到所有序列结束，kv_cache tensor 这块 83.5 GiB 的显存从未被释放或重新分配。整个过程 BlockManager 只是在逻辑层面更新 free_block_ids / used_block_ids / hash_to_block_id 三个 Python 数据结构，物理显存一动没动
  3. Prefix cache 的实际收益：这次实验中 seq_5 通过 REUSE 省掉了前 256 个 token 的 prefill 计算。按 §1 算过的 QKV 投影开销，这相当于省掉了 256 × 28 层 × 4 次 F.linear 的计算量——在 cu_seqlens_q=[0, 430, 607] 而不是 [0, 430, 863] 这个数字上直接体现出来
  4. 从代码到日志的完整对应：§3 里讲的每一个方法（warmup_model、allocate_kv_cache、allocate、compute_hash、prepare_prefill 里的 slot 循环、may_append、deallocate）都能在日志里找到对应的一行或几行。如果你想验证对源码的理解是否正确，最快的办法就是读自己跑出来的 DEBUG 日志

结语

写完这一篇，我自己对 KV Cache 和 Paged Attention 的理解比看了十几篇博客文章都要深。原因不是那些文章不好——它们大多讲得很清楚。而是只有一行一行读代码、一条一条看日志，才能建立起对”这套机制到底是怎么运行的”的真正直觉。

回过头看这一篇讲了什么：

从第一性原理出发，我们先想明白为什么需要 KV Cache（没它的话 decode 每步都要重算 O(L²) 的历史 K/V），然后看到 KV Cache 本身带来了新问题（显存占用巨大，还有碎片化），进而引出 PagedAttention 这套把操作系统分页思想搬到 GPU 显存上的解决方案。搞清楚理论之后，我们进入 nano-vllm 的代码：从三层架构的全局视角开始，逐层看 kv_cache tensor 的初始化（warmup 摸峰值、一次性预分配）、BlockManager 的内部设计（数据结构、三种分配路径、引用计数）、prefix caching 的链式 hash 机制（为什么要把前一个 block 的 hash 算进去、惰性回收让 REACTIVATE 成为可能）、以及 prepare_prefill 怎么把逻辑 block_table 翻译成 GPU kernel 能用的物理 slot_mapping。最后用一条真实跑出来的日志，把所有这些零件串成了一次完整的推理生命周期。

如果上一篇博客讲的是 nano-vllm 的”骨架”——LLMEngine / Scheduler / ModelRunner 的主循环、continuous batching 的调度逻辑，那这一篇讲的就是”血肉”——KV Cache 这块 83.5 GiB 的显存是怎么被有效管理、复用、回收的。骨架决定了什么时候做什么，血肉决定了每次做的时候有没有在浪费显存。

不过这一篇也留下了几个黑盒，需要后续专门的篇章来拆：

  • Attention 的内部数学：Q @ K^T → softmax → @V 的完整计算过程、causal mask、GQA 的共享结构。这部分和 rotary_embedding.py 一起，下一篇单独讲
  • FlashAttention 的内部实现：tiling、online softmax、SRAM vs HBM 的数据流。flash_attn_varlen_func 和 flash_attn_with_kvcache 内部到底在做什么
  • Triton kernel 的编程模型：store_kvcache_kernel 的代码我们这一篇贴过了，但没展开讲 tl.program_id / tl.load / tl.store 到底怎么映射到 GPU 的 SM 和线程。这部分和 CUDA kernel 基础一起讲
  • Tensor Parallel 下的 KV Cache 切分：多卡时 num_kv_heads 按 rank 切分，每个 rank 的 kv_cache 只存自己的那部分 KV head。这和 TP 通信机制一起讲

按写作进度，下一篇大概率先讲 Transformer 的单层 forward 和 Attention 数学——读懂 qwen3.py 的 Qwen3DecoderLayer.forward 和 attention.py 的三个 FlashAttention 分支。

最后想说一句：这篇文章里所有的代码引用、日志摘录、手算数字都可以复现。如果你跟着跑一遍 example.py、开 DEBUG 日志、自己对着 block_manager.py 和 model_runner.py 走一遍源码，应该能比只读这篇文章更深地理解 PagedAttention 的工作方式。这也是我建议每个想入门推理引擎的工程师都应该做一遍的事：找一个教学级实现、开 DEBUG 日志跑一次、把日志和源码对起来读。nano-vllm 是目前最适合做这件事的项目之一。

参考材料

nano-vllm 源码

  • nano-vllm GitHub 仓库 — Xingkai Yu 的教学级实现, MIT License. 本文所有代码引用均来自此项目
  • 本文重点涉及的文件：

    • nanovllm/engine/block_manager.py — §3.3, §3.4

    • nanovllm/engine/model_runner.py — §3.2, §3.5 的 warmup_model / allocate_kv_cache / prepare_prefill

    • nanovllm/layers/attention.py — §3.1 三层架构里的 store_kvcache 和 FlashAttention 分派

PagedAttention 论文

  • Kwon et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention”, SOSP 2023. arXiv:2309.06180
  • 本文引用的 Figure 3 (KV Cache memory waste)、Figure 6 (Block table translation) 和 Figure 7 均来自该论文
  • 如果你想了解 PagedAttention 的完整设计动机（尤其是碎片化问题的量化分析）和在 Parallel Sampling / Beam Search 等复杂场景下的应用，这篇论文是最好的一手资料

vLLM 官方博客

上一篇博客（如果你还没读）

致谢：再次感谢 nano-vllm 作者 Xingkai Yu 提供了这个教学级的推理引擎实现。如果没有 nano-vllm 这种”刚好能看懂”的代码量（约 1600 行）和精心设计的教学结构，直接读 vLLM 源码对入门者来说门槛会高得多。
另外感谢 vLLM 团队开源了论文和动图素材，让我能在这篇博客里直接引用——权威出处往往比博客作者的二次转述更可信。

avatar
文章
阅读
粉丝