Prefix Caching学习笔记

用户1318486753946
0 阅读8分钟

在部署大模型推理服务时，Prefix Caching（前缀缓存） 是近年工程上非常实用的一类优化：它把「多条请求里相同的前缀」对应的中间结果缓存起来，避免重复算一遍。

一、基础概念与背景

1.1 什么是 Prefix Caching？

定义：在自回归大语言模型推理中，对输入序列前缀部分的 Key/Value（KV）表示进行跨请求复用的技术。新请求若与历史请求（或同一会话内前文）共享同一段 token 前缀，可直接加载已算好的 KV，而不必对该前缀再做一遍完整的前向计算。

核心目标可以概括为三点：

  1. 减少重复计算：相同前缀只算一次（或少量增量）。
  2. 降低首 token 延迟（TTFT）：前缀越长、命中率越高，收益越明显。
  3. 提升吞吐：同样算力下能服务更多 QPS，或在固定 SLA 下承载更大并发。

1.2 为什么需要 Prefix Caching？

LLM 推理的成本结构

推理大致分为两阶段（后面第二部分会细讲）：

  • Prefill：一次性处理整段 prompt，计算量大（近似与 prompt 长度成超线性关系，取决于实现与并行策略），主要吃 算力。
  • Decode：逐 token 生成，每步算子规模小但步数多，往往更受 内存带宽 与 KV 显存占用 制约。

因此，长 prompt、多轮对话、RAG、带长 system 的 Agent 等场景里，Prefill 与 KV 存储都是明显瓶颈；若能跨请求复用前缀 KV，相当于把「已经付过账」的那一段直接从缓存里「取出来」。

典型场景里的「重复前缀」
场景重复内容通常是什么
多轮对话system、工具说明、历史前文（若拼接方式固定）
RAG检索到的文档块、模板化的 query 包装
Few-shot固定的若干示例 + 格式说明
批量评测 / 数据标注同一指令模板 + 不同样本后缀

只要多条请求的 token 序列从开头起有一段完全一致，这段就适合作为「可缓存前缀」的候选。

传统 KV Cache 的局限性

单请求推理里，KV Cache 只在本条请求内有效：Prefill 时按层、按位置写入 KV；Decode 时继续追加。请求结束后，显存中的 KV 通常被释放。

因此，另一条新请求若带着相同的前缀再来一遍，系统会重新做 Prefill，无法自动「继承」上一条的 KV。Prefix Caching 要做的，就是在多请求 / 多会话维度上，把「前缀 → KV 块」做成可查找、可复用、可淘汰的共享资源。

二、技术原理解析

2.1 KV Cache 机制回顾

Transformer 注意力里的 K/V

在每一层、每个注意力头中，对位置 (i) 的 hidden state 会线性映射得到 Query QiQ_i、Key KiK_i、Value ViV_i。自注意力需要当前步的 (Q) 与此前所有位置的 (K,V) 做点积与加权。若每步都把历史再算一遍，复杂度会爆炸。

KV Cache 的做法是：对已经处理过的位置，把各层的 (K,V) 存下来；后续步只算当前 token 的 (Q,K,V)，再与历史 (K,V) 拼接参与注意力。这样 Decode 每步的计算量与上下文长度相关，而不是每步重算整段历史。

Prefill 与 Decode 中的 KV
flowchart LR
  subgraph Prefill["Prefill（整段 prompt）"]
    P1["逐 token / 分块前向"]
    P2["写入各层 KV"]
  end
  subgraph Decode["Decode（逐 token 生成）"]
    D1["只算新 token 的 QKV"]
    D2["KV 追加到 cache"]
  end
  Prefill --> Decode
  • Prefill：输入整段 prompt，生成整段前缀的 KV（可并行度高，算力密集）。
  • Decode：每生成一个新 token，KV 长度 +1；计算形态以「读大量历史 KV + 写当前步」为主。
KV 的内存占用（直觉量级）

粗略地，单层单 token 的 KV 体量与 hidden_size × num_heads（或 head 维拆分方式） 相关；总显存约随 2 × 层数 × 序列长度 × 每 token KV 大小 增长（系数依张量并行、量化等变化）。因此长上下文 + 高并发下，KV 既是算力问题也是显存预算问题；前缀缓存还要在「复用收益」与「缓存占用」之间做权衡。

Prefix Caching 核心原理

前缀匹配：为何常用「块 + 哈希」

工程实现里很少按「单个 token」做细粒度一致性校验（开销大），多采用 PagedAttention 风格的固定大小块（block/page）：

  • 将序列按块切分，例如每块 16 个 token。
  • 对每个块的内容（或块内 token id 序列）算 hash，在全局表里查找是否已有相同块的 KV。

命中条件：从序列开头起，连续多个块的 hash 均命中，则这些块对应的 GPU 上 KV 可被复用；第一个未命中块之后的内容仍需正常 Prefill。

flowchart LR
  R["新请求 token 序列"]
  R --> M{"块级 hash 前缀匹配"}
  M -->|连续命中| H["加载已缓存 KV 块"]
  M -->|某块未命中| N["从该块起做 Prefill"]
  H --> N
  N --> G["进入 Decode 并更新缓存元数据"]

缓存粒度：Token 级 vs Block 级

