【infra之路】10-PagedAttention 与 KV Cache 管理

学习目标

理解为什么 KV Cache 是 LLM 推理的核心数据结构，以及 PagedAttention 如何通过虚拟内存思想解决显存碎片化问题，让 vLLM 的吞吐量比 HuggingFace 高 2-4 倍。

1. 为什么需要 KV Cache？

1.1 Decode 阶段的重复计算问题

回顾上一课，Decode 阶段每生成一个 token 都要做一次 Self-Attention：

第 t 步：输入 token_t Q_t = W_q · x_t # 当前 token 的 Query K_t = W_k · x_t # 当前 token 的 Key V_t = W_v · x_t # 当前 token 的 Value # Attention 计算：Q_t 和所有历史 token 的 K, V 做点积 scores = Q_t · [K_1, K_2, ..., K_t]^T / √d # [1, t] attn = softmax(scores) output = attn · [V_1, V_2, ..., V_t] # [1, d]

如果没有 KV Cache，每生成一个新 token 都要重新计算所有历史 token 的 K 和 V。生成长度为 N 的序列，总计算量是 O(N² · d)，和 Prefill 一样慢。

1.2 KV Cache 的作用

把每层算过的 K 和 V 缓存下来，下次直接用：

# 伪代码kv_cache={layer_idx:{"K":[],"V":[]}forlayer_idxinrange(num_layers)}defdecode_step(token_t,step):x=embedding(token_t)# [1, d]forlayerinrange(num_layers):x_norm=rmsnorm(x)Q,K,V=project(x_norm)# 各 [1, d]# 追加到 cachekv_cache[layer]["K"].append(K)# 现在是 [t, d]kv_cache[layer]["V"].append(V)# 现在是 [t, d]# Attention：Q 和整个 cache 做点积scores=Q @ kv_cache[layer]["K"].T/sqrt(d)attn=softmax(scores)attn_out=attn @ kv_cache[layer]["V"]x=x+attn_out# 残差x=x+ffn(rmsnorm(x))# FFN + 残差logits=lm_head(rmsnorm(x))# [1, vocab_size]returnsample(logits)

这样每步只需计算新 token 的 Q, K, V，历史 token 的 K, V 直接从 cache 读取。计算量从 O(N² · d) 降到 O(N · d)，Decode 阶段变成了 memory-bound（瓶颈在于读取 KV Cache 的显存带宽）。

2. KV Cache 的显存问题

2.1 KV Cache 有多大？

以 LLaMA-7B 为例，计算一个请求的 KV Cache 大小：

模型参数: num_layers = 32 hidden_size = 4096 num_kv_heads = 32（MHA）或 8（GQA） head_dim = 128 每个 token 的 KV Cache（每层）: K: num_kv_heads × head_dim × 2 bytes (FP16) = 32 × 128 × 2 = 8 KB V: 同上 = 8 KB 每层每 token: 16 KB 整个模型的 KV Cache（每 token）: 32 层 × 16 KB = 512 KB/token 对于 seq_len = 2048 的请求: 512 KB × 2048 = 1 GB

如果是 GQA（8 个 KV head），KV Cache 缩小 4 倍，变成 256 MB。但如果 batch_size = 32，仍然是 8 GB。

2.2 显存碎片化：预分配的浪费

朴素做法是为每个请求预分配最大长度的 KV Cache：

# 假设 max_seq_len = 2048forrequestinbatch:request.kv_cache=torch.zeros(num_layers,2,max_seq_len,num_kv_heads,head_dim)

问题：大多数请求的实际长度远小于 max_seq_len。比如一个请求只生成了 100 个 token，但你分配了 2048 的空间，浪费了 95%。

更糟的是，不同请求的长度差异很大（有的 50 token，有的 2000 token），预分配要么浪费显存（按最大长度分配），要么限制 batch size（按平均长度分配但无法处理长请求）。

这就是显存碎片化问题——大量显存被预分配但未使用，导致 batch size 上不去，GPU 利用率低。

