Transformer 大模型架构深度解析（7）KV Cache 与 PD 分离

目录

Decode-only 推理过程

Decode-only Transformer 是一个自回归模型（Autoregressive model） —— 假定我们输入一个文本，模型就会输出一个长度为 N 的回答。那么，实际上模型过程中执行了 N 次推理过程，且一次推理只输出一个 token。当前轮输出的 token 会与之前输入的 tokens 拼接在一起，并作为下一轮的输入 tokens，这样不断重复直至遇见 EOS 终止符或生成的 token 数目达到设置的 max_new_token 为止停下。

Prefill 和 Decode 阶段

具体而言，自回归模型的推理过程会被划分为 2 个阶段：

1. Prefill（预填充阶段），又称为 Prompt phase（提示词处理阶段）：模型接受到用户输入的 Prompt（Is tomato a fruit?），根据输入 Tokens（Is, tomato, a, fruit, ?） 生成第一个输出 Token（Yes）。
2. Decode（增量解码阶段），又称为 Token-generation phase（token 生成阶段）：从生成第一个输出 Token（Yes） 之后开始，把 Prompt 以及已生成的输出 tokens 组成新的模型输入，采用自回归方式一次生成一个 Token，直到结束。
   

可见，自回归模型的生成模式是存在 2 个阶段的原因，其中的 KV cache 是主要的优化方向。

KV Cache

存在 KV Cache（KV 矩阵缓存）的原因是：Decode 阶段自回归生成时，历史 token 的 Key/Value 实际上是可以复用的，因此还可以省掉大量重复的计算。因此 KV 矩阵有必要缓存起来，实际上是用显存空间换来了计算和延迟的节省。

具体而言，KV Cache 只应用于推理场景中的 Decode 阶段的 MMHA（掩码注意力层）计算中，只要目的是加速 Q×K^T×V 的两次矩阵相乘操作。通过缓存 K 和 V 矩阵的值，模型可以减少计算量，从而提高推理速度。但同时，KV cache 需要占用额外的显存空间。

下图中 Decode 的红框就是 KV cache 生效的地方。

KV Cache 的基本原理

下面是一个具体的例子：

第一次推理: 输入=[BOS]新；输出=年 第二次推理: 输入=[BOS]新年；输出=大 第三次推理: 输入=[BOS]新年大；输出=吉 第四次推理: 输入=[BOS]新年大吉；输出=[EOS]

1. 输入 “新”，输出 “年"。
   

2. 将 “年” 拼接到 “新” 的后面作为新的输入，即本次推理的输入为 “新年”，预测得到 “快”。
   

3. 将 “快” 拼接到 “新年” 的后面作为新的输入，即本次推理的输入为 “新年快”，预测得到 “乐”。
   

上述例子可见，Decode 每轮生成输出 token 时，都可以复用前面轮次 token 的 KV 矩阵。如果没有 KV cache 的话，那么每轮都要重新算一次所有历史 token 的 KV 矩阵，复杂度为 O(n^2)。这显然是没必要的。

并且，实际上一次前向计算中的多个位置都存在 KV 的冗余计算，如下图所示。

1. embedding 操作
2. KV 矩阵生成操作
3. Q×K^T 操作
4. softmax 的注意力得分与 V 相乘操作

KV Cache 的生成流程

1. 输入 prompt “新年快”。此时会把 “新年快” 的 KV 矩阵计算出来，存储在 KV Cache 中。然后输出 token “乐”。
   

2. 下一轮，只需要计算 “乐” 的 KV 矩阵值，前面轮次的 “新年快” KV 矩阵直接重 KV Cache 获得，然后和新的 KV “拼接” 到一起。然后输出 token “万”。
   

下图为流程总结。

注意，KV cache 只是只是适用于 Decode-only 此类的 casual mask model（因果模型），即：每一个 token 的输出只依赖于它自己以及之前的输入，与之后的输入无关。因为casual mask 可以让每一层历史 token的 KV 不会被新 token 改变，这才是能够做 KV Cache 的根本原因。而 BERT 等类 Encoder 模型不满足这一性质。

KV Cache 的数学原理

如下图，通过缓存 KV 矩阵然后再 “拼接” 新生成矩阵的结果，和正常输入全序列的结果是等价的。

KV Cache 的容量计算

KV Cache 是空间换时间的一种实践，所以需要了解 KV Cache 到底占用了多少显存空间。

单层 KV Cache 存储空间 = 2 × B × S × H × D × P × N

• 2：代表 Key/Value 两个向量，每层都需存储。
• B：代表 batch size。
• S：代表 total sequence length（输入序列+输出序列）。

• H：代表 number of head。
• D：代表 hidden size of head，每个 head 的维度。
• P：代表 KV 的数据格式，比如 FP16 为 2Byte。
• N：带宽 MMHA 的层数

