Month 2 · PagedAttention ★

皇冠章节。这一个月学透，后面所有内容都顺。

这是 vLLM 之所以是 vLLM 的核心 idea。论文（Kwon et al., SOSP'23）的副标题是 "Efficient Memory Management for Large Language Model Serving with PagedAttention"。 注意是 Memory Management——这是一篇系统论文，不是机器学习论文。 作者用 OS 虚拟内存做类比，把 KV cache 的利用率从 ~40% 拉到 ~96%，端到端 throughput 提升 2–4×。 这一章把这套思路完整重建：从 KV cache 的数学，到 OS 分页的原理，到 vLLM 怎么把两者拼到一起。

01本章导航

这一章很长，因为内容是 vLLM 的核心。按下面的顺序走：

1. §02 KV cache 的数学 · 先把"KV cache 到底有多大"算清楚。
2. §03 朴素实现为什么失败 · 把驱动问题量化。
3. §04 OS 虚拟内存复习 · 这是 vLLM 的"母概念"。
4. §05 PagedAttention 设计 · 一对一对照映射。
5. §06 代码深读 · KVCacheManager、BlockPool、attention kernel。
6. §07 Prefix caching 深读 · "page cache"的现代版。
7. §08 CPU offload / Swap · 当显存真的不够。
8. §09 同类设计的对比 · SGLang RadixAttention、H2O。
9. §10 动手实验 · 跑出真实数字。
10. §11 常见误解 · 把"以为懂"挤出来。
11. §12 本月 PR 候选
12. §13 自检

02KV cache 到底是什么、有多大

读 vLLM 源码前，先把这个对象的大小算清楚。这一步省了，后面看 block_size 选 16 的讨论会觉得云里雾里。

2.1 数学定义

Transformer 的每一层都有自己的 K 和 V。对一个 token，每层产生一份 K 向量和一份 V 向量（每个 attention head 各一份）：

     模型每一层：

     输入 token 的 embedding x  ─→  Q = x @ W_q       (这一步不需要 cache)
                                ─→  K = x @ W_k       ← 这个要 cache
                                ─→  V = x @ W_v       ← 这个要 cache

     未来 token i 算 attention：

        attn = softmax(Q_i @ K_{:i}^T / √d) @ V_{:i}
                              ─────┬─────         ┬──
                              过去所有 K        过去所有 V    ← 必须从 cache 里取

所以单个 token 在整个模型占的 KV cache 大小是：

per_token_kv_bytes = 2 × num_layers × num_kv_heads × head_dim × dtype_size
                     │   │            │              │          │
                     │   │            │              │          └ fp16 → 2 字节
                     │   │            │              └ 每个 head 的维度（通常 128）
                     │   │            └ KV 头数（GQA 模型 ≠ 注意力头数）
                     │   └ 层数（Llama 3.1 8B = 32）
                     └ K 和 V 各一份

累计到 seq_len 个 token + 一个序列时：

seq_kv_bytes = per_token_kv_bytes × seq_len

2.2 实际数字 · Llama 3.1 8B (GQA)

把 Llama 3.1 8B 代入，立刻能算：

per_token_kv_bytes = 2 × 32 × 8 × 128 × 2 = 131072 bytes ≈ 128 KB / token

把它代入"我能装多长 / 多少并发"的问题：

2.3 70B / 长 context 的麻烦

同样算法，对 Llama 3.1 70B（80 层、8 KV head、head_dim 128）：

per_token_kv_bytes = 2 × 80 × 8 × 128 × 2 = 327680 bytes ≈ 320 KB / token

一个 32K 的 prompt 单序列 KV ≈ 10 GB。一张 A100 80GB 减掉权重（140 GB / 2 卡 TP = 70 GB / 卡，所以放不下；实际要 8 卡 TP，每卡承担 ~18 GB 权重），KV 预算就更紧。
所以"KV cache 利用率"这种听起来抽象的指标，直接决定了你能 serve 几个用户。

03朴素实现为什么失败（把驱动问题量化）

把 §02 的数字应用到"朴素 contiguous KV"，每个问题都能算出具体百分比：

3.1 内部碎片 · 量化