3. PagedAttention：虚拟内存思想

3.1 核心思想

vLLM 论文的关键洞察：KV Cache 的管理和操作系统管理虚拟内存极其相似。

不再预分配连续的 KV Cache，而是把显存分成固定大小的Block，每个 Block 存储固定数量的 token 的 KV。通过Block Table记录每个请求的 KV 分布在哪些 Block 中。

3.2 Block 结构

GPU 显存布局: Block 0: [token_0_KV, token_1_KV, ..., token_15_KV] ← 16 个 token 的 KV Block 1: [token_16_KV, ..., token_31_KV] Block 2: [空闲] Block 3: [token_0_KV, ..., token_15_KV] ← 另一个请求的 KV ... Block N: [空闲] 每个 Block 大小固定，比如 block_size = 16 tokens

Block 在物理显存中可以不连续，通过 Block Table 映射逻辑位置。

3.3 Block Table

每个请求维护一个 Block Table，记录它的 KV 分布在哪些物理 Block 中：

请求 A（已生成 50 个 token，block_size=16）: 逻辑 Block 0 (token 0-15) → 物理 Block 3 逻辑 Block 1 (token 16-31) → 物理 Block 7 逻辑 Block 2 (token 32-47) → 物理 Block 1 逻辑 Block 3 (token 48-50) → 物理 Block 12（部分填充） Block Table: [3, 7, 1, 12]

当请求需要更多空间时，只需分配一个新的空闲 Block，追加到 Block Table 中：

请求 A 生成第 51 个 token: 逻辑 Block 3 已满（16 个 token） 分配物理 Block 5 Block Table: [3, 7, 1, 12, 5]

3.4 Attention 计算时的 Block 访问

Decode 阶段计算 Attention 时，根据 Block Table 访问 KV Cache：

defpaged_attention(Q_new,block_table,kv_cache_blocks):scores=[]forblock_idxinblock_table:block=kv_cache_blocks[block_idx]# 读取物理 BlockK_block=block.K# [block_size, d]scores_block=Q_new @ K_block.T/sqrt(d)scores.append(scores_block)scores=concat(scores)# [total_tokens]attn=softmax(scores)output=0offset=0forblock_idxinblock_table:block=kv_cache_blocks[block_idx]V_block=block.V# [block_size, d]output+=attn[offset:offset+block_size]@ V_block offset+=block_sizereturnoutput

从 GPU kernel 的角度看，这只是一次普通的 Attention 计算，只不过 K 和 V 的内存地址不连续——kernel 通过 Block Table 找到每个 Block 的物理地址，依次读取。现代 GPU 的显存访问对非连续地址的容忍度很高（L2 cache 会帮忙），所以性能损失很小。

4. PagedAttention 的收益

4.1 显存利用率提升

显存浪费减少意味着同样的 GPU 显存可以容纳更大的 batch size。batch size 是 Decode 阶段吞吐量的直接决定因素（GPU 并行处理多个请求的 Attention）。

4.2 吞吐量提升

vLLM 论文的实验结果：相比 HuggingFace Transformers（预分配方式），PagedAttention 在相同 GPU 上实现了2-4 倍的吞吐量提升。

核心原因不是 Attention 计算更快了，而是batch size 更大了。显存利用率从 20-40% 提升到 96%+，同样的显存可以塞进更多请求，GPU 的计算单元被充分利用。

4.3 灵活的请求调度

Block Table 还带来了一些额外好处：

Copy-on-Write 共享前缀：多个请求如果共享相同的 prompt 前缀（比如 system prompt），可以让它们的 Block Table 指向同一组物理 Block，不需要复制 KV Cache。这就是 SGLang 的 RadixAttention 的基础。

Beam Search 共享：Beam search 的多个候选序列共享相同的 KV Cache 前缀，只需复制 Block Table 而不是实际的 KV 数据。

Preemption 和 Swapping：如果显存不够，可以把低优先级请求的 Block 换出到 CPU 内存，需要时再换回来——和 OS 的 swap 机制完全一样。