假设：总上下文 100K，60 层，8 个头，128 的嵌入维度，使用 BF16 存储，则 KV Cache 大小为：

可见，KV Cache 由输入输出序列决定，和模型权重参数没有关系。随着推理的 batch size、sequencxe length 变大变长，那么 KV Cache 占用的存储空间很可能超过模型本身。所以变长数据的 KV cache 存储是一个比较大的难题。

KV Cache 的计算量计算

由于 KV Cache 的存在，Mask Attention 计算的时候从矩阵乘法降为矩阵乘向量，计算量的比较如下：

• No KV Cache：24 b s h 2 + 4 b s 2 h 24bsh^2 + 4bs^2h24bsh2+4bs2h
  

• KV Cache：24 b h 2 + 4 b s h 24bh^2 + 4bsh24bh2+4bsh
  

可见，对于单次运算，KV cache 模型下的计算量减少了 s 倍。

PD 分离

具体而言，Prefill 和 Decode 的任务特性存在较大的差异，Preffill 强关联 TTFT、Decode 强关联 TPOT。因此将 P、D 分开部署在不同的设备上，一方面消除了二者之间的干扰。另一方面也更有利于差异优化，使得两个阶段都能更快地达到各自的 SLO（Service Level Objective，服务等级目标）。

• Prefill 阶段：并行计算 Prompt 的所有输入 token，然后生成第一个输出 token，同时生成用于后续 Decode 的 KV Cache。因为输入序列的长度很可以长，所以并行计算开销很大，很小的 batch size 就可以把 GPU 打满，进而导致 TTFP 增大。所以 Prefill 阶段是计算密集型的任务。

• Decode 阶段（开启 KV Cache）：使用先前的 KV Cache 以自回归方式逐步生成新的 token。每个自回归步仅生成一个 token，所以计算负载较小。但是 Decode 需要反复读取 KV Cache，故而显存 I/O 开销很大。所以 Decode 是访存密集型任务，并且因为 “内存墙” 的原因，所以 Decode 阶段的 GPU 算力（FLOPs）利用率较低。

值得注意的是，PD 分离也存在一定的劣势：

• 多加载了模型副本（耗显存）
• 涉及到 GPU 间的 KV Cache 传输（耗时间）。

因此，需要结合实际应用场景来进行选用。

PD 分离架构

如上图所示，推理框架的演进经过了 3 个时期，现在 PD 分离已经成为了主流。

• 阶段 1：No KV Cache 时期
• 阶段 2：KV Cache 时期
• 阶段 3：PD 分离时期

如上图所示，PD 分离的处理流程主要包括 3 个阶段：

1. Prefill：当一个请求进入系统后，首先发送到一个 Prefill Instance 执行 Prefill 运算，得到第一个 token 以及 prompt 对应的 KV Cache。
2. Migrate：将 Prefill 处理后的 token、KV Cache 迁移至 Decode Instance；
3. Decode： Decode Instance 对迁移的请求执行 Decode 运算。

如上图所示，KV Cache 的存储和传输就成为了 PD 分离架构的关键之一。

• Prefill 初始化 KV Cache 存储；
• Decode 检索和更新 KV Cache 存储。

这个 KV Cache 存储可能是 P 或 D 设备上的存储，也可能是一个外挂的第三方专用存储设备。

推理性能指标

评价模型推理过程的性能指标主要有：

1. TTFT（Time to first token，首 token 时延，单位：时间）：是用户发出请求后，模型输出第一个 token 所需的时间，它主要受 Prefill 阶段计算和缓存加载影响，决定了响应的 “开口速度”；
2. TPOT（Time per output token，每 token 生成时延，单位：时间）：是模型在开始生成后，平均每个 token 的输出耗时，它反映了 Decode 阶段的效率，决定了生成过程的 “语速流畅度”。
3. TPS（输出 token 吞吐量，单位：tok/s）
   

PD 分离差异化配置

Prefill 和 Decode 的特性很不同，所以参数配置上也有很多的不同。例如：批策略，量化类型，TP 大小，PP 大小，任务超时时间，重试时间，显存申请策略，RDMA 软硬件队列大小，是否可回退执行等。

注 1：为了充分发挥 Decode 的能力，应该配置更多的 Prefill 数量，和更少的 Decode 数量，以此来提供更多的 token 到 Decode 进行连续批处理。
注 2：W4A8 表示权重使用 4-bit 量化，而激活值保留较高的 16-bit 精度。激活值更接近原始精度，通常用于 FP16 或 BF16。

KV Cache 传输技术 —— Mooncake

Mooncake 是一个以 KV Cache 存储为中心的推理引擎组件，主要负责 KV Cache 的存储和传输。在 SGLang 推理引擎中，Mooncake 作为一种 KV connector backend 实现。

