PagedAttention:KV cache 的页表革命

先看传统做法为啥浪费

传统 KV cache 管理像租房:你进来时得按最长租期一次性买下全部家具——

假设 max_seq_len=2048,一条请求的 KV 显存按 2048 预留。
实际使用:
  请求 A 实际生成 200 tokens      —— 浪费 1848 格
  请求 B 实际生成 50 tokens       —— 浪费 1998 格
  请求 C 实际生成 1800 tokens     —— 浪费 248 格

综合利用率常常只有 20-40%。

PagedAttention 的核心比喻:虚拟内存

操作系统解决内存碎片靠的是分页(paging):把物理内存切成固定大小的 page,进程看到的是逻辑连续的虚拟地址,通过页表映射到物理页——物理页可以乱序、可以按需分配、可以在进程间共享。

PagedAttention 照搬:

一张图看清楚映射关系

序列 A 逻辑视角:
  token[0..15]   ──┐
  token[16..31]  ──┼─→ block_table = [7, 2, 11]
  token[32..47]  ──┘

物理显存 (大池子,共 128 个 block):
  ... [blk2] ... [blk7] ... [blk11] ...
       ↑        ↑          ↑
       序列A    序列A       序列A
       的逻辑   的逻辑      的逻辑
       block1  block0      block2

同时还有其他序列的 block 混在池子里:
  [blk3]=序列B  [blk5]=序列C  [blk9]=free

Block 不需要连续,任意跳跃都能查表到位。新 token 生成时只追加一个 block——哪里有空塞哪里。

按需分配 demo

初始:请求进来,prompt=20 tokens
  ─ 需要 ceil(20/16) = 2 个 block
  ─ block_manager 分配 [blk0, blk1]

生成到第 32 个 token:
  ─ 需要 ceil(32/16) = 2 个 block  → 还够

生成到第 33 个 token:
  ─ 需要 ceil(33/16) = 3 个 block  → 分配 [blk2]
  ─ block_table = [blk0, blk1, blk2]

序列结束:
  ─ free [blk0, blk1, blk2] 回池子
  ─ 下一条请求接着用

从此再也不会出现 "预留了 2048,实际用 50" 的浪费——只按实际 token 数申请 block。

前缀共享:Copy-on-Write

同一个 system prompt 被 N 个用户共用是常见场景。PagedAttention 允许多条序列的 logical block 指向同一个 physical block:

system_prompt = "你是一位客服助手..."  (长 80 token,占 5 个 block)

用户请求 A、B、C 各自有不同的问题:
  序列A:  block_table = [P, P, P, P, P, A1, A2]
  序列B:  block_table = [P, P, P, P, P, B1]
  序列C:  block_table = [P, P, P, P, P, C1, C2, C3]
                         ─ 共享 ─

物理显存里 P0-P4 这 5 个 block 只占一份,3 条请求共用。
各自的 A/B/C block 才是私有的。

Block Size 选多少

block_size 默认 16,可以调。两头的 trade-off:

block_size	优点	缺点
小(8)	内部碎片小,短序列省	block_table 大,查表开销高
中(16)	平衡点,默认	—
大(32)	查表快,CUDA kernel 友好	短序列浪费多

启动参数改:

vllm serve meta-llama/Llama-3-8B-Instruct \
  --block-size 16

GPU 上 KV cache 池子有多大

vLLM 启动时会根据 --gpu-memory-utilization(默认 0.9)预先把剩余显存切成 block 池。用公式估算:

KV_cache_per_block
  = block_size × 层数 × 头数 × 头维度 × 2(K+V) × dtype_size

Llama-3-8B (32 层, 32 头, 128 头维度, bf16):
  16 × 32 × 32 × 128 × 2 × 2
  = 8 MB / block

A100-80GB,模型占 16GB,KV 池 ≈ 56GB:
  可用 block 数 = 56*1024 / 8 ≈ 7168

每 block 装 16 token → 可同时缓存约 11 万 token。

vLLM 日志里会打印 GPU KV cache size: XXX tokens,这是你实际能支撑的并发总 token。

Swap 到 CPU:兜底溢出

block 池满了怎么办?vLLM 会把低优先级序列整体 swap 到 CPU 内存(--swap-space 4 默认 4GB):

1. 新请求要 block,池子没空——触发 preemption
2. 挑最低优先级的那条序列,所有 block 拷贝到 CPU swap 区
3. 释放这些 block 给新请求用
4. 被 swap 出去的序列排队,等有空再换回来继续 decode

比直接 OOM 或拒绝请求优雅得多,但 swap 回来要重算 attention——频繁 swap 会拖慢吞吐,生产环境看到日志里常报 swap 说明 pool 偏小,该加 GPU 或减 --max-num-seqs。

PagedAttention CUDA Kernel

传统 attention 假设 K / V 是连续的大 tensor,kernel 直接 gemm。PagedAttention kernel 多一步:按 block_table 把散落的 physical block gather 到寄存器再算。

# 伪代码,体会一下
def paged_attn(q, block_tables, block_pool_k, block_pool_v):
    out = []
    for seq_idx in range(batch):
        blocks = block_tables[seq_idx]        # [P0, P1, P2]
        K = gather(block_pool_k, blocks) # 拼回连续 K
        V = gather(block_pool_v, blocks)
        out.append(softmax(q[seq_idx] @ K.T) @ V)
    return stack(out)

真正的 kernel 用 warp-level 并行把 gather + gemm 融合,开销几乎可以忽略——这就是 vLLM 能做到"又省又快"的关键。

和其他方案的对比

方案	KV 管理	显存利用率
HF Transformers	预留 max_seq_len 连续显存	20-40%
TGI v1	padding 对齐 batch	40-60%
vLLM PagedAttention	block 分页 + 共享	96%+
SGLang RadixAttention	前缀树 + 自动共享	96%+ 且前缀重用率更高

本章小结