Nano-vLLM 源码解读 - 3. PagedAttention

nano-vllm 用千行代码拆解 vLLM 核心，是读懂大模型推理最快的捷径。

这是「Nano-vLLM 源码解读」第 3 讲。承接前两讲——L1 给出系统全景、L2 把 Sequence 这个数据结构吃透——这一讲专攻 nano-vllm 也是 vLLM 系列的灵魂：PagedAttention。它是把 KV Cache 当成"操作系统的分页内存"来管理的一套机制，决定了引擎的吞吐上限和可调度性。

读完这一讲，你能：

• • 说清楚朴素连续 KV Cache 的三大痛点，以及它们各自为什么严重
• • 解释 PagedAttention 的两级映射：物理块、逻辑序列、block_table、slot_mapping 各自的职责
• • 看懂 model_runner.py 里 allocate_kv_cache / prepare_prefill / prepare_decode 是怎么把这套机制串起来的
• • 评估 block_size = 256 这个超参数在元数据成本和内部碎片之间的取舍

1. 朴素 KV Cache 的三大痛点

如果完全按"训练框架的思路"分配 KV Cache，最自然的做法是：

# 朴素版：每条序列预留 max_seq_len 的连续 KVkv_cache = torch.empty(    num_seqs, max_seq_len,    num_layers, num_kv_heads, head_dim,    dtype=torch.bfloat16,)

这样一来三个问题立刻浮上来：

痛点 1：必须按 max_seq_len 预留。 一条 prompt 最终可能只产出 200 个 token，但你不得不为它预留 8192 个 token 的空间——因为没法预知。这意味着实际的有效利用率经常只有 5%–20%。

痛点 2：碎片化无法重用。 当 seq A 跑到第 100 个 token 就结束了，理论上它后面 8092 个 token 的空间应该空出来给别人。但如果别的序列已经在它"右边"开始预留连续空间，这块空闲就成了不可用的间隙。即使没人挤进来，这块空闲也只够装下"长度 ≤ 8092 的新序列"——一个 16k 的长 prompt 完全装不下，即使总空闲量够。

痛点 3：抢占（preempt）无法实现。 当 GPU 显存压力大、需要把某个低优先级序列踢出去时，你想做的事是"释放它的 KV"。如果 KV 在大数组里是一段连续的内存，释放等于在中间挖一个洞——其它人也用不上。换句话说，连续布局让"释放" 和"碎片"变成同一个问题。

这三个痛点合起来意味着：朴素布局的最大并发数 ≈ 显存 / max_seq_len 的 KV 体积。在 16k context、bf16、num_layers=32 的设定下，单条序列的 KV 占 ~2GB，一张 80GB 卡装不到 40 条并发——而真实的请求平均长度可能只有几百 token，理论上能装上千条。

2. PagedAttention 的核心思想：操作系统启发

vLLM 的设计灵感来自虚拟内存：

核心动作：把"按序列连续分配"换成"按固定大小的块分配，由 block_table 间接寻址"。 序列对自己的 KV 仍然有"逻辑连续"的视角（token 0 → token 1 → ...），但物理上每 256 个 token 是一块，块和块之间在显存里可能完全分散。

这一改动直击三大痛点：

• • 痛点 1（按 max_seq_len 预留）→ 解决：只在序列真的写到第 N 个 token 时才分配第 ⌈N/256⌉ 个块
• • 痛点 2（碎片）→ 解决：所有空闲块都来自同一个池，释放后任何序列都能取来用
• • 痛点 3（抢占）→ 解决：抢占 = 把这个序列的所有块归还池子，逻辑上简单，物理上彻底干净

下面这张图把"为什么 paged"一次性讲完——上边是朴素布局逐条暴雷的三大痛点，中间是 block_table 的两级寻址，下边是三条序列共享同一物理池后的清爽世界：

3. 物理块 vs 逻辑序列：两级映射

engine/model_runner.py:allocate_kv_cache 一次性给整个进程分配 KV 池：

