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

学习目标理解为什么 KV Cache 是 LLM 推理的核心数据结构以及 PagedAttention 如何通过虚拟内存思想解决显存碎片化问题让 vLLM 的吞吐量比 HuggingFace 高 2-4 倍。1. 为什么需要 KV Cache1.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):xembedding(token_t)# [1, d]forlayerinrange(num_layers):x_normrmsnorm(x)Q,K,Vproject(x_norm)# 各 [1, d]# 追加到 cachekv_cache[layer][K].append(K)# 现在是 [t, d]kv_cache[layer][V].append(V)# 现在是 [t, d]# AttentionQ 和整个 cache 做点积scoresQ kv_cache[layer][K].T/sqrt(d)attnsoftmax(scores)attn_outattn kv_cache[layer][V]xxattn_out# 残差xxffn(rmsnorm(x))# FFN 残差logitslm_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 32MHA或 8GQA 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如果是 GQA8 个 KV headKV Cache 缩小 4 倍变成 256 MB。但如果 batch_size 32仍然是 8 GB。2.2 显存碎片化预分配的浪费朴素做法是为每个请求预分配最大长度的 KV Cache# 假设 max_seq_len 2048forrequestinbatch:request.kv_cachetorch.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虚拟内存页PageKV Cache Block页表Page TableBlock 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 tokensBlock 在物理显存中可以不连续通过 Block Table 映射逻辑位置。3.3 Block Table每个请求维护一个 Block Table记录它的 KV 分布在哪些物理 Block 中请求 A已生成 50 个 tokenblock_size16: 逻辑 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 Cachedefpaged_attention(Q_new,block_table,kv_cache_blocks):scores[]forblock_idxinblock_table:blockkv_cache_blocks[block_idx]# 读取物理 BlockK_blockblock.K# [block_size, d]scores_blockQ_new K_block.T/sqrt(d)scores.append(scores_block)scoresconcat(scores)# [total_tokens]attnsoftmax(scores)output0offset0forblock_idxinblock_table:blockkv_cache_blocks[block_idx]V_blockblock.V# [block_size, d]outputattn[offset:offsetblock_size] V_block offsetblock_sizereturnoutput从 GPU kernel 的角度看这只是一次普通的 Attention 计算只不过 K 和 V 的内存地址不连续——kernel 通过 Block Table 找到每个 Block 的物理地址依次读取。现代 GPU 的显存访问对非连续地址的容忍度很高L2 cache 会帮忙所以性能损失很小。4. PagedAttention 的收益4.1 显存利用率提升方案显存浪费原因预分配 max_seq_len60-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-offBlock Size优点缺点大如 64Block Table 短管理开销小最后一个 Block 浪费多小如 8显存浪费少Block Table 长管理开销大vLLM 默认使用block_size 16这是在实验中发现的较优平衡点。每个 Block 的显存占用LLaMA-7B, FP16block_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 tokens6. 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 memoryCUDA 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 BatchingPagedAttention 让请求随时加入/退出成为可能