PagedAttention学习笔记

逻辑部分

一、基础概念与背景

1.1 什么是 PagedAttention？

定义 ：PagedAttention 是一种借鉴操作系统虚拟内存 / 分页 思想的注意力 KV 缓存管理方案：不再为每个请求预留一整段连续 的 KV 显存，而是把 KV 切成固定大小的页（Block） ，通过块表（Block Table）把「逻辑上的第 i 个块」映射到「物理块 ID」，从而在物理显存上实现非连续存储。

核心目标：

1. 降低内部碎片：只为已生成的 token 分配块，未用到的槽位不占用物理块。
2. 缓解外部碎片 ：所有请求使用统一块大小，空闲块可放入全局池复用，避免「大块装不下、小块又浪费」的锯齿状空洞。
3. 支撑共享：同一前缀（系统提示、多采样分支等）可映射到同一组物理块，配合引用计数与写时复制策略复用显存。

1.2 为什么需要 PagedAttention？

瓶颈：自回归推理中，KV Cache 随序列长度线性增长，且与层数、注意力头数、batch 成正比；大模型服务中 KV 往往超过权重本身的常驻显存。

传统「按最大长度连续预分配」的问题：

问题	含义
内部碎片	为 max_model_len 预留一整条连续 KV，实际序列可能很短，大量槽位永远不用。
外部碎片	若按「各请求实际长度」分别 malloc 不同大小，释放后易产生无法合并的小空洞，大连续区难以再分配。
吞吐受限	显存浪费 → 同卡上能并发的请求数变少 → GPU 算力空转。

与 OS 分页的类比（帮助记忆）：

flowchart LR subgraph 逻辑地址空间 L0["逻辑块 0"] L1["逻辑块 1"] L2["逻辑块 2"] end subgraph 物理显存 P7["物理块 7"] P2["物理块 2"] P15["物理块 15"] end L0 -->|"块表"| P7 L1 --> P2 L2 --> P15

逻辑上连续，物理上任意离散；注意力核通过块表做间接寻址加载 K/V。

---

二、技术原理解析

2.1 PagedAttention 核心原理

分页 ：设 block_size = B（每块存 B 个 token 的 K/V 槽位）。序列长度 L 需要 $\lceil \; L\; /\; B\; \rceil \; \backslash lceil\; L\; /\; B\; \backslash rceil$⌈L/B⌉ 个逻辑块 ；每个逻辑块对应一个物理块 ID（在 vLLM 实现里由分配器给出）。

逻辑块 vs 物理块：

• 逻辑块 ：第 i 块对应 token 下标区间 $[\; i\; B\; ,\; min\; \; (\; (\; i\; +\; 1\; )\; B\; ,\; L\; )\; )\; [iB,\backslash \; \backslash min((i+1)B,\backslash \; L))$[iB, min((i+1)B, L))，是「用户视角」的顺序切片。
• 物理块 ：显存大池中的第 block_id 号槽位，大小固定；多个请求的逻辑块 0 可以都映射到不同物理块 ，也可以在前缀缓存命中时映射到同一物理块。

块表 ：对每个请求维护 block_table[row] = [phys_0, phys_1, ...]。注意力核根据 (query 位置 → 逻辑块索引 → 物理块 ID → 块内偏移) 取 K/V。

按需分配：只有当序列增长跨过块边界时才向块池申请新物理块；未写满的块中尾部槽位虽在块内但可视为「尚未使用」，无需为未到长度提前占满所有块。

共享与写时复制（概念层）：

• 并行采样 / 束搜索：多个序列在分叉前共享同一前缀 → 多块表指向同一批物理块。
• 写时复制 ：当某条序列要在「仍被他人引用」的块上就地改写 时，应复制出新物理块再写（避免脏写共享数据）。工程上常与「新 token 只追加在新块」组合出现；前缀缓存命中时也会通过 ref_cnt 与分配路径保证安全共享。

引用计数 ：每个 KVCacheBlock 维护 ref_cnt；引用为 0 且非占位块时，块可回到空闲队列供 LRU 类策略驱逐或再分配。

stateDiagram-v2 [*] --> 空闲池: 初始化 空闲池 --> 已分配: allocate / touch(ref++) 已分配 --> 空闲池: free(ref--) 且 ref==0 已分配 --> 前缀缓存索引: 块满且写入 hash 前缀缓存索引 --> 空闲池: 驱逐 evict

与 FlashAttention 的关系

二者正交：vLLM 常在 Paged KV 布局之上调用融合/分块注意力核（名称随版本演进，思想仍是「按块取 K/V」）。

• FlashAttention ：在单次 注意力计算内做 tiling / 重排，减少 HBM↔SRAM 流量，是 kernel 级数值实现优化。
• PagedAttention ：在多次 forward 之间管理 KV 驻留位置与生命周期 ，是 系统级显存管理。

维度	FlashAttention	PagedAttention
优化目标	单次注意力算子的 I/O 与中间态	跨步 KV 的分配与布局
主要手段	Tiling、在线 softmax 等	固定块、块表、块池
是否改变 attention 公式	否（数值等价）	否（存取路径改变）

---

三、算法流程与数据结构

sequenceDiagram participant Sched as 调度器 / KVCacheManager participant Pool as 物理块池 BlockPool participant BT as BlockTable(CPU/GPU) participant Kern as Attention Kernel Sched->>Pool: 为新 token 或新块申请物理块 Sched->>BT: 更新每请求的 block_ids Sched->>BT: compute_slot_mapping(req_idx, pos) BT->>BT: commit_block_table / slot_mapping 到 GPU Kern->>BT: 按 slot 或 block 表间接加载 K/V Kern->>Kern: softmax(QK^T)V 等