假设朴素策略：每个请求来时按 max_model_len = 4096 预分配。如果用户实际只生成 100 token：

分配槽位 = 4096 × 128 KB = 512 MB / 请求
实际使用 = 100 × 128 KB =  12.5 MB / 请求
利用率 = 12.5 / 512 = 2.4%

论文 Figure 2 显示真实 workload 下利用率长期在 20–40%。大部分显存在等永远不会来的 token。

3.2 外部碎片

请求 A 用完 512 MB 段释放回内存池。下一个请求 B 想要 max_model_len=8192，需要 1 GB 连续段。即使总自由内存有 2 GB，分散在 4 个 512 MB 小段里，依然装不下 B。这是经典 buddy allocator 痛点。

3.3 无法共享 prefix

"你是一个客服" 这种 system prompt 在 100 个并发请求里被算 100 次 KV，存 100 份。如果能共享，节省 = 100 × (system prompt 长度 × per_token_kv)。 system prompt 100 token、并发 100 时，能省 100 × 12.5 MB = 1.25 GB，几乎一整张消费级 GPU 的容量。

3.4 不能动态适配 batch 长度

同一个 batch 里请求长度参差不齐（一个 100 token、一个 3000 token）。朴素做法 = padding 到最长。3000 token 的 batch 矩阵：

(batch_size × 3000 × hidden_size) 的输入
其中 99% 是 padding，浪费的 FLOPs ≈ 99%

结论：四个问题都是数量级浪费，加起来轻松 60-80% 资源被烧。这就是 PagedAttention 之前 LLM serving 的现实。

04OS 虚拟内存复习（必要的母概念）

如果你完全没接触过虚拟内存，下面这段速通；学过的扫一眼当复习。后面的对应才能立刻生效。

4.1 为什么发明分页

1960 年代的计算机：每个程序看到的是物理内存的连续段。运行多个程序很难，因为它们要竞争同一块物理 RAM 的不同区段。解法 = 把"程序看到的地址" 和 "实际物理位置" 解耦：

4.2 关键带走的 4 点

• 固定大小的页（典型 4KB）让分配/释放变简单 —— 永远只对齐到 page boundary，外部碎片消失。
• 页表是 indirection 层 —— 应用看到的连续地址不是真的连续。
• 页可以共享 —— 两个进程的页表项指向同一个物理页（典型场景：共享库、COW fork）。
• 页可以延迟分配 (demand paging) —— 程序声明 4GB 虚拟空间，OS 只在你真访问到某页时才给物理页。Page fault 触发分配。

4.3 TLB · 页表的代价

查页表本身要访问内存，每次内存访问变两次（一次查表、一次取数据）。为了不爆炸，CPU 有 TLB（Translation Lookaside Buffer）——硬件 cache，存最近用过的 V→P 映射，~64-512 项，几个 cycle 命中。TLB miss → 重新走页表 → 慢 ~100 cycle。

这一点要记住：indirection 是有成本的。后面我们会看到 vLLM 在 GPU 上"模拟 TLB"的方式（其实是直接把整个 block table 塞进 kernel 的常量参数）。

4.4 COW Fork · prefix caching 的祖宗

Unix fork() 复制整个进程。复制 4GB 地址空间太慢，所以现代实现是 copy-on-write：父子进程的页表初始指向同一份物理页，标记只读；任一方写时才真的复制。

这就是 prefix caching 的原型：两个请求 prompt 相同时，它们的 block table 指向同一份物理 KV，谁开始生成不同的 token 谁分裂出新 block。我们在 §07 详细看。

05PagedAttention：把上面这套搬到 KV cache

vLLM 的做法结构上完全平行：

5.1 把 OS 概念逐行映射到 vLLM

5.2 一个 physical block 物理上是什么

vLLM 把整个 GPU 显存的 KV 区域看成一个大 tensor，按 block 切：

kv_cache = torch.zeros(
    num_blocks,          # 池里有多少块
    block_size,          # 一块装多少 token（16）
    num_kv_heads,        # 8
    head_dim,            # 128
    dtype=torch.float16,
)
# K 和 V 各一份，所以模型每层有两个这样的 tensor
# shape 在新版可能略不同（layout: NHD vs HND），但大小一样