粒度优点缺点
Token 级理论上复用最细元数据与查找开销大，碎片严重
Block 级与分页 KV 一致，易与 allocator 结合仅块内全同才能命中；块边界可能浪费少量 token

主流推理框架（如 vLLM）与 Prefix Caching 结合时，块级是自然选择。

命中后的执行路径（概念上）

  1. 请求到达，调度器拿到 token id 序列。
  2. 自前缀起按块查表：命中则记录「物理块 → 逻辑前缀」映射；未命中处开始算子执行 Prefill。
  3. GPU 上为未命中后缀分配新块，必要时把命中块的 KV 引用到本请求的 block table（见下文引用计数）。
  4. TTFT 主要受益于跳过的 Prefill 计算量；后续 Decode 与无缓存时一致（仍要读整段 KV）。

2.2 数据结构与算法

Radix Tree（基数树）

当需要支持动态插入、按前缀查找、合并公共前缀时，基数树很合适：每条从根到叶的路径对应一段 token（或块）序列；共享前缀对应共享路径。插入新序列时，可能在某节点分裂；淘汰或过期时可能剪枝。

与「纯 hash 块链」相比，Radix 更利于表达层次化的前缀关系和可变长度的共享边界（具体实现因框架而异）。

flowchart LR
  root["根"]
  root --> A["块: You are..."]
  A --> B["块: 工具定义..."]
  A --> C["块: 另一分支..."]
  B --> L1["...会话 1 后缀"]
  C --> L2["...会话 2 后缀"]
LRU（或近似 LRU）淘汰

缓存容量有限时，需要驱逐策略。常见做法：

  • 维护块的最后访问时间或访问频次；
  • LRU：最近最少使用的块先被淘汰；
  • 实际系统可能还有引用计数：正在被某进行中的请求引用的块不可释放。
块分配与引用计数

一致性要点：

  • 多个请求共享同一物理 KV 块时，用 ref_count 表示引用数；请求结束 ref_count 减一。
  • refcount 归零且 LRU 需要空间时，块可回收到空闲池。
  • 新写入的块在 Prefill/Decode 过程中逐步填充，并加到全局「块 hash → 物理块」的映射上。

这样可以在安全复用与显存回收之间取得平衡。

三、主流框架实现对比

vLLM：Automatic Prefix Caching

vLLM 的 KV 以 PagedAttention 的块为单位管理。开启 Automatic Prefix Caching 后，会在块粒度上对前缀做 hash 与全局缓存映射。

实现要点（与上文概念对应）：

  • 块池 / 空闲队列：物理块在空闲与已缓存之间流转。
  • 缓存块映射：内容 hash → 物理块（命中时把该块链入当前请求的 block table）。
  • 与连续批处理、抢占等机制叠加时，需要保证同一物理块在并发请求间生命周期正确。

配置：命令行或 API 中启用前缀缓存（具体参数名随版本演进，常见为 --enable-prefix-caching 或配置项 enable_prefix_caching）。升级版本时建议查阅对应 release 的说明，确认默认值与已知限制（例如与某些 speculative 特性、部分并行模式的组合）。

适合理解的重点：vLLM 的路线是 「PagedAttention 块 + hash 匹配」，与 allocator 强绑定，工程上落地快、与现有 KV 分页一致。

SGLang：RadixAttention

SGLang 在前缀复用上强调 RadixAttention：用 Radix Tree 组织已出现过的前缀，节点与 KV 块对应，随请求动态插入、分裂、合并。

特点：

  • 多会话共享：适合大量短请求、长 system、树状分叉的对话与 Agent 图。
  • 节点演化：新 prompt 插入时可能分裂已有节点以区分后缀；淘汰时合并或删除子树以省显存。
  • 自洽性采样 / 多路复用：同一 prompt 多分支采样时，公共前缀在树中只存一份，各分支共享根到分叉点的 KV，降低重复 Prefill。

与 vLLM 的对比直觉：

维度vLLM Automatic Prefix CachingSGLang RadixAttention
组织方式块 hash + 映射表为主Radix 树 + 块
强项与 PagedAttention 一体、易集成批调度分叉多、多路采样、前缀树状演化
学习曲线相对贴近「分页 + 哈希表」需理解树操作与节点生命周期

四、实验

4.1 基准测试

  • 实验设置：不同提示词长度、并发请求数。
  • 性能指标：TTFT、吞吐量、缓存命中率。
  • 结果分析：Prefix Caching的收益边界。

小结

  • Prefix Caching 解决的是：跨请求复用相同前缀的 KV，从而省 Prefill、降 TTFT、提吞吐。
  • 工程上多为 块级 hash 匹配 + 全局块池与淘汰（LRU + refcount）；高阶实现用 Radix Tree 管理分叉与前缀共享。
  • vLLM 与 SGLang 分别代表了「分页块缓存」与「基数树注意力」两条主流落地路径，选型时可结合业务前缀重复形态（线性长前缀 vs 多分叉）与框架生态评估。

若你正在调某一版 vLLM/SGLang，建议把 块大小、缓存容量、并发下命中率 与 TTFT/P99 放在同一 dashboard 上看；Prefix Caching 的收益高度依赖负载是否真有稳定前缀，需要先有数据再决定是否默认开启。

avatar
文章
阅读
粉丝