关键步骤简述：

1. 元数据：为 batch 中每个请求维护逻辑块 → 物理块 ID 列表（可能多组，对应 hybrid 模型的不同 KV group）。
2. Slot mapping：把「本步参与的每个 (request, position)」映射到线性 KV tensor 中的槽位（或 -1 表示本 rank 不负责，见分布式交错）。
3. 内核：读取 Q 与块表，按块 gather K/V，完成注意力。

3.1 关键数据结构

概念	作用
物理块	固定容量，存放一段连续 token 的 K/V；有 block_id。
块表	每请求一行，逻辑块序 → block_id 序列。
KV 管理器	分配、释放、前缀命中、抢占/交换策略的协调者。
空闲块队列	可驱逐候选的顺序（如 LRU）；与 ref_cnt 联动。

代码实现部分

以vLLM为例，

• 核心 ：调度侧 KVCacheManager / SingleTypeKVCacheManager 等与 KVCacheBlock + FreeKVCacheBlockQueue + 块表 配合；Worker 侧 BlockTable 维护 GPU 可见的块表与 slot_mapping。
• 前缀缓存 ：BlockHashToBlockMap 等结构把「块内容哈希」映射到可复用物理块（跨请求）。
• Hybrid 模型 ：kv_cache_utils 中按 KVCacheGroupSpec 将多层归为若干组，每组一张块表 ，组内层共享映射；物理块每页字节数需对齐（见 unify_hybrid_kv_cache_specs 等逻辑）。

关键的模型自顶向下为：

块表与槽位映射：BlockTable

文件：vllm/v1/worker/block_table.py

• block_table ：[max_num_reqs, max_num_blocks_per_req] 的 int32 缓冲（CPU 填、再 commit_block_table 拷到 GPU）。
• num_blocks_per_row：每请求当前已挂接的块数。
• compute_slot_mapping ：由 (req_indices, positions) 得到每个 token 在「扁平化 KV 布局」中的 slot；普通单卡下为
  block_id = block_table[req, pos // B]，slot = block_id * B + (pos % B)。
  DCP/PCP 时使用 virtual_block_size 与 mask，把不属于本 rank 的项写为 -1。
• 混合块 ：当「显存分配块大小」与「kernel 要求的 block_size」不一致时，map_to_kernel_blocks 把管理器层的一个大块拆成多个核侧小块。

这与提纲中的「间接寻址」「元数据传入内核」一致：块表 + slot_mapping 即运行期页表。

物理块元数据与空闲队列：KVCacheBlock / FreeKVCacheBlockQueue

文件：vllm/v1/core/kv_cache_utils.py

• KVCacheBlock ：block_id、ref_cnt、可选 _block_hash（前缀缓存）、双向链表指针 prev_free_block / next_free_block（供 LRU 队列 O(1) 摘除）。
• FreeKVCacheBlockQueue ：注释中说明为何不用 deque------避免额外小对象分配，并在队列中间摘除块。

块池与缓存索引：block_pool.py

文件：vllm/v1/core/block_pool.py

• BlockHashToBlockMap ：block_hash → KVCacheBlock（或冲突时的多块字典），服务于跨请求前缀复用；与论文式 PagedAttention 的「共享物理页」对应。

分组块表与 hybrid 模型

文件：vllm/v1/core/kv_cache_utils.py（如 _get_kv_cache_groups_uniform_page_size）

• 同一 物理页大小 的前提下，将不同 attention 类型的层分组 ，每组重复若干层 → 减少块表套数、保证块对齐。文内英文注释给出了「10 full + 20 sliding → 3 层 pattern × 10 组」的示例。

组件关系（总览图）

flowchart TB subgraph SchedulerCPU["调度侧 CPU"] KM["KVCacheManager / SingleType..."] BP["BlockPool + ref_cnt + LRU 队列"] BH["BlockHashToBlockMap 前缀缓存"] end subgraph WorkerGPU["Worker GPU"] BT["BlockTable: block_table + slot_mapping"] KERN["Attention / FA 等内核"] end KM --> BP KM --> BH KM -->|"序列的 block_ids"| BT BT --> KERN

实验部分

内存利用率评估

目标 ：对比「按 max_model_len 为每条请求连续预留 KV」与「PagedAttention 按需块分配」的有效显存占用或可并发请求数。

指标建议：

• 每请求 KV 占用：平均 / P95 序列长度下，实际分配块数 × 每块字节数。
• 碎片率 proxy ：1 - (所有已用槽位数 / (已分配块数 × B)) 在稳态负载下的分布（Paged 应接近 0，预分配则随长度分布变差）。

可视化：对固定时间窗内「每物理块已写入 token 数」做热力图（横轴块 ID，纵轴时间），可直观看到填充是否饱满。

吞吐量与延迟

目标：在相同 P99 延迟约束下，对比最大可持续 RPS；或相同 RPS 下的延迟尾。

图表：

• 吞吐--延迟曲线：横轴 RPS，纵轴 mean / P99 latency。
• Batch size ：观察 vLLM 在动态 batch 下实际 num_batched_tokens 分布是否更大。

实验提示：控制变量（同 GPU、同模型精度、同输入分布）；PagedAttention 的收益在长序列、高并发、输入长度差异大时通常更明显。