从这个 layout 你可以反算单块大小：

block_bytes = block_size × num_kv_heads × head_dim × dtype_size × 2_for_KV
            = 16 × 8 × 128 × 2 × 2 = 65536 bytes = 64 KB / block / layer
全模型一块：
            = 64 KB × 32 layers = 2 MB / block / sequence_token-coverage_16

所以 vLLM 启动时打的日志 GPU KV cache size: X tokens，X = num_blocks × block_size。从此你看到这个数字会立刻能解读。

5.3 Block size 为什么是 16

这是 PagedAttention 论文 §6.4 实验出来的 sweet spot。trade-off：

论文图 18 显示 4/8/16/32 几乎相当，64+ 开始掉。16 是最甜的点，对 16-token Llama tokenizer 的 BPE 分布也对得齐（平均一个英文词 4 token，约 4 个词 = 一句话短语）。

5.4 attention kernel 怎么用 block table

读到这里你应该好奇：朴素 attention 一个 torch.matmul(Q, K.transpose(-1,-2)) 就完事，要求 K 是连续 tensor。如果 K 散在 16 块物理 block 里，怎么矩阵乘？

答案：kernel 内部 gather。论文叫 "PagedAttention kernel"。伪代码：

# 每个 query token 一个 thread block
# 任务：算这个 token 对过去所有 KV 的 attention

q = load_query()                       # 一次性加载（query 只有 1 个 token）
running_max = -inf
running_sum = 0
out = 0

