【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": []} for layer_idx in range(num_layers)}

def decode_step(token_t, step):
    x = embedding(token_t)                    # [1, d]
    for layer in range(num_layers):
        x_norm = rmsnorm(x)
        Q, K, V = project(x_norm)             # 各 [1, d]
        
        # 追加到 cache
        kv_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]
    return sample(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 = 2048
for request in batch:
    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 的管理和操作系统管理虚拟内存极其相似。

操作系统	PagedAttention
虚拟内存页（Page）	KV Cache Block
页表（Page Table）	Block Table
按需分页（Demand Paging）	按需分配 Block
物理内存	GPU 显存

不再预分配连续的 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：

def paged_attention(Q_new, block_table, kv_cache_blocks):
    scores = []
    for block_idx in block_table:
        block = kv_cache_blocks[block_idx]    # 读取物理 Block
        K_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 = 0
    offset = 0
    for block_idx in block_table:
        block = kv_cache_blocks[block_idx]
        V_block = block.V                     # [block_size, d]
        output += attn[offset:offset+block_size] @ V_block
        offset += block_size
    
    return output

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

---

4. PagedAttention 的收益

4.1 显存利用率提升

方案	显存浪费	原因
预分配 max_seq_len	60-80%	大多数请求远未达到最大长度
PagedAttention	< 4%	按需分配，只有最后一个 Block 可能有浪费

显存浪费减少意味着同样的 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：

Block Size	优点	缺点
大（如 64）	Block Table 短，管理开销小	最后一个 Block 浪费多
小（如 8）	显存浪费少	Block Table 长，管理开销大

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 设计的虚拟内存方案，更高效。

---

总结

概念	要点
KV Cache	缓存历史 token 的 K, V，避免 Decode 阶段重复计算
显存碎片化	预分配 max_seq_len 导致 60-80% 显存浪费
PagedAttention	虚拟内存思想，固定大小的 Block + Block Table 映射
核心收益	显存浪费 < 4%，batch size 增大 2-4 倍，吞吐量提升 2-4 倍
Block 管理	按需分配，最后一个 Block 可能部分浪费
与 Continuous Batching	PagedAttention 让请求随时加入/退出成为可能

---