上图是 Mooncake 的软件架构图，包括：

• KVCache-centric Conductor：根据当前的 KV Cache 分布和工作负载分派请求。
• Prefill Pool：处理用户输入的 Prompt，目标在 Prefill 阶段也尽可能多的重用 KV Cache，以避免冗余计算，进而优化 TTFT。
• Decoding Pool：自回归式的流式输出，影响 TBT（time-between-tokens）。
• KV Cache Pool：利用每台 GPU 服务器上的主存、SSD 磁盘等存储空间组成一个 KVCache Pool 来进行全局的 Prefix Cache。进而，全局 Prefix Cache 通过全局的调度能够大幅度提升复用率从而提升总吞吐。

Prefill Pool

Prefill Pool 实现了以下关键技术：

1. CPP（Chunked pipeline parallelism）：通过 Chunk 分块流水线并行机制来扩展单个请求的处理跨多个节点。并且以流水线的模式结合 KV Cache 的异构传输来支持 Prefill 和 Decode 流程的 overlap。
   

2. Layer-Wise Prefill：分层完成 Prefill。每个 MMHA 层的注意力计算开始之前，模型会等待该层的 KV Cache 的异步加载完成，并触发下一层的异步 KV Cache 加载。在注意力计算完成后，会启动该层 KV Cache 的异步存储。
   

3. Multi-Node Prefill：分块流水线并行，对于每个请求，其输入标记被分成不同的 Chunk 块，每个块不超过预填充块。同一请求的不同块可以由不同节点同时处理，从而实现处理的并行化并减少 TTFT。

KV Cache Pool

使用 KV Cache Pool 后，在 P 和 D 之间增加了一个中间存储，Prefill 节点先将 KV Cache 写到中间存储，Decode 节点从中间存储读。

KV Cache Pool 实现了以下技术：

1. Chunk 块管理：KV Cache 以 Chunk 为单位进行组织，每个 Chunk 块都附有一个哈希值，该哈希值由其自身的哈希值和 Prefix 决定。
2. 支持多会话的共享缓存：全局范围内避免重复存储和计算。
3. 支持分层存储管理：高频 KV Cache 驻留 CPU 主存，低频数据下沉至 SSD 磁盘。
4. 支持 RDMA 传输：KV Cache 在 CPU 和 GPU 之间的传输由一个支持 GPUDirect / RDMA 的 Messenger 模块处理。
5. 缓存驱逐算法：如 LRU（最近最少使用）、LFU（最不频繁使用）或基于请求特征的缓存淘汰算法，算会会将 KV Cache 转入 CPU 内存。

工作流程

1. 对于每个新请求，Conductor 根据请求的 seq_len 和 prefix_len 估计相应的执行时间。（因实例而异）
2. 将请求的估计等待时间相加，以获取该实例上的预计 TTFT。
3. Conductor 将请求分配给 TTFT 最短的实例，并相应地更新该实例的缓存和队列时间。
4. 至此，Conductor 选择了一对 Prefill 节点和一个 Decode 节点。
5. Conductor 根据 Prefix Cache 的 Chunk ID 将 Prefix Cache 从远程 CPU 主存加载到 GPU 显存，以启动请求（KV Cache 重用）。
6. Prefill 节点（或组）使用 Prefix Cache 完成 Prefill 阶段，并将新生成的增量 KV Cache 存储回 CPU 主存。如果未命中缓存的输入 token 数量超过了一定阈值（prefill_chunk），则将 Prefill 阶段分成多个 Chunk 并以流水线方式执行。
7. Messanger 将每个模型层生成的 KV Cache 流式传输到目标 Decode 节点的 CPU 主存，此 KV Cache 的传输是异步执行的，并与上述增量预填充步骤重叠，这样可以减少等待时间。
8. 当 Decode 节点的 CPU 主存中接收到所有的 KV Cache 后，Mooncake 会以连续批处理的方式加入下一批请求。
9. Conductor 根据 Decode 节点当前的负载预先选择 Decode 节点，以保证 TBT SLO 要求。

SGLang 集成

SGLang 的 KV connector 都有四个角色：

1. KVManager：是 KV connector 的管理器，Prefill 和 Decode 都有。负责初始化内存、BootstrapServer 以及发送 KV Cache。
2. KVBootstrap / MooncakeKVBootstrapServer：只用于 Prefill，记录 Prefill 和 Decode 交互所用的信息，可以有很多种，Decode 会请求该信息和 Prefill 进行连接。
3. KVSender：专属于 Prefill，用于发送 KV Cache。
4. KVReceiver：专属于 Decode，用于和 Prefill 握手，获取 BootstrapServer 的交互信息，也用于接收 KV Cache。