5. Block 大小的选择

Block size 是一个 trade-off：

vLLM 默认使用block_size = 16，这是在实验中发现的较优平衡点。

每个 Block 的显存占用（LLaMA-7B, FP16）：

block_size = 16 tokens 每层每 Block: 16 × 16 KB = 256 KB 32 层: 32 × 256 KB = 8 MB/Block 80GB 显存的 A100： 模型权重: ~14 GB (FP16) 可用: ~66 GB Block 数量: 66 GB / 8 MB ≈ 8400 个 Block 可存储 token 总数: 8400 × 16 ≈ 134K tokens

6. vLLM 架构概览

理解了 PagedAttention，再来看 vLLM 的整体架构就清晰了：

┌─────────────────┐ │ Scheduler │ 决定哪些请求可以执行 │ (调度器) │ 基于可用 Block 数量 └────────┬────────┘ │ ▼ ┌─────────────────┐ │ Block Manager │ 管理物理 Block 分配 │ (Block 表管理) │ 维护 Block Table └────────┬────────┘ │ ▼ ┌───────────────────┼───────────────────┐ │ │ │ ┌────▼────┐ ┌────▼────┐ ┌────▼────┐ │ Worker 0│ │ Worker 1│ │ Worker 2│ Tensor Parallel │ (GPU 0) │ │ (GPU 1) │ │ (GPU 2) │ └─────────┘ └─────────┘ └─────────┘ │ │ │ └───────────┬───────┴───────────────────┘ │ PagedAttention Kernel (根据 Block Table 访问 KV Cache)

Scheduler维护一个等待队列，根据当前可用的物理 Block 数量决定哪些请求可以进入 batch。如果显存快满了，可以抢占（preempt）低优先级请求，把它的 Block 换出。

Block Manager负责分配和回收物理 Block，维护每个请求的 Block Table。

Worker是实际执行推理的 GPU 进程，多 Worker 之间通过 Tensor Parallelism 协作。

7. 与 Static Batching 的对比

Static Batching（传统方式）

Batch: [请求A(100 tokens), 请求B(500 tokens), 请求C(1000 tokens)] 问题：请求 A 生成完后必须等 B 和 C 都生成完，才能开始新请求 GPU 在等待期间算力浪费

PagedAttention + Continuous Batching

时刻 1: Batch = [A, B, C] 时刻 2: A 完成，立即加入 D → Batch = [B, C, D] 时刻 3: B 完成，立即加入 E → Batch = [C, D, E] ... 请求完成即退出，新请求随时加入 GPU 始终满载

PagedAttention 让 Continuous Batching 成为可能——因为新请求只需分配几个 Block，不需要预分配一大块连续显存。

8. 面试高频问题

Q1: PagedAttention 的 Block Table 存在哪里？

Block Table 本身很小（每个请求几十个 int），存在 GPU 显存中。Attention kernel 启动时作为参数传入。

Q2: 非连续的 KV Cache 会不会影响性能？

影响很小。现代 GPU 的 L2 cache 很大（A100 有 40MB），可以缓存热点 Block。而且 Attention kernel 对非连续访问做了优化（比如 FlashDecoding 的设计）。实测性能损失 < 5%。

Q3: PagedAttention 和 FlashAttention 是什么关系？

两者解决不同问题。FlashAttention 优化单次 Attention 的计算（通过 tiling 减少 HBM 访问），PagedAttention 优化 KV Cache 的显存管理。它们可以组合使用——vLLM 的 Attention kernel 内部就用类 FlashAttention 的 tiling 策略。

Q4: 为什么不用 CUDA 的 unified memory？

CUDA unified memory 可以在 GPU 和 CPU 之间自动迁移数据，但粒度太粗（4KB page），管理开销大，而且无法实现 Block 级别的精细调度。PagedAttention 是专门为 KV Cache 设计的虚拟内存方案，更高效。

总结