self.kv_cache = torch.empty(2,                              # K 和 V 两份    hf_config.num_hidden_layers,    # 每层一份    config.num_kvcache_blocks,      # 全部块数（启动时根据剩余显存计算）self.block_size,                # 每块容纳 256 个 token    num_kv_heads,                   # TP 切分后本 rank 的 head 数    head_dim,)layer_id = 0for module inself.model.modules():ifhasattr(module, "k_cache") andhasattr(module, "v_cache"):        module.k_cache = self.kv_cache[0, layer_id]        module.v_cache = self.kv_cache[1, layer_id]        layer_id += 1

注意几件事：

• • 这是全进程共享的一个大张量，所有序列写入和读取都从这里面切片
• • 每个 attention 层的 k_cache / v_cache 都是这个池的视图（shared view），没有数据拷贝
• • 池的大小（num_kvcache_blocks）是启动时根据剩余显存反推的——下一讲（L6）会拆这个公式

而每个 Sequence 自己持有的只是一个"逻辑块号 → 物理块号"的映射：

seq.block_table = [7, 12, 3]   # 这条序列的第 0/1/2 个逻辑块分别对应物理块 #7 #12 #3

要找到 token at position p 对应的具体内存位置，需要走两步映射：

第一步整除 p // B 把逻辑位置切到"哪个逻辑块"，再借 block_table 翻成物理块号；第二步取模 p % B 算出"在这块里第几个 slot"，加上块的起始地址 block_id × B 就是池里的最终下标。两步都是 O(1) 计算，attention 内核每读写一个 token 都要走一遍。

这套间接寻址是 PagedAttention 的全部秘密。一旦理解，剩下的复杂度都是工程细节——块从哪来、什么时候归还、怎么把 slot_mapping 喂给内核——但核心机制就这一个。

4. slot_mapping：告诉 GPU 内核每个 token 该写到哪里

它解决了什么问题

把 KV Cache 切成块以后，attention 内核（在 GPU 上跑的算子，FlashAttention 或自写的 Triton kernel）撞上一个新麻烦：它没法自己算出每个 token 的 K/V 应该写到显存的什么位置。

朴素连续布局下不存在这个问题——内核拿到 (batch, seq_len) 形状的输入张量，第 i 个 token 的地址就是 base + i × per_token_size，一个乘加就够。分块以后，"逻辑上第 p 个 token 物理上在哪儿"必须先查 seq.block_table 才知道（第 3 节那张两级映射图）。但 block_table 是 Python 端 Sequence 对象的字段，GPU 内核 forward 过程中没法"中断一下让 CPU 帮我查字典"——它只看得到喂进来的张量。

怎么解决的

nano-vllm 的解法很直接：把这层翻译提前到调度阶段，由 CPU 把"每个 token 该写到哪个 slot"算好，结果打成一个张量喂进 GPU。这个张量就叫 slot_mapping。

它的"协议"由三件事定义：

• • 形状：一维 int32 张量，长度 = 本 step 在 batch 里要算 K/V 的 token 总数（prefill 序列贡献多个、decode 序列各贡献 1 个，按 batch 顺序拼起来）。
• • 语义：slot_mapping[i] 表示"输入张量第 i 个 token 的 K/V 应该写到 KV 池的第几格"。
• • 代价：CPU 端多算一次小循环（开销可忽略）；GPU 内核原本要做的"查 block_table → 乘加算地址"被压成 pool[slot_mapping[i]] = K_i, V_i——一次数组下标读取就到位。

换句话说，slot_mapping 是 CPU 调度器和 GPU 内核之间的地址翻译协议：调度器在分块世界里查表，把结果用一维数组打平交给内核；内核完全不需要懂"块"这件事，按表对号入座写入就行。

下图把上面两段拆成 ① CPU 阶段 + ② GPU 阶段两半，串成一条数据流。简单起见 batch 里只放一条序列，3 个 token，这条序列 block_table = [2]：

① 灰色的 CPU 半区里，调度器拿到 seq.block_table = [2]，对每个本步要算的 token 位置 i=0,1,2 套公式 slot[i] = bt[i//B] × B + (i%B)，得到 slot_mapping = [8, 9, 10]——这就是 CPU 提前算好的"地址翻译表"。蓝色箭头表示这张表作为张量被上传到 GPU。

② 浅色的 GPU 半区里，attention 内核同时拿到两份输入：上一行是按 batch 顺序排好的输入张量（蓝色，3 个 token 的 K/V 待写入），下一行是从 CPU 接收到的 slot_mapping（黄色）。内核执行 pool[slot_mapping[i]] = K_i, V_i——三条橙色箭头落到 KV 池里物理块 #2 的前 3 个 slot（编号 8、9、10）。整个过程内核完全不需要懂"块"概念，只看映射表对号入座。

这是最简单的情形：本步要算的 3 个 token 集中在同一个新块里，slot 编号也连续。下一节就来看 chunked prefill 怎么把这种"集中"打破。

Prefill 路径

prefill 的麻烦在于：输入 batch 在 token 维度上是连续的，但它们在物理 slot 上未必连续。一条做 chunked prefill 的序列，本步只啃 prompt 的中间一段——这一段在逻辑上是连续的位置区间 [start, end)，但翻译到 KV 池里可能跨多个物理块，并且首尾两个块都要做边界处理：首块只从某个 offset 开始写（前面的 slot 已经被上一 step cache 过了），末块只写到某一格就停（后面的位置在本 step 还没轮到算）。

下图是这两种边界情况的具体例子。一条 6-token 的 prompt，前 3 个已经在上一个 step cache 完成，本步预算 3 个把剩下的 prefill 完——block_size=4，block_table=[1, 6]：

上半的逻辑视角里，灰条纹的前 3 个 token 已经在前面的 step cache 过；橙色的后 3 个是本步要算的。这 3 个 token 在逻辑上分属两个块：position 3 是逻辑块 0 的最后一格，positions 4、5 是逻辑块 1 的开头两格。下半的物理视角里，逻辑块 0 经 block_table[0]=1 翻成物理块 #1，逻辑块 1 经 block_table[1]=6 翻成物理块 #6——本步真正要写的就是 #1 的最后一格（slot 7） + #6 的前两格（slot 24、25）。

engine/model_runner.py:prepare_prefill 里这段是核心：

for seq in seqs:    start = seq.num_cached_tokens                # 这次开始处理的 token 起点    seqlen_q = seq.num_scheduled_tokens    end = start + seqlen_q    ...    start_block = start // self.block_size    end_block = (end + self.block_size - 1) // self.block_sizefor i inrange(start_block, end_block):        slot_start = seq.block_table[i] * self.block_sizeif i == start_block:            slot_start += start % self.block_size       # 首块要跳过已 cache 的 offsetif i != end_block - 1:            slot_end = seq.block_table[i] * self.block_size + self.block_sizeelse:            slot_end = seq.block_table[i] * self.block_size + end - i * self.block_size  # 末块只到 end        slot_mapping.extend(range(slot_start, slot_end))

逐行翻译：

• • 算出这次要写哪些 token，对应哪些"逻辑块"（start_block 到 end_block）
• • 对每个逻辑块，查 block_table[i] 拿到物理块号，乘 block_size 得到这一块在池里的起始 slot
• • 处理两个边界：首块从 start % block_size 开始（前面已经 cache 过的部分跳过），末块只到 end 那一格（后面还没分配的不写）
• • 把这一段 slot 编号 extend 进 slot_mapping

套到上图那条序列上：start=3、end=6、start_block=0、end_block=2，循环跑两次，slot_mapping 在这条序列贡献的部分就是 [7, 24, 25]——i=0 的 token 写 slot 7（首块只写一格），i=1、i=2 的 token 写 slot 24、25（末块只写两格）。这正是图里三条箭头的终点。

Decode 路径

prepare_decode 简单得多——每个 seq 一次只算 1 个 token，永远是最后一个位置：

for seq in seqs:    input_ids.append(seq.last_token)    positions.append(len(seq) - 1)    context_lens.append(len(seq))    slot_mapping.append(        seq.block_table[-1] * self.block_size + seq.last_block_num_tokens - 1    )

公式 last_block * block_size + last_block_num_tokens - 1 等价于"最后一块物理位置 + 序列在最后一块里已用 slot 数 - 1"——刚好定位到刚 append 的那个 token 的 slot。

这个数组喂给谁

写 KV：自写的 Triton kernel store_kvcache_kernel，直接拿 slot_mapping[i] 索引。L11 会拆这个 kernel。

读 KV：FlashAttention 的两个 API（flash_attn_varlen_func / flash_attn_with_kvcache），它们接受 block_table 自行算地址。同样在 L11 详谈。

5. 块大小的取舍

config.kvcache_block_size 默认 256，源码里还有 assert self.kvcache_block_size % 256 == 0——必须是 256 的倍数。这个"256"是个折中。

下图把这个折中的两端画出来：

横轴是 block_size（对数刻度）。两条虚线分别画两类成本——红色的元数据/调度成本随块变小而上升（块越多 → block_table 越长、hash 越频繁、调度的物件越多），蓝色的内部碎片/对齐损失随块变大而上升（一条短序列被迫吃一整块、preempt 整块还、prefix-cache 对齐粒度变粗）。两者相加就是黑实线那条 U 形总成本曲线，最低点对应"两类成本平衡得最好的那个 block_size"——大约就在 256 这一档。下面分别看两端各自的代价。

取小（如 16/32）的代价

• • block_table 变长 N 倍，每次 attention 的元数据传输（prepare_block_tables 里那个 int32 张量 cuda 拷贝）成本变高
• • prefix cache 的 hash 计算更频繁——每个块都要算一次 xxhash（L5 详谈），变成 16 倍负担
• • 物理块数变多 → BlockManager 的 free_block_ids 队列、hash_to_block_id 字典都变大，CPU 调度成本上升

取大（如 1024/4096）的代价

• • 内部碎片回归——一条 200-token 的序列只用一块的 1/5，回到了"按 max_seq_len 预留"那种浪费
• • preempt 粒度变粗——抢占必须释放整块，可能误伤未来要用的部分
• • prefix cache 的命中粒度变粗——共享前缀必须正好对齐 1024 的倍数才能复用，命中率掉

为什么 256

• • 和 FlashAttention 内核的 chunk 偏好对齐
• • 真实 prompt 长度的 P50 在 100–500 区间，单条序列大多用 1–2 个 256 块——内部碎片的最坏情况是 255 个 token 被浪费，可控
• • xxhash 在 256 个 token（即 256 × 4 字节 = 1KB 上下）规模上 CPU 开销 < 1μs，频率合适

简单结论：256 是个"够大但不太大"的甜蜜点，也是和 FlashAttention 内核协同设计的产物。如果你的工作负载特别短（agent 类、平均 < 100 token），可以考虑往下调；特别长（长文档、code review），可以考虑往上调，但优先级在 prefix cache 等其它机制之后。

6. 视频：多序列共享物理 KV 池

下面这段动画展示三条同时存活的请求是怎么共享一个物理 KV 池的——它们独立成长、独立释放，物理块在它们之间动态流转：

注意几个值得反复看的画面：

• • 三条序列的 block_table 永远互相独立，但它们指向的物理块都来自同一个池
• • 当某条序列结束（FINISHED）时，它的物理块整体归还池子，不留任何"占用残留"
• • 同一块物理 slot 在不同时刻可能服务于不同的 sequence——这是 PagedAttention 让显存利用率变高的根本原因

7. 下一讲

接下来 L4 把镜头拉近 BlockManager 这个类本身：

• • free_block_ids 是 deque、used_block_ids 是 set，为什么这样选数据结构？
• • can_allocate / allocate / can_append / may_append / deallocate 这五个 API 各自的责任边界是什么？
• • 引用计数 ref_count 为什么存在——它是 L5 prefix cache 的伏笔。

理解了 PagedAttention 的"为什么"（L3）之后，L4 讲"怎么实现"——一段精炼的内存分配器。