for logical_block_idx in range(my_seq_len // 16):
    # 关键的间接寻址：查 block table
    physical_block_idx = block_table[seq_id, logical_block_idx]

    # 从物理块加载 16 个 token 的 K 和 V
    K_block = kv_cache_K[physical_block_idx]   # shape (16, n_kv_heads, head_dim)
    V_block = kv_cache_V[physical_block_idx]

    # 算这块的 attention scores
    scores = q @ K_block.T / sqrt(d)           # shape (16,)

    # online softmax 累加（FlashAttention 的核心 trick，见 M4）
    block_max = max(scores)
    new_max = max(running_max, block_max)
    correction = exp(running_max - new_max)
    running_sum = running_sum * correction + sum(exp(scores - new_max))
    out = out * correction + (exp(scores - new_max) @ V_block)
    running_max = new_max

return out / running_sum

关键就是 block_table[seq_id, logical_block_idx] 这一行 indirection。它把"逻辑上连续的 KV"翻译成"物理上分散的 block 号"。这就是 OS 页表在 GPU 上的样子。

06代码深读 · KV cache 子系统

下面的 anchor 是这个月的主菜。按顺序读，期间随时回头看上面的图和数字。5 段加起来 ~10-12 小时，分 3-4 天完成。

6.1 第一遍 · 整体结构

6.2 第二遍 · BlockPool

关键发现：free 不等于丢弃。一个 block 引用计数到 0、回到 free_queue 之后，如果还没被覆盖，它依然在 cached_block_hash_to_block 里——下次同 prefix 进来命中，可以从 free 列表里"救回"它。这就是 prefix cache 在 free 状态下仍然能命中的原因。

6.3 第三遍 · Hash 怎么算

Prefix cache 的命中靠 hash。同样的 token 序列必须 hash 到同一个 key。关键代码：

# 类似实现（具体函数名以源码为准，搜 hash_block_tokens）
def hash_block_tokens(parent_block_hash, token_ids):
    """链式 hash —— 一块的 hash 依赖上一块的 hash。

    这样保证: 只有"从开头到这块"完全相同的两个序列才会撞 hash。
    避免: 只是"中间一段相同"被错误命中。
    """
    return hash((parent_block_hash, tuple(token_ids)))

这是chain hash，跟 Merkle tree 的思路一样。每一块的指纹包含了它的父块指纹，保证 prefix 一致性。

6.4 第四遍 · 每序列的"翻页"逻辑

# 伪代码
def append_token(request, new_token):
    request.token_ids.append(new_token)
    # 当前 token 数已经填到第几块的第几格？
    num_blocks_needed = ceil(len(request.token_ids) / block_size)
    while len(request.block_ids) < num_blocks_needed:
        new_block = block_pool.get_new_blocks(1)[0]
        request.block_ids.append(new_block.id)

6.5 第五遍 · 走到 GPU 端

6.6 attention kernel 端绑定

07Prefix Caching · "page cache" 的现代版

OS 里 page cache 是："读过的磁盘块缓存在内存，下次再读直接命中"。 vLLM 里 prefix cache 是："算过的 prompt 的 KV 缓存在显存，下个共享 prompt 的请求直接命中"。

7.1 命中的时机线

请求 1 prompt:  [system: 你是一个客服] + [user: 怎么退货?]
                ─────────────── 16 tokens ───────────────  +  ─── 5 tokens ───
请求 2 prompt:  [system: 你是一个客服] + [user: 怎么改地址?]
                ─────────────── 16 tokens ───────────────  +  ─── 6 tokens ───

   前 16 tokens 的 KV 完全相同
   →  请求 1 算完后，把 P3 这块标记为可共享 (cache_blocks 写入 hash 表)
   →  请求 2 进来 hash 一对：命中！直接引用 P3，跳过 prefill 这一段
   →  延迟降低（少算 16 token 的 prefill）、显存少占（少存 16 token 的 KV）

7.2 三种状态

每个 block 在生命周期里有三种状态：

关键设计：EVICTABLE 状态让 prefix cache 的"命中率"不依赖于请求并发。即使两个共享 prefix 的请求错开到达（请求 1 已经结束、请求 2 后来），只要中间没把 P3 覆盖掉，请求 2 还能命中。

7.3 Eviction 策略

显存有限。free_queue 是双向链表，head 是最近释放的、tail 是最久没用的。取新块从 head 取（LRU 的 fresh-first）；如果不够，从 tail 强制 evict（覆盖掉旧的 EVICTABLE block）。

这和 OS page cache 完全是同一个算法（Linux 用的是 active/inactive LRU 链表，思路相同）。差别只在粒度（4KB vs 16 token KV）。

7.4 OS page cache vs vLLM prefix cache · 完整对照

7.5 限制：只能命中 block-aligned prefix

vLLM 的 prefix cache 是整块粒度。如果两个 prompt 在第 17 个 token 才分叉（第 16 个 token 之后），只能共享前 16 token 的那一块（因为 chain hash 是按 16-token 块算的）。第 17 token 那块的 hash 已经因为该 token 不同而完全不同。

SGLang 的 RadixAttention 用细粒度 radix tree 解了这个问题——但代价是 lookup 更慢、实现更复杂。当 prompt 共享是动态分叉（agent 多分支推理）时 SGLang 优；当 prompt 共享是固定 system prompt（90% serving 场景）时 vLLM 的简单方案足够。

08CPU Offload / Swap · 当显存真的不够

分页解决了"碎片"，prefix 共享解决了"重复"。但如果新请求来了、所有可 evict 的 block 都已经被 evict、仍然不够呢？两条路：

1. Preempt（M3 主题）—— 把某个 RUNNING 请求挤出去，让新请求/高优先级请求进来
2. Swap（本节）—— 把不活跃的 block 从 GPU 显存搬到 host RAM 或 NVMe，腾出空间

8.1 Swap 的工程现实

看似还行。但 swap 是全模型那么大：一个序列 4K context 大约 4096/16 = 256 块 × 2 MB = 512 MB。swap 一个 4K 序列 ~16 ms，**比一次 decode step (~10ms) 还慢**。所以 swap 必须提前异步，否则 decode 直接 stall。

8.2 代码入口

09同类设计的对比 · 横向看 KV 管理

把 PagedAttention 放回大家族里：

把这张表收藏 —— 看新论文时先定位"它在哪一行"。Prompt 看不出来? 看它怎么处理 KV cache，方向立刻明白。

10动手 · 复现碎片对比

这是本月的硬核作业。跑通它，PagedAttention 就刻进肌肉里了。

10.1 实验设计

1. 租 A10 (24 GB)，跑 Llama 3.1 8B fp16，max_model_len = 4096。
2. 启动 vLLM 时打开日志：能看到 GPU KV cache size: X tokens，Maximum concurrency: Y。把这两个数字记下来。
3. 用 benchmarks/benchmark_serving.py 跑三组 workload：
   • A · 短输出：100 个请求，prompt 50 tokens，max_tokens=10
   • B · 长输出：100 个请求，prompt 50 tokens，max_tokens=2000
   • C · 共享 prompt：100 个请求，prompt = 同一个 200 token system prompt + 不同 20 token 短问 ，max_tokens=100
4. 对每组记录：throughput (tok/s)、TTFT (time to first token p50/p99)、平均显存利用率（看 nvidia-smi 或 vLLM metrics）。
5. 再跑一遍 C 但加 --no-enable-prefix-caching，看 throughput 差多少。

10.2 命令样例

# Terminal 1: 起 server
vllm serve meta-llama/Llama-3.1-8B-Instruct \
  --max-model-len 4096 \
  --gpu-memory-utilization 0.85 \
  --enable-prefix-caching

# Terminal 2: benchmark
python benchmarks/benchmark_serving.py \
  --backend vllm \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --dataset-name sonnet \
  --dataset-path benchmarks/sonnet.txt \
  --num-prompts 100 \
  --request-rate 10 \
  --save-result

10.3 期待的数字

10.4 进阶 · 看 metrics

vLLM 暴露 Prometheus metrics（默认开在 :8000/metrics）：

curl http://localhost:8000/metrics | grep vllm
# vllm:num_requests_running
# vllm:num_requests_waiting
# vllm:gpu_cache_usage_perc
# vllm:cpu_cache_usage_perc
# vllm:num_preemptions_total
# vllm:prefix_cache_hits_total
# vllm:prefix_cache_queries_total

这些指标值得每个一看。prefix_cache_hits_total / queries_total = 命中率，这是你下次调参的关键数字。

11常见误解 · 把"以为懂"挤出来

12本月 PR 候选方向

Month 2 的 KPI 是一个内存子系统相关的 PR。难度梯度：

12.1 入门（500-行内，1-2 周）

• 给 BlockPool 加一个 metric（block 池碎片率、平均 block 寿命）暴露到 Prometheus。
• 给某条错误信息加上上下文（"failed to allocate X blocks，pool has Y free, Z evictable" 比 "OOM" 有用 10×）。
• 给 kv_cache_manager.py 写测试覆盖未测函数（搜 --cov-report 看 coverage gap）。
• 修文档：架构图示，docs/source/design/ 下加一个 prefix cache 的内部图。

12.2 中等（一两周 + 设计讨论）

• 修一个 prefix cache hash 的 edge case（搜 issue tracker prefix-cache label）。常见 case：tokenizer 加 special token 时 hash 没考虑。
• 给 BlockPool 加一个新策略选项（如 ARC、Clock-PRO），跟 LRU 对照 benchmark。
• 改进 swap_out 的 IO（用 cudaMemcpyAsync + pinned memory，看能不能降低延迟）。

12.3 高难（vLLM RFC 级，1+ 月）

• 实现一个细粒度 prefix cache 变种（接近 SGLang RadixAttention 但更轻量）。需要 RFC 提案。
• 把 prefix cache 扩展到跨进程/跨 host（Mooncake 风格 KV 共享）。

13本页自检 · 这是个长清单

14延伸阅读

• PagedAttention 原论文 · Kwon et al., SOSP 2023 · arxiv.org/abs/2309.06180 — §3 设计动机、§4 系统设计、§6.4 block size 消融。
• vLLM 官方架构 blog · blog.vllm.ai — 看最新的 v1 架构文章。
• SGLang / RadixAttention · arxiv.org/abs/2312.07104 — 对照学。
• H2O · Heavy Hitter Oracle · arxiv.org/abs/2306.14048 — KV cache 的 working set 视角。
• Mooncake KV Cache disaggregation · arxiv.org/abs/2407.00079 — Moonshot 的分布式 KV cache。
• OSTEP Ch 18 / 19 / 21 / 22 · pages.cs.wisc.edu/~remzi/OSTEP/ — OS 虚拟内存权威教材。
• MLA · DeepSeek-V2 · arxiv.org/abs/2405.04434 — KV 压缩到 latent 维度，长 context 杀器。